Protein sequence and inverse folding

Tied vocabulary head and masked/causal token losses

MediumRaw PyTorchAbout 20 min3 public examples plus private tests

ESM-1bESM-2AntiBERTyAbLang

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Your task

Look up token embeddings and apply the supplied masked-language head to external hidden states. Reuse the embedding table as the vocabulary decoder, then score eligible targets. Return embeddings, vocabulary logits, target-aligned token losses, and their requested scalar sum or mean.

Function signature

masked_lm_head_loss(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    mode: str = 'masked',
    reduction: str = 'mean',
    ignore_index: int = -100,
    eps: float = 1e-05,
) -> dict[str, torch.Tensor]

Worked example

The tied decoder scores the next selected target

Inspect the decoder/loss stage after the specified affine, exact GELU and LayerNorm stages have produced z=[[1,-1],[0,1]] in one batch. Let E=[[1,0],[0,1]], vocab_bias=[0,0], tokens=[[1,0]], labels=[[0,1]], and both token_mask and selected be all True. Lookup returns embeddings=[[[0,1],[1,0]]]; the same $E$ gives logits $zE^{\mathsf T}$ , namely [[[1,-1],[0,1]]]. In mode="causal", label 1 at target position 1 is scored by logits [1,-1] at position 0. Its probability and loss are

p=\frac{1}{1+e^2},\qquad -\log p=\log(1+e^2)\approx2.126928.

Thus token_loss=[[0,2.126928]] to six decimal places. With one eligible target, mean and sum both equal approximately 2.126928.

Parameters

Parameter	Shape and meaning
`activations`dict	Supplied intermediate tensors.
`parameters`dict	Supplied model parameters, keyed as shown below.
`masks`dict	Supplied masks. Each field below defines what True includes or permits.
`structure`dict	Fields supplied together in this dictionary.
`mode`strDefault `'masked'`	Choose "masked" to score a target at its own position, or "causal" to score target t using logits at t-1. Selection always refers to the target position.
`reduction`strDefault `'mean'`	Choose "sum" or "mean" over eligible targets. Mean divides by the eligible-target count; no epsilon is added.
`ignore_index`intDefault `-100`	Integer target value outside the vocabulary range; excluded from the loss even when its selection mask is True.
`eps`floatDefault `1e-05`	Positive value added to final-axis population variance inside each LayerNorm square root. It is not a mask/count denominator or an output offset.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Implement an embedding lookup, tied vocabulary head and selected-token negative log likelihood. $B$ , $L$ , $C \ge 1$ and $V \ge 2$ . All floating tensors share dtype/device and are finite. tokens $[B,L]$ are integer IDs in $[0,V)$ ; labels $[B,L]$ are IDs in $[0,V)$ or ignore_index.
token_mask and selected are bool $[B,L]$ , with True meaning valid token and selected target respectively. Callers exclude BOS/EOS/padding from scoring with these masks. With $E=\mathrm{embedding}[V,C]$ , return $\mathrm{embeddings}=E[\mathrm{tokens}]$ .
The supplied hidden $[B,L,C]$ already comes from an external encoder; do not add embeddings to hidden. Compute $u=\mathrm{hidden}\mathbin{@}\mathrm{head\_weight}+\mathrm{head\_bias}$ ; $a=\operatorname{GELU}(u)$ , with exact $\operatorname{GELU}(x)=x\cdot (1+\operatorname{erf}(x/\sqrt{2}))/2$ .
Let $m=\mathrm{mean\_C}(a)$ and $v=\mathrm{mean\_C}((a-m)^{2})$ . $z=(a-m)/\sqrt{v+\varepsilon}\cdot \mathrm{norm\_weight}+\mathrm{norm\_bias}$ . $\mathrm{logits}=z\mathbin{@}E^{\top}+\mathrm{vocab\_bias}$ . The same $E$ input is used in lookup and output projection, without cloning or an independent decoder parameter.
All affine weights use $[\mathrm{input},\mathrm{output}]$ orientation except $E$ , which uses $[\mathrm{vocabulary},\mathrm{channel}]$ .
Config: mode is 'masked' (default) or 'causal'; reduction is 'mean' (default) or 'sum'; $\varepsilon$ is positive (default $10^{-5}$ ) and goes inside the square root; ignore_index is an integer outside $[0,V)$ (default $-100$ ). No other options.
In masked mode, target $t$ is eligible iff token_mask $[t]$ and selected $[t]$ and $\mathrm{labels}[t]\ne \mathrm{ignore\_index}$ , and is predicted by logits $[t]$ .
In causal mode, target $t\ge 1$ is predicted by logits $[t-1]$ ; eligibility additionally requires token_mask $[t-1]$ . selected and labels always refer to the TARGET position, not the predictor. Target $t=0$ is never scored causally.
For an eligible label $y$ , $\mathrm{loss} = \operatorname{log}(\sum_{v} \operatorname{exp}(\mathrm{logit\_v})) - \mathrm{logit\_y}$ , evaluated with a maximum shift for stability. Return token_loss $[B,L]$ aligned with target positions; ineligible entries are exactly zero.
Scalar loss sums these entries or divides by the number of eligible targets, not the number of sequences. No eligible targets (including causal $L=1$ ) must yield a finite, graph-connected scalar zero. logits and embeddings are returned for all positions without masking. The return contains outputs, not residual updates.

Why this matters

The input embedding table also serves as the vocabulary projection, so learning signals from the readout accumulate into the same supplied parameter.

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def masked_lm_head_loss(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    mode: str = 'masked',
    reduction: str = 'mean',
    ignore_index: int = -100,
    eps: float = 1e-05,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    embeddings: (B, L, C).
    logits: (B, L, V).
    token_loss: (B, L).
    loss: Scalar tensor, shape ()."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement masked_lm_head_loss")

Ready to implement it?

3 progressive hints, a solution walkthrough, and grading that checks values, shapes and gradients are waiting in the workspace.

Write and run your solution Sign in

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Sources and further reading

A primitive exercise, not an end-to-end reproduction of the named models. No affiliation or endorsement is implied.