Protein sequence and inverse folding

Projected Layer Mixture

MediumRaw PyTorchAbout 20 min1 public example plus private tests

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

Your task

Normalize and project each supplied representation layer, then combine layers at every valid token using one shared softmax weight vector.

Function signature

projected_layer_mixture(
    states: torch.Tensor,
    layer_logits: torch.Tensor,
    mask: torch.Tensor,
    norm_scale: torch.Tensor,
    norm_bias: torch.Tensor,
    projection: torch.Tensor,
    eps: float = 1e-05,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

Normalize each layer before mixing

For one valid token, take layer vectors [1,3] and [4,0], layer_logits [0,log(3)], norm_scale [1,1], norm_bias [0,0], projection [[1],[2]], and illustrative eps=1. The normalized layers are $[-1,1]/\sqrt{2}$ and $[2,-2]/\sqrt{5}$ ; their scalar projections are $1/\sqrt{2}$ and $-2/\sqrt{5}$ . Weights [1/4,3/4] give mixture $\frac{1}{4\sqrt{2}}-\frac{3}{2\sqrt{5}}\approx-0.494044$ . Normalize each original layer, then mix the projected values.

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]$ True=valid.
`norm_scale`torch.Tensor	$[C]$ Per-channel multiplier for each layer’s normalization.
`norm_bias`torch.Tensor	$[C]$ Per-channel offset for each layer’s normalization.
`projection`torch.Tensor	$[C, D]$ Bias-free projection of each normalized layer into output channels.
`eps`floatDefault `1e-05`	Add inside sqrt(population channel variance + eps), independently for each layer and token before projection.

Return

tuple[torch.Tensor, torch.Tensor]

Return these tensors in the stated order.

Required behavior

states has shape $[B,N,L,C]$ ; layer_logits is (L), shared across every batch/token. Normalize the logits with stable softmax over L, not over channels or tokens.
First apply supplied affine population LayerNorm over C independently at each layer/token, then projection $[C,D]$ . eps (default 1e-5) belongs inside the variance square root.
Sanitize masked tokens before arithmetic and zero their final mixture. Sum weighted layers along L, preserving batch/token and output-channel axes. Return the layer weights as well as the mixture.
Use CPU float32/float64 tensors with one shared floating dtype/device. Preserve strided layouts and do not mutate supplied tensors.
Use raw tensor operations, not a ready-made implementation of the requested primitive. Supplied parameters and random variates must not be reinitialized.

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