← 全部论文 / 2019

ALBERT: A Lite BERT for Self-supervised Learning of Language Representations

2019 TPU、加速器与 Transformer 时代 被引用 984 次
下载 PDF

ALBERT: A Lite BERT for Self-supervised Learning of Language Representations - 中文验证版

英文原文卡片：albert_2019.md

状态：已翻译。

元数据

  • Slug： albert_2019
  • 年份： 2019
  • 会议： ICLR
  • 作者： Zhenzhong Lan et al.
  • 阅读状态： read complete
  • 计算范式： TPU、加速器与 Transformer 时代
  • 主要来源： PDF、抽取文本

计算设置

论文明确给出了加速器设置：训练在 Cloud TPU v3 上完成，TPU 数量从 64 到 512 不等，具体取决于模型规模。所有模型更新使用 batch size 4096 和 LAMB 优化器，学习率为 0.00176。除非另有说明，模型训练 125,000 steps，最大输入长度 512，并有 10% 的概率使用较短的随机序列。这是一个大 batch TPU 训练范式，而非小规模 GPU 复现。

模型配置显示了硬件内存和通信为何是核心问题。BERT-base 有 108M 参数，BERT-large 有 334M。ALBERT-base 有 12M 参数，ALBERT-large 有 18M，ALBERT-xlarge 有 60M，ALBERT-xxlarge 约 235M 参数，同时使用 hidden size 4096 和 embedding size 128。来源文本还报告了等时比较：在相同数量的 TPU 上，BERT-large 用 34 小时训练 400K steps，而 ALBERT-xxlarge 用 32 小时训练 125K steps。

瓶颈

瓶颈是缩放 Transformer 编码器的参数和通信成本。论文指出，更大的模型会遇到 GPU/TPU 内存限制，而分布式训练因通信开销与模型参数数量成正比而受阻。标准的内存节省方法如 checkpointing 和可逆层以牺牲速度来减少内存，但它们并不直接减少参数通信。BERT 式的缩放还将 WordPiece embedding size 与 hidden size 绑定，因此增加 H 会使词汇 embedding 矩阵膨胀，该矩阵包含许多在训练中稀疏更新的参数。

ALBERT 的核心计算主张是表示容量与参数数量应当解耦。较大的隐藏层、注意力头和深度可以改善下游性能，但简单增加它们会产生一个更难装入加速器内存且在 TPU 芯片之间同步成本更高的模型。因此瓶颈不仅仅是 FLOPs，而是携带许多冗余或低效用参数所带来的内存和 all-reduce 成本。

方法适配

ALBERT 通过两种参数缩减手段来适配 BERT 式预训练。首先，分解 embedding 参数化将词汇到隐藏层的投影分解为两个较小的矩阵。模型不是将 one-hot WordPiece 向量直接投影到 hidden size H，而是先投影到较小的 embedding size E，然后再投影到 H。在主要配置中 E 固定为 128，因此隐藏层可以增长而不会使词汇 embedding 矩阵爆炸。这专门针对稀疏 embedding 的内存和通信问题。

其次，ALBERT 跨层共享参数。默认设置在所有层之间共享所有 Transformer 参数，但论文也消融了仅共享注意力或前馈参数的情况。共享使参数数量不会随深度线性增长，因此深度和重复计算可以增加，而通信的参数状态保持较小。从计算角度看，ALBERT 以重复使用相同权重来处理较大隐藏表示，从而换取参数内存和同步的节省。

论文还用 sentence-order prediction（SOP）替换了 BERT 的 next-sentence prediction。SOP 使用两个连续片段作为正例，交换其顺序作为负例。这主要不是为了降低硬件成本，而是改善句子间连贯性的训练信号，使得较小的参数预算不会被浪费在更容易的主题预测目标上。结合起来，该方法将 Transformer 编码器适配到 TPU 内存和通信限制，同时保持下游容量。

证据

论文的参数表是主要的计算证据。ALBERT-large 有 18M 参数，而 BERT-large 有 334M，在可比的大型配置下缩小了 18 倍。ALBERT-xxlarge 使用 hidden size 4096，仅约 235M 参数，约为 BERT-large 参数量的 70%，尽管其隐藏层更宽。作者报告 ALBERT-xxlarge 在 SQuAD v1.1、SQuAD v2.0、MNLI、SST-2 和 RACE 上优于 BERT-large，其中包括文本中描述的 RACE 提升 8.4 分。

等时实验对加速器预算尤其相关。在几乎相同的训练时间和 TPU 数量下，BERT-large 在 400K steps 和 34 小时后，所列任务的平均得分为 87.2。ALBERT-xxlarge 在 125K steps 和 32 小时后达到 88.7。RACE 对于 BERT-large 为 77.3，对于 ALBERT-xxlarge 为 82.5；SQuAD v2.0 分别为 86.9/84.3 和 88.3/85.3。这支持了以下主张：在固定 wall-clock 预算下，减少参数通信同时允许更大的隐藏结构可以改善结果。

最终报告的 state-of-the-art 数字也很重要。论文指出 GLUE 89.4，RACE 准确率 89.4，SQuAD 2.0 F1 92.2。这些是在大规模 TPU 范式下取得的，因此结果不是“小模型在小硬件上”；而是“大规模加速器训练变得更具参数效率”。

历史影响

ALBERT 将部分 Transformer 缩放从原始参数数量转向了参数高效的加速器利用。在 BERT/RoBERTa 时期，更大通常意味着更多参数、更多内存和更多分布式同步。ALBERT 表明可以将计算投入到更大的隐藏表示和重复的层应用上，同时保持参数状态较小。

这在历史上很重要，因为它分离了经常被混为一谈的三个量：参数数量、激活/FLOP 成本和下游容量。ALBERT 并没有让编码器预训练变得便宜，但它清楚地表明内存和通信可以主导 TPU 和 GPU 上可行的设计空间。其分解和共享思想成为高效 Transformer 缩放更广泛词汇的一部分。

局限

主要局限是较少的参数并不意味着较低的计算量。论文明确指出 ALBERT-xxlarge 比 BERT-large 有更少的参数和更好的结果，但由于其更大的结构，计算开销更高。Hidden size 4096 仍然产生较大的激活张量、注意力计算和前馈 FLOPs。参数共享减少了内存和通信，但不消除反复应用宽 Transformer block 的成本。

作者将稀疏注意力和块注意力指向为加速训练和推理的未来工作方向。这个限定条件很重要：ALBERT 更直接地解决了参数内存和同步问题，而不是序列长度的注意力成本。它也是一种用于自然语言理解基准的 encoder-only 预训练方法，而不是自回归生成接口。其历史教训是在 TPU 约束下实现高效容量，而非更便宜服务的通用配方。

链接