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

TL;DR

Rodimus and Rodimus+ introduce a novel, efficient attention mechanism for large language models that significantly reduces computational costs while maintaining or improving performance, challenging the traditional accuracy-efficiency trade-off.

Contribution

The paper presents Rodimus and Rodimus+ with innovative attention mechanisms that enhance efficiency and accuracy in LLMs, including a data-dependent tempered selection and hybrid attention techniques.

Findings

Rodimus+-1.6B outperforms larger models trained on more data.

Achieves superior downstream performance with reduced memory usage.

Open-sourced code and models for community use.

Abstract

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.