[論文レビュー] 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

ひとこと要約

本論文は、状態空間モデル（SSM）の一般化された直交基底射影フレームワークを導入し、S4モデルが長距離依存関係を捉える成功の背後にある数学的解釈を明確にした。ハイポ（HiPPO）行列を指数的歪みを加えたレジェンドル多項式への分解として再解釈することで、S4-FouT（フーリエ基底）を含む新たなSSM変種が導出された。これにより、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行列の役割を解消すること。
• 元のS4定式化を超える直交基底関数を用いたSSM構築の一般化され、数学的に整合性のあるフレームワークを提供すること。
• S4における時間定数パラメータΔの役割と適切な初期化方法を明確にすること。このパラメータは、捉えるべき依存関係の長さを制御する。
• フーリエ変換や局所畳み込みを自然に組み込む新たなSSM変種（例：S4-FouT）を導出すること。これにより、CNNやSTFTが一般化される。
• 関数再構成および記憶タスクにおける理論的フレームワークの実証的妥当性を検証し、理論と性能の整合性を示すこと。

提案手法

• 著者らは、重み付き測度（特に指数的減衰測度）の下で、SSMを直交基底関数への射影として一般化することで、HiPPOフレームワークを拡張した。
• 彼らは、元のS4状態行列Aが指数的歪みを加えたレジェンドル多項式への直交射影に対応することを証明し、これが長期間の系列モデル化において有効であることを説明した。
• このフレームワークにより、他の基底関数（例：切り捨てられたフーリエ基底関数）を選択することで新たなSSMが導出可能であり、これによりS4-FouT変種が得られた。
• 時間定数パラメータΔは、カーネルの有効長を制御すると解釈され、所望のコンテキスト窓に応じた明確な初期化ルールが導かれた。
• スケーリング、シフト、ユニタリ変換の下で閉じた性質を持つ、一般化された構造的状態空間モデル（TOSSMs）のクラスが導入され、正規直交性と安定性が保たれた。
• 理論的結果は、パデ近似および連分数に基づき、SSMカーネルが有理関数を介して指数関数を近似することを示した。

実験結果

リサーチクエスチョン

• RQ1 S4におけるHiPPO行列を、経験的初期化としてではなく、数学的に厳密に解釈することは可能か？
• RQ2 S4が長距離依存関係を捉える理論的メカニズムは何か？そして、他の基底関数へ一般化可能か？
• RQ3 S4における時間定数パラメータΔの正確な役割と正しい初期化方法は何か？
• RQ4 フーリエ変換や局所畳み込みを自然に組み込む新たなSSM変種を導出可能か？
• RQ5 基底関数の整合性に関する理論的予測は、系列モデルタスクにおける実験的性能とどの程度一致するか？

主な発見

• 元のS4モデルは、指数的減衰測度の下で指数的歪みを加えたレジェンドル多項式への射影として理論的に解釈され、長期間コンテキスト処理における優れた性能が説明された。
• 切り捨てられたフーリエ基底関数に基づく新規S4-FouT変種は、Path-Xタスクで96%の正確性を達成し、記憶タスクおよび連続信号タスクにおいて他のSSMを上回る最先端の性能を示した。
• 一般化されたフレームワークにより、時間定数パラメータΔの原理的初期化が可能となり、SSMがモデル化する依存関係の有効長を直接制御できるようになった。
• 基底関数の整合性に関する理論的予測は、実験的性能を正確に予測しており、S4-FouTはスパイク関数の符号化や局所パターン検出を要するタスクで優れた性能を示した。
• 最も挑戦的なLong Range Arenaベンチマークでは、元のS4モデルが平均86%の正確性を達成し、長距離モデル化におけるその強靭性が裏付けられた。
• フレームワークのスケーリング、シフト、ユニタリ変換の下での閉じる性質により、安定で分散を保つSSMパラメータ化が可能となり、信頼性の高い学習が可能になった。