Foundations · INTEGRATION

Stable logsumexp, log-softmax and masked softmax

MediumRaw PyTorchAbout 20 min3 public examples plus private tests

OpenFoldESM-2ESM-1v

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

Your task

Convert included scores into exponential weights along dim. Return each weight’s normalized share (softmax), its natural log (log_softmax), or the log of the unnormalized weight sum (logsumexp), according to variant.

Function signature

stable_masked_softmax(
    logits: torch.Tensor,
    mask: torch.Tensor,
    dim: int = -1,
    variant: str = 'softmax',
    sentinel: float = 0.0,
    keepdim: bool = False,
) -> torch.Tensor

Worked example

One normalizer, three outputs

For logits=[0,1,100], mask=[True,True,False], dim=0 and sentinel=-9, the included exponential sum is 1+e. softmax returns approximately [0.268941,0.731059,0]; log_softmax returns approximately [-1.313262,-0.313262,-9]. logsumexp returns log(1+e) ≈ 1.313262 (a scalar when keepdim=False). The third score participates in none of these calculations.

Parameters

Parameter	Shape and meaning
`logits`torch.Tensor	$[\mathrm{d0}, \ldots, d(R-1)]$ Finite floating tensor $[\mathrm{d0},\ldots,d(R-1)]$ , $R\ge 1$ ; normalization axis may have length zero.
`mask`torch.Tensor	Any shape broadcastable to `logits.shape`.Bool tensor broadcastable to logits.shape; `True` means valid, `False` excludes numerator and normalizer.
`dim`intDefault `-1`	One integer normalization axis; negative indices count from the end.
`variant`strDefault `'softmax'`	Choose "softmax", "log_softmax", or "logsumexp". The first two preserve the logits shape; the last reduces the selected axis.
`sentinel`floatDefault `0.0`	Finite value for masked log probabilities and empty log-sum-exp slices. It is a reporting convention, not log(0).
`keepdim`boolDefault `False`	Only logsumexp uses this setting: retain its reduced axis with size one when True. Softmax and log_softmax ignore it.

Return

torch.Tensor · Same as logits for softmax/log_softmax. For logsumexp, remove dim, or replace it by 1 when keepdim=True.

Return Same-shaped probabilities/log probabilities, or logsumexp with dim reduced according to keepdim.

Required behavior

Normalize one integer axis of logits. Broadcast the Boolean mask to exactly logits.shape; True means include. Logits on nonempty valid slices are finite and have representable output differences.
On each nonempty valid slice, let $a=\max_{i\in V} x_i$ , $Z=\sum_{i\in V}\exp(x_i-a)$ , and $L=a+\log Z$ . Subtract a valid maximum before exponentiating.
For softmax return $\exp(x_i-a)/Z$ on valid entries and zero elsewhere. For log_softmax return $x_i-a-\log Z$ on valid entries and the finite sentinel elsewhere.
For logsumexp return $L$ . Remove the normalized axis unless keepdim=True; keepdim has no effect on softmax or log_softmax.
An all-masked slice returns all zeros for softmax, all sentinel values for log_softmax, or one sentinel for logsumexp. These empty-support extensions have zero gradients.
A zero-length normalization axis gives empty softmax/log_softmax tensors, or sentinel values at retained logsumexp positions. Do not exponentiate masked padding.
Use no epsilon, probability floor, clipping, or post-mask renormalization. No output may be NaN or infinity in the stated domain.
Softmax and log_softmax are unchanged by adding a slice-constant. Nonempty logsumexp increases by that constant; an empty slice keeps its sentinel. Accept arbitrary valid input strides without changing input values.

Why this matters

Subtracting a valid maximum controls exponent size; an explicit empty-row branch prevents an all-masked input from becoming NaN or a uniform distribution.

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 stable_masked_softmax(
    logits: torch.Tensor,  # (d0, ..., d(R-1))
    mask: torch.Tensor,  # Any shape broadcastable to logits.shape.
    dim: int = -1,
    variant: str = 'softmax',
    sentinel: float = 0.0,
    keepdim: bool = False,
) -> torch.Tensor:
    """Return a tensor with shape Same as logits for softmax/log_softmax. For logsumexp,
        remove dim, or replace it by 1 when keepdim=True."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement stable_masked_softmax")

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.