UltraBERT: Beyond the Limit of BERT Pre-Training

AI summary

Overview

Research area: Natural Language Processing — bidirectional encoder pre-training (masked language modeling), with evaluation on classification and information retrieval.

Technical level: Intermediate. The paper assumes familiarity with BERT-style masked language modeling, knowledge distillation, and bi-encoder retrieval fine-tuning.

Scope: An empirical study of encoder pre-training design choices at trillion-token scale, culminating in the UltraBERT-150M and UltraBERT-380M models.

What This Paper Is About

Bidirectional encoders have been modernized with architectural tricks borrowed from decoder-only LLMs (RoPE, FlashAttention-2, long contexts, trillions of tokens), but the actual pre-training recipe — especially the masked language modeling (MLM) objective and its masking ratio — is still inherited from small-scale BERT-era conventions. The result is that no single encoder performs competitively on both classification and retrieval: models tend to specialize. The paper re-examines core pre-training choices at saturation scale, finds that classification and retrieval prefer opposite masking ratios, and builds a single encoder family that reaches state of the art on both.

Key Contributions

1. A saturation result for encoders. Training encoders up to 8.8T tokens across 150M and 1.3B scales, the authors show downstream performance saturates around 1–2T tokens even while training and validation losses keep decreasing — unlike decoders, whose performance keeps improving. This establishes an upper bound that more pre-training tokens cannot raise.

2. A re-examination of masking ratio at saturation. Rather than at 10B–100B token budgets, the paper compares masking ratios at convergence and finds classification consistently favors lower ratios across 150M, 380M, and 1.3B scales, while retrieval consistently favors 50%. This contradicts prior work reporting that larger ratios suit larger models.

3. LM head decoupling. The authors show that encoders trained at 6% and 50% masking learn nearly identical input embeddings but distinctly different LM head geometries, and that sharing one head across mixed masking ratios causes interference. Decoupling the head per ratio lets a single batch mix 5% and 50% masked sequences (at a 4:1 ratio in favor of 5%) and recovers the best of both regimes.

4. A revised distillation rule. Pruning and logit-based distillation from a 1.3B teacher improves classification at 15% masking (90.2 to 90.8 on GLUE) but degrades retrieval at 50% masking (47.2 to 46.5 on BEIR). UltraBERT therefore applies KL-divergence distillation only to the low-ratio sequences and standard MLM to the high-ratio ones.

Main Findings

• Downstream saturation: Encoders saturate on GLUE, SQuAD 2.0, and BEIR around 1–2T tokens; BEIR shows no significant improvement beyond the first checkpoint at 275B tokens. Losses continue to fall, so fine-tuned downstream performance is not directly correlated with loss as it is for decoders.

• Divergence from decoders: After instruction tuning on the multiple-choice QA format, the 1.3B encoder saturates while a 1.7B SmolLM2 decoder keeps improving, suggesting an encoder-specific ceiling.

• Classification prefers low masking: GLUE and SQuAD 2.0 degrade consistently as the masking ratio rises, across all three model sizes. At 150M, GLUE falls from 89.5 at 6% masking to 87.5 at 50%.

• Retrieval prefers high masking: BEIR nDCG@10 is highest at 50% masking for all model scales: 45.0 (50%) vs. 43.5 (6%) for the 150M encoder. Effective rank of mean-pooled representations increases with masking ratio.

• Regularization cannot fix the trade-off: Applying the KoLeo regularizer at 6% masking raises BEIR but less than plain 50% masking does, and it also harms GLUE.

• RTD's advantage peaks at 15%: Comparing MLM with replaced token detection (RTD) at matched corruption levels on 150M encoders, the gap peaks around 15% and shrinks at both extremes — because at 50% the generator produces trivial replacements, and at 3% the difference from [MASK] is negligible. Since classification already favors low ratios, encoders can keep MLM and avoid RTD's retrieval degradation.

• Naive ratio mixing collapses to the average: Uniform sampling from 5% to 50% (average 27.5%) performs almost identically to a single 30% ratio (88.0 vs. 88.3 GLUE, 44.2 BEIR for both), and its LM head geometry matches the 30% model. Decaying from 50% to 6% or 15% gives no GLUE improvement over a single ratio. Mixing {5%, 50%} with decoupled heads reaches 89.4 GLUE and 45.1 BEIR, matching the best single ratios of 6% (89.5) and 50% (45.0).

• Headline results: UltraBERT-380M reaches 91.0 GLUE and 47.7 BEIR; UltraBERT-150M reaches 89.5 GLUE and 45.2 BEIR. The 150M model matches RoBERTa-large on classification with 2.3× fewer parameters and outperforms Ettin-encoder-400M on retrieval at 2.6× smaller size.

• Largest gains are in code and math retrieval: 7.3 and 8.3 points over the second-best base model on RAR-b Code and Math, and 8.4 and 7.7 points at medium scale.

• Long context: Both UltraBERT sizes natively support 32K context, and both achieve the best LongEmbed score in Table 3 (72.0 at 152M, 74.9 at 380M); the base model outperforms other families' medium models. The one benchmark where UltraBERT-380M is not best is MLDR, where LFM2.5-Encoder-355M scores higher (44.5 vs. 44.1).

Methodology in Plain English

The authors first run controlled experiments at three sizes (150M, 380M, 1.3B) using a shared architecture — rotary position embeddings, SwiGLU, alternating local and global attention (every third layer global), untied LM head, and the SmolLM2 tokenizer with a 49,152-token vocabulary. All models are 24 layers deep and scaled by width (hidden size 768, 1,024, and 2,048).

Because masking ratio and convergence speed are normally confounded, they deliberately train each ratio to saturation — 1T tokens at 50% and 30%, 2T at 15%, and 5T at 6% and 3% — so that speed no longer explains the result. They then sweep masking ratios and measure both classification (GLUE, SQuAD 2.0) and retrieval (fine-tune on MS MARCO, evaluate nDCG@10 on BEIR), and inspect the geometry of the input embedding matrix and LM head using principal components and singular value spectra.

Diagnosing head-level interference, they test several mixing strategies and settle on putting 5% and 50% masked sequences in the same batch at a 4:1 ratio with separate LM heads. Finally they take the best 1.3B encoder as a teacher, prune it to 380M and then to 150M by width only (keeping all 24 layers), and distill using KL divergence on the 5% sequences plus standard MLM on the 50% sequences. The final models train on 2T tokens, with context extended to 32K over the last 100B tokens, on TPU pods of 256 chips.

Why This Matters

The paper reframes encoder task specialization as an artifact of the pre-training recipe rather than an inherent property of the architecture. That is a substantive claim: it suggests a single backbone can serve as both a classifier and a retriever, simplifying systems that currently maintain separate models.

Real-world applications:

• Retrieval-augmented generation: a single encoder that both embeds documents for retrieval and scores/classifies text removes a model from the RAG pipeline.
• Pre-training data quality filtering: encoders are used to classify and score web-scale corpora before training larger models; better classifiers directly improve the filtered data.
• Reward modeling for reinforcement learning: UltraBERT-380M scores 78.8 on RewardBench versus 76.7 for the next-best model listed, which matters for preference-based training.
• Code and math search: the largest margins in the paper are on RAR-b Code and Math, where the base model gains 7.3 and 8.3 points — relevant to developer tooling and technical documentation retrieval.

Industry relevance: Encoders are cheap at inference, and the paper's 150M model matches a much larger RoBERTa-large on classification while beating a 400M retrieval model. For teams serving high-volume embedding and scoring traffic, that is a direct cost reduction. The compute-allocation finding — that spending tokens past roughly 1–2T is wasted on encoders — is also an actionable budgeting result for anyone pre-training an encoder.

Future Directions

• Extending the saturation bound. The paper identifies a ceiling that more pre-training tokens cannot raise, but does not report a method that lifts it beyond the recipe changes it introduces. Whether other levers (data composition, objectives, architecture) raise the ceiling remains open.

• Multi-stage data mixing. The appendix notes that the authors explored adjusting mixture ratios across training stages, as Ettin and EuroBERT do, but the provided text is truncated mid-sentence and no results are reported here.

• Alternative distillation signals. The paper restricts itself to logit-based distillation because attention-based objectives require O(L²) memory and are incompatible with FlashAttention. Whether attention-based distillation behaves differently at low versus high masking ratios is not reported.

• Generalization beyond the evaluated benchmarks. The paper evaluates English classification, retrieval, code, math, and long-context tasks. Behavior on multilingual data, on other masking-ratio pairings besides 5% and 50%, and on the 4:1 batch ratio's sensitivity are not reported.

Target Audience

Researchers and engineers working on encoder pre-training, sentence embedding models, or retrieval systems who want a rigorous account of how masking ratio interacts with downstream task type. It is also valuable for practitioners deciding how to allocate pre-training compute, and for anyone building RAG or data-filtering pipelines who needs one encoder to do several jobs. Readers without background in masked language modeling or bi-encoder retrieval will find the empirical sections dense, though the high-level findings are stated plainly.

Authors’ abstract

Recent bidirectional encoders have been modernized with innovations from causal large language models (LLMs). However, the pretraining recipe itself remains inherited from small-scale conventions, leaving encoders underoptimized, with no single model performing competitively across both classification and retrieval tasks. In this work, we investigate how far encoder performance can be pushed by revisiting core design choices in training, revealing three key findings: (1) Encoders saturate on downstream tasks beyond a few trillion tokens, indicating that additional compute past this point does not help. (2) Across model sizes, lower masking ratios benefit classification while higher ones benefit retrieval, contradicting prior work suggesting larger ratios for larger models. (3) Distillation with pruning improves classification at low masking ratios but degrades retrieval at high ones. Guided by these findings, we present UltraBERT, a single encoder that assigns separate LM heads to low and high masking ratios within a batch and restricts the distillation objective to the low ratio. To our knowledge, UltraBERT is the first encoder to achieve state of the art on both classification and retrieval tasks, reaching 91.0 on GLUE and 47.7 on BEIR. Together, these results show how far encoder performance can be pushed and provide a principled basis for compute allocation, masking ratio design, and knowledge transfer in bidirectional models.