Add Lecture 2 flow matching and diffusion guidance examples

Standalone PyTorch examples, synthetic sequence data, setup instructions, mathematical notes, and offline tests.

.gitignore

@@ -0,0 +1,20 @@
+# Local Python environments and caches.
+.venv*/
+venv/
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.DS_Store
+.esm2_cache/
+
+# Generated results and downloaded/trained weights.
+outputs/
+esm2_flow_outputs/
+esm2_diffusion_outputs/
+*.pt
+*.pth
+*.safetensors
+
+# Local secrets must never be committed.
+.env
+.env.*

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Chatterjee Lab
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,3 +1,349 @@
 ---
 license: mit
+language:
+  - en
+tags:
+  - education
+  - pytorch
+  - flow-matching
+  - diffusion
+  - esm2
+  - protein-generation
+  - guidance
 ---
+
+# CIS 6270 · Generative Modeling
+
+### Lecture 2: Continuous Generative Models
+
+**Chatterjee Lab | Executable teaching examples in PyTorch**
+
+This repository accompanies the CIS 6270 lectures. We connect the mathematics
+of continuous generative models to working code, using frozen ESM-2 protein
+representations to compare **flow matching** and **DDPM diffusion**.
+
+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.
+
+> This is a source-code repository. The included dataset is synthetic, and no
+> trained generator weights are distributed here. These examples teach the
+> mechanics of guidance; they do not establish protein activity or solubility.
+
+## Start here
+
+| Example | What the model learns | Sampling |
+| --- | --- | --- |
+| [Flow matching](lecture_02/esm2_flow_guidance.py) | A velocity field between noise and clean ESM-2 latents | Euler integration from noise to data |
+| [Diffusion](lecture_02/esm2_diffusion_guidance.py) | The noise added to clean ESM-2 latents | A 1,000-step DDPM reverse chain |
+| [Mathematical and implementation notes](lecture_02/GUIDANCE_NOTES.md) | Normalization, guidance conventions, decoding, and limitations | Companion reading for both scripts |
+
+The scripts are independent. Neither imports the other, and each contains the
+complete example in one file. The same decoder appears in both so either can
+run on its own.
+
+## 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_02/esm2_flow_guidance.py --epochs 200 --samples 8
+python lecture_02/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_02/.esm2_cache/`. Later runs reuse those downloaded files. No
+Hugging Face login is required to read this public repository or checkpoint.
+
+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_02/esm2_flow_guidance.py --epochs 2 --samples 2 --output outputs/flow_smoke
+python lecture_02/esm2_diffusion_guidance.py --epochs 2 --samples 2 --output outputs/diffusion_smoke
+```
+
+Two epochs check execution, not sample quality.
+
+### 3. Find the generated sequences
+
+By default, outputs are written beside the scripts:
+
+```text
+lecture_02/
+├── 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.
+
+Repeating a command with the same output directory replaces those output files.
+Choose a new `--output` directory to keep separate experiments. The scripts do
+not automatically resume training or reload saved generators.
+
+## What happens in each example?
+
+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. A higher surrogate prediction need not survive decoding.
+
+## Where do the rewards come from?
+
+The [example CSV](lecture_02/esm2_example.csv) contains **64 synthetic sequences
+of length 24**. The labels are transparent composition calculations:
+
+| 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 |
+
+The first reward is not a full pH-dependent charge calculation. The second is
+not an experimental solubility measurement. The synthetic sequences are
+teaching data, not validated natural proteins or functional peptides.
+
+### 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_02/esm2_flow_guidance.py --data my_sequences.csv --output outputs/my_flow
+python lecture_02/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. The compact models do not implement variable-length masking.
+- 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. The code trains a differentiable latent-space surrogate; it
+  does not differentiate through the original measurement or sequence scorer.
+- 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.
+
+For a research experiment, split and sequence-cluster data before fitting
+normalization statistics or models. The bundled demonstration uses all 64
+sequences for training and contains no held-out performance claim.
+
+## Normalize first, then express preferences
+
+Each objective is standardized independently:
+
+```text
+r̃ₘ(s) = [rₘ(s) − μₘ] / max(σₘ, 10⁻⁶)
+```
+
+Here μₘ and σₘ are the mean and standard deviation of property m in the training
+data. Normalization prevents the numerical units of one property from
+determining its importance in a weighted sum. The saved training statistics
+must also be used for new examples; do not refit them on generated samples.
+
+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.
+
+## Compare the three guidance modes
+
+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, 20% of class labels are replaced
+by null class index 2. Sampling combines 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. CFG alone does not use the reward predictor.
+
+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. Hardware and library differences can still
+change numerical results.
+
+### How the guidance enters the 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, `κ(t)=4ηt(1−t)` is a chosen steering schedule. It is not an
+automatically derived probability-flow ODE coefficient. For DDPM, the sign and
+scale follow the score-to-noise relation `s=−ε/√(1−ᾱₖ)`.
+
+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.
+
+These are **surrogate-based steering examples**. A predicted expected property
+is not the exact log likelihood or log exponential-reward expectation needed
+for exact conditional or reward-tilted sampling. Neither example establishes
+such an exact sampling guarantee.
+
+## A hard constraint at 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}
+```
+
+Starting from the highest-logit residue at each position, the decoder counts
+members of 𝒫. If more are needed, it selects the smallest logit-cost substitutions
+into that set. This exactly maximizes the sum of the fixed per-position logits
+under the minimum-count constraint. It does not require differentiation through
+the final amino-acid choices.
+
+```bash
+# Require at least 16 selected residues.
+python lecture_02/esm2_flow_guidance.py --min-polar 16
+
+# Remove the minimum-count constraint.
+python lecture_02/esm2_diffusion_guidance.py --min-polar 0
+```
+
+Choose a minimum between zero and the sequence length. A hard composition rule
+is a guarantee about the output sequence's count. It is not a guarantee of
+physical solubility, folding, stability, safety, or biological function.
+
+## Tests and limitations
+
+Run the offline unit checks after installing dependencies:
+
+```bash
+python -m unittest discover -s tests -v
+```
+
+These check dataset annotations, scalarization weights, reward input gradients,
+DDPM schedule indexing, and the constrained decoder against exhaustive solutions
+on small examples. They do not download ESM-2 or evaluate biological quality.
+
+The teaching scripts were also run end to end with the actual ESM-2 checkpoint,
+200 training epochs, and all three guidance modes. All 48 decoded sequences in
+that smoke test met the hard count constraint. Some single- and multi-objective
+outputs decoded identically, and surrogate improvements did not consistently
+translate into improved decoded properties.
+
+Important boundaries:
+
+- Tiny synthetic data and narrow MLPs are intended for instruction.
+- The frozen ESM-2 head is not an exact inverse of the encoder; generated
+  latents may be off the encoding manifold.
+- Increasing guidance can exploit surrogate errors. Validate the predictor
+  and re-evaluate generated sequences using the original objective.
+- Constrained substitutions can change other properties and do not preserve
+  an exact conditioned generative distribution.
+- No biological or clinical claims follow from these examples.
+
+## 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:** the bundled experiment is an execution demo.
+  Use suitable data, validate the decoder and predictor, and assess held-out
+  quality before interpreting guidance effects.
+
+## 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](LICENSE), matching the
+repository's license setting. ESM-2 weights are downloaded separately and remain
+subject to their original distribution terms.

lecture_02/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_02/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_02/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_02/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()

requirements.txt

@@ -0,0 +1,3 @@
+# Tested teaching environment: Python 3.11.
+torch==2.9.1
+transformers==4.57.6

tests/test_examples.py

@@ -0,0 +1,97 @@
+"""Offline checks for the Lecture 2 examples; no model download is required."""
+import csv
+import itertools
+from pathlib import Path
+import sys
+from types import SimpleNamespace
+import unittest
+
+import torch
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT / "lecture_02"))
+import esm2_flow_guidance as flow
+import esm2_diffusion_guidance as diffusion
+
+
+class GuidanceTests(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        torch.set_num_threads(2)
+
+    def test_example_annotations(self):
+        with (ROOT / "lecture_02" / "esm2_example.csv").open(newline="") as handle:
+            rows = list(csv.DictReader(handle))
+        self.assertEqual(len(rows), 64)
+        self.assertEqual({len(row["sequence"]) for row in rows}, {24})
+        labels = torch.tensor([[float(row["r1"]), float(row["r2"])] for row in rows])
+        for module in (flow, diffusion):
+            computed = module.composition_proxies([row["sequence"] for row in rows])
+            torch.testing.assert_close(computed, labels, atol=1e-7, rtol=0)
+        for row in rows:
+            self.assertEqual(int(row["c"]), int(float(row["r1"]) > 0))
+
+    def test_scalarization_weights(self):
+        for module in (flow, diffusion):
+            with self.subTest(module=module.__name__):
+                weights = module.normalize_weights((7, 3)).cpu()
+                torch.testing.assert_close(weights, torch.tensor([0.7, 0.3]))
+                for invalid in ((0, 0), (-1, 2), (1,), (float("nan"), 1)):
+                    with self.assertRaises(ValueError):
+                        module.normalize_weights(invalid)
+
+    def test_reward_gradient_without_weight_updates(self):
+        for module in (flow, diffusion):
+            with self.subTest(module=module.__name__):
+                torch.manual_seed(3)
+                reward = module.RewardModel(4, 5).to(module.DEVICE)
+                reward.eval().requires_grad_(False)
+                state = torch.randn(2, 4, 5, device=module.DEVICE)
+                time = torch.tensor([0.3, 0.7], device=module.DEVICE)
+                weights = module.normalize_weights((0.7, 0.3))
+                gradient = module.reward_gradient(reward, state, time, weights)
+                delta = torch.zeros_like(state)
+                delta[0, 1, 2] = 1e-3
+                difference = reward(state + delta, time) - reward(state - delta, time)
+                estimate = (difference * weights).sum() / (2e-3)
+                torch.testing.assert_close(gradient[0, 1, 2], estimate, atol=3e-5, rtol=3e-3)
+                self.assertFalse(gradient.requires_grad)
+                self.assertTrue(torch.isfinite(gradient).all())
+                self.assertTrue(all(p.grad is None for p in reward.parameters()))
+
+    def test_ddpm_schedule(self):
+        self.assertEqual(diffusion.alpha_bars[0].item(), 1.0)
+        self.assertLess(diffusion.alpha_bars[-1].item(), 1e-4)
+        self.assertEqual(diffusion.posterior_variances[1].item(), 0.0)
+        self.assertTrue((diffusion.posterior_variances >= 0).all())
+        self.assertTrue(torch.isfinite(diffusion.posterior_variances).all())
+
+    def test_constrained_decoder_against_exhaustive_search(self):
+        torch.manual_seed(11)
+        scores = torch.randn(3, 4, 20)
+        tokenizer = SimpleNamespace(convert_tokens_to_ids=lambda tokens: list(range(20)))
+        esm = SimpleNamespace(lm_head=torch.nn.Identity())
+        stats = {"z_mean": torch.zeros(1, 1, 20), "z_std": torch.ones(1, 1, 20)}
+        for module in (flow, diffusion):
+            polar = [i for i, a in enumerate(module.AMINO_ACIDS) if a in module.POLAR_RESIDUES]
+            nonpolar = [i for i in range(20) if i not in polar]
+            for minimum in range(5):
+                sequences = module.decode(scores, esm, tokenizer, stats, min_polar=minimum)
+                for batch, sequence in enumerate(sequences):
+                    count = sum(a in module.POLAR_RESIDUES for a in sequence)
+                    self.assertGreaterEqual(count, minimum)
+                    actual = sum(float(scores[batch, j, module.AMINO_ACIDS.index(a)])
+                                 for j, a in enumerate(sequence))
+                    optimal = max(
+                        sum(float(scores[batch, j, polar if flag else nonpolar].max())
+                            for j, flag in enumerate(assignment))
+                        for assignment in itertools.product((0, 1), repeat=4)
+                        if sum(assignment) >= minimum
+                    )
+                    self.assertAlmostEqual(actual, optimal, places=5)
+            with self.assertRaises(ValueError):
+                module.decode(scores, esm, tokenizer, stats, min_polar=5)
+
+
+if __name__ == "__main__":
+    unittest.main()