Blame for lecture_3/latent_diffusion.py

"""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()