[论文解读] Simple linear attention language models balance the recall-throughput tradeoff

Simran Arora, Sabri Eyuboglu+7 more|arXiv (Cornell University)|Feb 28, 2024

Topic Modeling被引用 4

一句话总结

本文提出 Based，一种混合式语言模型架构，通过结合线性注意力与滑动窗口注意力，在召回性能与推理效率之间实现平衡。通过调节窗口大小与特征维度，Based 在显著降低显存使用的同时，实现与标准注意力相当的召回率，且推理吞吐量较 FlashAttention-2 提高 24 倍，相较于 Mamba 等 SOTA 模型，在召回密集型任务上准确率提升 6.22 个百分点。

ABSTRACT

Recent work has shown that attention-based language models excel at recall, the ability to ground generations in tokens previously seen in context. However, the efficiency of attention-based models is bottle-necked during inference by the KV-cache's aggressive memory consumption. In this work, we explore whether we can improve language model efficiency (e.g. by reducing memory consumption) without compromising on recall. By applying experiments and theory to a broad set of architectures, we identify a key tradeoff between a model's state size and recall ability. We show that efficient alternatives to attention (e.g. H3, Mamba, RWKV) maintain a fixed-size recurrent state, but struggle at recall. We propose BASED a simple architecture combining linear and sliding window attention. By varying BASED window size and linear attention feature dimension, we can dial the state size and traverse the pareto frontier of the recall-memory tradeoff curve, recovering the full quality of attention on one end and the small state size of attention-alternatives on the other. We train language models up to 1.3b parameters and show that BASED matches the strongest sub-quadratic models (e.g. Mamba) in perplexity and outperforms them on real-world recall-intensive tasks by 6.22 accuracy points. Implementations of linear attention are often less efficient than optimized standard attention implementations. To make BASED competitive, we develop IO-aware algorithms that enable 24x higher throughput on language generation than FlashAttention-2, when generating 1024 tokens using 1.3b parameter models. Code for this work is provided at: https://github.com/HazyResearch/based.

研究动机与目标

• 解决自回归语言模型中召回性能与推理效率之间的权衡问题。
• 识别并遍历不同架构下模型状态大小与召回能力之间的帕累托前沿。
• 设计一种模型，在保持标准注意力高召回能力的同时，降低显存消耗，类似高效模型（如 Mamba、RWKV）的实现方式。
• 开发 I/O 友好的优化策略，使所提架构在真实世界吞吐量基准测试中具备竞争力。
• 证明结合线性注意力与滑动窗口注意力可实现优于任一组件单独使用时的性能权衡表现。

提出的方法

• 提出一种混合架构 Based，在同一模型层中结合线性注意力与滑动窗口注意力。
• 通过可配置的滑动窗口大小（如 64、128）与线性注意力的特征维度，控制状态大小与性能权衡行为。
• 应用 I/O 友好的算法优化，减少生成过程中的内存访问开销，提升吞吐量。
• 使用标准训练方案（Adam 优化器、余弦衰减学习率、BFloat16 精度）训练最大达 1.3B 参数的模型。
• 滑动窗口组件使用旋转位置编码，线性组件不使用位置编码。
• 实施块独占的多项式分析，从理论上证明为解决关联召回任务，模型深度必须足够。

Figure 1: Based overview. Combining linear attention with tiny sliding window softmax attention (e.g., 64 or 128 tokens in width) enables improved recall accuracy with limited efficiency overhead vs. smaller tile sizes. ( Left ) Time to execute Cutlass GEMMs ( $y$ ) vs. sliding window attention size

实验结果

研究问题

• RQ1 模型能否在保持 Mamba 等高效模型的小状态尺寸的同时，达到标准注意力的召回性能？
• RQ2 在不同架构类别中，模型状态大小与召回能力之间的根本权衡关系是什么？
• RQ3 结合线性注意力与滑动窗口注意力能否形成帕累托最优解，优于各自独立组件？
• RQ4 如何设计 I/O 高效算法，使混合注意力模型在真实世界吞吐量上可与 FlashAttention 竞争？
• RQ5 线性注意力与局部注意力机制的结合，是否能实现优于任一机制单独使用时的长程与精确标记交互能力？

主要发现

• 在真实世界的召回密集型任务上，Based 相较 Mamba 准确率提升 6.22 分，证明其在更小状态尺寸下仍具备优越的召回能力。
• 在 1.3B 参数规模下，Based 在自回归生成 1024 个 token 时，吞吐量较 FlashAttention-2 提高 24 倍。
• 该混合架构成功遍历了召回-显存权衡的帕累托前沿，在仅增加适度效率开销的前提下，实现了与标准注意力相当的召回水平。
• 仅使用线性注意力在关联召回任务上失败，原因在于局部标记比较缺乏精度；而滑动窗口注意力受限于窗口大小。
• 理论分析表明，一个与数据无关的 BaseConv 模型至少需要 ⌊log(2p)⌋ 层才能解决关联召回任务，验证了此类模型对足够深度的需求。
• I/O 友好的优化使 Based 在吞吐量上可匹配或超越优化后的注意力实现，即使在使用混合注意力组件的情况下。

Figure 2: Throughput (memory) - recall tradeoff. $x$ -axis shows state size (bytes) during generation; $y$ -axis shows accuracy on the MQAR recall task [ 1 ] . For each architecture, we train several models varying hyperparameters that affect the recurrent state size ( e.g. model dimension). The plo