行列の直交化がリカレントモデルの記憶を改善

行列の直交化がリカレントモデルの記憶を改善 2026年6月30日 この研究はParadigmの資金提供を受けました。 トランスフォーマーは、驚くべき結合想起（AR）能力を示します。アテンションは、各トークンに先行するトークンへの直接アクセスを提供し、これはリカレントニューラルネットワーク（RNN）のような他のアーキテクチャでは達成が困難でした。しかし、一部のドメインでは、トランスフォーマーの二次アテンションのオーバーヘッドは許容できません。一例として、Dreamerのような長期間RLが挙げられます。このようなアプリケーションでは、リカレントニューラルネットワークを機能させる必要がありますが、結合想起を諦めたくありません。 結合想起に最適なRNNとして知られているのはmLSTMです。mLSTMはLSTMの変種で、行列記憶を保持します。mLSTMは、MQARというベンチマークでベースラインと比較して想起能力を大幅に向上させました。しかし、純粋な想起だけではリカレント性能を測るのに十分ではないかもしれません。環境の遷移がノイズの多い分野では、有用な代理テストとしてノイズの多い結合想起（NAR）があります。MQARはNARを測定しないため、MADのノイズの多いARタスクスイートを調べることができます。タスクの例を以下に示します。 0 9 3 10 12 13 15 14 0 9 5 8 2 9 ここで、キー0は値9に、キー3は値10に、というようにマッピングされます。MADジェネレータは、キー、値、ディストラクタに異なるトークン範囲を使用します。したがって、キーが0〜5の場合、トークン12〜15はディストラクタです。NARに優れたモデルは、最初に0 -> 9を見た後、挿入されたディストラクタトークンを無視して、10番目の位置に9を予測するはずです。 では、リカレントNARをどのように改善するのでしょうか？言語モデリングで非常に成功しているオプティマイザであるMuonからいくつかのアイデアを借りることができます。Muonはそのモーメンタを直交化し、表現された方向のイコライザーとして機能します。これは、いくつかの強い方向が更新を支配するのを防ぎ、弱い方向を持ち上げます。特に重要なのは、Muonが末端の結合記憶学習においてAdamを上回ることを示す最近の研究です。このアイデアは、このイコライゼーションが弱い記憶が押し出されるのを防ぐというものです。 これに触発されて、mLSTM記憶行列を読み出し時に直交化し、この追加プロセスでトレーニングすることで、NAR性能が向上するかどうかをテストすることにしました。私たちは、MADノイズARサンプルを使用して、次トークン予測におけるmLSTMベースラインとその直交化バリアントを比較します。トレーニングと評価にはMAD noisy-recallを使用し、さまざまな語彙サイズとシーケンス長でfrac_noiseを0.8に設定しました。すべてのモデルはAdamW（ベータ=0.9, 0.999, weight_decay = 0.01）を使用して、バッチサイズ64で2kステップトレーニングしました。学習率は、各タスク設定で3e-4, 1e-3, 3e-3, 1e-2をスイープして選択しました。各ステップで新しいバッチをトレーニング用に生成し、実験ごとに個別の固定検証セットを維持しました。直交化には、フロベニウスノルムで正規化し（eps = 1e-6）、5回のニュートン・シュルツ反復を適用しました。勾配はこのプロセスを通じて流れるようにしました。重要なことに、直交化された記憶を書き戻すことはありませんでした。これは性能を低下させることが分かったためです。私たちはそれを読み出しにのみ使用します。実験の完全に再現可能なコードはここにあります。 図1. 直交化されたmLSTMアルゴリズム。 図2. トレーニングステップ対検証精度。 $$ \small \begin{array}{lccc} \hline \text{Regime} & \text{Orthogonalized} & \text{Baseline} & \Delta \\ \hline \text{vocab 80, len 512} & 87.5 \pm 12.4\ (20/24) & 69.1 \pm 17.8\ (17/24) & +18.4 \pm 18.1 \\ \text{vocab 80, len 768} & 91.7 \pm 11.4\ (22/24) & 75.9 \pm 12.0\ (13/24) & +15.7 \pm 16.8 \\ \text{vocab 80, len 1024} & 98.5 \pm 2.4\ (23/24) & 83.3 \pm 13.6\ (19/24) & +15.2 \pm 14.0 \\ \text{vocab 96, len 768} & 62.4 \pm 18.4\ (14/24) & 22.0 \pm 14.4\ (4/24) & +40.4 \pm 17.6 \\ \text{vocab 96, len 1024} & 68.5 \pm 18.3\ (16/24) & 23.1 \pm 15.3\ (4/24) & +45.4 \pm 18.6 \\ \hline \end{array} $$ 図3. MADノイズ想起の結果。エントリは2kステップでの最終検証精度、24シードに対する平均+/- 95%信頼区間。括弧内は80%以上のシード数。デルタはシードごとにペアリング。パラメータ数はvocab 80で77,716、vocab 96で80,740。直交化された実行ではLR 3e-3を使用。ベースラインはvocab 80, seq len 768の場合のみLR 1e-2を使用し、それ以外はLR 3e-3を使用。 私たちは、直交化が成功率と平均精度を全体的に向上させることを発見しました。興味深いのは、vocab-96の領域に入るとギャップが広がるように見えることで、これは直交化が、生のmLSTMが苦戦する困難なNARタスクに最も役立つことを示唆しています。後者の2つのケース（vocab 96, seq len 768/1024）では、直交化によりmLSTMが失敗寸前（24シード中4つ解決）から、大幅に信頼性の高い性能（14〜16シード解決）へと引き上げられました。これは、小さな介入を意図していたものとしては驚くべきことです。ニュートン・シュルツ法は、パラメータ数を固定したままで追加のゲインをもたらしますが、追加のFLOPsと実時間とのトレードオフがあります。これらの結果を過大に解釈しないよう注意する必要があります。これらは小さなモデルの領域で有効であり、NARは合成タスクです。NARのゲインが、より大きなモデルの実世界のベンチマークでゲインに変換されるかどうかを調査する価値があるでしょう。 この投稿の執筆中にフィードバックと提案をくれたDan Robinson、Alpin Yukseloglu、Glen Taggartに感謝します。