中文

Rodimus*: 突破准确率-效率之间的权衡 with 高效注意力机制

计算与语言 2025-05-20 v2

摘要

最近的 Transformer-based 大型语言模型 (LLM) 进展已制定了自然语言处理领域的新标准。然而,经典的 softmax 注意力机制导致计算成本高昂,导致每 token 生成的时间复杂度为 O(T)O(T),其中 TT 表示上下文长度。本文通过引入 Rodimus 及其增强版本 Rodimus++,探索在保持性能的同时降低 LLM 的复杂度。Rodimus 采用创新的数据依赖温度选择 (DDTS) 机制,基于线性注意力的纯递归框架,实现了显著的准确率,同时大幅降低了递归模型通常associated的内存使用。该方法通过保持固定大小隐藏状态中的关键输入信息来实现语义压缩。构建于此之上,Rodimus++ 将 Rodimus 与创新的滑动窗口共享键注意力 (SW-SKA) 结合,采用混合方法,有效利用语义、token 和 head 压缩技术的互补性。我们的实验表明,Rodimus++ -1.6B 模型在训练 1 万亿 token 时,相较于训练更多 token 的 Qwen2-1.5B 和 RWKV6-1.6B 模型,在下游任务上实现了优异的性能,凸显了其在 LLM 中重新定义准确率-效率平衡的潜力。模型代码和预训练模型已开源于 https://github.com/codefuse-ai/rodimus。

关键词

引用

@article{arxiv.2410.06577,
  title  = {Rodimus*: Breaking the Accuracy-Efficiency Trade-Off with Efficient Attentions},
  author = {Zhihao He and Hang Yu and Zi Gong and Shizhan Liu and Jianguo Li and Weiyao Lin},
  journal= {arXiv preprint arXiv:2410.06577},
  year   = {2025}
}

备注

Accepted by ICLR 2025. Camera-ready Version