UltraMemV2: Масштабирование сетей памяти до 120 миллиардов параметров с превосходным обучением на длинных контекстах

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

August 26, 2025

Авторы: Zihao Huang, Yu Bao, Qiyang Min, Siyan Chen, Ran Guo, Hongzhi Huang, Defa Zhu, Yutao Zeng, Banggu Wu, Xun Zhou, Siyuan Qiao

cs.AI

Аннотация

Хотя модели Mixture of Experts (MoE) достигают впечатляющей эффективности за счет активации только подмножеств параметров, они страдают от высоких затрат на доступ к памяти во время вывода. Архитектуры с использованием слоев памяти предлагают привлекательную альтернативу с минимальным количеством обращений к памяти, но предыдущие попытки, такие как UltraMem, лишь соответствовали производительности моделей MoE с 2 экспертами, значительно уступая современным конфигурациям с 8 экспертами. Мы представляем UltraMemV2 — переработанную архитектуру с использованием слоев памяти, которая устраняет этот разрыв в производительности. Наш подход включает пять ключевых улучшений: интеграцию слоев памяти в каждый блок трансформера, упрощение расширения значений с помощью одиночных линейных проекций, использование обработки значений на основе FFN из PEER, внедрение принципиальной инициализации параметров и перебалансировку соотношения вычислений между памятью и FFN. В ходе обширной оценки мы демонстрируем, что UltraMemV2 достигает паритета производительности с моделями MoE с 8 экспертами при одинаковых вычислительных затратах и количестве параметров, но с значительно меньшим количеством обращений к памяти. Примечательно, что UltraMemV2 показывает превосходную производительность на задачах, интенсивно использующих память, с улучшениями на +1,6 балла в запоминании длинных контекстов, +6,2 балла в многораундовом запоминании и +7,9 балла в обучении в контексте. Мы подтверждаем наш подход на масштабных моделях с активированными параметрами до 2,5 млрд из общего числа 120 млрд параметров и устанавливаем, что плотность активации оказывает большее влияние на производительность, чем общее количество разреженных параметров. Наша работа выводит архитектуры с использованием слоев памяти на уровень производительности современных моделей MoE, предлагая убедительную альтернативу для эффективных разреженных вычислений.

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.