Protein sequence and inverse folding · INTEGRATION

Learned language-model layer mixtures

EasyRaw PyTorchAbout 10 min2 public examples plus private tests

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

Your task

Form one learned probability distribution over model layers and use it to combine their token representations. Apply the selected normalization/projection variant and return layer weights and mixed features.

Function signature

learned_layer_mixture(
    states: torch.Tensor,
    layer_logits: torch.Tensor,
    mask: torch.Tensor,
    norm_scale: torch.Tensor | None = None,
    norm_bias: torch.Tensor | None = None,
    projection: torch.Tensor | None = None,
    variant: str = 'raw',
    eps: float = 1e-05,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

Give unequal weight to two layers

For one valid token, let the two layer vectors be [1,3] and [5,7], with layer_logits [0,log(3)]. Their exponential weights are [1,3], so the shared probabilities are [1/4,3/4]. The raw mixture is (1/4)*[1,3]+(3/4)*[5,7]=[4,6]. Those same layer probabilities apply at every token, while a masked token’s mixture is zero.

Parameters

Parameter	Shape and meaning
`states`torch.Tensor	$[B, N, L, C]$ All hidden layers including an embedding layer if supplied.
`layer_logits`torch.Tensor	$[L]$ One learned scalar per supplied layer.
`mask`torch.Tensor	$[B, N]$ Validity mask; True includes the indicated observation.
`norm_scale`torch.Tensor \| NoneDefault `None`	$[C]$ Required for projected `variant`.
`norm_bias`torch.Tensor \| NoneDefault `None`	$[C]$ Required for projected `variant`.
`projection`torch.Tensor \| NoneDefault `None`	$[C, D]$ Bias-free projection required for projected `variant`.
`variant`strDefault `'raw'`	Choose "raw", "projected". The `'projected'` `variant` uses $z_{bnl}=\operatorname{LN}(\mathrm{states}_{bnl};\mathrm{norm\_scale},\mathrm{norm\_bias})\,\mathrm{projection}$ , independently on each last $C$ axis before mixing (ESMFold2).
`eps`floatDefault `1e-05`	Projected variant only: add inside the population-variance square root for each layer/token normalization.

Return

tuple[torch.Tensor, torch.Tensor]

Return these tensors in the stated order.

Required behavior

Mix supplied language-model layers with learned, globally shared softmax weights. states $[B,N,L,C]$ , layer_logits $[L]$ and mask $[B,N]$ are required, with $B,N,L,C\ge 1$ and True meaning valid. All floating tensors share float32 or float64 dtype. Set $\mathrm{w\_l}=\exp(\mathrm{a\_l}-\max(a))/\sum_{k}\,\exp(a_{k}-\max(a))$ .
The 'raw' default uses $z_{bnl}=\mathrm{states}_{bnl}$ (ESMFold's layer mixture). The 'projected' variant uses $z_{bnl}=\operatorname{LN}(\mathrm{states}_{bnl};\mathrm{norm\_scale},\mathrm{norm\_bias})\,\mathrm{projection}$ , independently on each last $C$ axis before mixing (ESMFold2).
$\operatorname{LN}(x)=((x-\operatorname{mean}_{C}(x))/\sqrt{\operatorname{mean}_{C}((x-\operatorname{mean}_{C}(x))^2)+\varepsilon})\cdot \mathrm{norm\_scale}+\mathrm{norm\_bias}$ . Population variance, eps default $1\times10^{-5}$ , strictly positive. projection $[C,D]$ has no bias.
Return weights $[L]$ and mixture $[B,N,C\,\mathrm{or}\,D]$ = $\sum_{l}\,\mathrm{w\_l}\cdot z_{bnl}$ . Masked state payloads, including NaN, are sanitized before any arithmetic and the final mixture is exact zero at those tokens.
Weights are GLOBAL and still returned for an entirely masked example; the mask does not renormalize weights over layers. Config variant is 'raw' or 'projected'; eps is the only other key. $L=1$ gives weight $1$ and zero logit derivative. $C=1$ in projected mode leaves only normalization bias before projection.
$B,N,L,C,D$ must be positive; logits and parameters are finite; valid states are finite. Normalization/reduction/projection intermediates, including centered squares, must remain finite in the common dtype, and eps must remain strictly positive when represented in that dtype.
These bounds exclude arbitrary-magnitude finite values and epsilon that rounds to zero. Invalid variant/nonpositive eps raises ValueError. Frozen source-language-model execution and downstream pair construction are not performed.

Why this matters

Each layer contributes a view of the same token. Global layer weights choose their relative importance; projected mode first normalizes and projects each layer so that mixing happens in the requested output feature space.

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 learned_layer_mixture(
    states: torch.Tensor,  # (B, N, L, C)
    layer_logits: torch.Tensor,  # (L,)
    mask: torch.Tensor,  # (B, N)
    norm_scale: torch.Tensor | None = None,  # (C,)
    norm_bias: torch.Tensor | None = None,  # (C,)
    projection: torch.Tensor | None = None,  # (C, D)
    variant: str = 'raw',
    eps: float = 1e-05,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    weights: (L,).
    mixture: (B, N, C) or (B, N, D)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement learned_layer_mixture")

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.

Related exercises

Sources and further reading

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