行列型再帰ユニット (MRU) が Transformer に匹敵する長文脈処理を実現──線形計算量で連想記憶を保持
RNN の隠れ状態を行列化する MRU 手法。Attention なしで線形計算量を維持しつつ、長文脈依存性と連想記憶能力を大幅に向上。(原題: Matrix Recurrent Units: An Attention Alternative)
リリース: 2024-12-16 · 読了 5 分RNN の隠れ状態を従来のベクトルから行列 H (dxd) に拡張し、情報の書き込みと読み出しを行列演算として定義した。
Associative Recall(連想想起)タスクにおいて、従来の LSTM や GRU が失敗する 1000 トークン以上の文脈でも 99% 以上の高い正解率を維持した。
計算量はシーケンス長 L に対して線形 O(L) であり、推論時のメモリ使用量はシーケンス長に依存しない O(1) に抑えられる。
Long Range Arena (LRA) の ListOps タスクにおいて、標準的な Transformer を 5.2 ポイント上回る 63.4% の精度を達成した。
Transformer の計算コスト O(L^2) を回避しつつ、RNN の弱点だった複雑な情報の保持(連想記憶)を克服している。この発見を知らないと、長文脈処理には Attention が必須であるという誤った前提で設計を進めてしまう。
Mamba などの状態空間モデル (SSM) とは異なる「行列演算による再帰」というアプローチで線形計算量を達成しており、アーキテクチャ設計の新たな選択肢となる。
長文脈を扱うモデルの推論コスト削減を目指す開発者は、Attention の代替として MRU を検討すべき。特にメモリ制約の厳しいエッジデバイスでのデプロイにおいて、定数メモリで動作しつつ高い表現力を持つ MRU は実装の有力候補になる。
国内固有の追加文脈は限定的(汎用的に有用)。
| Metric | Score | Δ |
|---|---|---|
| Associative Recall (L=1000) | 99.8 | +99.8 |
| LRA ListOps | 63.4 | +5.2 |