UltraMemV2: Redes de Memória Escalando para 120 Bilhões de Parâmetros com Aprendizado Superior de Contexto Longo

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

August 26, 2025

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

cs.AI

Resumo

Embora os modelos de Mistura de Especialistas (MoE) alcancem eficiência notável ao ativar apenas subconjuntos de parâmetros, eles sofrem com altos custos de acesso à memória durante a inferência. As arquiteturas de camadas de memória oferecem uma alternativa atraente com muito poucos acessos à memória, mas tentativas anteriores, como o UltraMem, só conseguiram igualar o desempenho de modelos MoE com 2 especialistas, ficando significativamente aquém das configurações state-of-the-art com 8 especialistas. Apresentamos o UltraMemV2, uma arquitetura de camadas de memória redesenhada que fecha essa lacuna de desempenho. Nossa abordagem introduz cinco melhorias principais: integração de camadas de memória em cada bloco do transformer, simplificação da expansão de valores com projeções lineares únicas, adoção do processamento de valores baseado em FFN do PEER, implementação de inicialização de parâmetros fundamentada e reequilíbrio das proporções de computação entre memória e FFN. Por meio de avaliação extensiva, demonstramos que o UltraMemV2 alcança paridade de desempenho com modelos MoE de 8 especialistas sob a mesma computação e parâmetros, mas com acesso à memória significativamente menor. Notavelmente, o UltraMemV2 mostra desempenho superior em tarefas intensivas em memória, com melhorias de +1,6 pontos em memorização de contexto longo, +6,2 pontos em memorização de múltiplas rodadas e +7,9 pontos em aprendizado em contexto. Validamos nossa abordagem em escala com modelos de até 2,5 bilhões de parâmetros ativados de um total de 120 bilhões de parâmetros, e estabelecemos que a densidade de ativação tem maior impacto no desempenho do que a contagem total de parâmetros esparsos. Nosso trabalho leva as arquiteturas de camadas de memória à paridade de desempenho com os modelos MoE state-of-the-art, apresentando uma alternativa convincente para computação esparsa eficiente.

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.