15 天前
如何训练你的HiPPO:基于广义正交基投影的状态空间模型
Albert Gu, Isys Johnson, Aman Timalsina, Atri Rudra, Christopher Ré

摘要
线性时不变状态空间模型(Linear Time-Invariant State Space Models, SSM)是工程学与统计学中的经典模型,近年来在机器学习领域展现出巨大潜力,这主要得益于结构化状态空间序列模型(Structured State Space sequence model, S4)。S4的核心组件之一是将状态矩阵初始化为一种特定矩阵——HiPPO矩阵,这一初始化策略在实践中被证明对S4处理长序列能力至关重要。然而,S4所采用的具体矩阵实际上源自先前研究中针对特定时变动力系统的设计,而将其用于时不变SSM并无已知的数学解释。因此,S4模型如何建模长程依赖关系的理论机制至今仍不明确。本文提出了一种更一般且更具直观意义的HiPPO框架形式化方法,首次从数学上清晰地揭示了S4的本质:其可被理解为在指数加权勒让德多项式(exponentially-warped Legendre polynomials)基上的分解。这一解释阐明了S4为何能够有效捕捉长距离依赖关系。我们的推广框架不仅为SSM构建了一个理论丰富的新型类,还使得我们能够自然推导出适用于其他基函数(如傅里叶基)的S4变体,并进一步解释了S4训练中的若干关键细节,例如如何合理初始化其核心的时间尺度参数。基于这些理论洞见,我们改进后的S4模型在Long Range Arena基准测试中取得了86%的准确率,其中在最具挑战性的Path-X任务上达到了96%的优异表现,显著提升了模型性能。