Transformer en boucle efficace en mémoire : découpler le calcul de la mémoire dans les modèles de langage bouclés

Résumé

Les architectures LLM récurrentes sont devenues une approche prometteuse pour améliorer le raisonnement, car elles permettent un calcul multi-étapes dans l'espace de plongement sans générer de jetons intermédiaires. Des modèles tels que Ouro effectuent un raisonnement en mettant à jour itérativement les représentations internes tout en conservant un cache standard de clés-valeurs (KV) à travers les itérations, ce qui entraîne une consommation de mémoire qui croît linéairement avec la profondeur de raisonnement. Par conséquent, augmenter le nombre d'itérations de raisonnement peut conduire à une utilisation prohibitive de la mémoire, limitant ainsi l'évolutivité pratique de telles architectures. Dans ce travail, nous proposons le Transformateur Bouclé à Mémoire Efficace (MELT), une architecture novatrice qui dissocie la profondeur de raisonnement de la consommation de mémoire. Au lieu d'utiliser un cache KV standard par couche et par boucle, MELT maintient un seul cache KV par couche qui est partagé entre les boucles de raisonnement. Ce cache est mis à jour au fil du temps via un mécanisme de portail apprenable. Pour permettre un entraînement stable et efficace sous cette architecture, nous proposons d'entraîner MELT en utilisant un entraînement par blocs dans une procédure en deux phases : transition interpolée, suivie d'une distillation alignée sur l'attention, toutes deux à partir du modèle de départ LoopLM vers MELT. Empiriquement, nous montrons que les modèles MELT affinés à partir des paramètres pré-entraînés d'Ouro surpassent les LLM standard de taille comparable, tout en maintenant une empreinte mémoire comparable à celle de ces modèles et nettement inférieure à celle d'Ouro. Dans l'ensemble, MELT réalise un raisonnement itératif à mémoire constante sans sacrifier les performances de LoopLM, en utilisant seulement une procédure de post-entraînement légère.

English

Recurrent LLM architectures have emerged as a promising approach for improving reasoning, as they enable multi-step computation in the embedding space without generating intermediate tokens. Models such as Ouro perform reasoning by iteratively updating internal representations while retaining a standard Key-Value (KV) cache across iterations, causing memory consumption to grow linearly with reasoning depth. Consequently, increasing the number of reasoning iterations can lead to prohibitive memory usage, limiting the practical scalability of such architectures. In this work, we propose Memory-Efficient Looped Transformer (MELT), a novel architecture that decouples reasoning depth from memory consumption. Instead of using a standard KV cache per layer and loop, MELT maintains a single KV cache per layer that is shared across reasoning loops. This cache is updated over time via a learnable gating mechanism. To enable stable and efficient training under this architecture, we propose to train MELT using chunk-wise training in a two phase procedure: interpolated transition, followed by attention-aligned distillation, both from the LoopLM starting model to MELT. Empirically, we show that MELT models fine-tuned from pretrained Ouro parameters outperform standard LLMs of comparable size, while maintaining a memory footprint comparable to those models and dramatically smaller than Ouro's. Overall, MELT achieves constant-memory iterative reasoning without sacrificing LoopLM performance, using only a lightweight post-training procedure.