我要投稿

BERT新版本：ModernBERT -- Smarter, Better, Faster, Longer

发布日期：2024-12-25 12:11:48 浏览次数： 2595

作者：鸿煊的学习笔记

微信搜一搜，关注“鸿煊的学习笔记”

一、背景

自 BERT 发布后，仅编码器的 Transformer 模型在自然语言处理（NLP）的检索和分类任务中占据主导地位。尽管大语言模型（LLMs）不断发展，但仅编码器模型因其推理要求适度，在处理大规模文档检索和判别性任务时仍广泛应用。然而，当前的应用大多依赖旧模型，如原始 BERT，存在诸多局限性，包括序列长度受限、模型设计和词汇量不理想、架构效率低以及训练数据有限且领域狭窄等问题。

二、创新点

（一）架构改进

1. Transformer 特性集成

偏置项调整：除最终解码器线性层外，所有线性层的偏置项被禁用，层归一化中的偏置项也被去除，使更多参数可用于线性层。
位置嵌入优化：采用旋转位置嵌入（RoPE）替代绝对位置嵌入，因其在短、长上下文语言模型中表现出色，且便于上下文扩展，在多数框架中实现高效。
归一化策略改进：使用预归一化块和标准层归一化，在嵌入层后添加层归一化，并对注意力层中的归一化进行优化，有助于稳定训练。
激活函数升级：采用基于 Gated-Linear Units（GLU）的 GeGLU 激活函数，相比原始 BERT 的 GeLU 激活函数有经验性改进。

2. 效率提升措施

交替注意力机制：每三层使用全局注意力（theta 为 160,000），其余层采用局部注意力（128 token 滑动窗口，theta 为 10,000），提高计算效率。
无填充技术应用：在训练和推理中去除填充令牌，通过 Flash Attention 的可变长度注意力和 RoPE 实现，减少计算浪费，提升性能。
先进的注意力算法：在全局和局部注意力层分别使用 Flash Attention 3 和 Flash Attention 2，提供内存和计算高效的注意力内核。
编译优化加速：利用 PyTorch 的内置编译功能，提升训练效率，且编译开销可忽略不计。

（二）数据与训练优化

1. 大规模数据训练

基于多种数据源的 2 万亿令牌进行训练，采用现代数据混合策略，通过实验确定最佳数据混合方式，使模型学习到更丰富的语言知识。

2. 分词器优化

使用改进的 OLMo 分词器，相比多数重用原始 BERT 分词器的编码器，其在代码相关任务中令牌效率和性能更高，且词汇量设置为 50,368（64 的倍数），包含 83 个未使用令牌以支持下游应用，优化了 GPU 利用率。

3. 序列打包技术

采用贪婪算法进行序列打包，确保训练批次内序列长度的均匀性，避免因无填充导致的批量大小方差问题。

3. 训练设置精细调整

改进的 MLM 设置：遵循 MosaicBERT 的 MLM 设置，去除 Next-Sentence Prediction 目标以减少开销，同时将掩码率从 15% 提高到 30%，提升训练效果。
优化器选择：使用 StableAdamW 优化器，通过 Adafactor-style 的更新裁剪实现自适应学习率调整，改进了 AdamW 优化器，使训练更稳定，在下游任务中表现更优。
学习率调度策略：采用修改后的梯形学习率调度（Warmup-Stable-Decay），结合 1 - sqrt LR 衰减，相比线性和余弦衰减效果更好，且能避免冷启动问题，实现持续训练。
批量大小动态调整：从较小梯度累积批量开始，随时间增加到全批量大小，如 ModernBERT-base 从 768 增加到 4,608，ModernBERT-large 从 448 增加到 4,928，通过不均匀的令牌调度确保每个批量大小有相同的更新步骤，加速训练。
权重初始化与继承：ModernBERT-base 使用 Megatron 初始化，ModernBERT-large 从 ModernBERT-base 的权重初始化，加速大型模型的初始训练。
上下文长度扩展训练：先在 1024 序列长度上训练，然后通过调整 RoPE theta 扩展到 8192 序列长度，并在扩展后的长度上进行额外训练，采用特定的学习率策略，平衡了不同任务的性能。

（三）高效模型设计

通过广泛的小规模实验，旨在优化模型在常见 GPU 上的推理效率。ModernBERT 的base model和large model分别设计为 22 层和 28 层，对应的参数总量分别为 1.49 亿和 3.95 亿。在隐藏层大小和 GLU 扩展比例方面进行精细调整，确保在不同 GPU 配置下，特别是目标 GPU 中的张量核心和平铺效率，从而在下游任务性能和硬件利用率之间实现最佳平衡。这种设计策略使得模型在实际应用中能够充分利用硬件资源，同时保持良好的性能表现。

三、模型训练

（一）数据处理

1. 数据混合

模型在 2 万亿令牌的大规模数据上进行训练，数据来源广泛，涵盖网络文档、代码和科学文献等。通过一系列实验确定最终的数据混合策略，以确保模型能够学习到多样化的语言表达和知识。

2. 分词器

采用改进的 OLMo 分词器，与多数重用原始 BERT 分词器的编码器不同，该分词器在处理代码相关任务时表现出更高的令牌效率和性能。词汇量设定为 50,368，这是 64 的倍数，其中包含 83 个未使用令牌，专门用于支持下游应用，同时优化了 GPU 的内存利用效率。

3. 序列打包

为避免因无填充操作导致训练批次内序列长度的高度变化，采用贪婪算法进行序列打包。这一技术确保了训练过程中批量大小的高度一致性，从而稳定训练过程，提高训练效率。

（二）训练设置

见训练设置精细调整

四、下游任务评估

（一）评估设置

1. 自然语言理解（NLU）

使用 GLUE 基准测试，这是衡量模型在句子或句对理解任务（如情感检测、语言蕴含等）上性能的标准测试。通过在每个 GLUE 子集中进行超参数搜索，确保与其他模型的可比性。

2. 文本检索

单向量检索：在 DPR 设置下，模型使用单个向量表示文档，通过对比学习进行微调，以计算查询与文档间的相似度（如余弦相似度）。在 MSMARCO 数据集上进行训练，并使用挖掘的硬负例增强训练效果。
多向量检索：采用 ColBERT 模型，通过每个文档的所有单个令牌向量表示文档，使用 MaxSim 算子计算查询与文档的相似度。在 MS-Marco 数据集上训练，并通过知识蒸馏从教师模型中获取信息。

3. 长上下文检索

由于缺乏标准化长上下文基准，在 MLDR（包含 200,000 多个长文档）的英语子集中进行评估，设置了单向量（域外和域内）和多向量（域外）三种评估方式。

4. 代码检索

使用 CodeSearchNet 和 StackOverflow-QA 基准测试评估模型在编程相关任务中的性能，将其视为单向量检索任务，并重用之前确定的最佳超参数进行训练。

（二）评估结果

1. 下游性能总体表现

ModernBERT 在 BASE 和 LARGE 模型尺寸上均表现卓越，在所有评估类别中超越原始 BERT 和 RoBERTa 模型，实现了全面的性能提升。

2. 短上下文检索

在 BEIR 评估套件的 DPR 和 ColBERT 设置下，ModernBERT 均优于现有编码器，包括专门为检索设计的 GTE-en-MLM 和 NomicBERT 模型。ModernBERT-large 在参数较少的情况下仍领先于 GTE-en-MLM-large。

3. 长上下文检索

单向量检索
在 DPR 设置下的 MLDR 任务中，ModernBERT 表现出色，但在未进行长上下文微调时，虽优于短上下文模型和 NomicBERT，但落后于 GTE-en-MLM；域内评估时，两者性能接近，表明 ModernBERT 可处理长上下文，但可能需更适配的调优。
多向量检索
在 ColBERT 设置下，长上下文模型（GTE-en-MLM、NomicBERT 和 ModernBERT）表现优异，ModernBERT 领先其他长上下文模型，可能得益于长预训练和局部注意力与 ColBERT 式检索的协同作用。

4. 自然语言理解

在 GLUE 测试中，ModernBERT 模型展现出卓越的 NLU 能力。ModernBERT-base 超越所有现有基础模型，成为首个超越 DeBERTaV3-base 的 MLM 训练模型；ModernBERT-large 是第二好的大型编码器，参数少但处理速度快，接近 DeBERTaV3-large 的性能。

5. 代码任务

在编程任务中，ModernBERT 在代码到文本（CodeSearchNet）和长上下文混合设置（StackQA）中表现出色，因其是唯一在训练数据中包含编程数据的评估编码器，表明其在代码理解上的提升未影响自然文本处理能力。

五、模型效率

（一）评估设置

为衡量模型在不同序列长度下的推理效率，创建了 4 组合成文档集，包括固定短上下文（512 令牌）、固定长上下文（8192 令牌）以及两种可变长度（分别以 256 和 4096 令牌为中心的正态分布）。在 NVIDIA RTX 4090 上评估模型每秒处理的令牌数，对 GTE-en-MLM 模型进行了原生和使用 xformers 库增强两种设置的测试，所有测试均重复十次取平均值。

（二）评估结果

ModernBERT 在效率方面表现卓越。在短上下文处理上，虽比原始 BERT 和 RoBERTa 稍慢，但比其他近期编码器快；在长上下文处理上，显著快于所有竞争编码器，处理速度分别比同类模型中第二快的编码器快 2.65 倍（BASE）和 3 倍（LARGE）。在可变长度输入方面，ModernBERT 比 GTE-en-MLM 更高效，尤其在长上下文时优势明显，这得益于其局部注意力机制。同时，ModernBERT 在内存利用上也表现出色，在两种模型尺寸下均是最内存高效的模型，如 ModernBERT-base 在处理长上下文输入时，可处理的最大批量大小是其他模型的两倍，ModernBERT-large 在短上下文输入时虽略逊于原始 BERT-large，但在处理长上下文输入时，其可处理的批量大小比其他大型模型大至少 60%。

五、结论

ModernBERT 是仅编码器模型领域的重大突破，在广泛的分类和检索任务上超越现有模型，实现了更快的处理速度和更高的内存效率。通过采用现代架构改进，如 GeGLU 层、RoPE 位置嵌入和交替局部 - 全局注意力机制，结合大规模训练数据和优化的训练策略，模型在多个基准测试中取得领先成绩。其原生支持 8192 令牌的长序列处理，为长上下文应用提供了有力支持。尽管仍存在一些局限性，如仅适用于英语、可能存在数据偏差、有潜在的有害内容生成风险、训练目标单一以及未充分探索模型参数缩放等，但整体上推动了编码器技术的显著发展，为自然语言处理任务提供了更强大、高效的工具。

53AI，企业落地大模型首选服务商

产品：场景落地咨询+大模型应用平台+行业解决方案

承诺：免费场景POC验证，效果验证后签署服务协议。零风险落地应用大模型，已交付160+中大型企业