稀疏增量记忆：通过稀疏性扩展线性RNN的状态 Sparse Delta Memory: Scaling the State of Linear RNNs through Sparsity

SDM的完整方法流程包含四个主要步骤，每步都有明确的输入输出和具体操作。第一步是稀疏键选择，从输入 $x_t \in \mathbb{R}^d$ 通过学习的线性投影 $W_k, W_q \in \mathbb{R}^{d \times 2\sqrt{N}}$ 生成写键和读查询。每个投影向量被分成两半用于键和查询，它们的外和产生N个分数，每个对应一个内存槽位。对写分数应用top-W，对读分数应用top-R，从N个可能槽位中选择W个写索引和R个读索引，计算复杂度为 $O(\sqrt{N} \times d + W^2 + R^2)$。第二步是带门控的delta写入，对于每个选定的写槽位，首先应用遗忘门 $\tilde{M}_t[i] \leftarrow \alpha_t \cdot M_{t-1}[i]$，其中 $\alpha_t = \exp(-A \cdot \text{softplus}(W_a x_t + b_d^t))$ 是每头的遗忘门。然后应用delta更新 $M_t[i] \leftarrow \tilde{M}_t[i] + \beta_t \cdot k_t^{(i)} \cdot (v_t - \tilde{M}_t[i])$，其中 $\beta_t = \sigma(W_b x_t)$ 是输入门。第三步是稀疏读取，通过R个选定读槽位的加权和从内存检索 $y_t = M_t^\top q_t = \sum_{i \in I_r^t} q_t^{(i)} \cdot M_t[i]$。第四步是标准化、门控和头部混合，检索的内存通过RMS-Norm标准化，用门控向量进行元素级门控，最后通过投影混合所有SDM头的输出产生最终层输出。

SDM的技术新颖性体现在多个方面。首先，它是首个将GDN的delta规则成功稀疏化的工作，证明了防止记忆干扰的机制可以与稀疏访问模式兼容。其次，SDM采用了动态加静态的混合内存设计：内存既可以通过delta规则动态更新存储上下文信息，又可以通过学习初始状态存储预训练知识。第三，SDM在严格的isoFLOP约束下实现了内存容量数千倍的增长，这在稀疏神经网络领域是罕见的。第四，SDM的自适应内存访问模式显示了智能的行为：写分布比读分布更集中，读分布更均匀，并且读写分布会相互适应。最后，SDM的训练效率优化也很新颖，使用chunk-wise并行（通过GDN/FLA的WY表示）和内存高效的反向传播，使得大规模训练成为可能。当 $H=1$ 时，SDM的内存大小按照公式 $M_{\text{size}} = \frac{(d_{\text{total}}^{qk}/2H)^2 \cdot d_{\text{total}}^v}{4H^2}$ 缩放，这表明通过调整头数H可以在不增加FLOPs的情况下控制状态大小。

SDM存在几个可以改进的弱点。首先，内核效率问题严重限制了训练和推理速度，当前MFU比GDN低一个数量级，这意味着即使理论FLOPs相同，实际运行时间也显著更长。改进方向包括开发专门的GPU内核优化，利用稀疏访问模式的特性，探索异步内存传输策略，以及优化内存访问模式以提高缓存命中率。其次，在multikey_2等复杂多键检索任务上性能仍然远低于FullAttn，这可能是因为稀疏访问模式无法捕获某些复杂的多键关联。改进方向包括增加读写槽位的数量，探索更智能的键选择策略，或者设计混合架构在关键层使用全注意力。第三，内存开销对于资源受限环境可能过大，8B模型的7.963B状态与参数相当。改进方向包括研究状态压缩技术，开发分层或分区的内存架构，或者设计动态内存分配策略只激活任务需要的部分状态。第四，当前的自适应内存访问虽然表现出一定智能，但仍是基于softmax的静态分布，可能无法充分适应不同任务的内存需求模式。改进方向包括探索强化学习驱动的动态读写策略，设计任务感知的内存访问机制，或者引入可学习的读写预算分配。

作者提出的未来工作方向包括设计更高效的内核以进一步扩展SDM模型。基于成果可延伸的研究方向包括：探索SDM在视频处理和多模态任务中的应用，特别是长视频理解需要处理数百万帧的场景；研究分层或分区的内存架构，以更好地平衡内存容量和访问效率；开发任务感知的动态内存分配策略，使模型能够根据任务需求自适应调整读写模式；探索SDM与检索增强生成的结合，将外部知识检索与内部稀疏内存相结合；研究SDM在持续学习和增量学习场景中的应用，利用其delta更新规则的自然适应性；探索更复杂的键选择机制，如基于注意力权重的自适应选择或基于学习的索引策略；研究SDM在不同模态之间的知识共享机制；开发更高效的训练算法；探索SDM在边缘设备上的部署；研究SDM的可解释性，分析内存槽位中学到的知识表示和访问模式。

SDM的可复现性评估表明该项目具有较好的开放性和可复现性。作者已在GitHub上开源了代码，这为复现提供了基础。论文详细描述了所有架构配置和超参数设置，包括缩放阶梯的具体配置，每个规模的参数数量、token数量、学习率、状态大小等关键信息。实验使用了多样化的文本数据进行训练，最高规模的训练达到8B模型和超过1万亿token，这对计算资源提出了很高的要求。论文报告了详细的缩放定律分析和各个规模的性能指标，包括1.4B和8B模型的完整结果。消融研究设计合理，覆盖了学习初始状态、内存大小、读写配置等关键超参数。然而，复现8B规模可能需要大量的GPU资源，这对大多数研究团队来说是不可及的。较小的模型应该更容易复现。总体而言，虽然大规模复现具有挑战性，但开源代码和详细配置使得较小规模的验证和概念验证成为可能。