MoM: Linear Sequence Modeling with Mixture-of-Memories

要約

線形注意、状態空間モデリング、線形RNNなどの線形シーケンスモデリング方法は、トレーニングと推論の複雑さを減らすことにより、大幅な効率改善を提供します。
ただし、これらの方法は通常、入力シーケンス全体を単一の固定サイズメモリ状態に圧縮し、リコール集約型のダウンストリームタスクの最適ではないパフォーマンスにつながります。
神経科学、特に「記憶干渉」を緩和しながら堅牢な長期記憶を維持する脳の能力からインスピレーションを得て、混合物（MOM）と呼ばれる新しいアーキテクチャを紹介します。
Momは複数の独立したメモリ状態を利用しており、ルーターネットワークが特定のメモリ状態に入力トークンを向けています。
このアプローチは、メモリ干渉を最小限に抑えながら、全体的なメモリ容量を大幅に向上させます。
その結果、ママはリコール集約型のタスクで非常にうまく機能し、既存の線形シーケンスモデリング手法を上回ります。
複数のメモリ状態を組み込んでいるにもかかわらず、各メモリ状態の計算は複雑さが直線的なままであるため、ママはトレーニング中に線形複雑さの利点を保持することができますが、推論中は一定の複雑さがあります。
私たちの実験結果は、MOMが下流の言語タスク、特にリコール集約型タスクの現在の線形シーケンスモデルを大幅に上回ることを示しており、トランスモデルに匹敵するパフォーマンスを達成することさえあります。
このコードはhttps://github.com/opensparsellms/momでリリースされ、https：//github.com/opensparsellms/linear-moeの一部としてリリースされます。

要約(オリジナル)

Linear sequence modeling methods, such as linear attention, state space modeling, and linear RNNs, offer significant efficiency improvements by reducing the complexity of training and inference. However, these methods typically compress the entire input sequence into a single fixed-size memory state, which leads to suboptimal performance on recall-intensive downstream tasks. Drawing inspiration from neuroscience, particularly the brain’s ability to maintain robust long-term memory while mitigating ‘memory interference’, we introduce a novel architecture called Mixture-of-Memories (MoM). MoM utilizes multiple independent memory states, with a router network directing input tokens to specific memory states. This approach greatly enhances the overall memory capacity while minimizing memory interference. As a result, MoM performs exceptionally well on recall-intensive tasks, surpassing existing linear sequence modeling techniques. Despite incorporating multiple memory states, the computation of each memory state remains linear in complexity, allowing MoM to retain the linear-complexity advantage during training, while constant-complexity during inference. Our experimental results show that MoM significantly outperforms current linear sequence models on downstream language tasks, particularly recall-intensive tasks, and even achieves performance comparable to Transformer models. The code is released at https://github.com/OpenSparseLLMs/MoM and is also released as a part of https://github.com/OpenSparseLLMs/Linear-MoE.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google