MoM: Lineare Sequenzmodellierung mit Mixture-of-Memories

MoM: Linear Sequence Modeling with Mixture-of-Memories

February 19, 2025

Autoren: Jusen Du, Weigao Sun, Disen Lan, Jiaxi Hu, Yu Cheng

cs.AI

Zusammenfassung

Lineare Sequenzmodellierungsmethoden wie lineare Aufmerksamkeit, Zustandsraummodellierung und lineare RNNs bieten erhebliche Effizienzsteigerungen, indem sie die Komplexität von Training und Inferenz reduzieren. Diese Methoden komprimieren jedoch typischerweise die gesamte Eingabesequenz in einen einzigen festen Speicherzustand, was zu suboptimaler Leistung bei erinnerungsintensiven Downstream-Aufgaben führt. Inspiriert von der Neurowissenschaft, insbesondere der Fähigkeit des Gehirns, ein robustes Langzeitgedächtnis zu erhalten und gleichzeitig „Gedächtnisinterferenz“ zu minimieren, führen wir eine neuartige Architektur namens Mixture-of-Memories (MoM) ein. MoM nutzt mehrere unabhängige Speicherzustände, wobei ein Router-Netzwerk Eingabetoken zu bestimmten Speicherzuständen leitet. Dieser Ansatz erhöht die gesamte Speicherkapazität erheblich und minimiert gleichzeitig Gedächtnisinterferenzen. Dadurch schneidet MoM bei erinnerungsintensiven Aufgaben außergewöhnlich gut ab und übertrifft bestehende lineare Sequenzmodellierungstechniken. Obwohl mehrere Speicherzustände integriert sind, bleibt die Berechnung jedes Speicherzustands linear in der Komplexität, sodass MoM den Vorteil der linearen Komplexität während des Trainings beibehält, während die Inferenz konstant komplex bleibt. Unsere experimentellen Ergebnisse zeigen, dass MoM aktuelle lineare Sequenzmodelle bei Downstream-Sprachaufgaben, insbesondere erinnerungsintensiven Aufgaben, deutlich übertrifft und sogar eine Leistung erzielt, die mit Transformer-Modellen vergleichbar ist. Der Code ist unter https://github.com/OpenSparseLLMs/MoM und auch als Teil von https://github.com/OpenSparseLLMs/Linear-MoE veröffentlicht.

English

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.