[论文解读] When Attention Meets Fast Recurrence: Training Language Models with Reduced Compute

Tao Lei|arXiv (Cornell University)|Feb 24, 2021

Topic Modeling被引用 6

一句话总结

本文提出 SRU++，一种高效序列模型，通过结合快速循环机制与最小注意力机制，在将训练计算量减少 3x–10x 的同时，实现了最先进的语言建模性能。通过用自注意力模块替代 SRU 中的输入投影，SRU++ 保持了并行计算能力，并在 enwik8、Wiki-103 和 Billion Word 数据集上，以显著更少的 GPU 训练时间，实现了优于顶尖 Transformer 模型的每字符位数（bits-per-character）和困惑度（perplexity）。

ABSTRACT

Large language models have become increasingly difficult to train because of the growing computation time and cost. In this work, we present SRU++, a highly-efficient architecture that combines fast recurrence and attention for sequence modeling. SRU++ exhibits strong modeling capacity and training efficiency. On standard language modeling tasks such as Enwik8, Wiki-103 and Billion Word datasets, our model obtains better bits-per-character and perplexity while using 3x-10x less training cost compared to top-performing Transformer models. For instance, our model achieves a state-of-the-art result on the Enwik8 dataset using 1.6 days of training on an 8-GPU machine. We further demonstrate that SRU++ requires minimal attention for near state-of-the-art performance. Our results suggest jointly leveraging fast recurrence with little attention as a promising direction for accelerating model training and inference.

研究动机与目标

• 为应对训练大语言模型日益增长的计算成本，探索纯注意力机制 Transformer 的高效架构替代方案。
• 研究结合快速循环与最小注意力是否能在大幅降低训练与推理成本的同时，保持高建模能力。
• 证明循环与注意力在序列建模中是互补组件，其中循环机制可在更低计算成本下实现优异性能。
• 提供一种计算高效的 Transformer 替代方案，仅需极少注意力机制且无需位置编码，但仍能实现接近最先进水平的结果。

提出的方法

• SRU++ 通过用自注意力机制替换 SRU 的输入投影，扩展了简单循环单元（SRU），同时保持了架构的并行计算能力。
• 模型使用单层自注意力组件，在应用 SRU 的门控循环机制前，计算上下文感知的输入表征。
• 该架构避免使用位置编码与多头注意力，转而依赖单头注意力组件，以减少计算开销。
• 模型保持了 SRU 的高效参数化与计算方式，支持快速前向传播与更低的内存占用。
• 训练采用标准语言建模范式，结合混合精度训练与标准优化设置。
• 实现代码以 PyTorch 开源，支持可复现性与未来架构扩展。

实验结果

研究问题

• RQ1 结合快速循环与最小注意力的模型是否能在标准语言建模基准上实现最先进性能，同时显著降低训练成本？
• RQ2 在困惑度与每字符位数方面，采用有限注意力的循环模型相较于完整注意力的 Transformer 模型表现如何？
• RQ3 在序列模型中，注意力机制可被最小化到何种程度而不损失建模能力？
• RQ4 当与快速循环结合时，缺乏位置编码与多头注意力是否会损害模型性能？

主要发现

• 在 enwik8 数据集上，SRU++ 仅用 8-GPU 机器训练 1.6 天，即达到 1.017 的最先进每字符位数，优于训练成本仅为 1/8 的 Transformer-XL。
• SRU++ 在 Wiki-103 与 Billion Word 数据集上，困惑度达到或优于顶尖 Transformer 模型，同时训练计算量减少 3x–10x。
• 模型仅使用两层注意力机制即实现接近最先进性能，证明在结合快速循环时，最小注意力已足够。
• 由于减少了矩阵乘法与张量转置开销，特别是相对注意力组件中的开销，SRU++ 的每批次推理速度比 Transformer-XL 快 4–5 倍。
• 尽管无需位置编码与多头注意力，模型仍保持强大性能，表明循环机制可有效建模长距离依赖。
• SRU++ 的 PyTorch 实现已开源，支持可复现性，并推动高效序列建模架构的进一步研究。