Back Home

GitHub Repo

Transformers Proposes Native UltraBERT Integration with Dual-Mask Training for Classification and Retrieval

A Transformers integration proposal submitted on October 14 adds native support for UltraBERT, and the research team published a model announcement the same day. The English encoder uses separate language-modeling heads for two masking ratios, but the integration has not been merged, and its performance in Chinese has not been evaluated.

Bert Verhoeff for Anefo · CC0 · Image source
zh-Hant

On October 14, Transformers received a proposal for native UltraBERT integration. DMIS Lab, the research team, published a model introduction the same day, offering English encoders with approximately 150 million and 380 million parameters, along with ready-to-use embedding models. The series is licensed under Apache 2.0 and supports a 32K context window. The latest developments are the model announcement and framework integration; the related paper was submitted on October 10. Model announcement, paper

UltraBERT’s main change is its pretraining objective. The authors found in experiments that a lower masking ratio benefits classification, while a higher ratio benefits retrieval. They therefore mix sequences with 5% and 50% masking ratios at a four-to-one ratio in the same batch, using a shared encoder but separate language-modeling output heads. The low-mask branch is distilled from a 1.3-billion-parameter teacher, while the high-mask branch retains standard masked language modeling, to avoid distillation harming retrieval performance. Training method

The architecture alternates between local and global attention, which use different RoPE bases. The native integration proposal therefore needs to manage multiple sets of rotary position embeddings and make the model code conform to the Transformers attention-backend interface. The proposal author includes output comparisons for eager, SDPA, and FlashAttention 2. In the listed test cases, the hidden states of valid tokens match those from the original remote code. These are still limited tests and do not establish that all precisions and sequence lengths have been validated. Integration proposal

The authors report that the 380M model achieves GLUE 91.0 and BEIR 47.7, but these results are from task-specific fine-tuning. The base checkpoint still needs to be trained for classification or retrieval, so it should not be treated as a general-purpose embedding model ready for deployment. For engineering teams, the question worth testing is whether one compact backbone can support both document classification and search, reducing the cost of maintaining multiple models. Model card

As of this review, the native integration remains open for review, and the existing model card requires trust_remote_code=True. Follow-up should track the proposal’s merge and inclusion in an official release, and measure long-document latency, memory use, and retrieval quality on your own corpus. The model was trained on English data, and the available scores are not enough to infer its performance in Chinese. Proposal status, usage limitations

Sources

  1. Integrate the UltraBERT architecture — Transformers PR #49596
  2. UltraBERT: One Encoder for Classification and Retrieval
  3. UltraBERT: Beyond the Limit of BERT Pre-Training
  4. UltraBERT-150M 模型卡