改进初始化方法,让状态空间模型更擅长记忆和回溯。
Mimetic Initialization Helps State Space Models Learn to Recall
- 设计结构化初始化,让状态空间层模仿注意力机制
- 在复制和联想回忆任务上性能显著提升
- 适合研究高效序列建模与记忆机制的读者
近期研究表明,像 Mamba 这样的状态空间模型在基于记忆的任务上远不如 Transformer,原因在于其状态大小不随输入序列长度变化。但实践中状态空间模型的状态尺寸较大,我们推测其实际记忆能力应优于现有报告。本文探讨其较差的复制与回忆表现是否源于训练困难而非固有容量限制。通过观察其“注意力”图谱,提出一种结构化初始化方法,使状态空间层更易模仿注意力机制。在多种架构设置下,该初始化显著降低了 Mamba 从零开始学习复制和关联回忆的难度。
原文摘要 · Abstract (English)
Recent work has shown that state space models such as Mamba are significantly worse than Transformers on recall-based tasks due to the fact that their state size is constant with respect to their input sequence length. But in practice, state space models have fairly large state sizes, and we conjecture that they should be able to perform much better at these tasks than previously reported. We investigate whether their poor copying and recall performance could be due in part to training difficulties rather than fundamental capacity constraints. Based on observations of their "attention" maps, we propose a structured initialization technique that allows state space layers to more readily mimic attention. Across a variety of architecture settings, our initialization makes it substantially easier for Mamba to learn to copy and do associative recall from scratch.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。