UltraMemV2: Reti di Memoria che Scalano fino a 120 Miliardi di Parametri con un Apprendimento di Contesti Lunghi Superiore

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

August 26, 2025

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

cs.AI

Abstract

Mentre i modelli Mixture of Experts (MoE) raggiungono un'efficienza notevole attivando solo sottoinsiemi di parametri, soffrono di elevati costi di accesso alla memoria durante l'inferenza. Le architetture a strati di memoria offrono un'alternativa interessante con pochissimi accessi alla memoria, ma tentativi precedenti come UltraMem hanno eguagliato solo le prestazioni di modelli MoE a 2 esperti, rimanendo significativamente al di sotto delle configurazioni all'avanguardia a 8 esperti. Presentiamo UltraMemV2, un'architettura a strati di memoria riprogettata che colma questo divario prestazionale. Il nostro approccio introduce cinque miglioramenti chiave: integrazione degli strati di memoria in ogni blocco transformer, semplificazione dell'espansione dei valori con proiezioni lineari singole, adozione dell'elaborazione dei valori basata su FFN da PEER, implementazione di un'inizializzazione dei parametri basata su principi e riequilibrio dei rapporti di calcolo tra memoria e FFN. Attraverso una valutazione estensiva, dimostriamo che UltraMemV2 raggiunge la parità prestazionale con i modelli MoE a 8 esperti a parità di calcolo e parametri, ma con un accesso alla memoria significativamente inferiore. In particolare, UltraMemV2 mostra prestazioni superiori nei task intensivi di memoria, con miglioramenti di +1,6 punti nella memorizzazione a contesto lungo, +6,2 punti nella memorizzazione multi-round e +7,9 punti nell'apprendimento in-context. Convalidiamo il nostro approccio su larga scala con modelli fino a 2,5 miliardi di parametri attivati su un totale di 120 miliardi, e stabiliamo che la densità di attivazione ha un impatto maggiore sulle prestazioni rispetto al conteggio totale dei parametri sparsi. Il nostro lavoro porta le architetture a strati di memoria alla parità prestazionale con i modelli MoE all'avanguardia, presentando un'alternativa convincente per il calcolo sparso efficiente.

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.