Protein sequence and inverse folding

Weighted Layer Mixture

EasyRaw PyTorchAbout 10 min1 public example plus private tests

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

Your task

Compute global softmax weights from the supplied layer logits, then combine representation layers independently at each valid token. Return the layer weights followed by the weighted token features.

Function signature

weighted_layer_mixture(
    states: torch.Tensor,
    layer_logits: torch.Tensor,
    mask: torch.Tensor,
) -> 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]$ True=valid.

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.
Mix the original layer features directly; no normalization or projection is applied.
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.

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 weighted_layer_mixture(
    states: torch.Tensor,  # (B, N, L, C)
    layer_logits: torch.Tensor,  # (L,)
    mask: torch.Tensor,  # (B, N)
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    weights: (L,).
    mixture: (B,N,C)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement weighted_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.