MesaNet: Sequence Modeling by Locally Optimal Test-Time Training

要約

シーケンスモデリングは現在、SoftMaxの自己触媒を使用する因果変圧器アーキテクチャによって支配されています。
広く採用されていますが、変圧器はスケーリングメモリを必要とし、推論中に直線的に計算します。
最近の作業ストリームは、SoftMax操作を線形化し、Deltanet、Mamba、XLSTMなどの一定のメモリと計算コストを備えた強力な再発性ニューラルネットワーク（RNN）モデルをもたらしました。
これらのモデルは、それらの再発層のダイナミクスがすべて、オンライン学習ルールを通じてほぼ最適化されたコンテスト内回帰目標から派生できることに注意することで統合できます。
ここでは、この作業ラインに参加し、最近提案されたMESA層（Von Oswald et al。、2024）の数値的に安定した塊状の並列化可能なバージョンを導入し、10億パラメータースケールで言語モデリングで研究します。
この層は再びコンテキスト内の損失に起因しますが、高速コンジュゲート勾配ソルバーを使用して、すべての時点で最適化されるようになりました。
広範な一連の実験を通じて、最適なテスト時間トレーニングにより、特に長いコンテキストの理解を必要とするタスクで、以前のRNNよりも低い言語モデリングの困惑と下流のベンチマークパフォーマンスに到達することができることを示しています。
このパフォーマンスのゲインは、推論時間中に追加のフロップを費やした費用がかかります。
したがって、私たちの結果は、テスト時間計算の増加の最近の傾向に興味深く関連しています。ここでは、ニューラルネットワーク自体内の連続的な最適化問題を解決するために計算を費やすことにより。

要約(オリジナル)

Sequence modeling is currently dominated by causal transformer architectures that use softmax self-attention. Although widely adopted, transformers require scaling memory and compute linearly during inference. A recent stream of work linearized the softmax operation, resulting in powerful recurrent neural network (RNN) models with constant memory and compute costs such as DeltaNet, Mamba or xLSTM. These models can be unified by noting that their recurrent layer dynamics can all be derived from an in-context regression objective, approximately optimized through an online learning rule. Here, we join this line of work and introduce a numerically stable, chunkwise parallelizable version of the recently proposed Mesa layer (von Oswald et al., 2024), and study it in language modeling at the billion-parameter scale. This layer again stems from an in-context loss, but which is now minimized to optimality at every time point using a fast conjugate gradient solver. Through an extensive suite of experiments, we show that optimal test-time training enables reaching lower language modeling perplexity and higher downstream benchmark performance than previous RNNs, especially on tasks requiring long context understanding. This performance gain comes at the cost of additional flops spent during inference time. Our results are therefore intriguingly related to recent trends of increasing test-time compute to improve performance — here by spending compute to solve sequential optimization problems within the neural network itself.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google