揭示状态空间模型学习方向由记忆能力决定,提出优化初始化与固定权重的新策略。
Memory Determines Learning Direction: A Theory of Gradient-Based Optimization in State Space Models
- 通过分析输入在状态中的存储机制,发现记忆长度与精度存在权衡。
- 实验验证长记忆训练困难,初始参数对成功学习至关重要。
- 固定循环权重可加速收敛且性能更优,适用于长序列任务。
状态空间模型(SSMs)因其潜在超越Transformer的性能而受到关注,但其学习动态缺乏充分的理论解释。本研究提供了相关理论解释并提出改进训练策略。通过考察输入时间序列在当前状态中的存储方式,可评估SSMs的记忆容量,揭示了记忆精度与长度之间的权衡关系,并证明了结构化状态空间序列模型(S4)与带对角循环权重的简化S4在理论上等价。该理论基础使我们阐明了学习动态,证明初始参数的重要性:成功学习要求初始记忆结构尽可能长,即使记忆精度下降或梯度丢失教师信息亦然。在需要长记忆的任务上进行的实验确认,延长记忆极为困难,凸显初始化的关键作用。此外,我们发现固定循环权重比自适应权重更具优势,能在更快收敛下实现相当甚至更高的性能。研究成果为SSMs提供了新的理论基础,可能催生新型优化策略。
原文摘要 · Abstract (English)
State space models (SSMs) have gained attention by showing potential to outperform Transformers. However, previous studies have not sufficiently addressed the mechanisms underlying their high performance owing to a lack of theoretical explanation of SSMs' learning dynamics. In this study, we provide such an explanation and propose an improved training strategy. The memory capacity of SSMs can be evaluated by examining how input time series are stored in their current state. Such an examination reveals a tradeoff between memory accuracy and length, as well as the theoretical equivalence between the structured state space sequence model (S4) and a simplified S4 with diagonal recurrent weights. This theoretical foundation allows us to elucidate the learning dynamics, proving the importance of initial parameters. Our analytical results suggest that successful learning requires the initial memory structure to be the longest possible even if memory accuracy may deteriorate or the gradient lose the teacher information. Experiments on tasks requiring long memory confirmed that extending memory is difficult, emphasizing the importance of initialization. Furthermore, we found that fixing recurrent weights can be more advantageous than adapting them because it achieves comparable or even higher performance with faster convergence. Our results provide a new theoretical foundation for SSMs and potentially offer a novel optimization strategy.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。