UltraBERT: Dual Heads for Contrasting Masking Regimes

Why People Wanted It

A single masking ratio pushes an encoder toward one kind of representation. The reported comparisons show the tension: lower masking favors classification, while higher masking favors retrieval. Mixing ratios in one batch sounds like a way to serve both, but it also asks shared parameters to satisfy two different prediction tasks. The real question is not merely how to put both kinds of examples in a batch; it is whether they should also share the same prediction head.

The Core Move

Mix lightly and heavily masked sequences in each batch, but give each masking ratio its own language-model head. The stated setup combines 5% and 50% masking at a 4:1 sequence ratio; the paper gives no update equation for this method, so the mechanism can be shown without adding one.

Why Mixing Ratios Is a Multitask Problem

Now the version for readers who want more of the machinery: each masking ratio defines a distinct training task, and the shared encoder is optimized against both tasks' losses. The reported setup places 5%-masked and 50%-masked sequences in the same batch, with four times as many 5%-masked sequences because their gradients are sparser and converge more slowly. This is not a claim that every task gets equal weight: the sequence counts are deliberately unequal.

A uniform sweep across masking ratios is a different choice. In the reported comparison, sampling uniformly from 5% to 50% behaves much like using a single 30% ratio; mixing only the two extremes is the more targeted way to retain both regimes. Why might even that targeted mixture still interfere?

Why the LM Heads Need Different Geometries

Here is the part that gets technical: the input embedding matrices for models trained at 6% and 50% masking are described as almost identical, while their LM heads take clearly different shapes. The 6% head has more concentrated singular values, indicating that it occupies a lower-dimensional subspace. The encoder receives gradients through its LM head, so sharing a head across these different prediction geometries may create gradient-level interference.

Decoupling means using one LM head for each masking ratio while continuing to train the mixed batch. The reported 5% and 50% decoupled heads have distinct geometries. In the 150M comparison, the mixed setup without decoupling improves retrieval relative to the single 15% result but does not recover the low-mask classification result; with decoupling, the reported GLUE and BEIR scores reach the respective best single-ratio scores. Those are results for that comparison, not a guarantee for every model or training setup. What is gained by separating the heads, and what shared capacity remains a possible source of conflict?

Minimal Causal Code

This plain loop makes the routing visible: each example's masking ratio selects its own prediction head. It omits the paper's full encoder, masking, and loss implementation.

import torch
import torch.nn.functional as F

# encoder, head_5, and head_50 are ordinary PyTorch modules.
# Each batch item carries a ratio label: "5" or "50".
def train_step(batch, optimizer):
    tokens, labels, ratio = batch
    hidden = encoder(tokens)
    loss = torch.zeros((), device=tokens.device)

    for name, head in [("5", head_5), ("50", head_50)]:
        selected = ratio == name
        if selected.any():
            logits = head(hidden[selected])
            loss = loss + F.cross_entropy(logits, labels[selected])

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss

What You Gain, and What You Lose

Separate heads let the two masking regimes fit different output geometries instead of forcing one LM head to serve both. The reported mixed-and-decoupled setup reaches the best single-ratio GLUE and BEIR scores in the comparison, while the shared-head mixture falls short on classification.

The cost is another prediction head and a more structured training batch; the encoder is still shared, so decoupling does not remove every possible source of task conflict. The result is a targeted compromise, not proof that arbitrary masking mixtures will work equally well. A separate module covers how masking choices relate to downstream task preferences.