[论文解读] How to Train Your HiPPO: State Space Models with Generalized Orthogonal Basis Projections

Albert Gu, Isys Johnson+3 more|arXiv (Cornell University)|Jun 24, 2022

Blind Source Separation Techniques被引用 11

一句话总结

本文提出了一种广义正交基投影框架用于状态空间模型（SSMs），从数学上严谨地解释了S4模型在捕捉长程依赖关系方面成功的原因。通过将HiPPO矩阵重新解释为在指数加权勒让德多项式上的分解，作者推导出新的SSM变体（如基于傅里叶基的S4-FouT），在Long Range Arena上达到86%的性能，在Path-X上达到96%的性能，同时为时间尺度参数提供了更清晰的初始化规则。

ABSTRACT

Linear time-invariant state space models (SSM) are a classical model from engineering and statistics, that have recently been shown to be very promising in machine learning through the Structured State Space sequence model (S4). A core component of S4 involves initializing the SSM state matrix to a particular matrix called a HiPPO matrix, which was empirically important for S4's ability to handle long sequences. However, the specific matrix that S4 uses was actually derived in previous work for a particular time-varying dynamical system, and the use of this matrix as a time-invariant SSM had no known mathematical interpretation. Consequently, the theoretical mechanism by which S4 models long-range dependencies actually remains unexplained. We derive a more general and intuitive formulation of the HiPPO framework, which provides a simple mathematical interpretation of S4 as a decomposition onto exponentially-warped Legendre polynomials, explaining its ability to capture long dependencies. Our generalization introduces a theoretically rich class of SSMs that also lets us derive more intuitive S4 variants for other bases such as the Fourier basis, and explains other aspects of training S4, such as how to initialize the important timescale parameter. These insights improve S4's performance to 86% on the Long Range Arena benchmark, with 96% on the most difficult Path-X task.

研究动机与目标

• 为解决S4在建模长序列方面成功背后的理论理解不足问题，特别是HiPPO矩阵的作用。
• 提供一个通用且数学上一致的框架，用于构建SSMs，使用原始S4公式之外的正交基函数。
• 明确时间尺度参数Δ在S4中的作用及其正确初始化方式，该参数控制所捕捉依赖关系的长度。
• 推导出新的SSM变体（如S4-FouT），自然编码基于傅里叶的变换和局部卷积，从而推广CNN和STFT。
• 在函数重构和记忆任务上实证验证理论框架，表明理论与性能之间的一致性。

提出的方法

• 作者通过将SSMs表述为在加权测度下（特别是指数衰减测度）对正交基函数的投影，推广了HiPPO框架。
• 他们证明原始S4状态矩阵A对应于在指数加权勒让德多项式上的正交投影，从而解释了其在长上下文建模中的有效性。
• 该框架允许通过选择其他基函数（如截断傅里叶基函数）推导出新的SSM，从而导出S4-FouT变体。
• 时间尺度参数Δ被解释为控制核的有效长度，其初始化规则基于期望的上下文窗口。
• 该方法引入了一类广义的结构化状态空间模型（TOSSMs），具备缩放、平移和酉变换下的封闭性，保持正交性和稳定性。
• 理论结果基于Pade逼近和连分数，表明SSM核通过有理函数逼近指数函数。

实验结果

研究问题

• RQ1 S4中HiPPO矩阵能否以数学上严谨的方式解释，而非仅作为经验初始化？
• RQ2 S4通过何种理论机制捕捉长程依赖关系？该机制能否推广到其他基函数？
• RQ3 S4中时间尺度参数Δ的精确作用及其正确初始化方式是什么？
• RQ4 能否推导出能自然编码常见信号处理操作（如STFT或局部卷积）的新SSM变体？
• RQ5 关于基函数对齐的理论预测在序列建模任务上的经验性能中，其匹配程度如何？

主要发现

• 原始S4模型被理论解释为在指数衰减测度下对指数加权勒让德多项式进行投影，解释了其在长上下文任务中的优异表现。
• 基于截断傅里叶基函数的新变体S4-FouT在Path-X任务上达到96%的准确率，为当前最优性能，在记忆和连续信号任务上优于其他SSM。
• 广义框架实现了对时间尺度参数Δ的合理初始化，该参数直接控制SSM所建模的有效依赖长度。
• 关于基函数对齐的理论预测准确预测了经验性能，S4-FouT在需要脉冲函数编码和局部模式检测的任务中表现优异。
• 在最具挑战性的Long Range Arena基准上，原始S4模型仍为最佳性能者，平均准确率达86%，验证了其在长程建模中的鲁棒性。
• 该框架在缩放、平移和酉变换下具备封闭性，确保SSM参数化稳定且保持方差，从而支持可靠训练。