UltraMemV2: Redes de Memoria que Escalan hasta 120B Parámetros con un Aprendizaje de Contexto Largo Superior

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

Resumen

Si bien los modelos de Mezcla de Expertos (MoE) logran una eficiencia notable al activar solo subconjuntos de parámetros, sufren de altos costos de acceso a memoria durante la inferencia. Las arquitecturas de capas de memoria ofrecen una alternativa atractiva con muy pocos accesos a memoria, pero intentos previos como UltraMem solo han igualado el rendimiento de modelos MoE de 2 expertos, quedándose significativamente por debajo de las configuraciones de última generación con 8 expertos. Presentamos UltraMemV2, una arquitectura de capas de memoria rediseñada que cierra esta brecha de rendimiento. Nuestro enfoque introduce cinco mejoras clave: integrar capas de memoria en cada bloque del transformador, simplificar la expansión de valores con proyecciones lineales únicas, adoptar el procesamiento de valores basado en FFN de PEER, implementar una inicialización de parámetros fundamentada y reequilibrar las proporciones de cálculo entre memoria y FFN. A través de una evaluación exhaustiva, demostramos que UltraMemV2 alcanza un rendimiento equivalente a los modelos MoE de 8 expertos con el mismo cómputo y parámetros, pero con un acceso a memoria significativamente menor. Notablemente, UltraMemV2 muestra un rendimiento superior en tareas intensivas en memoria, con mejoras de +1.6 puntos en memorización de contexto largo, +6.2 puntos en memorización de múltiples rondas y +7.9 puntos en aprendizaje en contexto. Validamos nuestro enfoque a gran escala con modelos de hasta 2.5 mil millones de parámetros activados de un total de 120 mil millones, y establecemos que la densidad de activación tiene un mayor impacto en el rendimiento que el recuento total de parámetros dispersos. Nuestro trabajo lleva las arquitecturas de capas de memoria a un nivel de rendimiento equivalente a los modelos MoE de última generación, presentando una alternativa convincente para el cómputo disperso 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.