ESM-2 token-dropout scaling

Your task

Function signature

esm2_token_dropout(
    tokens: torch.Tensor,
    table: torch.Tensor,
    token_dropout: bool = True,
    padding_id: int = 0,
    mask_id: int = 1,
) -> dict[str, torch.Tensor]

Worked example

Different mask fractions, different sequence scales

import torch

tokens = torch.tensor([[2, 1, 0, 3], [1, 1, 2, 0]], dtype=torch.int64)
table = torch.tensor([[50.0, 60.0], [70.0, 80.0], [1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]], dtype=torch.float64, requires_grad=True)
token_dropout = True
padding_id = 0
mask_id = 1

result = esm2_token_dropout(tokens=tokens, table=table, token_dropout=token_dropout, padding_id=padding_id, mask_id=mask_id)

Expected return

{
    'output': torch.tensor([[[1.3199999999999998, 2.6399999999999997], [0.0, 0.0], [0.0, 0.0], [3.9599999999999995, 5.279999999999999]], [[0.0, 0.0], [0.0, 0.0], [2.6399999999999997, 5.279999999999999], [0.0, 0.0]]], dtype=torch.float64),
    'mask_fraction': torch.tensor([0.3333333333333333, 0.6666666666666666], dtype=torch.float64),
    'scale': torch.tensor([1.3199999999999998, 2.6399999999999997], dtype=torch.float64),
    'padding_mask': torch.tensor([[False, False, True, False], [False, False, False, True]], dtype=torch.bool),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Apply ESM-2 token dropout to a supplied embedding table. Inputs are tokens [B,L], an int64 tensor, and finite table [V,C], with common CPU placement and float32 or float64 table dtype. Require B,L,C\ge1 and V\ge2. Every token ID must lie in [0,V-1]. Do not mutate either input; noncontiguous layouts are valid.
• Config padding_id defaults to 0, mask_id to 1, and token_dropout to True. The IDs are distinct in-range integers, not Booleans; token_dropout must be Boolean. These configurable IDs do not assert ESM-2's alphabet numbering.
• Return Boolean padding_mask [B,L], true exactly at padding, plus mask_fraction and scale [B] and output [B,L,C], all in the table's floating dtype. Let n_b count nonpadding tokens and m_b count mask tokens. Every nonpadding special token, including any supplied beginning/end token, contributes to n_b.
• Define r_b=m_b/n_b for n_b>0, and r_b=0 for an all-pad row. Return r as mask_fraction, even when token dropout is disabled. With token dropout enabled, use s_b=\begin{cases}(1-0.12)/(1-r_b),&n_b>m_b,\\0,&n_b=m_b.\end{cases} Return scale s.
• For a token that is neither padding nor mask, set y_{b\ell c}=s_b E_{t_{b\ell},c}; all padding/mask positions are exact zeros. A row with no masks still uses s_b=0.88, not one. All-pad and all-mask rows have zero scale and connected-zero output: avoid evaluating a zero denominator and then hiding NaNs.
• These two finite boundary definitions extend the upstream formula, which is singular there. With token_dropout=False, every row has scale one, mask-token embeddings are retained, and padding is still zero.
• No random corruption, ordinary elementwise dropout, embedding-width scaling, alphabet conversion, BOS/EOS insertion or Transformer layer is performed. Incompatible ranks/dtypes, out-of-range tokens and invalid ID/flag configuration raise ValueError; nonfinite table values and overflowed arithmetic are outside the contract.

Why this matters

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 esm2_token_dropout(
    tokens: torch.Tensor,  # (B, L)
    table: torch.Tensor,  # (V, C)
    token_dropout: bool = True,
    padding_id: int = 0,
    mask_id: int = 1,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    output: (B, L, C).
    mask_fraction: (B,).
    scale: (B,).
    padding_mask: (B, L)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement esm2_token_dropout")

Sources and further reading

• facebookresearch/esm · esm2.py
• www.science.org