DAMI:面向门控线性注意力长上下文检索增强的双轴记忆初始化
DAMI: Dual-Axis Memory Initialization for Enhanced Long-Context Retrieval in Gated Linear Attention
Zhangchi Feng · Harry Jiang · Yangyifei Luo · Ray Zhang · Richong Zhang
中文摘要
线性注意力随序列长度线性扩展,但其有限状态记忆会限制长上下文检索。在门控线性递归中,记忆保持取决于状态衰减,而状态衰减由门参数 A_log、dt_bias 与输入激活共同塑造。在受控的 Gated DeltaNet(GDN)实验中,这两个参数在训练后仍与初始值高度相关,凸显了初始化的重要性。为此,提出双轴记忆初始化(Dual-Axis Memory Initialization,DAMI),在训练前联合配置这两个参数。DAMI 为不同头分配几何时间尺度以覆盖不同检索距离,并在层间采用低-高-低的输入敏感度剖面。少量激活样本即可校准 GDN 系列实现,使其同时匹配两个行为目标,且不增加可训练参数。在 GDN、KDA、GDN2、Mamba-2 和 RWKV-7 上,DAMI 相比对应的随机初始化基线,将 RULER 风格检索的平均准确率提升 5.29 至 16.71 个百分点。在语言建模、常识推理和 2K 真实世界检索上的表现与这些基线大体相当。代码见 https://speedai.com/research/dami。
关键要点
- 01线性注意力的有限状态记忆限制长上下文检索,而门控递归中状态衰减由 A_log、dt_bias 与输入激活共同决定。
- 02受控 GDN 实验显示 A_log 与 dt_bias 训练后仍与初始值高度相关,说明初始化对记忆保持至关重要。
- 03DAMI 在训练前联合配置两参数:头间几何时间尺度覆盖不同检索距离,层间采用低-高-低输入敏感度。
- 04在 GDN、KDA、GDN2、Mamba-2、RWKV-7 上,RULER 风格检索平均准确率较随机初始化提升 5.29 至 16.71 个百分点。
- 05语言建模、常识推理与 2K 真实检索表现与基线大体相当,且不增加可训练参数。
解读
尚无解读。
原始英文摘要
arXiv:2610.41994v1 Announce Type: cross Abstract: 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.