Add a readout head

Note

The contracts on this page are proposed until the task registry ships. The signatures match Tasks; once src/emma/ lands the subclass below runs as written.

Subclass ReadoutHead to map frozen-backbone representations to a new task’s outputs. The head is the only trainable part of the pipeline; the backbone weights stay frozen. This is the SUPERB (Speech processing Universal PERformance Benchmark) contract: one frozen backbone, many lightweight readouts, so the score reflects representation quality, not head capacity.

Goal

Implement a custom readout head for a task whose output shape the built-in heads do not cover, wire it into a TaskSpec, and confirm it runs under the OOD (out-of-distribution) protocol.

Prerequisites

  • You have a TaskSpec whose output shape is decided (regression vector, class distribution, ranked list, or text).

  • You have a backbone that satisfies the FrozenBackbone protocol, so encode returns a tensor of known dimension.

  • You have read Add a task; a head is one step of a task proposal.

Steps

  1. Subclass ReadoutHead. Implement the two abstract methods: build, which initializes head parameters from the backbone’s output dimension, and forward, which maps representations to task outputs.

    import torch
from torch import nn
from emma.tasks import ReadoutHead, FrozenBackbone

class AzimuthRegressionHead(ReadoutHead):
    """Regression head for an angle-of-arrival (AoA) task.

    Predicts a continuous azimuth; the backbone is frozen.
    """

    def __init__(self, hidden_dim: int = 256) -> None:
        self.hidden_dim = hidden_dim

    def build(self, backbone: FrozenBackbone) -> None:
        # Probe the backbone once to read its output dimension. The probe
        # matches the scene input shape (num_rx, 2, N); the recipe's
        # ArraySpec fixes num_rx, and N is the capture length.
        probe = torch.zeros(8, 2, 1024)
        out_dim = backbone.encode(probe).shape[-1]
        self.net = nn.Sequential(
            nn.Linear(out_dim, self.hidden_dim),
            nn.ReLU(),
            nn.Linear(self.hidden_dim, 1),
        )

    def forward(self, representations):
        return self.net(representations)

    AoA is angle-of-arrival. The head owns its parameters; the backbone does not.

  2. Respect the frozen-backbone contract. The head must not mutate backbone weights. The Evaluator shares one backbone instance across every task and re-fits the head per OOD fold, so the head’s parameters are constructed fresh in build for each fold.

  3. Keep the head lightweight. The backbone carries the representational load; a head with more capacity than the task needs inflates the score without measuring representation quality. Prefer a small MLP (multilayer perceptron) over a deep network.

  4. Wire the head into the TaskSpec. The spec references the head class, the Metric, and the OOD axis; the harness instantiates the head through the registry.

  5. Run the head under the protocol to confirm it produces outputs the metric accepts.

    $ emma eval --task E-LOC-AOA --model ./checkpoints/my-fm --scenes ./scenes --out ./runs/head-check

Expected result

The head builds from any FrozenBackbone, produces outputs the task metric consumes, and trains only its own parameters. The eval run completes with a metric value, confirming the head is wired correctly.

See Also