RotRNN: Modelling Long Sequences with Rotations

技術要約：RotRNN – 回転を用いた長距離シーケンスのモデリング

問題提起
長距離シーケンスのモデリングは、長期的な依存関係を捉えることの難しさから、機械学習における困難な領域であり続けています。Transformerは計算量が二次関数的にスケールし、誘導バイアス（inductive bias）に乏しいという課題があり、非線形RNNは勾配消失や勾漸爆発の問題に直面しています。一方で、近年の状態空間モデル（SSM）や線形回帰ユニット（LRU）は、最先端の性能を実現しています。しかし、これらのモデルには以下のような重大な欠点があります：

1. 複雑な初期化: S4のようなSSMは、最適多項式投影に基づくHiPPO初期化などの複雑な理論的導出に依存しており、これらは学習プロセス中において十分に理解されていません。
2. 実装と理論の乖離: LRUは概念的にはより単純ですが、その理論的な動機を実用面で完全には反映できていません。具体的には、実数行列の固有値が共役複素数ペアであることを強制できておらず、実数値の出力を保証するための制約を欠いており、また、学習された正規化係数が学習全体を通じて一貫した隠れ状態の大きさを保証しないという問題があります。
3. 安定性: 既存の手法は、異なる深さやランダムシードに対して堅牢ではない、複雑な正規化スキームや初期化手順を必要とすることがよくあります。

手法
著者らは、再帰状態行列 $A$ を回転行列としてパラメータ化する線形回帰モデルである RotRNN (Rotational Recurrent Neural Network) を提案します。この手法は、以下の3つの柱に基づいています。

1. 回転行列によるパラメータ化:
   制約のある行列を直接学習する代わりに、RotRNNは一般的な重み行列 $M$ を学習し、反対称行列の行列指数関数を用いてそれを特殊直交群 $SO(N)$へと写像します（$A = \exp(M - M^\top)$）。これにより、$A$ が学習を通じて有効な回転行列であり続けることが保証されます。
   線形回帰の畳み込み形式に必要な行列累乗の効率的な計算を可能にするため、著者らは回転行列 $A$ をブロック対角形式へと分解します：$A = P \Theta P^\top$。ここで、$P$ は直交行列であり、$\Theta$ は角度 $\theta_i$ で定義される $2\times2$ の回転ブロックで構成されます。この構造により、高コストな密行列演算を行うことなく、$A^k$ を高速に計算できます。

2. 原理的な正規化:
   回転行列の直交性（固有値が単位円上に存在すること）を利用して、著者らは訓練を通じて期待される隠れ状態のノルムが一定（$E[\|x_t\|^2] = c$）に保たれる正規化手順を導出しています。
   再帰は $x_t = \gamma A x_{t-1} + \xi B u_t$ と定義されます。ここで、$\gamma$ は学習可能なスカラー減衰因子であり、$\xi$ は一定のノルム条件を満たすように導出された再スケーリング係数です。無限シーケンスの極限に基づいて正規化を行うLRUとは異なり、RotRNNの正規化はすべてのタイムステップにおいて成立するため、正規化定数が $t$ 乗されることで発生する初期入力の指数関数的な減衰を防ぐことができます。

3. マルチヘッド・アーキテクチャ:
   単一のスカラー減衰因子 $\gamma$ の限界（異なるメモリホライゾンを必要とするタスクへの汎用性の欠如）に対処するため、RotRNNは $H$ 個の独立した低次元ヘッドを採用しています。各ヘッドは、独自の回転行列 $A^{(h)}$、入力行列 $B^{(h)}$、および減衰因子 $\gamma^{(h)}$ を持ちます。出力は結合され、線形混合されることで、モデルが複数のスケールで依存関係を捉えることを可能にします。

主な貢献

• 新規な線形RNNレイヤー: 再帰状態行列を回転としてパラメータ化することで、アルゴリズムが学習を通じてその理論的動機に忠実であることを保証する、新しい再帰レイヤーの提案。
• 効率的な計算: パラメトリックな回転をブロック対角形式に分解することで、行列累乗を効率的に計算する手法を提供し、高速な線形回帰を可能にした。
• 堅牢な正規化: 漸近的な収束に依存する従来の研究とは異なり、一定の期待隠れ状態の大きさを保持する、原理に基づいた正規化手順の導出。
• 理論的洞察: LRUの特殊なケース（固有値が共役ペアを形成する場合）とマルチヘッドRotRNNとの間の数学的な等価性を示す導出を行い、先行研究におけるLRUの内部メカニズムの解明に寄与した。

結果
著者らは、いくつかの長距離シーケンスモデリングベンチマークを用いてRotRNNを評価しました。

• Long Range Arena (LRA): RotRNNは、最先端の線形回帰モデル（S4, S5, LRU）に対して競争力のある性能を示しました。ListOps、Text (IMDB)、Retrievalといった離散入力タスクにおいて、ベースラインの中で最高スコアを達成しました。しかし、ピクセルレベルの画像タスク（Path-X, CIFAR）では、一部のベースラインを下回りました。
• Speech Commands: 35クラスのSpeech Commands分類タスクにおいて、RotRNNは（再現された）LRUと同等の精度（95.2%）を、同程度のパラメータ数で達成し、より複雑なディープ状態空間モデルに対しても競争力を維持しました。
• 安定性分析: ListOpsデータセットにおける隠れ状態のノルムの経験的分析により、RotRNNは学習、層の深さ、およびランダムシードに関わらず、ほぼ一定の大きさを維持することが明らかになりました。対照的に、LRUは隠れ状態のノルムに激しい変動を示し、安定した大きさに収束するまでに大幅に長い時間を要しました。

意義と主張
本論文は、RotRNNが既存の線形回帰モデルに対して、概念的に単純で効率的、かつ数学的に原理に基づいた代替案を提供すると主張しています。その主要な意義は、理論的な動機と実用的な実装の間の溝を埋めることにあります。回転行列を使用することで、RotRNNはSSMのような複雑な初期化スキームや、LRUのような実装上の不整合を回避しています。著者らは、自らのモデルが「ブラックボックス」的な初期化に頼ることなく、安定性を保証する堅牢な正規化手順を提供し、標準的なベンチマークにおいて競争力のある性能を達成していると断言しています。さらに、本研究は、回転ベースのブロック構造という観点からLRUを捉えることが、先行モデルにおける線形回帰のメカニズムを明確にする助けとなる可能性を示唆しています。