dmis-lab / UltraBERT-380M

UltraBERT-380M

UltraBERT is a family of English bidirectional encoders with a native 32K context that reach state of the art on both classification and retrieval, two task families on which previous encoders had to choose a side.

For more information about UltraBERT, please check out our paper .

✨ Highlights

• Classification and retrieval at once. Earlier encoders each excel in one area: Ettin on GLUE, RoBERTa on NER, LFM2.5-Encoder on BEIR. UltraBERT is the first encoder to take the top score on both classification and retrieval at the same size.
• Native 32K context. The model is trained at 2K tokens and extended to 32K context length. On the LongEmbed passkey test it keeps 99% accuracy at every length up to 32K, where 8K-context encoders drop sharply.
• Decoupled LM heads. Each pre-training batch mixes sequences masked at 5% and at 50%, and each masking ratio has its own LM head. Low ratios favor classification and high ratios favor retrieval; separate heads let one encoder learn both without the interference a shared head causes.
• Pruning and distillation from a 1.3B teacher. The 380M model is width-pruned from a 1.3B encoder and the 150M model from the 380M model. Both are distilled from the 1.3B teacher with a KL loss, applied only to the 5% masked sequences, while the 50% sequences use the standard MLM loss.

🚀 Usage

The model needs transformers >= 5 (tested with 5.16.1) and trust_remote_code=True . FlashAttention 2 is recommended for long inputs; SDPA and eager attention are also supported.

import torch
from transformers import AutoModelForMaskedLM, AutoTokenizer

model_id = "dmis-lab/UltraBERT-380M"
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForMaskedLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",  # or "sdpa"
    device_map="cuda",
)

text = f"The capital of France is {tokenizer.mask_token}."
inputs = tokenizer(text, return_tensors="pt").to(model.device)
with torch.no_grad():
    logits = model(**inputs).logits

mask_index = (inputs.input_ids == tokenizer.mask_token_id).nonzero()[0, 1]
print(tokenizer.decode(logits[0, mask_index].argmax()))

Fine-Tuning

The checkpoint loads into the usual task classes:

from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    "dmis-lab/UltraBERT-380M", num_labels=2, trust_remote_code=True
)

We recommend starting from the settings used in the paper. For classification: AdamW (β1 0.9, β2 0.98, ε 1e-6), weight decay 0.1, a linear schedule with 5% warm-up, and a learning rate swept over {5e-6,8e-6,1e-5,3e-5,5e-5,8e-5}. For dense retrieval: mean pooling with a contrastive loss and a wider learning-rate sweep of {1e-5,3e-5,5e-5,8e-5,1e-4,2e-4,3e-4,5e-4,8e-4,1e-3}.

💡 Note : UltraBERT consistently prefers lower learning rates than other encoders, by up to 10x. For example, the best learning rates for UltraBERT-150M and UltraBERT-380M are 5e-5 and 3e-5, respectively, whereas other encoders typically prefer 3e-4 to 5e-4.

📊 Evaluation

Every model is fine-tuned per task with the same grid search. For the retrieval tasks, each encoder is fine-tuned on a subset of MS MARCO or MLDR under the same sweep, so these scores reflect the potential information retrieval capability of each encoder. Dedicated embedding models are usually trained further on hundreds of millions of pairs with much longer training schedules.

Base Size (at most 200M parameters)

Medium Size (at most 500M parameters)

Long-Context Passkey Retrieval

Passkey retrieval accuracy on LongEmbed as a function of document length, for (a) base and (b) medium encoders. Every baseline degrades once the document exceeds its context length: the 512-context models fail beyond 1K tokens and the 8K-context models drop sharply at 16K and 32K. UltraBERT is the only model that retains 99% accuracy at every length up to 32K, at both sizes.

🏛️ Model Description

UltraBERT is a pre-norm Transformer encoder with RoPE, SwiGLU feed-forward layers, RMSNorm and no bias terms. Two of every three layers use local attention restricted to 128 tokens on each side, and every third layer attends over the whole sequence. The depth is fixed at 24 layers for every size and only the width is scaled. The LM head is untied from the input embeddings.

What Is in This Checkpoint

• model.safetensors in float32 (the pre-training master weights), with config.dtype: float32 .
• The encoder and the MLM head.
• The MLM head is the one trained on the 5% masked sequences (the distilled head); the 50% head is not included.
• Remote code for transformers 5.x with UltrabertModel , UltrabertForMaskedLM , UltrabertForSequenceClassification , UltrabertForQuestionAnswering and UltrabertForTokenClassification .

⚙️ Pre-Training

Objective. Masked language modeling with two masking ratios in every batch, 5% and 50% at a 4:1 sequence ratio, each with its own LM head. The 5% head is trained with a KL divergence loss against the logits of a 1.3B teacher; the 50% head with the standard MLM loss.

Pruning and distillation. The 1.3B teacher is width-pruned to 380M and the distilled 380M model is pruned again to 150M. Importance of hidden channels, attention heads and FFN neurons is estimated from activations on 2,000 calibration documents, and the top-scoring units are kept. Both sizes are distilled from the 1.3B teacher directly.

Optimization. AdamW (β1 0.9, β2 0.95, ε 1e-8), weight decay 0.1, gradient clipping 1.0, peak learning rate 1e-4 with a warmup-stable-decay schedule (10B warm-up tokens, linear decay over the last 10%). Batch size 4M tokens, and 2T tokens in total; the context is extended to 32K, raising the global-attention RoPE base from 640K to 2.56M. Trained on a single TPU v4-512, v5e-256 or v6e-256 slice with data parallelism and FSDP.

Data. One fixed mixture for the whole run (Table 4 of the paper).

Limitations

• English only. The pre-training mixture is English web, code, academic, math and instruction text.
• This is a pre-trained encoder, not a ready-to-use classifier or embedding model: it needs fine-tuning for a task. For a ready-to-use text embedding model built on it, see UltraBERT-Embed-150M .
• The model can reproduce biases of its web-scale training data; evaluate it for your use case.

Citation

@article{park2026ultrabert,
  title   = {UltraBERT: Beyond the Limit of BERT Pre-Training},
  author  = {Park, Jungwoo and Lee, Taewhoo and Kang, Jaewoo},
  journal = {arXiv preprint arXiv:{ARXIV_ID}},
  year    = {2026}
}

Acknowledgments

The pre-training was primarily supported by Cloud TPUs from Google's TPU Research Cloud (TRC) program.