arXiv:2410.06577cs.CL2024-10ICLR被引 6

用新注意力机制突破大模型精度与效率的矛盾

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

  • 提出数据依赖的温度选择机制,实现线性复杂度的递归注意力
  • Rodimus+在1.6B参数下超越训练量更大的Qwen2和RWKV6模型
  • 适合追求高效推理且不牺牲性能的大模型应用开发者

基于Transformer的大语言模型虽推动自然语言处理发展,但传统softmax注意力带来$O(T)$的计算开销,其中$T$为上下文长度。本文提出Rodimus及其增强版Rodimus+,通过线性注意力与纯递归框架结合,引入数据依赖的温度选择(DDTS)机制,在保持高性能的同时显著降低内存占用。该方法通过固定大小隐状态实现语义压缩,有效保留关键输入信息。Rodimus+进一步融合滑动窗口共享键注意力(SW-SKA),综合运用语义、分词和头压缩技术。实验显示,使用1万亿tokens训练的Rodimus+-1.6B模型,在下游任务中表现优于训练数据更多的Qwen2-1.5B和RWKV6-1.6B模型,展现了重新定义大模型精度-效率平衡的潜力。模型代码与预训练权重已开源于https://github.com/codefuse-ai/rodimus。

原文摘要 · Abstract (English)

Recent advancements in Transformer-based large language models (LLMs) have set new standards in natural language processing. However, the classical softmax attention incurs significant computational costs, leading to a $O(T)$ complexity for per-token generation, where $T$ represents the context length. This work explores reducing LLMs' complexity while maintaining performance by introducing Rodimus and its enhanced version, Rodimus$+$. Rodimus employs an innovative data-dependent tempered selection (DDTS) mechanism within a linear attention-based, purely recurrent framework, achieving significant accuracy while drastically reducing the memory usage typically associated with recurrent models. This method exemplifies semantic compression by maintaining essential input information with fixed-size hidden states. Building on this, Rodimus$+$ combines Rodimus with the innovative Sliding Window Shared-Key Attention (SW-SKA) in a hybrid approach, effectively leveraging the complementary semantic, token, and head compression techniques. Our experiments demonstrate that Rodimus$+$-1.6B, trained on 1 trillion tokens, achieves superior downstream performance against models trained on more tokens, including Qwen2-1.5B and RWKV6-1.6B, underscoring its potential to redefine the accuracy-efficiency balance in LLMs. Model code and pre-trained checkpoints are open-sourced at https://github.com/codefuse-ai/rodimus.

大模型优化注意力机制高效推理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。