[论文解读] Single Headed Attention RNN: Stop Thinking With Your Head

Stephen Merity|arXiv (Cornell University)|Nov 26, 2019

Topic Modeling参考文献 17被引用 62

一句话总结

本文提出 SHA-RNN，一种具有单头注意力的 LSTM 模型，在仅在单个 GPU 上进行极少超参数调优的情况下，接近 enwik8 的字节级语言建模的状态-of-the-art，挑战多头注意力的必要性。它还讨论了分词敏感性和训练效率的实际考虑。

ABSTRACT

The leading approaches in language modeling are all obsessed with TV shows of my youth - namely Transformers and Sesame Street. Transformers this, Transformers that, and over here a bonfire worth of GPU-TPU-neuromorphic wafer scale silicon. We opt for the lazy path of old and proven techniques with a fancy crypto inspired acronym: the Single Headed Attention RNN (SHA-RNN). The author's lone goal is to show that the entire field might have evolved a different direction if we had instead been obsessed with a slightly different acronym and slightly different result. We take a previously strong language model based only on boring LSTMs and get it to within a stone's throw of a stone's throw of state-of-the-art byte level language model results on enwik8. This work has undergone no intensive hyperparameter optimization and lived entirely on a commodity desktop machine that made the author's small studio apartment far too warm in the midst of a San Franciscan summer. The final results are achievable in plus or minus 24 hours on a single GPU as the author is impatient. The attention mechanism is also readily extended to large contexts with minimal computation. Take that Sesame Street.

研究动机与目标

• 激发在语言建模中探索非 Transformer 为中心的架构，并展示基于低资源的 LSTM 营造的具有竞争力的结果。
• 引入一个简单、内存高效的注意力机制，与 LSTM（SHA-RNN）以及 Boom 前馈层集成。
• 表明单个注意力头就能带来大部分收益，并且内存/计算效率使在普通硬件上进行快速训练成为可能。
• 讨论分词、训练工具以及超越 Transformer 主导文献的更广泛研究方向的含义。

提出的方法

• 提出将嵌入层、一个或多个带单个注意力头的 SHA-RNN 层，以及权重绑定的 softmax 分类器结合的 SHA-RNN 架构。
• 使用单头的简化注意力机制，以及一次性写入的记忆（没有每个时间步的大规模矩阵乘法）。
• 引入 Boom 层，一种计算高效的前馈组件，替代传统下投影层的一部分。
• 应用 LAMB 优化器的最小信任版本，以在没有残差连接的网络上改善收敛。
• 在字节级 enwik8 上进行实验，并讨论与词级分词以及 WikiText 数据集的比较。

实验结果

研究问题

• RQ1 在 LSTM 之上单头注意力机制是否能在极少调优的情况下接近或达到字节级语言建模的最先进水平？
• RQ2 减少注意力头数量如何影响 enwik8 的性能与训练效率？
• RQ3 分词方案对困惑度及跨模型可比性的实际影响有哪些？
• RQ4 哪些架构组件（Boom 层、过参数化的静态向量）有助于训练稳定性和性能？

主要发现

• 单头注意力的 SHA-RNN 在性能提升上几乎可与多头配置相媲美。
• 在单个 GPU 上进行极少超参数调优的训练，在大约描述的训练时间内产出具有竞争力的结果（某些设置的 epoch 时间约为 1800 秒）。
• 在 enwik8 实验中，单头 SHA-LSTM 在效率方面取得显著提升，1 头模型的训练速度快于 4 头变体（约 30 分钟/ epoch 对比 ~67 分钟）。
• Boom 层相较于传统的下投影层在参数和计算上更高效，同时仍提供出色的性能。
• 分词选择会显著影响困惑度，凸显了可能的分词攻击以及当分词不同导致跨模型公平比较的难度。
• 该工作主张在 Transformer 之外保持模型多样性，提出 SHA-RNN 作为可行的替代方案，以及为蒸馏或更广泛的可获取性提供潜在基础。