[论文解读] N-Grammer: Augmenting Transformers with latent n-grams

Aurko Roy, Rohan Anil+14 more|arXiv (Cornell University)|Jul 13, 2022

Natural Language Processing Techniques被引用 6

一句话总结

N-Grammer 通过注意力头输出的乘积量化，引入可学习的离散潜在 n-gram，增强 Transformer 模型，实现更快的推理速度，且不损失性能。其在 C4 和 SuperGLUE 上的表现与更大的 Transformer 模型及 Primer 相当，同时减少了对昂贵密集矩阵运算的依赖。

ABSTRACT

Transformer models have recently emerged as one of the foundational models in natural language processing, and as a byproduct, there is significant recent interest and investment in scaling these models. However, the training and inference costs of these large Transformer language models are prohibitive, thus necessitating more research in identifying more efficient variants. In this work, we propose a simple yet effective modification to the Transformer architecture inspired by the literature in statistical language modeling, by augmenting the model with n-grams that are constructed from a discrete latent representation of the text sequence. We evaluate our model, the N-Grammer on language modeling on the C4 data-set as well as text classification on the SuperGLUE data-set, and find that it outperforms several strong baselines such as the Transformer and the Primer. We open-source our model for reproducibility purposes in Jax.

研究动机与目标

• 在不损害性能的前提下降低大型 Transformer 模型的推理成本。
• 探索受统计语言建模和离散潜在表示启发的高效架构改进。
• 通过学习的 n-gram 结构，用稀疏嵌入查找替换密集运算，实现更快的推理。
• 在 C4 和 SuperGLUE 等标准 NLP 基准上评估潜在 n-gram 的有效性。

提出的方法

• N-Grammer 层对每个注意力头的输出应用乘积量化（PQ），生成离散的潜在代码。
• 连续的潜在代码构成二元组，并通过查找表映射到可训练的 n-gram 嵌入。
• 最终表示为原始输入嵌入与 n-gram 嵌入之间的残差连接。
• 模型采用标准目标端到端训练，n-gram 表在推理时缓存。
• 该架构使用稀疏嵌入查找代替密集矩阵乘法，降低计算成本。
• 通过缓存词元到聚类的映射，支持高效部署，实现推理开销恒定。

实验结果

研究问题

• RQ1 从离散表示中学到的潜在 n-gram 是否能在不降低性能的前提下提升 Transformer 的效率？
• RQ2 在准确率和推理速度方面，N-Grammer 层与更大的 Transformer 模型及 Primer 模型相比如何？
• RQ3 所学的潜在聚类在多大程度上捕捉了有意义的语义分组（例如：体育、艺术、地点）？
• RQ4 能否用稀疏嵌入查找替代注意力机制中的密集运算，以降低推理延迟？
• RQ5 对注意力头应用乘积量化是否能生成有意义且语义连贯的 n-gram 表示？

主要发现

• N-Grammer 模型在 C4 上的困惑度为 64.98，优于标准 Transformer（66.74）和 Primer（62.20）。
• 在 SuperGLUE 上，N-Grammer 模型的平均得分为 68.27，超过 Transformer（65.38）和 Primer（63.46）。
• 使用 8K 个潜在聚类和 196K 个 n-gram 嵌入的模型，在性能上与更大模型相当，但推理速度显著更快。
• 聚类分析显示，潜在代码将语义相关的词元分组（如体育、艺术、地点），证明了有意义的表征学习。
• N-Grammer 层通过用快速稀疏嵌入查找替换昂贵的密集矩阵乘法，减少了对计算资源的依赖。
• 缓存词元到聚类的映射使推理期间的开销恒定，显著提升了模型在偏好稀疏运算的硬件上的可部署性。