← 返回 2026-10-14

UltraBERT:突破BERT预训练的极限 UltraBERT: Beyond the Limit of BERT Pre-Training

Jungwoo Park, Taewhoo Lee, Jaewoo Kang 📅 2026-10-10 👍 61 2026-11-11 18:30
BERT 掩码语言建模 文本嵌入 知识蒸馏 预训练

混合5%与50%双掩码率并解耦LM头,首个分类与检索同时SOTA的编码器

前置知识

掩码语言建模(MLM)与掩码率

MLM 是 BERT 式预训练目标:随机将输入中一定比例的 token 替换为 [MASK],让模型依据双向上下文还原被遮盖的词,损失为 $\mathcal{L}_{\text{MLM}}=-\sum_{i\in M}\log p(x_i\mid\tilde{x})$。掩码率指被遮盖 token 的占比,BERT 传统取 15%。比例越高,每条序列的训练信号越多,但上下文被破坏得越严重。

掩码率是本文的核心自变量:作者证明它决定编码器偏向分类还是检索,并据此提出双掩码率混合与 LM 头解耦这一核心方法。

双向编码器与因果解码器

编码器用双向注意力建模整段输入,输出每个位置的上下文表示,适合判别式任务和语义向量检索;解码器用因果掩码自回归预测下一个词。同参数量下,微调后的编码器在定义明确的任务上常能匹敌大得多的解码器,且推理成本低得多,是 RAG、数据质量过滤、奖励模型的主力。

本文的饱和结论和掩码率偏好都是相对解码器对照得出的,理解两类模型的差异才能体会『编码器下游性能与损失解耦』这一发现的不寻常之处。

知识蒸馏与 Cascade 剪枝蒸馏

知识蒸馏让小模型(学生)模仿大模型(教师)的输出分布,最常用的是 logits 上的 KL 散度损失 $\mathcal{L}_{\text{KD}}=\mathrm{KL}(p_T\Vert p_S)$。Cascade 蒸馏先依据激活统计量对大模型做宽度剪枝得到结构化初始化,再从原教师蒸馏,且可级联进行(如 1.3B→380M→150M),源自 LLM 蒸馏实践。

UltraBERT 的 150M 与 380M 两个尺寸正是这样得到的;论文还发现蒸馏收益强烈依赖掩码率,这是对『蒸馏总是有益』信念的重要修正。

有效秩(effective rank)

衡量一组表示向量实际使用了多少维空间。对协方差矩阵的特征值 $\lambda_1,\dots,\lambda_d$ 归一化后计算熵:$r_{\text{eff}}=\exp\big(-\sum_k \tilde{\lambda}_k\log\tilde{\lambda}_k\big)$,谱越平坦秩越高,表示信息在各维度分布越均匀(越各向同性)。

论文用它刻画不同掩码率塑造的表征几何差异,解释高掩码率为何利于检索,并用 KoLeo 正则做了因果性检验。

替换词检测(RTD / ELECTRA)

用一个小生成器把被遮位置替换成语义上说得通的词,判别器学习辨别每个 token 是原始词还是替换词。训练效率高,传统认为分类能力优于 MLM,但会显著损害检索性能;其优势通常被归因于训练效率而非学习信号本身。

作者用 RTD 做『上下文损坏』的对照实验,证明其优势峰值出现在 15% 掩码处且 CoLA 是混淆变量,这支撑了全文对掩码率与损坏程度的分析框架。

GLUE 与 BEIR 基准

GLUE 是 8 项句/句对分类任务(MNLI、QQP、RTE 等)的平均分,代表判别式语言理解;BEIR 是 15 个零样本检索数据集的平均 $\text{nDCG}@10$(先在 MS MARCO 上微调),代表稠密检索的泛化能力。两者分别对应编码器的两大应用场景。

全文的核心主张就是单一模型同时刷新两端:GLUE 91.0 与 BEIR 47.7,这两个数字是衡量本文贡献的锚点。

研究动机

尽管 ModernBERT、Ettin、NeoBERT 等新一代编码器已吸收 LLM 的架构进步(RoPE、FlashAttention-2、交替局部/全局注意力、数万亿 token 训练),它们的预训练配方——尤其是 MLM 的掩码率——仍沿用 BERT 时代小规模实验的结论。Wettig 等人(2023)与 Gisserot-Boukhlef 等人(2026)『模型越大掩码率应越高』的建议建立在 10B–100B token 的预算上,此时收敛速度与最终性能纠缠在一起,掩盖了掩码率的真实影响。其后果是编码器各有偏科:在本文统一评测协议下,Ettin Encoder-395M 在 GLUE 上达 90.8 但 BEIR 仅 45.9;ModernBERT-149M 检索较强(BEIR 42.6)但 GLUE 88.4;ELECTRA 式 RTD 模型分类强而检索崩塌——DeBERTaV3-large 分类平均 78.1、检索平均仅 25.8。没有任何单一模型能在分类与检索两端同时领先,使用者只能按下游任务二选一。

本文的目标是本文要回答两个递进的问题。其一,在数万亿 token 规模下,编码器预训练的核心设计选择——训练 token 量、掩码率、剪枝与蒸馏——究竟如何影响下游性能,哪些在小预算下得出的旧结论在饱和状态下依然成立?其二,能否据此造出一个单一编码器,同时刷新分类(GLUE、SQuAD 2.0、NER、RewardBench)与检索(BEIR、RAR-b、MLDR、LongEmbed)两端的纪录?作者的量化目标是让一个模型在两个任务家族的全部基准上同时达到 SOTA(GLUE 约 91、BEIR 约 47.7),并原生支持 32K 长上下文,最终提供 150M 与 380M 两个实用尺寸。

与已有工作不同的是,独特之处在于把『掩码率之争』从收敛速度的迷雾中解放出来:先证明编码器在 1–2T token 后下游性能饱和,再在饱和点上比较掩码率,此时收敛快慢不再混淆因果。在这个干净设定下得到与先前文献相反的结论,并把分类与检索的相反偏好重新表述为一个多任务学习问题——把每种掩码率当作一个任务。作者进一步通过表征几何(LM 头的主成分与奇异值谱)定位到任务冲突的根源在输出头而非输入嵌入,从而提出 LM 头解耦这一最小干预方案;再把蒸馏的适用边界与掩码率挂钩,发现蒸馏只在低掩码率下有益。『先消除混淆、再定位冲突、最后选择性蒸馏』这条研究路线在编码器文献中是首次出现。

核心方法

直觉上,两种掩码率塑造两种专长:低掩码率利于分类、高掩码率利于检索,与其折中,不如让一个模型同时经历两种训练信号,但把互不兼容的输出头分开。技术路线分四步:(1)以第 3 节受控实验中训练出的最强 1.3B 编码器为教师;(2)按 Cascade Distillation 做纯宽度剪枝,1.3B→380M,蒸馏后再 380M→150M,150M 直接从 1.3B 教师蒸馏;(3)每个训练 batch 按 4:1 混合 5% 与 50% 掩码序列,两个掩码率各配一个独立 LM 头,5% 序列用教师 logits 的 KL 蒸馏损失、50% 序列用标准 MLM 损失;(4)共训练 2T token(batch 4M token,WSD 学习率调度,峰值 $10^{-4}$),最后 100B token 把上下文从 2K 扩展到 32K(全局注意力 RoPE base 从 640K 增至 2.56M)。架构沿用现代组件:RoPE、SwiGLU、每三层一次全局注意力的局部/全局交替注意力、SmolLM2 分词器(49,152 词表)、固定 24 层、按宽度缩放(hidden size 768/1024/2048)。

核心创新一是 LM 头解耦。作者观察到 6% 与 50% 掩码训练出的编码器输入嵌入几乎相同,但 LM 头几何形态截然不同——6% 模型的头奇异值集中、占据低维子空间,与模型的有效秩一致(Figure 7、8)。由于编码器梯度经由 LM 头反传,共享头会让两种掩码机制在梯度层面互相拖拽。解耦后,{5%,50%} 混合在 150M 上同时达到 GLUE 89.4(逼近单 6% 的 89.5)与 BEIR 45.1(超过单 50% 的 45.0),而不解耦时 GLUE 只有 88.9(Table 1)。核心创新二是蒸馏的条件化:15% 掩码下从 1.3B 教师蒸馏把 GLUE 从 90.2 提到 90.8,但 50% 掩码下 BEIR 反而从 47.2 降到 46.5——教师在高掩码率下的重建信号过于平凡,因此 UltraBERT 把 KL 蒸馏严格限制在 5% 掩码序列上。这与『蒸馏总是有益』的普遍信念形成本质区别。

方法步骤详情

第一步(教师):在与第 3 节相同的现代架构下训练 1.3B 编码器(24 层、hidden 2048)作为教师。第二步(剪枝):仅沿宽度剪枝,用 2000 篇预训练文档打包成 1024 条 1024-token 序列做前向校准,按激活幅度在批内的 $L_2$ 范数估计隐藏通道、注意力头与 FFN 神经元的重要性,保留高分单元;注意力头维度 256→128 时保留每隔一个 RoPE 频率对并把 query 缩放 $1/\sqrt{2}$,value/output 投影用共享随机正交矩阵保证乘积期望不变。第三步(蒸馏训练):batch 含 4 份 5% 掩码序列与 1 份 50% 掩码序列(4:1 补偿低掩码率收敛慢);5% 头的损失为教师与学生 logits 的 KL 散度 $\mathcal{L}_{\text{KL}}\big(\,z_T\,\Vert\,z_S\,\big)$,50% 头用标准 MLM 交叉熵 $\mathcal{L}_{\text{MLM}}=-\sum_{i\in M}\log p(x_i\mid\tilde{x})$。第四步(优化):AdamW($\beta_1=0.9,\beta_2=0.95$、权重衰减 0.1、梯度裁剪 1.0),WSD 调度峰值学习率 $10^{-4}$,batch 4M token(2048 条 × 2048 token),共 2T token,最后 100B 扩展至 32K,在 256 芯片的 TPU(v4-512/v5e-256/v6e-256)上以 16 路数据并行 × 16 路 FSDP 训练。第五步(评估):分类任务逐任务网格搜索微调;检索在 MS MARCO 上用缓存多负例排序损失微调后评测 BEIR、RAR-b、MLDR、LongEmbed。

技术新颖性

新颖性体现在三个层面。实证层面:首次在饱和状态(1–2T token)下系统比较 3%–50% 掩码率,推翻了『模型越大掩码率应越高』的旧结论——在 150M/380M/1.3B 三个规模上分类都随掩码率单调下降,检索则单调上升;且用 KoLeo 正则人为抬高有效秩无法替代高掩码率(BEIR 提升不及 50% 掩码且 GLUE 受损),证明这是预训练目标的内在权衡而非表征统计的表面差异。方法层面:把掩码率混合建模为多任务学习并用 LM 头解耦化解冲突,不同于以往只调调度策略的尝试——Table 1 显示衰减式(50%→6%)与均匀采样 U(5%,50%) 都收敛到『平均掩码率』的效果,只有解耦头才真正两全其美。蒸馏层面:首次刻画 logit 蒸馏与掩码率的交互,发现高掩码率下蒸馏无益甚至有害,给出『蒸馏仅用于低掩码率』的可操作规则,并把 LLM 社区的 Cascade 蒸馏(Muralidharan 等;Liu 等)系统移植到编码器。

Geometry of input embeddings and LM heads.
Figure 7: Geometry of input embeddings and LM heads.
Relative singular values of (a) the input embedding matrix and (b) the LM head, for 150M encoders pre-trained at 6% (solid) and 50% (dashed) masking.
Figure 8: Relative singular values of (a) the input embedding matrix and (b) the LM head, for 150M encoders pre-trained at 6% (solid) and 50% (dashed) masking.
LM head geometry under masking ratio mixing.
Figure 9: LM head geometry under masking ratio mixing.
Overview of UltraBERT.
Figure 10: Overview of UltraBERT.

实验结果

发现一(饱和):150M 与 1.3B 编码器训练至 8.8T token,GLUE 与 SQuAD 2.0 在 1–2T 后不再提升,而训练与验证损失持续下降;BEIR 甚至在首个评测点 275B 后就饱和。对照实验显示 1.3B 编码器指令微调后 MCQA 平均停在 75.3%,而 1.7B 解码器 SmolLM2 指令微调后达 69.9% 且随预训练继续攀升,证明编码器存在 token 无法突破的下游上限。发现二(掩码率):分类上从 3% 到 50% 三个规模一致单调变差;检索上 50% 一贯最佳(150M 的 BEIR:6% 为 43.5,50% 为 45.0 nDCG@10);有效秩随掩码率升高;6% 加 KoLeo 正则虽改善 BEIR 但不及 50% 掩码且损害 GLUE。发现三(RTD 与蒸馏):匹配损坏程度后 RTD 对 MLM 的优势在 15% 处达峰、向两端衰减;从 50% MLM 模型继续 RTD 训练 100B,CoLA 从 61.6 跳至 71.5(+9.9)而其余 GLUE 仅 +0.5,证明 CoLA 是混淆变量。蒸馏在 15% 提 GLUE 90.2→90.8,在 50% 降 BEIR 47.2→46.5。主结果:UltraBERT-380M 在 GLUE 91.0、SQuAD 2.0 F1 90.9、MultiCoNER v2 63.5、Few-NERD 68.3、RewardBench 78.8、BEIR 47.7、RAR-b Code 47.0、RAR-b Math 61.7、LongEmbed 74.9 上除 MLDR(44.1,略低于 LFM2.5-355M 的 44.5)外全部第一;UltraBERT-150M 除 MultiCoNER v2 外全部第一,其 LongEmbed 72.0 甚至超过所有其他家族的 medium 模型;Passkey 任务在 256–32K 全长度保持 ≥99% recall@1。代码/数学检索提升最大:base 组较次优 +7.3/+8.3 分,medium 组 +8.4/+7.7 分。与 DeBERTaV3 相比检索差距达 24.6/22.9 分。

Comparison of masking ratio mixing for the 150M encoder.
Table 1: Comparison of masking ratio mixing for the 150M encoder.
Logit-based distillation from 1.3B to 380M at 15% and 50% ratios.
Table 2: Logit-based distillation from 1.3B to 380M at 15% and 50% ratios.
Comparison with existing encoders on classification and retrieval benchmarks.
Table 3: Comparison with existing encoders on classification and retrieval benchmarks.
Pre-training data mixture.
Table 4: Pre-training data mixture.
Effect of continued RTD training.
Table 6: Effect of continued RTD training.
Passkey retrieval accuracy on LongEmbed as a function of document length, for encoders in two size groups.
Figure 11: Passkey retrieval accuracy on LongEmbed as a function of document length, for encoders in two size groups.
查看结构化数据
任务指标本文基线提升
GLUE(8 任务分类平均) 开发集平均分 91.0(380M)/ 89.5(152M) Ettin Encoder 90.8(395M)/ 88.9(149M) +0.2 / +0.6
SQuAD 2.0(抽取式问答) F1 90.9 / 89.2 Ettin Encoder 90.4 / 87.4 +0.5 / +1.8
BEIR(15 数据集零样本检索) 平均 nDCG@10 47.7 / 45.2 LFM2.5-355M 47.3 / ModernBERT-149M 42.6 +0.4 / +2.6
RAR-b Code(代码检索) nDCG@10 47.0 / 39.1 次优:LFM2.5-355M 38.6 / Ettin-149M 31.8 +8.4 / +7.3
RAR-b Math(数学检索) nDCG@10 61.7 / 55.4 次优:Ettin-395M 54.0 / Ettin-149M 47.1 +7.7 / +8.3
RewardBench(偏好分类/奖励建模) 4 部分平均准确率 78.8 / 73.0 Ettin Encoder 76.7 / 70.3 +2.1 / +2.7
LongEmbed(长上下文嵌入) 6 任务平均 74.9 / 72.0(32K 上下文) LFM2.5-355M 71.7 / ModernBERT-149M 63.2(8K) +3.2 / +8.8

局限与改进

作者承认的局限:medium 尺寸在 MLDR 上 44.1 低于 LFM2.5-355M 的 44.5;base 尺寸在 MultiCoNER v2 上 57.9 非组内最佳(NomicBERT 60.1);medium CoLA 71.2 以 0.2 分之差低于 ModernBERT 的 71.4;有效秩与检索性能的因果方向作者自己也表示无法确定。我自己的观察:其一,5%/50% 两个端点与 4:1 配比是经验选择,论文未做配比敏感性分析;其二,双 LM 头增加约 $2\times|V|\times d$ 的参数(词表 49,152 时并非小数目),且 Table 3 的参数计数明确排除 LM 头,与基线比较时略不对等;其三,下游微调时该用哪个头(分类用 5% 头?检索用哪个?)论文没有明说,实践者需自行摸索;其四,『50% 蒸馏失败』只有教师信号平凡的类比性猜测,未做机理验证;其五,所有实验限于英文,训练数据 66.9% 为英文网页,结论对多语言是否成立未知;其六,饱和结论基于特定数据混合与架构,其他配方下饱和点可能移动。

独立分析的弱点

弱点一:双头设计的微调接口含糊——预训练有两个 LM 头,但 GLUE/NER 等判别微调和 MS MARCO 对比微调各自该用哪个头没有说明,若两个头都保留还会增加部署复杂度。改进方向:文档化默认用法(如冻结 50% 头),或研究头合并/蒸馏成单头的方案。弱点二:4:1 配比与 5%/50% 端点缺少消融,作者仅用脚注说明以 4 倍序列补偿 5% 的慢收敛,未探索 2:1、8:1 或加入 15% 构成三点混合。改进:网格搜索配比,或按两项损失的梯度冲突自动调节权重。弱点三:整条流水线依赖先训出的 1.3B 教师,且配方探索本身消耗了多组万亿 token 级训练(8.8T×2 规模),小团队无法重复这一探索过程。改进:开源教师权重(已承诺)并发布中间检查点可部分缓解。弱点四:高掩码率蒸馏失败的机理未验证,『教师信号平凡』只是类比 RTD 的猜测。改进:测量教师与学生 logits 的一致性、KL 梯度范数随掩码率的变化来定位原因。弱点五:评测全部集中在英文基准上。改进:在 EuroBERT 式多语言设定下复验掩码率结论与双头设计。弱点六:CoLA 被认定为目标混淆变量并从 MLM-vs-RTD 比较中剔除,但主表仍报告 CoLA——低掩码率的 UltraBERT 在 base 组 CoLA 65.2 低于 Ettin 的 66.9,这类代价可能被低估。

未来方向

作者提出的方向:结论章指出编码器的偏科『不是固有属性而是配方的结果』,暗示该配方应推广到其他规模与领域;附录 A.3 观察到多阶段数据混合对微调后性能几乎无影响、只改变零/少样本行为,作者认为值得研究暴露比例为何不改变微调所依赖的内在语言学知识。基于本文成果可延伸:其一,连续掩码率课程——把 {5%,50%} 两点扩展为随训练动态调度的连续分布,并与已证明失败的衰减式调度对照;其二,低掩码率 RTD 与高掩码率 MLM 的混合目标(作者已证明 RTD 优势集中于 15% 附近且低掩码率下与 MLM 无异);其三,研究 50% 头学到的表征在对比微调中保留多少,打通预训练-检索微调的配接;其四,对饱和现象做理论刻画——为何 MLM 编码器的下游性能与损失解耦而解码器持续改善;其五,多语言 UltraBERT 与多语言下的掩码率偏好验证;其六,把『按任务选择性施加辅助损失』的思想推广到句子级对比预训练等其他目标;其七,在 <100M 的边缘设备规模下检验结论是否保持。

复现评估

开源情况:作者承诺发布评测代码、模型权重与预训练数据配比,可复现性写作非常扎实——附录 A 给出完整架构、优化器、剪枝校准与蒸馏细节,附录 A.3 给出逐数据集 token 数的混合表,附录 B 给出每个基准的微调网格与数据来源,附录 E 给出全部逐任务分数;Table 7 还验证其 GLUE 协议能在 0.2 分内复现 Ettin 官方分数,评测可信度高。数据全部为公开数据集(DCLM、FineWeb-Edu、Stack-Edu、peS2o、MegaMath、FineMath、FLAN 等)。算力门槛:预训练在 256 芯片的 TPU v4-512/v5e-256/v6e-256 上进行,主训练 2T token,前置的配方探索更是训练了 150M 与 1.3B 两个规模共 8.8T token 的多组模型,完整端到端复现研究结论需要机构级 TPU 资源(估计数千芯片日)。但使用门槛低得多:权重发布后,150M 模型单卡即可微调,380M 亦属轻量;检索评测协议需 8 卡 GPU 做大 batch 对比学习。总体而言:微调与评测复现容易(高质量协议+公开数据),预训练复现难但方法描述完整,属于编码器文献中复现性最好的工作之列。

相似论文推荐