MoM: Моделирование линейных последовательностей с использованием смеси памяти

MoM: Linear Sequence Modeling with Mixture-of-Memories

February 19, 2025

Авторы: Jusen Du, Weigao Sun, Disen Lan, Jiaxi Hu, Yu Cheng

cs.AI

Аннотация

Методы линейного моделирования последовательностей, такие как линейное внимание, моделирование пространства состояний и линейные рекуррентные нейронные сети (RNN), предлагают значительное повышение эффективности за счет снижения сложности обучения и вывода. Однако эти методы обычно сжимают всю входную последовательность в одно фиксированное состояние памяти, что приводит к неоптимальной производительности на задачах, требующих интенсивного воспроизведения информации. Вдохновляясь нейробиологией, в частности способностью мозга поддерживать устойчивую долговременную память, минимизируя "интерференцию памяти", мы представляем новую архитектуру под названием Mixture-of-Memories (MoM). MoM использует несколько независимых состояний памяти, а сеть-маршрутизатор направляет входные токены в конкретные состояния памяти. Этот подход значительно увеличивает общую емкость памяти, минимизируя интерференцию. В результате MoM демонстрирует исключительную производительность на задачах, требующих интенсивного воспроизведения, превосходя существующие методы линейного моделирования последовательностей. Несмотря на использование нескольких состояний памяти, вычисление каждого состояния остается линейным по сложности, что позволяет MoM сохранять преимущество линейной сложности при обучении и постоянной сложности при выводе. Наши эксперименты показывают, что MoM значительно превосходит современные линейные модели последовательностей на языковых задачах, особенно на задачах, требующих интенсивного воспроизведения, и даже достигает производительности, сравнимой с моделями Transformer. Код доступен по адресу https://github.com/OpenSparseLLMs/MoM и также является частью https://github.com/OpenSparseLLMs/Linear-MoE.

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.