UltraBERT: Beyond the Limit of BERT Pre-Training

Jungwoo Park, Taewhoo Lee, Jaewoo Kang

arXiv 2610.15028 31 pages Read the paper

The short version

The problem: Text encoders are useful for labeling text and finding relevant documents, but a model trained to do one well often does the other poorly. Training it on more text does not necessarily fix the imbalance.
The idea: The paper finds that hiding only a few words at a time helps labeling, while hiding many words helps search. It trains one encoder in both ways at once, gives each way its own prediction head, and transfers a larger model’s knowledge only through the lightly masked examples—like teaching one student with two specialized practice sets while avoiding a conflicting lesson.
What it got: The resulting models perform strongly on both kinds of tasks rather than specializing in just one. The 380-million-parameter model scores 91.0 on GLUE, a suite of language-understanding tests, and 47.7 on BEIR, a suite of document-search tests.

Modules from this paper

UltraBERT: Dual Heads for Contrasting Masking Regimes

Mixing lightly and heavily masked sequences creates competing prediction geometries; UltraBERT gives each masking ratio its own language-model head to reduce that interference.

language-modelingmasked-language-modelingmulti-task-learningBERT

Low-Mask-Only Distillation in UltraBERT

Pruning makes a large encoder compact, but distillation can behave differently at different masking ratios; UltraBERT applies teacher-logit distillation only to low-masked examples and keeps standard masked-language-model training for high-masked examples.

distillationpruningmasked-language-modelingBERT

UltraBERT: Saturation and Masking Choices for Encoder Tasks

Encoder downstream scores can plateau while pretraining loss keeps falling; UltraBERT examines how to choose masking ratios differently for classification and retrieval once that plateau is reached.

encodersmasked-language-modelingpretrainingclassification