Request: pre-decay (stable-phase) checkpoints and the 50%-masking LM head, for continued pretraining
Hi DMIS team, thanks for UltraBERT and for the unusually thorough release!
I'm adapting UltraBERT-150M for a cybersecurity search and classification stack (a cross-encoder reranker distilled from an LLM teacher, plus document classification). We'd like to continue pretraining on domain text (security advisories, CVE/CWE/CAPEC, internal technical documentation) with general-data replay, before task fine-tuning.
Two requests, if you're able to share them:
- A checkpoint from the end of the stable LR phase, before the final linear decay, for 150M and 380M. Continuing from a fully decayed checkpoint means re-warming the LR, which the continual-pretraining literature (Gupta et al. 2023; Ibrahim et al. 2024) shows costs a loss spike and some forgetting. A pre-decay checkpoint would let us follow the BioClinical ModernBERT approach: continue at a stable LR, then run our own decay.
- The 50%-masking LM head. The model card notes only the 5% (distilled) head ships. With both heads, we could continue your decoupled-head recipe (5% + 50% at 4:1) during domain adaptation, rather than training a fresh 50% head from the 5% one, and keep both the classification and retrieval strengths you describe.
In return, here's what we've measured so far, in case it's useful:
- Fine-tuned as a cross-encoder reranker on ~90k LLM-graded pairs, UltraBERT-150M started well ahead of ModernBERT-base and converged with it after full training. That's consistent with your data-efficiency finding.
- Starting from the raw checkpoint matched starting from an MS MARCO-tuned one, once distilled on in-domain labels.
- For short task-adaptive pretraining before distillation, 5% masking beat 50% for the downstream cross-encoder, even though 50% gave much lower MLM loss on the domain text.
Happy to share more detail on any of it. Thanks again!
Hi, thank you for your interest in UltraBERT and for sharing your results!
- Unfortunately, due to several constraints, we are unable to release the intermediate checkpoints. We understand the concern about the loss spike and forgetting when re-warming from a decayed checkpoint, and we hope the effect turns out to be small for your domain adaptation.
- We have added the 50%-masking LM head. Please see here.
We were also curious that UltraBERT-150M converged with ModernBERT-base after full reranker training, as we thought the early gap might carry through. Could you tell us a bit more about it?