中文

MosaicBERT:面向快速预训练优化的双向编码器

计算与语言 2024-01-17 v2 机器学习

摘要

尽管BERT风格的编码器模型在自然语言处理研究中被广泛使用,但由于训练成本高昂,许多研究者并不从头预训练自己的BERT。在BERT首次崭露头角的过去五年中,其他Transformer架构和训练配置取得了许多进展,但这些进展尚未被系统地整合到BERT中。在此,我们介绍MosaicBERT,一种经验优化的BERT风格编码器架构和训练方案,旨在实现快速预训练。这种高效架构将FlashAttention、线性偏置注意力(ALiBi)、门控线性单元(GLU)、动态移除填充令牌的模块以及低精度层归一化整合到经典的Transformer编码器块中。训练方案包括掩码语言建模(MLM)目标30%的掩码比例、bfloat16精度、针对GPU吞吐量优化的词汇表大小,以及来自RoBERTa和其他编码器模型的最佳实践。在C4数据集上从头预训练时,该基础模型在8块A100 80 GB GPU上仅用1.13小时(成本约20美元)就达到了下游平均GLUE(开发集)得分79.6。我们绘制了广泛的准确率与预训练速度帕累托曲线,并表明MosaicBERT基础版和大模型在与竞争性的BERT基础版和大模型相比时始终处于帕累托最优。这种预训练速度的经验提升使研究人员和工程师能够以低成本预训练自定义的BERT风格模型,而不是在现有的通用模型上进行微调。我们开源了模型权重和代码。

关键词

引用

@article{arxiv.2312.17482,
  title  = {MosaicBERT: A Bidirectional Encoder Optimized for Fast Pretraining},
  author = {Jacob Portes and Alex Trott and Sam Havens and Daniel King and Abhinav Venigalla and Moin Nadeem and Nikhil Sardana and Daya Khudia and Jonathan Frankle},
  journal= {arXiv preprint arXiv:2312.17482},
  year   = {2024}
}

备注

10 pages, 4 figures in main text. 25 pages total