Activations and gated feed-forward transitions

Your task

Function signature

gated_feedforward(
    x: torch.Tensor,
    w_up: torch.Tensor | None = None,
    w_gate: torch.Tensor | None = None,
    w_down: torch.Tensor | None = None,
    b_up: torch.Tensor | None = None,
    b_gate: torch.Tensor | None = None,
    b_down: torch.Tensor | None = None,
    activation: str = 'silu',
    variant: str = 'feedforward',
    gated: bool = True,
) -> torch.Tensor

Worked example

A gate controls each hidden channel

Parameters

Return

torch.Tensor · Exactly x.shape in activation mode; [\ldots, \mathrm{Cout}] in feedforward mode.

Required behavior

• variant is 'feedforward' (default) or 'activation'. activation is 'silu' (default), 'gelu', 'gelu_tanh', 'gelu_sigmoid', or 'relu'.
• Define \operatorname{SiLU}(z)=z/(1+\operatorname{exp}(-z)), evaluated without overflow; exact \operatorname{GELU}(z)=0.5\cdot z\cdot (1+\operatorname{erf}(z/\sqrt{2})); and \operatorname{GELU}_{\tanh}(z)=\tfrac12 z\left[1+\tanh\left(\sqrt{2/\pi}(z+0.044715z^3)\right)\right]; \operatorname{GELU}_{\mathrm{sigmoid}}(z)=z/(1+\exp(-1.702z)); and \operatorname{ReLU}(z)=\operatorname{max}(z,0), with zero derivative at exactly zero.
• Sigmoid-GELU is Enformer's approximation, not \operatorname{tanh}-\operatorname{GELU} or exact GELU.
• In activation mode, x is any floating shape and \mathrm{output}=\mathrm{activation}(x).
• In feedforward mode x [\ldots,\mathrm{Cin}], w_up [\mathrm{Cin},H], w_down [H,\mathrm{Cout}], optional b_up [H] and b_down [\mathrm{Cout}] are supplied. Let u=x\mathbin{@}w_{\mathrm{up}}+b_{\mathrm{up}}.
• gated defaults to True: then w_gate [\mathrm{Cin},H] and optional b_gate [H] are supplied and \mathrm{hidden}=u\cdot \mathrm{activation}(x\mathbin{@}w_{\mathrm{gate}}+b_{\mathrm{gate}}), with multiplication elementwise. If gated=False, \mathrm{hidden}=\mathrm{activation}(u) and gate tensors are absent.
• The output is \mathrm{hidden}\mathbin{@}w_{\mathrm{down}}+b_{\mathrm{down}}, omitting missing biases. This is the update only: no residual, normalization or dropout. Leading axes and noncontiguous layout must work; Cin,H,Cout are positive. The gate branch, not the up branch, is activated in gated mode.
• All values and weights are supplied, and output keeps their dtype/device.

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 gated_feedforward(
    x: torch.Tensor,  # Any floating shape for variant="activation"; (..., Cin) for "feedforward".
    w_up: torch.Tensor | None = None,  # (Cin, H)
    w_gate: torch.Tensor | None = None,  # (Cin, H)
    w_down: torch.Tensor | None = None,  # (H, Cout)
    b_up: torch.Tensor | None = None,  # (H,)
    b_gate: torch.Tensor | None = None,  # (H,)
    b_down: torch.Tensor | None = None,  # (Cout,)
    activation: str = 'silu',
    variant: str = 'feedforward',
    gated: bool = True,
) -> torch.Tensor:
    """Return a tensor with shape x.shape in activation mode; (..., Cout) in feedforward
        mode."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement gated_feedforward")

Related exercises

• Builds onReLU
• Builds onExact GELU
• Builds onSiLU
• Builds onTanh GELU Approximation
• Builds onSigmoid GELU Approximation
• Builds onFeedforward Block
• Builds onGated Feedforward Block

Sources and further reading

• facebookresearch/esm · modules.py
• facebookresearch/esm · modules.py
• Biohub/esm · layers.py
• Biohub/esm · blocks.py
• google-deepmind/deepmind-research · enformer.py