ICLR 2025 | Rodimus*：兼顾性能与效率的混合注意力机制
仅用于站内搜索，没有排版格式，具体信息请跳转上方微信公众号内链接
来自：CodeFuse
近年来，基于Transformer架构的大语言模型（LargeLanguageModels，简称LLMs）在自然语言处理领域取得了显著成就。然而，这类模型采用的传统softmaxattention机制在推理阶段需要存储大量key-value缓存信息，导致生成每个token时的时空复杂度达到O(Td²)（其中T代表上下文长度，d代表特征维度）。如此高的空间复杂度严重制约了模型在长文本处理和计算资源受限场景中的实际应用。
在模型创新之外，我们还基于Rodimus架构开发并开源了轻量级代码生成模型Rodimus+-Coder。该模型提供1.6B和4B参数两个版本，在同规模模型对比中均实现了超越现有最优水平（SOTA）的性能表现。
为充分展示Rodimus*的应用潜力和实用价值，我们将首先分享Rodimus+-Coder的详细评测结果。
Huggingface：https ://huggingface.co/collections/codefuse-ai
github：https ://github.com/codefuse-ai/rodimus
图1：各Chat/Instruct模型性能对比，包括1.6B和4B两个尺寸
我们在7个代码任务榜单、4个通用任务榜单和2个数学任务榜单上评估Rodimus+-Coder的Base模型的整体性能，并与类似尺寸的Qwen2.5-Coder、Gemma系列模型进行了比较。
从整体上看，Rodimus+-Coder的Base模型在同尺寸下优于Qwen2.5-Coder、Gemma系列模型。甚至在代码任务上，Rodimus+-Coder-4B的平均代码能力达到了与Qwen2.5-Coder-7B持平的水准。另外，即使经过了大规模的代码垂域加训，Rodimus+-Coder在通用自然语言和数学任务上的表现仍然具有竞争力。
我们在13个代码任务榜单、4个通用任务榜单和2个数学任务榜单上评估Rodimus+-Coder的Chat模型的整体性能，并与不同尺寸的Qwen2.5-Coder、Gemma和Phi-4-Mini模型进行比较。
paper：https ://openreview.net/forum?id=IIVYiJ1ggK
核心问题：Transformer模型的强大性能来自于其softmaxattention机制，该机制通过存储所有历史token构成的Key-Value（KV）cache，确保上下文信息的完整性。然而，这种无上限存储的设计导致模型推理复杂度随上下文长度线性增长（O(T)），不仅推高了时间消耗，还对内存提出了高要求。
现有优化方法的三大方向概括：
语义压缩：通过linearattention或者state-spacemodel替代softmaxattention，利用固定大小的隐藏状态压缩历史信息，但有限的隐藏状态记忆容量无法容纳所有历史信息。
Token压缩：通过sparseattention机制，减少需要处理的token数量，但仍可能带来具有关键信息的tokens丢失。
Head压缩：通过在不同head间共享Key-Valuecache减少缓存带来的内存占用，但牺牲了部分MHA的表达能力，且仍无法解决O(T)的内存复杂度。
图2：(a)在WikiText-103数据集上测试，记录最佳困惑度（PPL），通过调整扩展因子改变内存占用；(b)使用MQAR任务评估模型的记忆召回能力。
Rodimus*系列模型包括Rodimus和Rodimus+，其架构如图3所示。Rodimus是一种纯循环模型，通过迭代方式将历史上下文压缩为固定大小的隐藏状态（即语义压缩），并利用该隐藏状态进行下一个标记的预测。得益于新提出的DDTS机制，Rodimus能有效过滤无关信息，从而在提升性能的同时减少隐藏状态的大小。而Rodimus+在Rodimus的基础上加入了所提出的SW-SKA技术，这一增强方法在不牺牲效率的前提下进一步提升性能，实现了语义压缩、标记压缩和头部压缩方法的无缝结合。
Rodimus块将这些门控机制集成，并结合门控线性单元（GLU）实现高效的标记和通道混合，并通过增强局部上下文聚合能力，相比传统Transformer具备更高的参数效率。在训练阶段，Rodimus块通过分块并行化算法保证了次二次复杂度；在推理阶段如下式，仅需维护固定大小的隐藏状态，从而实现O(1)的空间复杂度。这种设计使Rodimus能够在大规模序列建模中，以较低的内存开销实现高性能。
其中，是选择门，是温控门，表示作用在输入的输入门控，是可学习权重，表示的输出；、和分别表示当前token的query、key与value表示。
图4：注意力机制对比。MQA和GQA通过共享值实现有损压缩，而SKA保留了MHA的多值设置，仅在所有头部间共享键，既减少了内存占用，又保留了MHA的表达能力。
Rodimus+通过将Rodimus块与SW-SKA相结合，解决了Rodimus块中隐藏状态固定容量可能导致的性能限制问题。其中，Rodimus块提供全局上下文表示，而SWA（SlidingWindowAttention）强调对局部上下文的关注，专注于生成答案中最重要的邻近标记。这种结合在不增加计算复杂度的情况下，实现了全局与局部信息的互补。
此外，Rodimus+引入共享键注意力（SKA），一种无损的头部压缩方法，通过在所有头部之间共享一个键（Key）矩阵，同时保留每个头独立的值（Value）矩阵，如图4。在传统MHA中，每个头均有独立的键和值，而SKA保持值的多样性同时共享键，从而在减少内存开销的情况下，保留了MHA的表达能力。相比之下，MQA和GQA通过共享值实现压缩，但会导致信息损失，因此属于有损压缩。在Rodimus+中，SKA进一步替代SWA中传统的多头注意力（MHA），形成SW-SKA，对SWA的键值缓存（KVCache）进行了优化，进一步减少了内存开销。
表4：WikiText-103实验结果
在图5的scalinglaws实验中，Rodimus和Rodimus+随着参数数量增加展现了更好的可扩展性，尤其在大模型情况下性能领先。
图5：基于ScalingLaws的Scaling曲线
此外，在表5下游任务的零样本评估中，Rodimus超过了大部分同类模型，同时，Rodimus+通过整合滑动窗口注意力和共享键机制进一步提升了性能，展示了混合架构的潜力。
表5：下游任务实验结果
图6：NeedleBench实验结果
图7：不同序列长度MQAR召回任务的准确性
如果您想更快地获取到最新信息，欢迎加入我们的微信群。
备注：昵称-学校/公司-方向/会议(eg.ACL)，进入技术/投稿群
id：DLNLPer，记得备注呦
