UltraMemV2 : Réseaux de mémoire évolutifs jusqu'à 120 milliards de paramètres avec un apprentissage de contexte long supérieur

UltraMemV2: Memory Networks Scaling to 120B Parameters with Superior Long-Context Learning

August 26, 2025

papers.authors: Zihao Huang, Yu Bao, Qiyang Min, Siyan Chen, Ran Guo, Hongzhi Huang, Defa Zhu, Yutao Zeng, Banggu Wu, Xun Zhou, Siyuan Qiao

cs.AI

papers.abstract

Bien que les modèles de mélange d'experts (MoE) atteignent une efficacité remarquable en activant uniquement des sous-ensembles de paramètres, ils souffrent de coûts élevés d'accès à la mémoire lors de l'inférence. Les architectures à couches mémoire offrent une alternative attrayante avec très peu d'accès à la mémoire, mais les tentatives précédentes comme UltraMem n'ont égalé que la performance des modèles MoE à 2 experts, restant bien en deçà des configurations de pointe à 8 experts. Nous présentons UltraMemV2, une architecture à couches mémoire repensée qui comble cet écart de performance. Notre approche introduit cinq améliorations clés : l'intégration de couches mémoire dans chaque bloc de transformateur, la simplification de l'expansion des valeurs avec des projections linéaires uniques, l'adoption du traitement des valeurs basé sur FFN de PEER, la mise en œuvre d'une initialisation des paramètres raisonnée, et le rééquilibrage des ratios de calcul mémoire-FFN. Grâce à une évaluation approfondie, nous démontrons qu'UltraMemV2 atteint une parité de performance avec les modèles MoE à 8 experts pour un même calcul et nombre de paramètres, mais avec un accès à la mémoire significativement réduit. Notamment, UltraMemV2 montre une performance supérieure sur les tâches intensives en mémoire, avec des améliorations de +1,6 points sur la mémorisation de contexte long, +6,2 points sur la mémorisation multi-tours, et +7,9 points sur l'apprentissage en contexte. Nous validons notre approche à grande échelle avec des modèles allant jusqu'à 2,5 milliards de paramètres activés sur un total de 120 milliards, et établissons que la densité d'activation a un impact plus important sur la performance que le nombre total de paramètres épars. Notre travail amène les architectures à couches mémoire à une parité de performance avec les modèles MoE de pointe, offrant une alternative convaincante pour le calcul épars efficace.

English

While Mixture of Experts (MoE) models achieve remarkable efficiency by activating only subsets of parameters, they suffer from high memory access costs during inference. Memory-layer architectures offer an appealing alternative with very few memory access, but previous attempts like UltraMem have only matched the performance of 2-expert MoE models, falling significantly short of state-of-the-art 8-expert configurations. We present UltraMemV2, a redesigned memory-layer architecture that closes this performance gap. Our approach introduces five key improvements: integrating memory layers into every transformer block, simplifying value expansion with single linear projections, adopting FFN-based value processing from PEER, implementing principled parameter initialization, and rebalancing memory-to-FFN computation ratios. Through extensive evaluation, we demonstrate that UltraMemV2 achieves performance parity with 8-expert MoE models under same computation and parameters but significantly low memory access. Notably, UltraMemV2 shows superior performance on memory-intensive tasks, with improvements of +1.6 points on long-context memorization, +6.2 points on multi-round memorization, and +7.9 points on in-context learning. We validate our approach at scale with models up to 2.5B activated parameters from 120B total parameters, and establish that activation density has greater impact on performance than total sparse parameter count. Our work brings memory-layer architectures to performance parity with state-of-the-art MoE models, presenting a compelling alternative for efficient sparse computation.