Papers

DAMI:双轴记忆初始化提升门控线性注意力的长文本检索

DAMI: Dual-Axis Memory Initialization for Enhanced Long-Context Retrieval in Gated Linear Attention

Why we picked it

训练前改一下初始化,GDN、Mamba-2、RWKV-7 这些线性注意力模型的长文检索就提 5 到 17 个点,方法还挺省事。

线性注意力随序列长度线性扩展,但有限状态记忆会限制长上下文检索。论文提出 Dual-Axis Memory Initialization(DAMI),在训练前联合配置 Gated DeltaNet 的 A_log 和 dt_bias 两个参数:按注意力头分配几何时间尺度、按层设置低-高-低输入敏感度,且不增加可训练参数。在 GDN、KDA、GDN2、Mamba-2、RWKV-7 五种架构上,DAMI 将 RULER 式检索平均准确率比随机初始化基线提高 5.29 至 16.71 个百分点。语言建模、常识推理和 2K 真实检索任务上的表现与基线大致相当。

Original · arXiv cs.AI

DAMI: Dual-Axis Memory Initialization for Enhanced Long-Context Retrieval in Gated Linear Attention

Linear attention scales linearly with sequence length, but its finite-state memory can limit long-context retrieval. In gated linear recurrences, memory retention depends on state decay, which is shaped jointly by gate parameters A_log and dt_bias and by input activations. In controlled Gated DeltaNet (GDN) experiments, both parameters remain highly correlated with their initial values after training, highlighting the importance of initialization. Therefore, we introduce Dual-Axis Memory Initialization (DAMI), which jointly configures these parameters before training. DAMI assigns geometric timescales across heads for varied retrieval distances and a low-high-low input-sensitivity profile across layers. A small activation sample calibrates the GDN-family implementation to match both behavioral targets without adding trainable parameters. Across GDN, KDA, GDN2, Mamba-2, and RWKV-7, DAMI improves average RULER-style retrieval accuracy by 5.29 to 16.71 percentage points over the corresponding random-initialization baselines. Performance on language modeling, commonsense reasoning, and 2K real-world retrieval remains broadly comparable to those baselines. Our code is available at https://speedai.com/research/dami.