Add Lecture 2 MNIST code and checkpoint; place guidance materials in Lecture 3

lecture_2/README.md

@@ -0,0 +1,115 @@
+# Lecture 2 · MNIST Flow Matching
+
+This lecture contains one Python script:
+[flow_matching_unet_lecture.py](flow_matching_unet_lecture.py).
+It covers MNIST downloading and data loading, a simple U-Net class, conditional
+flow-matching training, and Euler sampling. Digit labels are ignored during
+training, so generation is unconditional.
+
+## Train and generate
+
+After installing the [repository requirements](../requirements.txt), run this
+command from the repository root:
+
+```bash
+python lecture_2/flow_matching_unet_lecture.py
+```
+
+The script downloads the 60,000-image MNIST training split into `data/`,
+normalizes pixel values to the range [-1, 1], trains for 20 epochs, and saves
+these files in `flow_matching_outputs/`:
+
+| File | Contents |
+| --- | --- |
+| `flow_unet_mnist.pt` | Learned model weights, loss history, channel width, and training seed |
+| `samples.png` | A grid of 16 generated images |
+| `trajectory.png` | Four samples progressing from noise to images |
+
+Paths are relative to the working directory. Change the constants in section 02
+to adjust epochs, batch size, sampling steps, or output paths. CUDA, Apple MPS,
+and CPU are selected automatically according to availability.
+
+## Saved checkpoint
+
+The bundled [flow_unet_mnist.pt](https://huggingface.co/ChatterjeeLab/CIS6270/resolve/main/lecture_2/flow_unet_mnist.pt?download=true)
+was trained using the script's model, training function, and sampler.
+
+| Setting | Saved run |
+| --- | --- |
+| Training data | All 60,000 MNIST training images |
+| Completed epochs | 5 |
+| Optimizer updates | 2,345 |
+| Batch size | 128 |
+| Optimizer | AdamW, learning rate 0.0002, weight decay 0 |
+| U-Net widths | 32, 64, 128 channels |
+| Trainable parameters | 471,265 |
+| Training seed | 7 |
+| Sampling | 100 forward Euler steps |
+
+The saved run used CPU execution with PyTorch 2.14.0 and TorchVision 0.29.0.
+The checkpoint contains `model`, `losses`, `base_channels`, `seed`, and
+`epochs_completed`. It has five completed epochs; the script's default for a
+new training run is twenty. Optimizer state is not included.
+
+## Generate from the checkpoint
+
+In Python or a notebook, with the repository root as the working directory:
+
+```python
+from pathlib import Path
+import torch
+from torchvision.utils import save_image
+from lecture_2.flow_matching_unet_lecture import DEVICE, FlowUNet, sample_images
+
+# Read the supplied weights and reconstruct the same U-Net.
+checkpoint = torch.load(
+    "lecture_2/flow_unet_mnist.pt", map_location="cpu", weights_only=True,
+)
+model = FlowUNet(base_channels=checkpoint["base_channels"]).to(DEVICE)
+model.load_state_dict(checkpoint["model"])
+
+# Generate new images directly from Gaussian noise.
+torch.manual_seed(8)
+samples, trajectory = sample_images(model, num_samples=16, steps=100)
+
+# Convert the generated values to display pixels and save the results.
+output = Path("flow_matching_outputs")
+output.mkdir(exist_ok=True)
+save_image(((samples.cpu() + 1) / 2).clamp(0, 1),
+           output / "checkpoint_samples.png", nrow=4)
+display_path = ((trajectory[:4] + 1) / 2).clamp(0, 1)
+save_image(display_path.flatten(0, 1), output / "checkpoint_trajectory.png",
+           nrow=trajectory.shape[1])
+```
+
+Importing the module does not start training or download MNIST. The images are
+generated entirely from the saved weights and freshly sampled Gaussian noise.
+Clipping occurs only when preparing the PNGs; intermediate ODE states remain
+unclipped.
+
+## Selected generated examples
+
+These 16 images were manually selected from 512 outputs of the five-epoch
+checkpoint. They illustrate the clearest generated digits in that batch.
+The sampling example above generates a fresh, unselected batch.
+
+![Selected MNIST generations](mnist_selected_examples.png)
+
+The following trajectories show the first four selected images at times
+0, 0.25, 0.5, 0.75, and 1:
+
+![Selected noise-to-image trajectories](mnist_selected_trajectories.png)
+
+## Script sections
+
+| Sections | Lecture material |
+| --- | --- |
+| 01–03 | Imports, configuration, and MNIST data loading |
+| 04 | Noise, interpolation time, intermediate image, and target velocity |
+| 05–07 | Convolution blocks, the U-Net, and its forward pass |
+| 08 | Loss, backpropagation, and optimizer updates |
+| 09 | Sampling by integrating the learned velocity |
+| 10–13 | Complete training run, checkpoint saving, and image export |
+
+Continue with [Lecture 3](../lecture_3/README.md) for flow matching, diffusion,
+and guidance using ESM-2 residue embeddings.

lecture_2/flow_matching_unet_lecture.py

@@ -0,0 +1,175 @@
+"""CIS 6270, Lecture 2: unconditional MNIST generation with flow matching.
+
+Install: python -m pip install torch torchvision
+Run: python flow_matching_unet_lecture.py
+Notation: x0 = prior noise, x1 = data, xt = state, ut = target, v_pred = prediction.
+"""
+
+# %% 01. Imports
+from pathlib import Path  # Dataset and output paths.
+
+import torch  # Tensors, gradients, and optimization.
+import torch.nn as nn  # Layers and model classes.
+import torch.nn.functional as F  # Resize feature maps in the decoder.
+from torch.utils.data import DataLoader  # Shuffle and batch real images.
+from torchvision import datasets
+from torchvision.transforms import v2  # Convert and normalize image pixels.
+from torchvision.utils import save_image  # Export tensors as PNG grids.
+
+
+# %% 02. Configuration
+BATCH_SIZE = 128  # Images per training batch.
+EPOCHS = 20  # Passes through the training set.
+LEARNING_RATE = 2e-4  # Optimizer step size.
+BASE_CHANNELS = 32  # U-Net widths: 32, 64, 128.
+SAMPLE_STEPS = 100  # Euler steps from t=0 to t=1.
+NUM_SAMPLES = 16  # Generate a 4-by-4 image grid.
+SEED = 7  # Seed for initialization and training.
+DATA_DIR = Path("data")  # MNIST download location.
+OUTPUT_DIR = Path("flow_matching_outputs")  # Checkpoint and generated image folder.
+
+DEVICE = torch.device(
+    "cuda" if torch.cuda.is_available()  # NVIDIA GPU.
+    else "mps" if torch.backends.mps.is_available()  # Apple GPU.
+    else "cpu"  # CPU fallback.
+)
+
+
+# %% 03. Load real images
+def make_loader():
+    transform = v2.Compose([
+        v2.ToImage(),  # Channel-first image: [1, 28, 28].
+        v2.ToDtype(torch.float32, scale=True),  # Integer pixels 0-255 -> floats 0-1.
+        v2.Normalize(mean=(0.5,), std=(0.5,)),  # Real data pixels: [0, 1] -> [-1, 1].
+    ])
+    dataset = datasets.MNIST(
+        root=DATA_DIR, train=True, download=True, transform=transform,  # Use the training split.
+    )
+    return DataLoader(
+        dataset, batch_size=BATCH_SIZE, shuffle=True,  # Batch shape: [B, 1, 28, 28].
+        num_workers=0, pin_memory=(DEVICE.type == "cuda"),  # Load in the main process.
+    )
+
+
+# %% 04. Sample the conditional path
+def sample_conditional_path(x1):  # x1 contains normalized real images.
+    x0 = torch.randn_like(x1)  # Independent Gaussian prior; same shape as x1.
+    t = torch.rand(x1.shape[0], device=x1.device)  # One random time per image; [B].
+    t_image = t[:, None, None, None]  # [B, 1, 1, 1]; share time across pixels.
+    xt = (1.0 - t_image) * x0 + t_image * x1  # Noise at t=0; real image at t=1.
+    ut = x1 - x0  # Derivative of the straight interpolation.
+    return xt, t, ut  # Network input, time, and target velocity.
+
+
+# %% 05. Define a simple convolution block
+def conv_block(in_channels, out_channels):
+    return nn.Sequential(
+        nn.Conv2d(in_channels, out_channels, 3, padding=1), nn.SiLU(),  # Keep H and W.
+        nn.Conv2d(out_channels, out_channels, 3, padding=1), nn.SiLU(),  # Refine features.
+    )
+
+
+# %% 06. Define the small U-Net
+class FlowUNet(nn.Module):
+    def __init__(self, base_channels=BASE_CHANNELS):
+        super().__init__()
+        c = base_channels  # Base width; c=32 by default.
+        self.encoder1 = conv_block(2, c)  # Image + time -> c channels.
+        self.encoder2 = conv_block(c, 2 * c)  # Encoder: c -> 2c.
+        self.middle = conv_block(2 * c, 4 * c)  # Bottleneck: 2c -> 4c.
+        self.decoder2 = conv_block(6 * c, 2 * c)  # Upsampled 4c + skip 2c -> 2c.
+        self.decoder1 = conv_block(3 * c, c)  # Upsampled 2c + skip c -> c.
+        self.output = nn.Conv2d(c, 1, 1)  # One signed velocity per pixel.
+        self.pool = nn.MaxPool2d(2)  # Halve the image height and width.
+
+    # %% 07. Compute the velocity field
+    def forward(self, xt, t):
+        t_image = t[:, None, None, None].expand_as(xt)  # Constant time channel.
+        x = torch.cat([xt, t_image], dim=1)  # [B, 2, 28, 28].
+        skip1 = self.encoder1(x)  # Keep full-resolution features.
+        skip2 = self.encoder2(self.pool(skip1))  # 28x28 -> 14x14.
+        x = self.middle(self.pool(skip2))  # 14x14 -> 7x7.
+        x = F.interpolate(x, size=skip2.shape[-2:], mode="nearest")  # 7x7 -> 14x14.
+        x = self.decoder2(torch.cat([x, skip2], dim=1))  # Restore encoder detail.
+        x = F.interpolate(x, size=skip1.shape[-2:], mode="nearest")  # 14x14 -> 28x28.
+        x = self.decoder1(torch.cat([x, skip1], dim=1))  # Restore fine detail.
+        v_pred = self.output(x)  # Predict v_theta(xt, t).
+        return v_pred  # [B, 1, 28, 28], matching xt.
+
+
+# %% 08. Train for one epoch
+def train_one_epoch(model, loader, optimizer):
+    model.train()  # Select training mode.
+    total_loss, num_images = 0.0, 0  # Accumulate an image-weighted epoch loss.
+    for x1, _ in loader:  # Ignore digit labels: unconditional generation.
+        x1 = x1.to(DEVICE)  # Move real images to the model device.
+        xt, t, ut = sample_conditional_path(x1)  # Fresh noise and times for this batch.
+        v_pred = model(xt, t)  # The model receives only xt and t.
+        squared_error = (v_pred - ut).square()  # Velocity error at every pixel.
+        loss = squared_error.flatten(1).sum(1).mean()  # Sum pixels, average images.
+        optimizer.zero_grad(set_to_none=True)  # Clear old gradients.
+        loss.backward()  # Compute gradients with respect to weights.
+        optimizer.step()  # Update the U-Net weights.
+        total_loss += loss.item() * x1.shape[0]  # Weight by the actual batch size.
+        num_images += x1.shape[0]  # Include the smaller final batch.
+    return total_loss / num_images  # Average training loss per image.
+
+# Save snapshots every quarter of the default 100-step run; keep ODE states unclipped.
+# %% 09. Generate images with Euler integration
+@torch.no_grad()  # Sampling needs no gradient tracking.
+def sample_images(model, num_samples=NUM_SAMPLES, steps=SAMPLE_STEPS):
+    if num_samples < 1 or steps < 1:
+        raise ValueError("num_samples and steps must be positive")
+    model.eval()  # Select evaluation mode.
+    device = next(model.parameters()).device  # Use the model device.
+    xt = torch.randn(num_samples, 1, 28, 28, device=device)  # Initial x0.
+    dt = 1.0 / steps  # Positive time increment.
+    snapshots = [xt.cpu().clone()]  # Keep the initial noise for plotting.
+    for step in range(steps):  # Integrate forward in time.
+        t = torch.full((num_samples,), step * dt, device=device)  # Shared time grid.
+        v_pred = model(xt, t)  # Reevaluate velocity at the current state.
+        xt = xt + dt * v_pred  # Euler: state += time increment * velocity.
+        if (step + 1) % max(steps // 4, 1) == 0 or step + 1 == steps:
+            snapshots.append(xt.cpu().clone())  # Save intermediate and final states.
+    return xt, torch.stack(snapshots, dim=1)  # Raw images; [B, snapshots, 1, 28, 28].
+
+
+# %% 10. Set up the complete run
+def main():
+    torch.manual_seed(SEED)  # Seed training randomness.
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)  # Create the output folder.
+    loader = make_loader()  # Download, normalize, and batch MNIST.
+    model = FlowUNet().to(DEVICE)  # Create the velocity network.
+    optimizer = torch.optim.AdamW(
+        model.parameters(), lr=LEARNING_RATE, weight_decay=0.0,  # Optimize the flow loss.
+    )
+    num_parameters = sum(p.numel() for p in model.parameters())  # Count weights and biases.
+    print(f"device={DEVICE} parameters={num_parameters:,}")
+
+    # %% 11. Train and save the learned parameters
+    losses = []  # Record the training loss after each epoch.
+    for epoch in range(1, EPOCHS + 1):  # Complete EPOCHS passes through the data.
+        loss = train_one_epoch(model, loader, optimizer)  # Run the training loop.
+        losses.append(loss)  # Keep the loss history.
+        print(f"epoch={epoch:02d} squared_norm_loss={loss:.3f}")
+    torch.save({
+        "model": model.state_dict(), "losses": losses,  # Weights and training curve.
+        "base_channels": BASE_CHANNELS,  # Width needed to reconstruct the model.
+        "seed": SEED,  # Record the training seed.
+    }, OUTPUT_DIR / "flow_unet_mnist.pt")
+
+    # %% 12. Sample and save images
+    torch.manual_seed(SEED + 1)  # Fix the starting noise for generation.
+    samples, trajectory = sample_images(model)  # Integrate the learned ODE.
+    display_samples = ((samples.cpu() + 1.0) / 2.0).clamp(0.0, 1.0)  # Clip for display only.
+    save_image(display_samples, OUTPUT_DIR / "samples.png", nrow=4)  # Final images.
+    display_path = ((trajectory[:4] + 1.0) / 2.0).clamp(0.0, 1.0)  # First four sample paths.
+    save_image(
+        display_path.flatten(0, 1), OUTPUT_DIR / "trajectory.png",  # One row per sample.
+        nrow=trajectory.shape[1],  # Time increases across columns.
+    )
+
+
+# %% 13. Run the script
+if __name__ == "__main__":  # Execute when launched directly.
+    main()  # Load data, train, sample, and save.

lecture_2/flow_unet_mnist.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:20373a7d33489c0b596a4bd2f54750795570eda7dd0bef1f064b51a1367bd2d7
+size 1893377

lecture_3/GUIDANCE_NOTES.md

@@ -0,0 +1,156 @@
+# ESM-2 flow matching and diffusion guidance
+
+Two independent teaching scripts, each with encoding, normalization, a small
+conditional generator, a time-conditioned reward predictor, training, three
+guidance examples, and constrained sequence decoding. Neither imports the other.
+
+## Run
+
+Install Python 3.10+ and the dependencies in your own environment:
+
+```sh
+pip install torch transformers==4.57.6
+python esm2_flow_guidance.py --epochs 200 --samples 8
+python esm2_diffusion_guidance.py --epochs 200 --samples 8
+```
+
+Keep both scripts and `esm2_example.csv` together. Their default paths are
+relative to the scripts, so the commands also work with absolute script paths.
+The first run downloads the public `facebook/esm2_t6_8M_UR50D` checkpoint into
+`.esm2_cache`. CUDA is used when available; otherwise the scripts use CPU.
+These examples were executed with PyTorch 2.9.1 and Transformers 4.57.6.
+
+## Motivation
+
+We will generate protein-sequence representations with a flow and a diffusion
+model, then guide generation toward specified properties. Using frozen ESM-2
+residue embeddings, we can compare classifier-free conditioning with
+single-objective and weighted multi-objective reward steering in the same latent
+space. After generation, we decode the latents into amino acids and enforce an
+explicit residue-count constraint. This separates learning the sequence
+distribution, expressing preferences, and guaranteeing a discrete output rule.
+
+## Where the rewards come from
+
+The bundled dataset contains 64 synthetic sequences of length 24. These are
+teaching examples, not natural proteins or experimentally validated peptides.
+
+For sequence s of length L, `composition_proxies()` computes:
+
+- r1 = (number of K and R minus number of D and E) / L.
+  This is a side-chain charge-count proxy, not a complete pH-dependent charge model.
+- r2 = number of residues in `DEHKNQRST` / L.
+  This is a defined polar/charged composition fraction, not measured solubility.
+- The bundled class label c is 1 when r1 > 0, and 0 otherwise.
+
+CSV columns are `sequence,c,r1,r2`. Both classes 0 and 1 must be present. If r1
+and r2 are omitted, the scripts compute the composition proxies directly.
+Supplied r1/r2 columns override the proxy calculation, allowing measured
+objectives or external predictor outputs. Orient each objective so larger is
+better before supplying it. For an undesirable quantity, negate it first.
+
+All sequences in one run must have the same length, between 1 and 128, using
+the 20 canonical amino acids. The default hard minimum is 12, so shorter
+sequences require a smaller `--min-polar` value.
+
+## Normalization and reward prediction
+
+- Preserve one ESM-2 embedding per residue, giving [N, L, 320]. Remove BOS/EOS.
+- Standardize each latent feature using its mean and standard deviation over
+  training sequences and residue positions.
+- Standardize each property separately: r_tilde = (r - mean) / std.
+- Save these statistics. Do not refit them on generated, validation, or test data.
+- Train a small predictor on intermediate latents and their clean-sequence
+  standardized property labels. Its two outputs estimate expected endpoint
+  rewards from the current state and time.
+
+The examples use all 64 sequences for training and do not claim held-out quality.
+For research, split and cluster data before fitting statistics or networks, then
+assess prediction and sequence reconstruction on held-out sequences.
+
+## Three sampling modes in each script
+
+1. CFG: `w=2`, `eta=0`. Combine unconditional and class-conditioned predictions
+   as F_uncond + w * (F_cond - F_uncond). Here w=0 is unconditional and w=1 is
+   ordinary conditional sampling. Drop the class to null index 2 in 20% of
+   training examples. CFG does not use the external reward predictor.
+2. Single objective: `w=0`, `eta=1`, `lambdas=(1,0)`.
+3. Multiple objectives: `w=0`, `eta=1`, `lambdas=(0.7,0.3)`.
+
+The weights are normalized to sum to one. Lambda controls the relative
+tradeoff in standardized property units; eta controls overall steering strength.
+Set both w and eta positive to combine CFG and reward steering. Edit the three
+calls in `main()` to change the demonstration settings.
+
+Flow training uses Z_0=noise, Z_1=data, a straight conditional path, and target
+velocity Z_1-Z_0. Sampling integrates from t=0 to t=1 with Euler steps. Reward
+steering adds kappa(t) times the reward gradient to the velocity, with the chosen
+schedule kappa(t)=4*eta*t*(1-t). This is heuristic velocity steering, not a claim
+of exact sampling from a reward-tilted density.
+
+DDPM training uses Z_0=data and predicts the Gaussian noise used to construct
+Z_k. Sampling runs k=1000 down to 1. The small noise predictor includes the
+Gaussian-reference skip sqrt(1-alpha_bar_k)*Z_k and learns an additive correction
+scaled by sqrt(alpha_bar_k). This is a parameterization of the noise predictor;
+the target and DDPM equations remain noise-prediction equations. It lets the
+small network pass through high-dimensional noise without reconstructing every
+coordinate through its narrow hidden layer.
+
+For DDPM steering, epsilon_guided = epsilon_CFG - eta *
+sqrt(1-alpha_bar_k) * gradient(R_lambda). This follows the score-to-noise
+conversion s = -epsilon / sqrt(1-alpha_bar_k). The posterior standard deviation
+used for the reverse random increment is a different quantity.
+
+In both scripts, gradient ascent uses a learned expected-reward predictor.
+It is not the exact log conditional likelihood or log exponential-reward
+expectation needed for exact conditional or reward-tilted sampling. The
+generator and reward weights are frozen during sampling; gradients are enabled
+only for the current latent. No gradients through a discrete sequence scorer
+are required. A research implementation must validate surrogate quality and
+recheck the true objectives after decoding.
+
+## One decoder at the end
+
+The identical `decode()` function appears in both files to keep each script
+standalone. Both samplers use this same final decoding operation.
+
+1. Undo latent standardization.
+2. Apply ESM-2's frozen language-model head and retain the 20 amino-acid logits.
+3. Take the highest-logit residue at each position.
+4. If fewer than M positions contain a residue in `DEHKNQRST`, replace exactly
+   the missing number using the lowest logit-cost changes to that set.
+
+For the default M=12, every decoded 24-residue output contains at least 12
+members of the selected set. `--min-polar 0` removes the constraint. The procedure
+maximizes the sum of the fixed per-position logits subject to this minimum
+count. It does not enforce the constraint along the latent trajectory, preserve
+an exact conditioned generative distribution, or guarantee physical solubility.
+
+The language-model head is a simple available decoder, not a mathematically
+exact inverse of ESM-2. Generated latents can be off the encoding manifold.
+For serious protein generation, validate decoding and consider training a
+dedicated sequence decoder. Property improvements in the surrogate need not
+survive discretization or the constrained substitutions.
+
+## Outputs and checks
+
+Each script writes `cfg.fasta`, `single.fasta`, `multi.fasta`, and `results.pt`
+to its own output directory. The tensor file contains standardized generated
+latents, both networks' state dictionaries, normalization statistics, ESM model
+name, latent shape, and constraint settings. The terminal prints example
+sequences and re-evaluated mean composition proxies.
+
+Both examples have been executed through actual ESM-2 encoding, 200 training
+epochs, all three guidance modes, and final decoding. Checks included finite
+outputs, reward input gradients, normalization, schedule indexing, and the
+hard constraint. For small test cases, constrained decoding was compared with
+exhaustive search over polar/nonpolar assignments. These checks establish code
+mechanics, not biological validity or reliable property optimization.
+
+## Primary references
+
+- ESM model: https://huggingface.co/facebook/esm2_t6_8M_UR50D
+- ESM implementation: https://huggingface.co/docs/transformers/model_doc/esm
+- Flow Matching: https://arxiv.org/abs/2210.02747
+- DDPM: https://arxiv.org/abs/2006.11239
+- Classifier-Free Diffusion Guidance: https://arxiv.org/abs/2207.12598

lecture_3/README.md

@@ -0,0 +1,310 @@
+# Lecture 3 · Continuous Generative Models
+
+These examples accompany Lecture 3 of CIS 6270. We implement flow matching and DDPM diffusion using frozen ESM-2
+residue embeddings, with shared property definitions and normalization to
+compare guidance across the two sampling processes.
+
+Each example includes data loading, latent and property normalization, a small
+conditional model, training, classifier-free guidance, single-objective reward
+steering, weighted multi-objective steering, and final sequence decoding with a
+hard residue-count constraint.
+
+## Examples
+
+| Example | What the model learns | Sampling |
+| --- | --- | --- |
+| [Flow matching](esm2_flow_guidance.py) | A velocity field between noise and clean ESM-2 latents | Euler integration from noise to data |
+| [Diffusion](esm2_diffusion_guidance.py) | The noise added to clean ESM-2 latents | A 1,000-step DDPM reverse chain |
+| [Mathematical and implementation notes](GUIDANCE_NOTES.md) | Normalization, guidance conventions, decoding, and limitations | Companion reading for both scripts |
+
+Each script contains the complete training and sampling implementation,
+including the final sequence decoder, and can run independently.
+
+## Quick start
+
+Use **Python 3.11**, or another compatible Python version at least 3.10. A fresh
+virtual environment is recommended. The dependencies are pinned to the versions
+used to test these examples: PyTorch 2.9.1 and Transformers 4.57.6.
+
+### 1. Download and install
+
+```bash
+git clone https://huggingface.co/ChatterjeeLab/CIS6270
+cd CIS6270
+
+python3 -m venv .venv
+source .venv/bin/activate
+python -m pip install --upgrade pip
+python -m pip install -r requirements.txt
+```
+
+On Windows PowerShell, use `python -m venv .venv` and activate with
+`.venv\Scripts\Activate.ps1`.
+
+### 2. Train and sample
+
+Run either example, or both:
+
+```bash
+python lecture_3/esm2_flow_guidance.py --epochs 200 --samples 8
+python lecture_3/esm2_diffusion_guidance.py --epochs 200 --samples 8
+```
+
+Each command trains a generator and a property predictor from scratch, then
+generates sequences using three guidance settings. The first run downloads the
+public [ESM-2 8M checkpoint](https://huggingface.co/facebook/esm2_t6_8M_UR50D)
+into `lecture_3/.esm2_cache/`, which subsequent runs reuse. Both the repository
+and checkpoint are publicly accessible.
+
+The scripts use CUDA when available and CPU otherwise. Apple Silicon currently
+uses the CPU path. Runtime depends on hardware and the first-run download.
+
+For a short end-to-end installation check:
+
+```bash
+python lecture_3/esm2_flow_guidance.py --epochs 2 --samples 2 --output outputs/flow_smoke
+python lecture_3/esm2_diffusion_guidance.py --epochs 2 --samples 2 --output outputs/diffusion_smoke
+```
+
+The two-epoch runs exercise encoding, training, sampling, and decoding.
+
+### 3. Find the generated sequences
+
+By default, outputs are written beside the scripts:
+
+```text
+lecture_3/
+├── esm2_flow_outputs/
+│   ├── cfg.fasta
+│   ├── single.fasta
+│   ├── multi.fasta
+│   └── results.pt
+└── esm2_diffusion_outputs/
+    ├── cfg.fasta
+    ├── single.fasta
+    ├── multi.fasta
+    └── results.pt
+```
+
+The FASTA files contain decoded sequences. `results.pt` contains generated
+standardized latents, generator and reward-predictor state dictionaries,
+normalization statistics, the ESM checkpoint name, latent dimensions, and
+constraint settings. The terminal also prints re-evaluated composition
+properties of the decoded sequences.
+
+Each run initializes new generator and reward-predictor parameters. Repeating
+a command with the same output directory replaces its saved files; choose a
+new `--output` directory to retain results from separate experiments.
+
+## Implementation
+
+1. **Encode sequences.** Use frozen ESM-2 to obtain one 320-dimensional vector
+   per residue, retaining position information and removing BOS/EOS tokens.
+2. **Normalize.** Standardize latent coordinates and each property using
+   training-set statistics.
+3. **Train.** Learn a conditional generative field and a time-conditioned
+   predictor of the clean sequence's standardized properties.
+4. **Guide generation.** Compare CFG, one reward, and a weighted reward sum.
+5. **Decode once at the end.** Undo latent normalization, apply the frozen ESM-2
+   language-model head, and enforce the selected residue-count constraint.
+6. **Re-evaluate.** Calculate the composition properties on the actual decoded
+   sequences so comparisons between guidance settings reflect the final
+   amino-acid outputs.
+
+## Sequence data and property labels
+
+For the teaching examples, we use [64 synthetic sequences of length
+24](esm2_example.csv). We calculate both property labels directly
+from residue counts and assign the conditioning class according to the first
+property:
+
+| Column | Definition | Interpretation |
+| --- | --- | --- |
+| `sequence` | A canonical amino-acid sequence | Input to frozen ESM-2 |
+| `r1` | `(count(K) + count(R) - count(D) - count(E)) / length` | A simple charge-count proxy |
+| `r2` | `count(residues in DEHKNQRST) / length` | A defined polar/charged fraction |
+| `c` | `1` if `r1 > 0`, otherwise `0` | The binary conditioning class |
+
+### Use your own data
+
+Supply a CSV with `sequence,c,r1,r2` columns:
+
+```csv
+sequence,c,r1,r2
+QYWSDSWWESQMMSPWYPMSPLSV,0,-0.08333333,0.41666667
+CKSEFQPPHLMGHDFFACEMRNFK,0,0.00000000,0.45833333
+WPYGEHMLADNNVVKKRLQQWCFI,1,0.04166667,0.41666667
+KVVGALPIESFYTAKMESIAVEVI,0,-0.04166667,0.33333333
+```
+
+```bash
+python lecture_3/esm2_flow_guidance.py --data my_sequences.csv --output outputs/my_flow
+python lecture_3/esm2_diffusion_guidance.py --data my_sequences.csv --output outputs/my_diffusion
+```
+
+- Include at least four sequences and both class labels, 0 and 1.
+- Use one fixed sequence length, at most 128 residues, and the 20 canonical
+  amino acids to match the fixed-length latent tensors in both models.
+- Supply finite property values and orient each objective so **larger is better**.
+  For a quantity to minimize, negate it before writing the CSV.
+- If both `r1` and `r2` are omitted, `composition_proxies()` calculates the
+  two example properties directly. Supplied reward columns override them.
+- Experimental measurements or scores from a separate predictor can fill the
+  reward columns. We fit a differentiable surrogate to these labels and
+  evaluate its gradients with respect to the intermediate latent during sampling.
+- The printed composition proxies always retain their defined count-based
+  meaning, even when custom reward labels are supplied. Re-evaluate custom
+  objectives with the corresponding assay or scorer after decoding.
+
+The bundled demonstration uses all 64 sequences for training. For studies with
+held-out evaluation, partition sequences by cluster before fitting normalization
+statistics or model parameters.
+
+## Property normalization and scalarization
+
+Each objective is standardized independently:
+
+```text
+r̃ₘ(s) = [rₘ(s) − μₘ] / max(σₘ, 10⁻⁶)
+```
+
+Here μₘ and σₘ are the training-set mean and standard deviation of property m.
+Standardization expresses each property in units of its training-set variation,
+so the scalarization weights specify relative preferences on a common scale.
+Reuse the saved training statistics for new examples to maintain that scale
+throughout evaluation and generation.
+
+The reward predictor estimates standardized clean-sequence properties from an
+intermediate latent and its time. During sampling, we combine its predictions:
+
+```text
+Rλ(z,t) = λ₁ r̂₁(z,t) + λ₂ r̂₂(z,t)
+λ₁ ≥ 0, λ₂ ≥ 0, λ₁ + λ₂ = 1
+```
+
+The nonnegative tradeoff weights λ are normalized to sum to one. Overall
+steering strength η is a separate parameter.
+
+## Guidance configurations
+
+Both scripts run these settings in `main()`:
+
+| Output | CFG strength `w` | Reward strength `eta` | Property weights `lambdas` |
+| --- | ---: | ---: | --- |
+| `cfg.fasta` | 2.0, class 1 | 0.0 | Unused |
+| `single.fasta` | 0.0 | 1.0 | `(1.0, 0.0)` |
+| `multi.fasta` | 0.0 | 1.0 | `(0.7, 0.3)` |
+
+**Classifier-free guidance.** During training, we replace 20% of class labels
+with null class index 2. At sampling time, we combine the conditional and
+unconditional predictions as `F_uncond + w * (F_cond - F_uncond)`.
+Thus `w=0` is unconditional, `w=1` is ordinary conditional sampling, and `w>1` amplifies the conditional
+difference. F is a velocity for flow matching and predicted noise for DDPM.
+
+**Reward steering.** Freeze both networks, enable gradients only for the
+current latent, calculate the scalarized reward, and differentiate it with
+respect to that latent. Set `(1, 0)` for the first objective alone or change
+the weights to express a tradeoff. This correction uses the property predictor,
+while CFG uses the conditional and unconditional generative predictions.
+
+To change these settings, edit the three `sample(...)` calls in `main()`.
+To combine CFG with reward steering, set both strengths:
+
+```python
+latent = sample(model, reward_model, n=8, c=1,
+                w=2.0, eta=1.0, lambdas=(0.7, 0.3))
+```
+
+Each call resets the sampling seed to compare methods using the same random
+draws for the same batch size.
+
+### Guidance in the sampling dynamics
+
+| | Flow matching | DDPM diffusion |
+| --- | --- | --- |
+| Clean-data endpoint | Z₁ | Z₀ |
+| Training target | Velocity Z₁ − Z₀ | Injected Gaussian noise ε |
+| Generation direction | t = 0 → 1 | k = 1000 → 1 |
+| Reward correction | Add κ(t)∇Rλ to velocity | Subtract η√(1−ᾱₖ)∇Rλ from predicted noise |
+| Numerical step | Euler ODE update | Reverse mean plus posterior Gaussian noise |
+
+For the flow, we choose `κ(t)=4ηt(1−t)` to taper reward steering near the noise
+and data endpoints. For DDPM, we convert the reward-gradient score correction
+into a noise-prediction correction using `s=−ε/√(1−ᾱₖ)`. Both implementations
+use heuristic gradient steering based on learned estimates of endpoint properties.
+
+The DDPM predictor includes a Gaussian-reference skip plus a learned correction
+so the small MLP can carry high-dimensional noise. The training target remains
+noise, and the sampler uses the standard DDPM reverse equations.
+
+## Constrained sequence decoding
+
+The default decoder requires **at least 12 of the 24 output residues** to
+belong to:
+
+```text
+𝒫 = {D, E, H, K, N, Q, R, S, T}
+```
+
+We begin with the highest-logit residue at each position and count the members
+of 𝒫. When the count falls below the specified minimum, we replace the required
+number of residues using the lowest-cost substitutions into 𝒫. This discrete
+decoding procedure maximizes the sum of the fixed per-position logits subject
+to the minimum-count constraint.
+
+```bash
+# Require at least 16 selected residues.
+python lecture_3/esm2_flow_guidance.py --min-polar 16
+
+# Remove the minimum-count constraint.
+python lecture_3/esm2_diffusion_guidance.py --min-polar 0
+```
+
+Set `--min-polar` between zero and the sequence length to specify the minimum
+number of selected polar or charged residues in each decoded sequence.
+
+## Verification
+
+Run the offline unit checks after installing dependencies:
+
+```bash
+python -m unittest discover -s tests -v
+```
+
+The five unit tests cover dataset annotations, scalarization weights, reward
+input gradients, DDPM schedule indexing, and constrained decoding. For the
+decoder test, we compare the selected sequences with exhaustive solutions on
+small examples. All unit tests run locally with the installed dependencies.
+
+We also checked both scripts end to end using the ESM-2 checkpoint, 200 training
+epochs, and all three guidance modes. All 48 decoded sequences in that run
+satisfied the specified minimum residue count.
+
+## Troubleshooting
+
+- **Download failure:** the first run needs internet access to the ESM-2
+  checkpoint. Cached subsequent runs can use `HF_HUB_OFFLINE=1` once all files
+  have been downloaded.
+- **Dependency conflicts:** use a clean virtual environment and the pinned
+  requirements. For a specific CUDA build, follow the official
+  [PyTorch installation instructions](https://pytorch.org/get-started/locally/).
+- **Input validation error:** check equal lengths, canonical residues, both
+  class labels, finite rewards, and a minimum count no greater than the length.
+- **Weak or repetitive samples:** inspect training convergence, evaluate
+  reconstruction through the ESM-2 head, and assess property-predictor accuracy
+  on held-out sequences before adjusting guidance strength. Compare guidance
+  settings using properties recalculated after sequence decoding.
+
+## References
+
+- [ESM-2 checkpoint](https://huggingface.co/facebook/esm2_t6_8M_UR50D)
+  and [Transformers ESM documentation](https://huggingface.co/docs/transformers/model_doc/esm).
+- [Flow Matching for Generative Modeling](https://arxiv.org/abs/2210.02747).
+- [Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239).
+- [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+
+## License
+
+Repository code is distributed under the [MIT License](https://huggingface.co/ChatterjeeLab/CIS6270/blob/main/LICENSE), matching the
+repository's license setting. ESM-2 weights are downloaded separately and remain
+subject to their original distribution terms.

lecture_3/esm2_diffusion_guidance.py

@@ -0,0 +1,264 @@
+"""Standalone ESM-2 diffusion guidance example for CIS 6270.
+Input CSV: sequence,c,r1,r2. Sequences have one fixed length; c is 0 or 1.
+Both objectives are oriented so larger is better. The bundled CSV is synthetic.
+Run: python esm2_diffusion_guidance.py --data esm2_example.csv --epochs 200
+Install: pip install torch transformers==4.57.6
+Outputs: guided residue latents, model weights, and decoded amino-acid sequences.
+"""
+import argparse
+import csv
+from pathlib import Path
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, TensorDataset
+from transformers import AutoTokenizer, EsmForMaskedLM
+
+ROOT = Path(__file__).resolve().parent
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ESM_NAME = "facebook/esm2_t6_8M_UR50D"
+AMINO_ACIDS = "ACDEFGHIKLMNPQRSTVWY"
+POLAR_RESIDUES = "DEHKNQRST"  # Operational polar/charged set; not a solubility assay.
+BATCH_SIZE, HIDDEN, LEARNING_RATE = 16, 128, 1e-3
+CONDITION_DROP = 0.2  # Drop c during training so the same model learns the null condition.
+
+
+def composition_proxies(sequences):
+    # Transparent teaching rewards; these are not measured activity or solubility.
+    return torch.tensor([
+        [(sum(a in "KR" for a in s) - sum(a in "DE" for a in s)) / len(s),
+         sum(a in POLAR_RESIDUES for a in s) / len(s)]
+        for s in sequences
+    ], dtype=torch.float32)
+
+
+# 1. Load annotated sequences and encode frozen ESM-2 residue vectors.
+@torch.no_grad()
+def load_data(path):
+    with Path(path).open(newline="") as handle:
+        rows = list(csv.DictReader(handle))
+    sequences = [row["sequence"].strip().upper() for row in rows]
+    if len(rows) < 4 or any(not s or set(s) - set(AMINO_ACIDS) for s in sequences):
+        raise ValueError("Supply at least four sequences using the 20 standard amino acids")
+    lengths = {len(s) for s in sequences}
+    if len(lengths) != 1 or max(lengths) > 128:
+        raise ValueError("This compact example requires one fixed sequence length, at most 128")
+    c = torch.tensor([int(row["c"]) for row in rows], dtype=torch.long)
+    # With sequence,c only, compute example rewards directly from residue counts.
+    # Optional r1/r2 columns override them with supplied labels, such as measurements.
+    if {"r1", "r2"}.issubset(rows[0]):
+        r = torch.tensor([[float(row["r1"]), float(row["r2"])] for row in rows])
+    else:
+        r = composition_proxies(sequences)
+    if set(c.tolist()) != {0, 1} or not torch.isfinite(r).all():
+        raise ValueError("Include both c=0 and c=1, with finite r1 and r2")
+
+    tokenizer = AutoTokenizer.from_pretrained(ESM_NAME, cache_dir=ROOT / ".esm2_cache")
+    esm = EsmForMaskedLM.from_pretrained(
+        ESM_NAME, cache_dir=ROOT / ".esm2_cache", use_safetensors=True
+    ).to(DEVICE).eval().requires_grad_(False)
+    encoded = []
+    for start in range(0, len(sequences), BATCH_SIZE):
+        tokens = tokenizer(sequences[start:start + BATCH_SIZE], return_tensors="pt")
+        tokens = {key: value.to(DEVICE) for key, value in tokens.items()}
+        hidden = esm.esm(**tokens).last_hidden_state
+        encoded.append(hidden[:, 1:-1].cpu())  # Remove BOS/EOS; retain all L residue positions.
+    z = torch.cat(encoded)                    # [N, L, D], with D=320 for this checkpoint.
+    z_mean = z.mean((0, 1), keepdim=True)
+    z_std = z.std((0, 1), correction=0, keepdim=True).clamp_min(1e-4)
+    r_mean, r_std = r.mean(0), r.std(0, correction=0).clamp_min(1e-6)
+    dataset = TensorDataset((z - z_mean) / z_std, c, (r - r_mean) / r_std)
+    stats = {"z_mean": z_mean, "z_std": z_std, "r_mean": r_mean, "r_std": r_std}
+    return dataset, esm, tokenizer, stats
+
+
+# 2. Small model classes: flattening lets each output depend on the entire sequence.
+class DiffusionModel(nn.Module):
+    def __init__(self, length, dim):
+        super().__init__()
+        self.length, self.dim = length, dim
+        self.time = nn.Sequential(nn.Linear(1, 32), nn.SiLU(), nn.Linear(32, 32))
+        self.condition = nn.Embedding(3, 16)    # Indices 0,1 are classes; 2 is the null class.
+        self.net = nn.Sequential(
+            nn.Linear(length * dim + 48, HIDDEN), nn.SiLU(),
+            nn.Linear(HIDDEN, HIDDEN), nn.SiLU(),
+            nn.Linear(HIDDEN, length * dim),
+        )
+
+    def forward(self, z, t, c):
+        time = self.time(t[:, None])
+        inputs = torch.cat([z.flatten(1), time, self.condition(c)], dim=1)
+        k = (t * K).round().long().clamp(0, K)
+        a = alpha_bars[k, None, None]
+        return (1 - a).sqrt() * z + a.sqrt() * self.net(inputs).reshape_as(z)
+        # Gaussian-reference noise prediction plus a learned correction.
+        # The skip carries all coordinates; the correction vanishes near pure noise.
+
+
+class RewardModel(nn.Module):
+    def __init__(self, length, dim):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(length * dim + 1, HIDDEN), nn.SiLU(),
+            nn.Linear(HIDDEN, HIDDEN), nn.SiLU(), nn.Linear(HIDDEN, 2),
+        )
+
+    def forward(self, z, t):
+        return self.net(torch.cat([z.flatten(1), t[:, None]], dim=1))
+        # Two predicted standardized endpoint objectives, given the current latent and time.
+
+
+# The same DDPM schedule as Section 7; k=0 denotes the clean latent.
+K = 1000
+betas = torch.cat([torch.zeros(1), torch.linspace(1e-4, 0.02, K)]).to(DEVICE)
+alphas = 1.0 - betas
+alpha_bars = alphas.cumprod(0)
+previous = torch.cat([torch.ones(1, device=DEVICE), alpha_bars[:-1]])
+posterior_variances = betas * (1 - previous) / (1 - alpha_bars).clamp_min(1e-20)
+
+
+# 3. DDPM: corrupt Z_0 at step k and regress the exact noise epsilon.
+def train(dataset, epochs):
+    _, length, dim = dataset.tensors[0].shape
+    loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
+    model = DiffusionModel(length, dim).to(DEVICE)
+    reward_model = RewardModel(length, dim).to(DEVICE)
+    optimizer = torch.optim.Adam(list(model.parameters()) + list(reward_model.parameters()), lr=LEARNING_RATE)
+    for epoch in range(epochs):
+        total = 0.0
+        for z0, c, r_tilde in loader:
+            z0, c, r_tilde = z0.to(DEVICE), c.to(DEVICE), r_tilde.to(DEVICE)
+            k = torch.randint(1, K + 1, (len(z0),), device=DEVICE)
+            t = k.float() / K                  # Normalized k is only the network's time input.
+            a = alpha_bars[k, None, None]
+            epsilon = torch.randn_like(z0)
+            zk = a.sqrt() * z0 + (1 - a).sqrt() * epsilon
+            dropped = c.masked_fill(torch.rand(len(c), device=DEVICE) < CONDITION_DROP, 2)
+            loss_diffusion = F.mse_loss(model(zk, t, dropped), epsilon)
+            loss_reward = F.mse_loss(reward_model(zk, t), r_tilde)
+            loss = loss_diffusion + loss_reward
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            optimizer.step()
+            total += loss.item()
+        if (epoch + 1) % 50 == 0 or epoch + 1 == epochs:
+            print(f"epoch {epoch+1}: combined training loss {total / len(loader):.4f}")
+    for network in (model, reward_model):
+        network.eval().requires_grad_(False)
+        for parameter in network.parameters():
+            parameter.grad = None
+    return model, reward_model
+
+
+# 4. R_lambda = sum_m lambda_m * r_tilde_m; differentiate the current latent only.
+def reward_gradient(reward_model, z, t, lambdas):
+    with torch.enable_grad():                  # Re-enable input gradients inside sampling.
+        state = z.detach().requires_grad_(True)
+        R_lambda = (reward_model(state, t) * lambdas).sum(dim=1)
+        grad = torch.autograd.grad(R_lambda.sum(), state)[0]
+    return grad.detach()                      # Do not retain graphs across sampling steps.
+
+
+def normalize_weights(lambdas):
+    values = torch.as_tensor(lambdas, dtype=torch.float32, device=DEVICE)
+    if values.shape != (2,) or not torch.isfinite(values).all() or (values < 0).any() or values.sum() <= 0:
+        raise ValueError("Use two finite, nonnegative weights with a positive sum")
+    return values / values.sum()              # lambda controls tradeoffs, eta controls strength.
+
+
+# 5. Sampling: CFG first; optional reward gradient then changes predicted noise.
+@torch.no_grad()
+def sample(model, reward_model, n=8, c=1, w=0.0, eta=0.0, lambdas=(1.0, 0.0)):
+    if c not in (0, 1) or w < 0 or eta < 0:
+        raise ValueError("Use c=0/1 and nonnegative w and eta")
+    lambdas = normalize_weights(lambdas)
+    torch.manual_seed(123)                     # Same starting and reverse noise across comparisons.
+    z = torch.randn(n, model.length, model.dim, device=DEVICE)
+    null = torch.full((n,), 2, dtype=torch.long, device=DEVICE)
+    condition = torch.full((n,), c, dtype=torch.long, device=DEVICE)
+    for k in range(K, 0, -1):
+        t = torch.full((n,), k / K, device=DEVICE)
+        eps_uncond = model(z, t, null)
+        eps = eps_uncond
+        if w != 0:
+            eps = eps_uncond + w * (model(z, t, condition) - eps_uncond)
+        sigma = (1 - alpha_bars[k]).sqrt()       # FORWARD corruption standard deviation.
+        if eta != 0:
+            grad_R = reward_gradient(reward_model, z, t, lambdas)
+            eps = eps - eta * sigma * grad_R    # s_guided=s_theta+eta*grad R; s=-epsilon/sigma.
+        mean = (z - betas[k] * eps / sigma) / alphas[k].sqrt()
+        z = mean + posterior_variances[k].sqrt() * torch.randn_like(z) if k > 1 else mean
+    return z                                   # No fresh noise at k=1.
+
+
+# 6. One decoder, used only AFTER the flow or diffusion trajectory is complete.
+@torch.no_grad()
+def decode(z, esm, tokenizer, stats, min_polar=12):
+    latent = z * stats["z_std"].to(z.device) + stats["z_mean"].to(z.device)
+    logits = esm.lm_head(latent)               # Frozen head produces one vocabulary distribution per residue.
+    aa_ids = torch.tensor(tokenizer.convert_tokens_to_ids(list(AMINO_ACIDS)), device=z.device)
+    logits = logits.index_select(-1, aa_ids)   # Restrict the vocabulary to the 20 amino acids.
+    if not 0 <= min_polar <= z.shape[1]:
+        raise ValueError("min_polar must lie between zero and the sequence length")
+    polar = torch.tensor([a in POLAR_RESIDUES for a in AMINO_ACIDS], device=z.device)
+    polar_ids = polar.nonzero().flatten()
+    best_scores, choices = logits.max(dim=-1)  # Start from unrestricted amino-acid argmax.
+    polar_scores, local = logits[..., polar_ids].max(dim=-1)
+    polar_choices = polar_ids[local]          # Best polar/charged amino acid at each position.
+    for i in range(len(z)):
+        already_polar = polar[choices[i]]
+        missing = max(0, min_polar - int(already_polar.sum()))
+        if missing:
+            cost = (best_scores[i] - polar_scores[i]).masked_fill(already_polar, float("inf"))
+            positions = cost.topk(missing, largest=False).indices
+            choices[i, positions] = polar_choices[i, positions]
+    return ["".join(AMINO_ACIDS[i] for i in row) for row in choices.cpu().tolist()]
+    # Exact maximum-logit decode subject to at least min_polar selected residues.
+    # This enforces composition, not measured solubility; no gradient through argmax.
+
+
+# 7. Run CFG, single-objective steering, and scalarized multi-objective steering in order.
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--data", type=Path, default=ROOT / "esm2_example.csv")
+    parser.add_argument("--epochs", type=int, default=200)
+    parser.add_argument("--samples", type=int, default=8)
+    parser.add_argument("--min-polar", type=int, default=12)
+    parser.add_argument("--output", type=Path, default=ROOT / "esm2_diffusion_outputs")
+    args = parser.parse_args()
+    if args.epochs < 1 or args.samples < 1:
+        parser.error("epochs and samples must be positive")
+    torch.manual_seed(7)
+    if DEVICE.type == "cpu":
+        torch.set_num_threads(2)
+    dataset, esm, tokenizer, stats = load_data(args.data)
+    if not 0 <= args.min_polar <= dataset.tensors[0].shape[1]:
+        parser.error("min-polar must lie between zero and the sequence length")
+    model, reward_model = train(dataset, args.epochs)
+    outputs = {
+        "cfg": sample(model, reward_model, args.samples, c=1, w=2.0),
+        "single": sample(model, reward_model, args.samples, eta=1.0, lambdas=(1.0, 0.0)),
+        "multi": sample(model, reward_model, args.samples, eta=1.0, lambdas=(0.7, 0.3)),
+    }
+    args.output.mkdir(parents=True, exist_ok=True)
+    for name, latent in outputs.items():
+        if not torch.isfinite(latent).all():
+            raise RuntimeError(f"Nonfinite {name} output; reduce guidance or check training")
+        sequences = decode(latent, esm, tokenizer, stats, args.min_polar)
+        assert all(sum(a in POLAR_RESIDUES for a in s) >= args.min_polar for s in sequences)
+        fasta = "".join(f">{name}_{i+1}\n{seq}\n" for i, seq in enumerate(sequences))
+        (args.output / f"{name}.fasta").write_text(fasta)
+        count = sum(a in POLAR_RESIDUES for a in sequences[0])
+        print(f"{name}: {sequences[0]}  polar/charged residues={count}")
+        print("  decoded mean composition proxies:", composition_proxies(sequences).mean(0).tolist())
+    torch.save({"standardized_latents": {k: v.cpu() for k, v in outputs.items()},
+                "model": model.state_dict(), "reward_model": reward_model.state_dict(),
+                "stats": stats, "esm_name": ESM_NAME,
+                "length": model.length, "dim": model.dim, "min_polar": args.min_polar,
+                "polar_residues": POLAR_RESIDUES}, args.output / "results.pt")
+    print(f"Saved latent tensors and FASTA sequences to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

lecture_3/esm2_example.csv

@@ -0,0 +1,65 @@
+sequence,c,r1,r2
+QYWSDSWWESQMMSPWYPMSPLSV,0,-0.08333333,0.41666667
+CKSEFQPPHLMGHDFFACEMRNFK,0,0.00000000,0.45833333
+WPYGEHMLADNNVVKKRLQQWCFI,1,0.04166667,0.41666667
+KVVGALPIESFYTAKMESIAVEVI,0,-0.04166667,0.33333333
+MEFHAYSLGPASRWSSRHGYFTNL,1,0.04166667,0.45833333
+TCYLFCLAQIVMFGRAYVGGYWAK,1,0.08333333,0.16666667
+VFSHIVACCEIETRDTCVWNHMYM,0,-0.08333333,0.41666667
+QRNQTEIRTLIYMPYGNKTYLRGS,1,0.12500000,0.54166667
+WDNQRIKTALAGHTGVDSGWMHNF,0,0.00000000,0.50000000
+PPVCRGWQMPRAALWVMRCRFPRQ,1,0.20833333,0.29166667
+HPNAKNYGETIFFGLWGNDRKNLV,1,0.04166667,0.45833333
+CMKEIYESFLDETTIWVIHMIAAR,0,-0.08333333,0.41666667
+RWDTGMWGKDCIHRPIQKKDPTWA,1,0.08333333,0.50000000
+WCCQLMAYPPAKMGHPTFDHPLPK,1,0.04166667,0.29166667
+QPVRDFVLEVGQSILESHSLKVPS,0,-0.04166667,0.50000000
+GREMRGFQKTPNATKKEGTYKPAP,1,0.16666667,0.54166667
+ENEWNGSRHYWCVRWLGTLKILTF,1,0.04166667,0.45833333
+AVSQIAPHARQQASSQENKHGETT,0,0.00000000,0.66666667
+NADTREKMPSNDLSWWLYNMFMPA,0,-0.04166667,0.45833333
+SLWFAIMPYVAKYVIYGHVEPTAT,0,0.00000000,0.25000000
+YGWHYTLLMWYDERCGFLGRGRAQ,1,0.04166667,0.33333333
+EEQEKNLDDMNSLDFSQNLSLFEY,0,-0.25000000,0.66666667
+ETWDHYPQYKCVKCKNNHCPHVAY,1,0.04166667,0.50000000
+YHIMRKVPSTNEFRPLAHNNCTLQ,1,0.08333333,0.54166667
+DFRMNQYGSKDYPKTTGAREHLFM,1,0.04166667,0.54166667
+IQWWSDVQTMQCMADEGMHHQYFV,0,-0.12500000,0.45833333
+WAQMMMVRQIDCRNFIRHTVYSNF,1,0.08333333,0.45833333
+DMTFGHSSDAILAKVYDSKNCPQL,0,-0.04166667,0.50000000
+DRSICNKIPPSQTIPFMGPFISQN,1,0.04166667,0.45833333
+QNHLSDAEIKYCIKWNVLNGKMLP,1,0.04166667,0.45833333
+RQYCPEPSWSQHTNKPACMKVEEC,0,0.00000000,0.54166667
+IWWWHVNCQCRDDHEHCPLCYAPI,0,-0.08333333,0.37500000
+SDKSVTHIMERDIWPIIAWLRDRV,0,0.00000000,0.50000000
+WSDAWDVYIWILAMNCEVGRMYYC,0,-0.08333333,0.25000000
+VAKELIRNYGDEFCWLVNSSMKED,0,-0.08333333,0.50000000
+VTWANPHINITFCDPMHIIRIQCQ,0,0.00000000,0.41666667
+GYYWVTHLVHMVDIWFSDMPIHYE,0,-0.12500000,0.33333333
+TRALRHDDVWVDYYFIWTALTYQF,0,-0.04166667,0.41666667
+CGNGSDFDTFIQMWICACGMCIDK,0,-0.08333333,0.33333333
+NGQMACTQENWAWSVFHFEWWCPK,0,-0.04166667,0.41666667
+GSRETKHEVTADCPFQWNGRALKK,1,0.08333333,0.58333333
+CHGKYTFTFGMEIRDMPAWYGHFV,0,0.00000000,0.33333333
+GSYDNPLAIKWQMVCNSSRMDWTM,0,0.00000000,0.45833333
+FRIAQIRSKIGYLQNIFVMSLLLR,1,0.16666667,0.37500000
+SCCDWRQMQDRVGPAMMKEWACGE,0,-0.04166667,0.41666667
+FFSELMKCHYHYCYAWRRPIAKDW,1,0.08333333,0.37500000
+AWGPVLDFNIQSDIFYPCRTDMGL,0,-0.08333333,0.33333333
+PHTCRRIMMMEQDLTNFPLAYCTQ,0,0.00000000,0.45833333
+DQINDGSAHENQYCQNPSDPFVWK,0,-0.12500000,0.58333333
+QMRSHFQYGCGWMITPMRFPKNMA,1,0.12500000,0.37500000
+EECYALAGDYDFKVRIQHAQCTWQ,0,-0.08333333,0.45833333
+ITKQMWVNDVWDTEVRHHSFMTLF,0,-0.04166667,0.54166667
+EHAKYVKNSNCYNKENNVDLFYCN,0,0.00000000,0.58333333
+SYYVPKTKQRLTTSRCGLKGVVKA,1,0.25000000,0.50000000
+TDAKGWHYNWHWPLHQFCHLPVHA,0,0.00000000,0.41666667
+YIVRMSDSFHSGHDHDKKNLIMYH,0,0.00000000,0.58333333
+CTHVQNPNSIETNCHEHYDFGTSM,0,-0.12500000,0.62500000
+RWWYQNHLMCYNSKDACYHPWINM,1,0.04166667,0.41666667
+RKSNGGKSLIQNAGITCVKWCYGV,1,0.16666667,0.41666667
+WWNQITTPHWCKWWASHWMQYSRW,1,0.08333333,0.45833333
+IDTICWQVCILVQESDAKDVFEGN,0,-0.16666667,0.45833333
+NFHEWVLWELTGQWRVWSRQHSKE,0,0.00000000,0.58333333
+VTYCMMAKFAGMDYEQKICVHVAC,0,0.00000000,0.29166667
+KVKTTLKVEREAPPYTRMSLSAWT,1,0.12500000,0.54166667

lecture_3/esm2_flow_guidance.py

@@ -0,0 +1,251 @@
+"""Standalone ESM-2 flow matching guidance example for CIS 6270.
+Input CSV: sequence,c,r1,r2. Sequences have one fixed length; c is 0 or 1.
+Both objectives are oriented so larger is better. The bundled CSV is synthetic.
+Run: python esm2_flow_guidance.py --data esm2_example.csv --epochs 200
+Install: pip install torch transformers==4.57.6
+Outputs: guided residue latents, model weights, and decoded amino-acid sequences.
+"""
+import argparse
+import csv
+from pathlib import Path
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, TensorDataset
+from transformers import AutoTokenizer, EsmForMaskedLM
+
+ROOT = Path(__file__).resolve().parent
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ESM_NAME = "facebook/esm2_t6_8M_UR50D"
+AMINO_ACIDS = "ACDEFGHIKLMNPQRSTVWY"
+POLAR_RESIDUES = "DEHKNQRST"  # Operational polar/charged set; not a solubility assay.
+BATCH_SIZE, HIDDEN, LEARNING_RATE = 16, 128, 1e-3
+CONDITION_DROP = 0.2  # Drop c during training so the same model learns the null condition.
+
+
+def composition_proxies(sequences):
+    # Transparent teaching rewards; these are not measured activity or solubility.
+    return torch.tensor([
+        [(sum(a in "KR" for a in s) - sum(a in "DE" for a in s)) / len(s),
+         sum(a in POLAR_RESIDUES for a in s) / len(s)]
+        for s in sequences
+    ], dtype=torch.float32)
+
+
+# 1. Load annotated sequences and encode frozen ESM-2 residue vectors.
+@torch.no_grad()
+def load_data(path):
+    with Path(path).open(newline="") as handle:
+        rows = list(csv.DictReader(handle))
+    sequences = [row["sequence"].strip().upper() for row in rows]
+    if len(rows) < 4 or any(not s or set(s) - set(AMINO_ACIDS) for s in sequences):
+        raise ValueError("Supply at least four sequences using the 20 standard amino acids")
+    lengths = {len(s) for s in sequences}
+    if len(lengths) != 1 or max(lengths) > 128:
+        raise ValueError("This compact example requires one fixed sequence length, at most 128")
+    c = torch.tensor([int(row["c"]) for row in rows], dtype=torch.long)
+    # With sequence,c only, compute example rewards directly from residue counts.
+    # Optional r1/r2 columns override them with supplied labels, such as measurements.
+    if {"r1", "r2"}.issubset(rows[0]):
+        r = torch.tensor([[float(row["r1"]), float(row["r2"])] for row in rows])
+    else:
+        r = composition_proxies(sequences)
+    if set(c.tolist()) != {0, 1} or not torch.isfinite(r).all():
+        raise ValueError("Include both c=0 and c=1, with finite r1 and r2")
+
+    tokenizer = AutoTokenizer.from_pretrained(ESM_NAME, cache_dir=ROOT / ".esm2_cache")
+    esm = EsmForMaskedLM.from_pretrained(
+        ESM_NAME, cache_dir=ROOT / ".esm2_cache", use_safetensors=True
+    ).to(DEVICE).eval().requires_grad_(False)
+    encoded = []
+    for start in range(0, len(sequences), BATCH_SIZE):
+        tokens = tokenizer(sequences[start:start + BATCH_SIZE], return_tensors="pt")
+        tokens = {key: value.to(DEVICE) for key, value in tokens.items()}
+        hidden = esm.esm(**tokens).last_hidden_state
+        encoded.append(hidden[:, 1:-1].cpu())  # Remove BOS/EOS; retain all L residue positions.
+    z = torch.cat(encoded)                    # [N, L, D], with D=320 for this checkpoint.
+    z_mean = z.mean((0, 1), keepdim=True)
+    z_std = z.std((0, 1), correction=0, keepdim=True).clamp_min(1e-4)
+    r_mean, r_std = r.mean(0), r.std(0, correction=0).clamp_min(1e-6)
+    dataset = TensorDataset((z - z_mean) / z_std, c, (r - r_mean) / r_std)
+    stats = {"z_mean": z_mean, "z_std": z_std, "r_mean": r_mean, "r_std": r_std}
+    return dataset, esm, tokenizer, stats
+
+
+# 2. Small model classes: flattening lets each output depend on the entire sequence.
+class FlowModel(nn.Module):
+    def __init__(self, length, dim):
+        super().__init__()
+        self.length, self.dim = length, dim
+        self.time = nn.Sequential(nn.Linear(1, 32), nn.SiLU(), nn.Linear(32, 32))
+        self.skip = nn.Linear(32, 1)           # Preserve full-dimensional state/noise through a time gate.
+        self.condition = nn.Embedding(3, 16)    # Indices 0,1 are classes; 2 is the null class.
+        self.net = nn.Sequential(
+            nn.Linear(length * dim + 48, HIDDEN), nn.SiLU(),
+            nn.Linear(HIDDEN, HIDDEN), nn.SiLU(),
+            nn.Linear(HIDDEN, length * dim),
+        )
+
+    def forward(self, z, t, c):
+        time = self.time(t[:, None])
+        inputs = torch.cat([z.flatten(1), time, self.condition(c)], dim=1)
+        return self.skip(time)[:, :, None] * z + self.net(inputs).reshape_as(z)
+        # A time-dependent linear part plus a learned nonlinear velocity correction.
+
+
+class RewardModel(nn.Module):
+    def __init__(self, length, dim):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(length * dim + 1, HIDDEN), nn.SiLU(),
+            nn.Linear(HIDDEN, HIDDEN), nn.SiLU(), nn.Linear(HIDDEN, 2),
+        )
+
+    def forward(self, z, t):
+        return self.net(torch.cat([z.flatten(1), t[:, None]], dim=1))
+        # Two predicted standardized endpoint objectives, given the current latent and time.
+
+
+# 3. Flow matching: Z_t=(1-t)Z_0+tZ_1; target velocity is Z_1-Z_0.
+def train(dataset, epochs):
+    _, length, dim = dataset.tensors[0].shape
+    loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
+    model = FlowModel(length, dim).to(DEVICE)
+    reward_model = RewardModel(length, dim).to(DEVICE)
+    optimizer = torch.optim.Adam(list(model.parameters()) + list(reward_model.parameters()), lr=LEARNING_RATE)
+    for epoch in range(epochs):
+        total = 0.0
+        for z1, c, r_tilde in loader:
+            z1, c, r_tilde = z1.to(DEVICE), c.to(DEVICE), r_tilde.to(DEVICE)
+            z0 = torch.randn_like(z1)          # Flow convention: Z_0=noise, Z_1=clean latent.
+            t = torch.rand(len(z1), device=DEVICE)
+            zt = (1 - t[:, None, None]) * z0 + t[:, None, None] * z1
+            dropped = c.masked_fill(torch.rand(len(c), device=DEVICE) < CONDITION_DROP, 2)
+            loss_flow = F.mse_loss(model(zt, t, dropped), z1 - z0)
+            loss_reward = F.mse_loss(reward_model(zt, t), r_tilde)
+            loss = loss_flow + loss_reward    # Independent parameter sets; one optimizer suffices.
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            optimizer.step()
+            total += loss.item()
+        if (epoch + 1) % 50 == 0 or epoch + 1 == epochs:
+            print(f"epoch {epoch+1}: combined training loss {total / len(loader):.4f}")
+    for network in (model, reward_model):
+        network.eval().requires_grad_(False)   # Freeze weights; still allow gradients of input z.
+        for parameter in network.parameters():
+            parameter.grad = None
+    return model, reward_model
+
+
+# 4. R_lambda = sum_m lambda_m * r_tilde_m; differentiate the current latent only.
+def reward_gradient(reward_model, z, t, lambdas):
+    with torch.enable_grad():                  # Re-enable input gradients inside sampling.
+        state = z.detach().requires_grad_(True)
+        R_lambda = (reward_model(state, t) * lambdas).sum(dim=1)
+        grad = torch.autograd.grad(R_lambda.sum(), state)[0]
+    return grad.detach()                      # Do not retain graphs across sampling steps.
+
+
+def normalize_weights(lambdas):
+    values = torch.as_tensor(lambdas, dtype=torch.float32, device=DEVICE)
+    if values.shape != (2,) or not torch.isfinite(values).all() or (values < 0).any() or values.sum() <= 0:
+        raise ValueError("Use two finite, nonnegative weights with a positive sum")
+    return values / values.sum()              # lambda controls tradeoffs, eta controls strength.
+
+
+# 5. Sampling: CFG first; optional reward gradient then changes the velocity.
+@torch.no_grad()
+def sample(model, reward_model, n=8, c=1, w=0.0, eta=0.0, lambdas=(1.0, 0.0), steps=200):
+    if c not in (0, 1) or steps < 1 or w < 0 or eta < 0:
+        raise ValueError("Use c=0/1, positive steps, and nonnegative w and eta")
+    lambdas = normalize_weights(lambdas)
+    torch.manual_seed(123)                     # Compare methods from the same initial noise.
+    z = torch.randn(n, model.length, model.dim, device=DEVICE)
+    null = torch.full((n,), 2, dtype=torch.long, device=DEVICE)
+    condition = torch.full((n,), c, dtype=torch.long, device=DEVICE)
+    dt = 1.0 / steps
+    for step in range(steps):
+        t = torch.full((n,), step * dt, device=DEVICE)
+        v_uncond = model(z, t, null)
+        v = v_uncond
+        if w != 0:
+            v = v_uncond + w * (model(z, t, condition) - v_uncond)
+        if eta != 0:
+            grad_R = reward_gradient(reward_model, z, t, lambdas)
+            kappa = eta * 4 * t[:, None, None] * (1 - t[:, None, None])
+            v = v + kappa * grad_R             # Chosen steering rule; not exact conditional transport.
+        z = z + dt * v                         # Integrate forward from t=0 to t=1.
+    return z
+
+
+# 6. One decoder, used only AFTER the flow or diffusion trajectory is complete.
+@torch.no_grad()
+def decode(z, esm, tokenizer, stats, min_polar=12):
+    latent = z * stats["z_std"].to(z.device) + stats["z_mean"].to(z.device)
+    logits = esm.lm_head(latent)               # Frozen head produces one vocabulary distribution per residue.
+    aa_ids = torch.tensor(tokenizer.convert_tokens_to_ids(list(AMINO_ACIDS)), device=z.device)
+    logits = logits.index_select(-1, aa_ids)   # Restrict the vocabulary to the 20 amino acids.
+    if not 0 <= min_polar <= z.shape[1]:
+        raise ValueError("min_polar must lie between zero and the sequence length")
+    polar = torch.tensor([a in POLAR_RESIDUES for a in AMINO_ACIDS], device=z.device)
+    polar_ids = polar.nonzero().flatten()
+    best_scores, choices = logits.max(dim=-1)  # Start from unrestricted amino-acid argmax.
+    polar_scores, local = logits[..., polar_ids].max(dim=-1)
+    polar_choices = polar_ids[local]          # Best polar/charged amino acid at each position.
+    for i in range(len(z)):
+        already_polar = polar[choices[i]]
+        missing = max(0, min_polar - int(already_polar.sum()))
+        if missing:
+            cost = (best_scores[i] - polar_scores[i]).masked_fill(already_polar, float("inf"))
+            positions = cost.topk(missing, largest=False).indices
+            choices[i, positions] = polar_choices[i, positions]
+    return ["".join(AMINO_ACIDS[i] for i in row) for row in choices.cpu().tolist()]
+    # Exact maximum-logit decode subject to at least min_polar selected residues.
+    # This enforces composition, not measured solubility; no gradient through argmax.
+
+
+# 7. Run CFG, single-objective steering, and scalarized multi-objective steering in order.
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--data", type=Path, default=ROOT / "esm2_example.csv")
+    parser.add_argument("--epochs", type=int, default=200)
+    parser.add_argument("--samples", type=int, default=8)
+    parser.add_argument("--min-polar", type=int, default=12)
+    parser.add_argument("--output", type=Path, default=ROOT / "esm2_flow_outputs")
+    args = parser.parse_args()
+    if args.epochs < 1 or args.samples < 1:
+        parser.error("epochs and samples must be positive")
+    torch.manual_seed(7)
+    if DEVICE.type == "cpu":
+        torch.set_num_threads(2)
+    dataset, esm, tokenizer, stats = load_data(args.data)
+    if not 0 <= args.min_polar <= dataset.tensors[0].shape[1]:
+        parser.error("min-polar must lie between zero and the sequence length")
+    model, reward_model = train(dataset, args.epochs)
+    outputs = {
+        "cfg": sample(model, reward_model, args.samples, c=1, w=2.0),
+        "single": sample(model, reward_model, args.samples, eta=1.0, lambdas=(1.0, 0.0)),
+        "multi": sample(model, reward_model, args.samples, eta=1.0, lambdas=(0.7, 0.3)),
+    }
+    args.output.mkdir(parents=True, exist_ok=True)
+    for name, latent in outputs.items():
+        if not torch.isfinite(latent).all():
+            raise RuntimeError(f"Nonfinite {name} output; reduce guidance or check training")
+        sequences = decode(latent, esm, tokenizer, stats, args.min_polar)
+        assert all(sum(a in POLAR_RESIDUES for a in s) >= args.min_polar for s in sequences)
+        fasta = "".join(f">{name}_{i+1}\n{seq}\n" for i, seq in enumerate(sequences))
+        (args.output / f"{name}.fasta").write_text(fasta)
+        count = sum(a in POLAR_RESIDUES for a in sequences[0])
+        print(f"{name}: {sequences[0]}  polar/charged residues={count}")
+        print("  decoded mean composition proxies:", composition_proxies(sequences).mean(0).tolist())
+    torch.save({"standardized_latents": {k: v.cpu() for k, v in outputs.items()},
+                "model": model.state_dict(), "reward_model": reward_model.state_dict(),
+                "stats": stats, "esm_name": ESM_NAME,
+                "length": model.length, "dim": model.dim, "min_polar": args.min_polar,
+                "polar_residues": POLAR_RESIDUES}, args.output / "results.pt")
+    print(f"Saved latent tensors and FASTA sequences to {args.output}")
+
+
+if __name__ == "__main__":
+    main()