Rodimus*: Breaking the Accuracy-Efficiency Trade-Off with Efficient Attentions

本文提出 Rodimus 和 Rodimus+ 模型，通过数据依赖温度选择（DDTS）和滑动窗口共享键注意力（SW-SKA）机制，在保持性能的同时显著降低大型语言模型的计算和内存复杂度，挑战了准确性与效率的权衡。

Large Language Model, Transformer, Efficiency, Reasoning, Multimodality

Zhihao He, Hang Yu, Zi Gong, Shizhan Liu, Jianguo Li, Weiyao Lin

上海交通大学, 蚂蚁集团

Generated by grok-3

Background Problem

大型语言模型（LLMs）基于 Transformer 的架构在自然语言处理中取得了显著进展，但其核心的 softmax 注意力机制带来了高计算成本，尤其是在长上下文生成中，表现为每 token 生成的复杂度为 O(T)，其中 T 为上下文长度。这种高复杂度导致内存使用和推理时间随着上下文长度线性增长，限制了模型在资源受限环境下的应用。为此，研究者探索了更高效的注意力机制替代方案，试图在保持性能的同时降低计算和内存需求。本文提出 Rodimus 和 Rodimus+ 模型，旨在通过语义压缩、token 压缩和头部压缩技术，解决 LLMs 在准确性和效率之间的权衡问题。

Method

本文提出了两种模型：Rodimus 和 Rodimus+，分别针对高效注意力机制的不同方面。

Rodimus：一个纯递归模型，基于线性注意力机制，通过固定大小的隐藏状态实现语义压缩。其核心创新是数据依赖温度选择（DDTS）机制，通过以下方式工作：
- 引入选择门控 $g_t$ 和温度门控 $\tau_t$ ，动态调整隐藏状态更新，决定保留历史信息还是纳入当前输入，公式为 $\alpha_t = \exp(-g_t \odot \tau_t)$ 和 $\hat{\alpha}_t = g_t^{\tau_t}$ 。
- 使用低秩矩阵设计 $\hat{\beta}_t$ ，减少参数量并过滤输入噪声。
- 整体架构结合了短卷积（ShortConv）和门控线性单元（GLU），增强局部上下文聚合和参数效率。关键点在于，Rodimus 在推理时仅需固定大小的隐藏状态，时间和空间复杂度为 O(1)，相比 softmax 注意力机制的 O(T) 具有显著优势。然而，DDTS 的复杂门控设计可能增加训练难度，尤其是在数据不足时可能导致过拟合风险。
Rodimus+：在 Rodimus 基础上，结合滑动窗口共享键注意力（SW-SKA），整合了语义、token 和头部压缩技术。其创新点包括：
- SW-SKA 使用滑动窗口注意力聚焦局部上下文，同时通过共享键（Shared-Key）机制减少 KV 缓存大小，理论上实现无损头部压缩。
- 采用两跳残差连接（Two-Hop Residual）设计，确保全局语义（Rodimus 块）和局部 token（SW-SKA 和 FFN）信息的有效融合。尽管 Rodimus+ 试图在性能和效率间取得更好平衡，但其混合架构增加了计算和内存开销，尤其是在 SW-SKA 部分，可能会抵消部分效率优势。此外，论文未充分讨论混合模型在训练稳定性上的潜在问题。

Experiment

实验在多个基准数据集（如 WikiText-103、Pile）上进行，涵盖语言建模、下游任务和召回能力评估，具体设置如下：

数据集与设置：在 WikiText-103 上训练约 44M 参数模型，在 Pile 上扩展到 125M 至 1.3B 参数规模，下游任务包括常识推理、阅读理解等，使用零样本评估。召回任务（如 NeedleBench、MQAR）测试模型在长上下文中的信息检索能力。
结果分析：
- Rodimus 在 WikiText-103 上以更小的隐藏状态大小（扩展因子 n=32 对比 Mamba2 的 n=128）实现了更低的困惑度（PPL），表明其在语义压缩上的效率优势。
- Rodimus+-1.6B 在下游任务中平均性能比 Qwen2-1.5B（训练于 7T tokens）高 0.31%，比 RWKV6-1.6B（训练于 1.4T tokens）高 2.3%，尽管其仅训练于 1T tokens。然而，这种比较可能存在不公平性，因为训练数据的质量和分布未详细披露，难以排除数据选择偏差。
- 在 NeedleBench 等召回任务中，Rodimus 和 Rodimus+ 甚至超越了基于 softmax 注意力机制的 Pythia，显示出在长上下文召回上的潜力，但 Pythia 在短上下文任务中仍占优势，表明递归模型在非外推区间内的局限。
实验设计的合理性与不足：实验设置覆盖了多种模型规模和任务类型，较为全面，但对训练数据量和质量的控制不足，可能导致结果的高估。此外，论文未提供足够的多设备推理延迟数据，仅强调理论复杂度优势，而实际应用中的 I/O 瓶颈未被充分探讨。总的来说，方法改进在效率上明显，但在性能提升的可持续性和泛化性上仍需更多验证。

Further Thoughts

Rodimus 和 Rodimus+ 的设计理念在高效注意力机制领域具有启发性，尤其是 DDTS 机制通过动态门控实现信息过滤，可能为其他递归模型（如 Mamba 系列）的改进提供新思路。然而，我对其实际应用中的可扩展性持谨慎态度，尤其是在超大规模模型（如 70B+ 参数）上的表现，因为论文未提供足够的大规模实验数据。此外，SW-SKA 的无损头部压缩理念值得进一步探索，是否能在其他混合架构中（如 Jamba 或 Samba）进一步优化 KV 缓存管理，是一个有趣的研究方向。另一个值得思考的点是，递归模型在长上下文任务中的召回能力提升是否会随着上下文长度进一步增加而衰减，这可能需要结合更先进的上下文扩展技术（如 ALiBi 或 RoPE）进行验证。总之，本文为高效 LLMs 的研究提供了有价值的视角，但需要在更广泛的实际场景和资源约束下进一步测试其鲁棒性。