CIS 6270 Course Code
This repository contains code examples for CIS 6270, organized by lecture. We will develop the implementations alongside the course material, which connect the mathematical definitions, training objectives, and sampling equations to the executable PyTorch code herein. Each lecture directory contains its examples, data, and implementation notes.
Lecture index
| Lecture | Topics | Code and documentation |
|---|---|---|
| 2 | Unconditional MNIST image generation with flow matching and a simple U-Net | Guide 路 Script 路 Saved checkpoint |
| 3 | Flow matching, diffusion, and guidance for ESM-2 residue embeddings | Guide 路 Flow matching 路 Diffusion |
| 4 | Discrete diffusion, masked and uniform corruption, block generation, and guidance | Guide 路 Training and generation 路 Slide code map |
| 5 | Discrete flow matching, Dirichlet and Fisher paths, Gumbel-Softmax, rectification, and multi-objective generation | Guide 路 Training and generation 路 Slide code map |
| 6 | Continuous, latent, categorical, posterior, expanding, and strong stochastic flow maps | Guide 路 Training and generation 路 Slide code map |
| 7 | Optimal transport, Sinkhorn, continuous and discrete Schr枚dinger bridges, matching, reward tilting, branching, and interacting control | Guide 路 Training and generation 路 Slide code map |
Installation
Use Python 3.11, or another compatible Python version at least 3.10, in a new
virtual environment. The shared requirements.txt pins PyTorch 2.9.1,
TorchVision 0.24.1, and Transformers 4.57.6. Lecture 2 uses PyTorch and
TorchVision; Lecture 3 also uses Transformers. Lectures 4 and 5 use PyTorch,
NumPy, and SciPy. Lecture 6 uses PyTorch and NumPy. Lecture 7 uses PyTorch, NumPy, and SciPy. These lecture folders also provide minimal requirements.
git clone https://huggingface.co/ChatterjeeLab/CIS6270
cd CIS6270
git lfs pull
python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install -r requirements.txt
On Windows PowerShell, create the environment with python -m venv .venv
and activate it with .venv\Scripts\Activate.ps1.
The saved checkpoint uses Git LFS. If Git LFS is unavailable, download
flow_unet_mnist.pt
directly and place it in lecture_2/. Training from scratch does not require
the saved checkpoint.
Lecture 2: MNIST flow matching
Lecture 2 contains one self-contained, commented script. It downloads MNIST, loads batches of normalized images, defines a small U-Net velocity network, trains it with conditional flow matching, and generates images from Gaussian noise using Euler integration.
From the repository root:
python lecture_2/flow_matching_unet_lecture.py
The default run trains for 20 epochs and writes samples.png, trajectory.png,
and a new flow_unet_mnist.pt to flow_matching_outputs/. The script selects
CUDA, Apple MPS, or CPU according to availability. Its numbered # %% sections
match the lecture walkthrough.
The bundled checkpoint has 5 completed epochs on all 60,000 MNIST training images. The Lecture 2 guide includes a short example for loading it and generating images without training, along with selected generated digits and their noise-to-image trajectories.
Lecture 3: Flow and diffusion guidance
From the repository root:
python lecture_3/esm2_flow_guidance.py --epochs 200 --samples 8
python lecture_3/esm2_diffusion_guidance.py --epochs 200 --samples 8
Both scripts train on ESM-2 residue embeddings and compare classifier-free guidance, single-objective reward steering, and scalarized multi-objective steering. We use the same property definitions, normalization, and final constrained decoder to compare the flow and diffusion implementations.
The first run downloads the public ESM-2 checkpoint. Each script then writes generated sequences and model parameters to its own output directory. The lecture guide describes the data format, training and sampling settings, property calculations, normalization, and residue-count constraint, with commands for using a custom dataset.
Lecture 4 - Discrete diffusion
Train small DNA denoisers and generate sequences with MDLM, UDLM, block diffusion, classifier-free guidance, exact and gradient-based classifier guidance, and a PepTune-style search. The guide includes each method's command and mathematical assumptions. The code map links the slide walkthroughs to their functions.
From the repository root, run the complete MDLM example.
python lecture_4/run.py --method mdlm --data lecture_4/data/dna_train.tsv --out lecture_4/outputs/mdlm
python lecture_4/run.py --method mdlm --mode sample --out lecture_4/outputs/mdlm
The script trains, saves a checkpoint, and writes generated DNA and loss logs. The bundled data are synthetic, and the guidance objectives are explicit toy properties. No pretrained model or external dataset is required.
Lecture 5 - Discrete flow matching
Start with Gat et al.'s discrete flow matching, then run Dirichlet, Fisher, Gumbel-Softmax, rectified flow, ReDi, MOG-DFM, and AReUReDi examples. Each method has a complete training and generation command in the guide.
python lecture_5/run.py --method gat --data lecture_5/data/dna_train.tsv --out lecture_5/outputs/gat
python lecture_5/run.py --method gat --mode sample --out lecture_5/outputs/gat
Both folders include numerical examples, mathematical tests, and saved results from seeded CPU runs. The guides explain finite endpoint approximations and classroom simplifications for each method.
Lecture 6 - Flow Maps
Learn finite-time motion after the local flows from Lecture 5. The new folder contains 16 complete examples covering flow-map matching and self-distillation, consistency, Shortcut, MeanFlow, learned latent representations, Flow Map Language Models, Categorical and Discrete Flow Maps, Diamond Maps, Meta Flow Maps, Expanding Flow Maps, and Strong Stochastic Flow Maps.
python lecture_6/run.py --method self-distill --out lecture_6/outputs/self-distill
python lecture_6/run.py --mode sample --out lecture_6/outputs/self-distill --sample-steps 1
python lecture_6/run_all.py --quick
The Lecture 6 guide includes all method commands, training objectives, data formats, and numerical assumptions. The slide code map links implementations to the Flow Maps presentation. Source notes identify the exact papers and inspected author-code revisions, including differences between paper pseudocode and released implementations. Verified examples contain actual training logs, generated samples, and checkpoint-reload checks.
Lecture 7 - Optimal Transport and Schr枚dinger Bridges
The final lecture includes 17 complete numerical and learned examples: OT, Sinkhorn, finite and continuous-time bridges, DSB, DSBM, SF2M, DDSBM, CSBM, TR2-D2, BranchSBM, and EntangledSBM. Each runner saves its solved distribution or learned parameters and can reload them to generate samples.
python lecture_7/run.py --method sf2m --out lecture_7/outputs/sf2m
python lecture_7/run.py --mode sample --out lecture_7/outputs/sf2m
python lecture_7/run_all.py --quick
The guide, mathematical notes, slide code map, and verified examples follow the earlier lecture structure. The synthetic demonstrations distinguish exact calculations, finite discretizations, and learned approximations. Lectures 6 and 7 also include spoken readings of their equations before the intuition.
Repository organization
| Location | Contents |
|---|---|
| Repository root | Course index, installation requirements, and license |
lecture_2/ |
One MNIST flow-matching script, guide, trained checkpoint, and selected example images |
lecture_3/ |
ESM-2 flow and diffusion guidance scripts, sequence data, guide, and mathematical notes |
lecture_4/ |
Seven discrete diffusion and guidance examples, synthetic DNA, slide code map, and verified outputs |
lecture_5/ |
Eight discrete and simplex flow examples, synthetic DNA, slide code map, and verified outputs |
lecture_6/ |
Sixteen flow-map implementations, text data, mathematical notes, slide links, and verified outputs |
lecture_7/ |
Seventeen OT and bridge examples, derivations, equation readings, slide links, tests, and verified outputs |
Installation instructions and the lecture index are maintained at the repository root. Lecture-specific commands, mathematical explanations, and references accompany the corresponding code.
Tests
python -m unittest discover -s tests -v
python -m unittest discover -s lecture_4/tests -v
python -m unittest discover -s lecture_5/tests -v
python -m unittest discover -s lecture_6/tests -v
python -m unittest discover -s lecture_7/tests -v
The Lecture 3 unit tests cover property annotations, scalarization weights, reward gradients, DDPM schedule indexing, and constrained decoding.
Lecture 4 tests check reverse KL losses and guidance calculations. Lecture 5
tests check the master equation, Fisher geometry, Gumbel path derivatives, and
MH detailed balance. Lecture 6 checks map identities and JVPs, categorical
teachers, GLASS conditioning, reward gradients, insertion clocks, and Brownian
composition. Its end-to-end runner also verifies checkpoint reloads. Lecture 7 checks OT duality, bridge marginals and conditional paths, generator identities, matching targets, and control geometry, with 17 seeded checkpoint-reload checks.
Run a short end-to-end check of every method from its
lecture folder with python run_all.py --quick.
License
The repository code is distributed under the MIT License. The Lecture 2 U-Net checkpoint is included under the same license. The Lecture 3 ESM-2 weights are downloaded from their original repository under the corresponding distribution terms.