Korea University Releases UltraBERT, Balancing Classification and Retrieval Performance

The DMIS Lab at Korea University has released the English bidirectional encoders UltraBERT-150M/380M and the embedding model UltraBERT-Embed-150M. These models support a 32K token context and are distributed under the Apache 2.0 license, allowing for commercial use.

Key Technologies and Architecture

Based on the finding that classification performs better with low masking ratios while retrieval performs better with high masking ratios, the approach adopts mixing 5% and 50% masked sentences within a single batch and using separate LM heads for each ratio. The architecture consists of a fixed 24 layers, with only the width adjusted to provide two sizes: 150M (hidden 768) and 380M (hidden 1024). The vocabulary uses the SmolLM2 tokenizer, and local and global attention are applied alternately.

Training followed the Cascade Distillation method, which involves pruning a 1.3B teacher model and then distilling directly. The 5% masked sentences trained the classification head using KL divergence against teacher logits, while the 50% masked sentences trained the retrieval head using standard MLM loss. The final model training utilized 2T tokens, with the context expanded to 32K during the final 100B tokens.

Performance and Benchmarks

UltraBERT-380M recorded GLUE 91.0 and BEIR 47.7, ranking first in 9 out of 10 benchmarks. Notably, on the code and math retrieval benchmark RAR-b, it achieved Code 47.0 and Math 61.7, outperforming the second-place model by 8.4 and 7.7 points, respectively. UltraBERT-150M achieved performance similar to RoBERTa-large in classification and superior performance to Ettin-400M in retrieval.

In the LongEmbed passkey task, a test for 32K context performance, it maintained 99% accuracy across the entire range from 256 tokens to 32K tokens. This contrasts with ModernBERT and Ettin, which suffer sharp drops at 16K/32K. The embedding model UltraBERT-Embed-150M recorded MTEB(eng) 67.24 and MTEB(Code) 73.30, ranking first among dense embedding models under 200M parameters.

Usage Notes

The models are exclusively for English and code; Korean is not supported. The backbone models (150M/380M) require fine-tuning, while the ready-to-use embedding model is UltraBERT-Embed-150M. When setting the learning rate, it must be set up to 10 times lower than for other encoders. For retrieval fine-tuning, the optimal values are 5e-5 for 150M and 3e-5 for 380M. Additionally, the trust_remote_code=True option is required, so pinning the revision is recommended for production environments.