"""AMP-Diffusion, the same chain run in the latent space of a protein language model. Gaussian corruption needs a state that can be added to, and an amino acid sequence has no such structure. A frozen encoder supplies one. The peptide becomes a tensor of residue embeddings, the diffusion runs there, and the language-model head maps the generated latent back to residues. Against the image model the state is the latent and a transformer replaces the U-Net. The network predicts the clean latent in place of the noise, and a mask marks which positions hold real residues. In Chapter 3 step k = 0 holds a clean latent and corruption grows with k, up to pure noise at k = K. Chapter 3, Sections 3.21, 3.29 and 3.46 of the notes. Trained on toy embeddings by default, and on ESM-2 embeddings of the course peptides with --dataset esm2. Prints the corrupted coordinate at k = 500, one reverse step there, the masked loss, and the polar fraction before and after the constraint. python lecture_3/latent_diffusion.py # 800 steps on offline embeddings python lecture_3/latent_diffusion.py --steps 20 # a few seconds, for a first look python lecture_3/latent_diffusion.py --dataset esm2 # real ESM-2 latents and its head """ # %% 01. Imports and configuration import math import torch import torch.nn as nn from cis6270.nets import count_parameters from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train MODEL_NAME = "facebook/esm2_t33_650M_UR50D" # Pretrained ESM-2 encoder. EMBED_DIM = 1280 # ESM-2 residue embedding width D. MAX_LEN = 32 # Token positions L per padded peptide. NUM_LAYERS = 6 # Transformer layers in the denoiser. NUM_STEPS = 1000 # Diffusion steps K. AMINO_ACIDS = "ACDEFGHIKLMNPQRSTVWY" # The 20 standard amino acids. POLAR_RESIDUES = set("STNQYCHKRDE") # Polar or charged side chains. POLAR_COLUMNS = [AMINO_ACIDS.index(a) for a in "STNQYCHKRDE"] # %% 02. The schedule tensors # # The image tensors of Section 3.28, rebuilt here so the file stands alone. Index k holds the # quantity at step k, and we prefix a zero so that abar_0 = 1 at the clean latent. betas = torch.cat([torch.zeros(1), torch.linspace(1e-4, 0.02, NUM_STEPS)]) # beta_k. alphas = 1.0 - betas # One-step retention alpha_k. alphas_cumprod = torch.cumprod(alphas, dim=0) # Cumulative retention abar_k. alphas_cumprod_prev = torch.cat([torch.ones(1), alphas_cumprod[:-1]]) # abar_{k-1}. posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod) # beta-tilde_k. # %% 03. Masked corruption # # Z_k = (sqrt(abar_k) Z_0 + sqrt(1 - abar_k) eps) * M. The reshape to [B, 1, 1] broadcasts one # coefficient per sequence across all L positions and all D coordinates, and the mask holds padded # positions at exactly zero so the denoiser spends no capacity on them. def q_sample_latent(z0, k, mask, noise=None): if noise is None: noise = torch.randn_like(z0) # [B, L, D] standard Gaussian. abar = alphas_cumprod.to(z0.device)[k].reshape(-1, 1, 1) z_k = abar.sqrt() * z0 + (1.0 - abar).sqrt() * noise return z_k * mask.unsqueeze(-1), noise # Padding stays exactly zero. def timestep_embedding(k, dim): half = dim // 2 # Half the coordinates sine, half cosine. freqs = torch.exp(-math.log(10000.0) * torch.arange(half, device=k.device).to(torch.get_default_dtype()) / half) args = k.to(torch.get_default_dtype())[:, None] * freqs[None, :] # [B, half]. return torch.cat([torch.sin(args), torch.cos(args)], dim=1) # [B, dim]. # %% 04. The latent denoiser and its loss # # We add the step embedding, a learned position embedding and the latent before the transformer # acts on them, so the step reaches every residue as one shared offset and the attention layers # mix the positions. The loss counts valid residues only. class LatentDenoiser(nn.Module): def __init__(self, embed_dim=EMBED_DIM, num_layers=NUM_LAYERS, num_heads=8, max_len=MAX_LEN): super().__init__() self.embed_dim = embed_dim self.time_mlp = nn.Sequential( # Same sinusoidal step embedding. nn.Linear(embed_dim, embed_dim), nn.SiLU(), nn.Linear(embed_dim, embed_dim), ) self.pos_embed = nn.Embedding(max_len, embed_dim) # Residue index. encoder_layer = nn.TransformerEncoderLayer( d_model=embed_dim, nhead=num_heads, dim_feedforward=4 * embed_dim, batch_first=True, ) self.transformer = nn.TransformerEncoder( encoder_layer, num_layers=num_layers, enable_nested_tensor=False, # Keep a padded batch on the ordinary attention path. ) self.out = nn.Sequential(nn.LayerNorm(embed_dim), nn.Linear(embed_dim, embed_dim)) def forward(self, z_k, k, attention_mask): t_emb = self.time_mlp(timestep_embedding(k, self.embed_dim)) positions = torch.arange(z_k.shape[1], device=z_k.device) h = (z_k + t_emb.unsqueeze(1) # Shared across positions. + self.pos_embed(positions).unsqueeze(0)) # Shared across batch. h = self.transformer(h, src_key_padding_mask=~attention_mask.bool()) return self.out(h) # Predicted clean latent Z0-hat. def latent_loss(z0_hat, z0, mask): error = (z0_hat - z0) ** 2 # [B, L, D] squared error. error = error * mask.unsqueeze(-1) # Zero out padded positions. return error.sum() / (mask.sum() * z0.shape[-1]) # Valid tokens only. # %% 05. Reverse latent sampling # # The network predicts the clean latent, so we rearrange the corruption equation to recover the # noise it implies. The reverse mean is then the image sampler's expression unchanged. @torch.no_grad() def p_sample_latent(model, z_k, k, mask): z0_hat = model(z_k, k, mask) # Predict the clean latent. abar = alphas_cumprod.to(z_k.device)[k].reshape(-1, 1, 1) eps_hat = (z_k - abar.sqrt() * z0_hat) / (1.0 - abar).sqrt() beta_k = betas.to(z_k.device)[k].reshape(-1, 1, 1) alpha_k = alphas.to(z_k.device)[k].reshape(-1, 1, 1) mean = (z_k - beta_k / (1.0 - abar).sqrt() * eps_hat) / alpha_k.sqrt() var = posterior_variance.to(z_k.device)[k].reshape(-1, 1, 1) xi = torch.randn_like(z_k) keep = (k > 1).to(torch.get_default_dtype()).reshape(-1, 1, 1) return (mean + keep * var.sqrt() * xi) * mask.unsqueeze(-1) @torch.no_grad() def sample_latents(model, num_sequences, length, width, device): mask = torch.ones(num_sequences, length, device=device) z = torch.randn(num_sequences, length, width, device=device) for step in range(NUM_STEPS, 0, -1): # Z_K down to Z_0. k = torch.full((num_sequences,), step, device=device, dtype=torch.long) z = p_sample_latent(model, z, k, mask) return z # %% 06. Decoding back to amino acids # # We unstandardize the generated latent first, then read it through the head that was trained to # score residues from exactly these embeddings. The head also scores special tokens and ambiguity # codes, so we restrict the columns to the twenty canonical amino acids before the argmax. @torch.no_grad() def decode(z, head, columns, z_mean, z_std): latent = z * z_std + z_mean # Undo normalization. logits = head(latent) # Full vocabulary logits. logits = logits.index_select(-1, columns) # 20 canonical residues. scores, choices = logits.max(dim=-1) # Argmax at each position. return logits, choices def composition_proxies(seq): """Net charge per residue and polar/charged fraction.""" L = max(len(seq), 1) charge = (seq.count("K") + seq.count("R") - seq.count("D") - seq.count("E")) / L polar = sum(1 for a in seq if a in POLAR_RESIDUES) / L return charge, polar # %% 07. The composition constraint at decoding # # Guidance makes a high polar fraction more likely without guaranteeing it. We impose the # requirement afterwards by substituting residues at the positions where a substitution costs the # least logit, so one projection at the end satisfies it exactly. def enforce_polar(logits, choices, minimum): cols = torch.tensor(POLAR_COLUMNS, device=logits.device) polar_scores, pick = logits.index_select(-1, cols).max(dim=-1) polar_choices = cols[pick] # best polar residue already = torch.isin(choices, cols) # positions already polar sequences = [] for b in range(choices.shape[0]): row = choices[b].clone() missing = minimum - int(already[b].sum()) if missing > 0: chosen = logits[b].gather(1, row[:, None]).squeeze(1) cost = chosen - polar_scores[b] # logit lost by swapping cost = cost.masked_fill(already[b], float("inf")) _, swap = cost.topk(missing, largest=False) row[swap] = polar_choices[b][swap] sequences.append("".join(AMINO_ACIDS[i] for i in row.tolist())) return sequences # %% 08. The latents, standardized # # ESM-2 coordinates have means and scales of their own, and a corruption that drove such a cloud # toward N(0, I) would have to destroy a systematic offset as well as the structure. So we center # and scale each coordinate over the valid residues before any diffusion runs, and we keep the two # statistics for decoding. def load_latents(dataset: str, device, seed: int): """Return the standardized latents, the mask, the decoding head and its residue columns.""" if dataset == "toy": from cis6270.data import standardize, toy_embeddings z0, mask = toy_embeddings(num_sequences=128, length=16, width=32, seed=seed) z0, z_mean, z_std = standardize(z0, mask) # Without ESM-2 we have no language-model head, so a fixed projection stands in. It makes # the decoding path and the constraint runnable offline, on residues with no biology. generator = torch.Generator().manual_seed(seed) weight = torch.randn(len(AMINO_ACIDS), z0.shape[-1], generator=generator) / 4.0 head = lambda latent: latent @ weight.t().to(latent) # noqa: E731 columns = torch.arange(len(AMINO_ACIDS)) else: from cis6270.data import esm2_embeddings, standardize z0, mask, tokenizer, head = esm2_embeddings( model_name=MODEL_NAME, max_length=MAX_LEN, device=device, ) # The head is ESM-2's masked language-model head, frozen with the encoder. z0, z_mean, z_std = standardize(z0, mask) columns = torch.tensor([tokenizer.convert_tokens_to_ids(a) for a in AMINO_ACIDS]) return z0.to(device), mask.to(device), z_mean.to(device), z_std.to(device), head, columns # %% 09. The report def main() -> None: parser = common_parser(__doc__.split("\n")[0]) parser.add_argument("--dataset", choices=["toy", "esm2"], default="toy") parser.add_argument("--minimum-polar", type=int, default=8, help="polar residues required") parser.set_defaults(steps=800, batch_size=16, lr=1e-3, samples=8) args = parser.parse_args() set_seed(args.seed) device = pick_device(args.device) z0_all, mask_all, z_mean, z_std, head, columns = load_latents(args.dataset, device, args.seed) num_sequences, length, width = z0_all.shape # [N, L, D]. denoiser = LatentDenoiser(embed_dim=width, num_layers=NUM_LAYERS, num_heads=8, max_len=length).to(device) if not args.quiet: print(f"device={device} dataset={args.dataset} latents={tuple(z0_all.shape)} " f"parameters={count_parameters(denoiser):,}") def loss_fn(step: int) -> torch.Tensor: rows = torch.randint(0, num_sequences, (args.batch_size,), device=device) # [B]. z0, mask = z0_all[rows], mask_all[rows] # [B, L, D] and [B, L]. k = torch.randint(1, NUM_STEPS + 1, (z0.shape[0],), device=device) # One step per item. z_k, _ = q_sample_latent(z0, k, mask) # Corrupt the clean latent. return latent_loss(denoiser(z_k, k, mask), z0, mask) # Valid residues only. losses = train(denoiser, loss_fn, args.steps, args.lr, quiet=args.quiet) generated = sample_latents(denoiser, args.samples, length, width, device) # [S, L, D]. logits, choices = decode(generated, head, columns.to(device), z_mean, z_std) # The argmax decoding before the constraint, and the same sequences after it. plain = ["".join(AMINO_ACIDS[int(i)] for i in row) for row in choices] constrained = enforce_polar(logits, choices, args.minimum_polar) # The chapter's coordinate, corrupted at k = 500 and moved back one reverse step. abar = float(alphas_cumprod[500]) a_500, sigma_500 = math.sqrt(abar), math.sqrt(1.0 - abar) # 0.2803 and 0.9599. z_500 = a_500 * 1.2 + sigma_500 * 0.4 # 0.7204. eps_hat = (z_500 - a_500 * 1.05) / sigma_500 # 0.4438 from the prediction z0_hat = 1.05. mean_500 = (z_500 - float(betas[500]) / sigma_500 * eps_hat) / math.sqrt(float(alphas[500])) # The nine substitution costs of the notes, with two positions still to fill. costs = torch.tensor([0.9, 1.4, 2.6, 3.1, 3.3, 4.0, 4.2, 5.5, 6.1]) window = min(50, args.steps) # A short run averages over everything it has. report = { "parameters": count_parameters(denoiser), f"loss, first {window} steps": sum(losses[:window]) / window, f"loss, last {window} steps": sum(losses[-window:]) / window, "valid residues per batch": int(mask_all[:args.batch_size].sum()), "a_500 and sigma_500": [round(a_500, 4), round(sigma_500, 4)], # 0.2803 and 0.9599. "z_500 from z_0=1.2, eps=0.4": z_500, # 0.7204. "eps implied by z0_hat=1.05": eps_hat, # 0.4438, against the draw 0.4. "reverse mean at k=500": mean_500, # 0.7193. "reverse standard deviation": math.sqrt(float(posterior_variance[500])), # 0.1002. "clean estimate at a=0.8 from eps=-0.8": (7.4 + 0.6 * 0.8) / 0.8, # 9.85. "noise recovered from it": (7.4 - 0.8 * 9.85) / 0.6, # -0.8. "polar fraction, plain decoding": sum( composition_proxies(s)[1] for s in plain) / len(plain), "polar fraction, constrained": sum( composition_proxies(s)[1] for s in constrained) / len(constrained), "net charge per residue, constrained": sum( composition_proxies(s)[0] for s in constrained) / len(constrained), "two cheapest substitutions": [round(float(v), 1) for v in costs.topk(2, largest=False)[0]], "logit surrendered": float(costs.topk(2, largest=False)[0].sum()), # 2.3. } print_report(f"AMP-Diffusion on {args.dataset} latents", report, constrained) save_run(args.out, model=denoiser, config=vars(args), report=report, losses=losses, samples=constrained) # %% 10. Run the script if __name__ == "__main__": main()