Mamba Rempli : Effondrement de l'État et Capacité de l'État de la Modélisation à Long Contexte Basée sur les RNN

Stuffed Mamba: State Collapse and State Capacity of RNN-Based Long-Context Modeling

October 9, 2024

Auteurs: Yingfa Chen, Xinrong Zhang, Shengding Hu, Xu Han, Zhiyuan Liu, Maosong Sun

cs.AI

Résumé

Un avantage essentiel des réseaux neuronaux récurrents (RNN) par rapport aux modèles de langage basés sur les transformateurs est leur complexité computationnelle linéaire en ce qui concerne la longueur de la séquence, ce qui les rend beaucoup plus rapides pour traiter de longues séquences lors de l'inférence. Cependant, la plupart des RNN disponibles publiquement (par exemple, Mamba et RWKV) sont entraînés sur des séquences de moins de 10 000 jetons, et leur efficacité dans des contextes plus longs reste largement insatisfaisante jusqu'à présent. Dans cet article, nous étudions la cause de l'incapacité des RNN à traiter de longs contextes et suggérons des atténuations critiques. Nous examinons deux préoccupations pratiques lors de l'application des RNN de pointe à de longs contextes : (1) l'incapacité d'extrapoler vers des entrées plus longues que la longueur d'entraînement et (2) la limite supérieure de la capacité mémoire. Pour répondre à la première préoccupation, nous étudions d'abord l'effondrement de l'état (SC), un phénomène qui entraîne une dégradation importante des performances sur des longueurs de séquence non rencontrées lors de l'entraînement. À l'aide d'expériences contrôlées, nous attribuons cela à un surajustement dû à la surparamétrisation de l'état récurrent pour la longueur d'entraînement. Pour la deuxième préoccupation, nous entraînons une série de modèles Mamba-2 sur de longs documents pour estimer empiriquement la capacité de l'état récurrent dans la modélisation du langage et la récupération de clés. Ensuite, trois méthodes d'atténuation de SC sont proposées pour améliorer la capacité de généralisation de la longueur de Mamba-2, permettant au modèle de traiter plus d'1 million de jetons sans SC. Nous constatons également que la capacité de l'état récurrent dans la récupération de clés évolue de manière exponentielle avec la taille de l'état, et nous entraînons empiriquement un Mamba-2 370M avec une précision de récupération de clés quasi parfaite sur une longueur de contexte de 256 000. Cela laisse entrevoir un avenir prometteur pour la modélisation de longs contextes basée sur les RNN.

English

One essential advantage of recurrent neural networks (RNNs) over transformer-based language models is their linear computational complexity concerning the sequence length, which makes them much faster in handling long sequences during inference. However, most publicly available RNNs (e.g., Mamba and RWKV) are trained on sequences with less than 10K tokens, and their effectiveness in longer contexts remains largely unsatisfying so far. In this paper, we study the cause of the inability to process long context for RNNs and suggest critical mitigations. We examine two practical concerns when applying state-of-the-art RNNs to long contexts: (1) the inability to extrapolate to inputs longer than the training length and (2) the upper bound of memory capacity. Addressing the first concern, we first investigate *state collapse* (SC), a phenomenon that causes severe performance degradation on sequence lengths not encountered during training. With controlled experiments, we attribute this to overfitting due to the recurrent state being overparameterized for the training length. For the second concern, we train a series of Mamba-2 models on long documents to empirically estimate the recurrent state capacity in language modeling and passkey retrieval. Then, three SC mitigation methods are proposed to improve Mamba-2's length generalizability, allowing the model to process more than 1M tokens without SC. We also find that the recurrent state capacity in passkey retrieval scales exponentially to the state size, and we empirically train a Mamba-2 370M with near-perfect passkey retrieval accuracy on 256K context length. This suggests a promising future for RNN-based long-context modeling.