UltraMemV2: Geheugennetwerken die schalen tot 120B parameters met superieur lang-context leren

Samenvatting

Hoewel Mixture of Experts (MoE)-modellen opmerkelijke efficiëntie bereiken door slechts subsets van parameters te activeren, kampen ze met hoge geheugentoegangskosten tijdens inferentie. Geheugenlaagarchitecturen bieden een aantrekkelijk alternatief met zeer weinig geheugentoegang, maar eerdere pogingen zoals UltraMem hebben alleen de prestaties van 2-expert MoE-modellen geëvenaard, wat aanzienlijk achterblijft bij state-of-the-art 8-expert configuraties. Wij presenteren UltraMemV2, een herontworpen geheugenlaagarchitectuur die deze prestatiekloof dicht. Onze aanpak introduceert vijf belangrijke verbeteringen: het integreren van geheugenlagen in elk transformerblok, het vereenvoudigen van waarde-expansie met enkele lineaire projecties, het overnemen van FFN-gebaseerde waarde-verwerking van PEER, het implementeren van principiële parameterinitialisatie, en het herverdelen van geheugen-naar-FFN berekeningsverhoudingen. Door uitgebreide evaluatie tonen we aan dat UltraMemV2 prestatiepariteit bereikt met 8-expert MoE-modellen onder dezelfde berekening en parameters, maar met aanzienlijk lagere geheugentoegang. Opmerkelijk is dat UltraMemV2 superieure prestaties laat zien op geheugenintensieve taken, met verbeteringen van +1,6 punten op lang-context memorisatie, +6,2 punten op multi-ronde memorisatie, en +7,9 punten op in-context leren. We valideren onze aanpak op schaal met modellen tot 2,5B geactiveerde parameters uit 120B totale parameters, en stellen vast dat activatiedichtheid een grotere impact heeft op prestaties dan het totale aantal sparse parameters. Ons werk brengt geheugenlaagarchitecturen naar prestatiepariteit met state-of-the-art MoE-modellen, en presenteert een overtuigend alternatief voor efficiënte sparse berekening.

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.

UltraMemV2: Geheugennetwerken die schalen tot 120B parameters met superieur lang-context leren

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

Samenvatting

Support