UltraMemV2: Memory Networks Scaling to 120B Parameters with Superior Long-Context Learning
Authors: Zihao Huang
Yu Bao
Qiyang Min
Siyan Chen
Ran Guo
Hongzhi Huang
Defa Zhu
Banggu Wu
Yutao Zeng
zhou Xun
Siyuan Qiao
Venue: ICLR 2026 Poster (2026)
Keywords: memory network
moe
pretrain
long context
Primary area: foundation or frontier models, including LLMs
TL;DR: Previous  Memory-layer attempts have only matched the performance of 2-expert MoE models, falling significantly short of state-of-the-art 8-expert configurations. We present UltraMemV2 to closes this performance gap.

Abstract: 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.

Published as a conference paper at ICLR 2026
ULTRAMEMV2: MEMORY NETWORKS SCALING TO 120B PARAMETERS WITH SUPERIOR LONG-CONTEXT LEARNING
Zihao Huang, Yu Bao; Qiyang Siyan Chen, Ran Guo, Hongzhi Huang, Defa Zhu, Yutao Zeng, Banggu Wu, Xun Zhou, Siyuan Qiao
Bytedance Seed (huangz hao. notabot, baoyu. 3302, mingi hinnee/eeeeegeeeeeddeeeeee...o
ABSTRACT
While Mixture of Experts (MoE) models achieve remarkable efficiency by activat- ing 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-an 8-expert configurations. We present UltraMem redesigned memory-layer ar- chitecture that closes this performance gap. Our approach introduces five key im- provements: integrating memory layers into every transformer block, simplifying value expansion with single linear projections, adopting FFN-based value process- ing from PEER, implementing principled parameter initialization, and rebalancing memory-to-Pii -F-NNN computation ratios. Through extensive evaluation, we demon- strate that UlraMemM2 achieves performance parity with 8-expert MoE models under same computation and parameters but significantly low memory access. Notably, UltraMem V2 shows superior performance on memory-ineensive 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 per- formance than total sparse parameter count. Our work brings memory- ar- chitectures to performance parity with state-of-the-art MoE models, presenting a compelling alternative for efficient sparse computation.
INTRODUCTION
Large language models (LLMs) have achieved remarkable success across NLP tasks, but their ponential growth in parameters and computational complexity presents significant challenges for resource-constraiied deployment. Mixture of Experts (MoE)(Fedus et al. 2022: He et al.. 2021 Liu et al. 2024 Muennighoff et al. 2024 Yuan et al.) 2025) have emerged as a promising solution by se-
lectively activating expert subsets, effectively decoupling parameter count from computational cost. Recent works (Mucnnighoff et al. 20244 Krajewski et al. (2024) show that MoE with 8 activated trade-offs, significantly outperforming configura-
experts achieves optimal performance-effictnn tions with fewer experts. However, MoE inference suffers from high memory access costs due to expert routing overhead, particularly problematic when only a small fraction of tokens activate all experts.
Memory-layer grchitectures(Lampee et al. 2019 Huang et al. 2024 Berges al. 2024) offer an al- ternative sparse model with significantly less memory access. Unlike MoE's s FFN-type expert, mem- ory layers activate embeddings from large parameter table, enabling extremely slowly lincar scaling of memory access with sequence length. The Over-tokenized Transformer (Huang al. (025) can also be viewed as a memory-layer architecture, where an n-gram router activates embeddings from a memory table that subsequently added to the word embeddings. While architectures like Ultra-
"Equal contribution, Corresponding authors

---

Published as a conference paper at ICLR 2026
Mem(Huang et al. 2024), current memory-layed State Of The Art(SOTA), demonstrate promising inference characteristics, they have only matched the performance of MoE with activated experts, falling short of SOTA 8-expert configurations by substantial margin.
This performance gap motivates our work. We introduce UltraMem V2. a redesigned memory-layer architecture that bridges the performance divide between embedding- based and expert-based sparse models. Our approach incorporates five key innovations: (1) architectural integration: tighter coupling between memory layers and Transformer blocks with memory layers in every block; (2) simplified value expansion: streamlined Implicit Value Expansion (IVE) using single linear pro- jections; (3) expert-like value processing: adoption of PEER's FFN- value computation (He)
2024| (4) optimized initialization: principled parameter initialization preventing training diver- gence; and (5) computational rebalancing: adjusted memory-to- FFN computation ratios.
Through comprehensive evaluation, we demonstrate that UltraMem V2 achieves performance parity with 8-expert MoE models while maintaining memory layer advantages. Notably, UltraMem V2 shows superior performance on memory-intensive tasks including long-context memorization (+1.6 points), multi-round memorization (+6.2 points), and in-context learning (+7.9 points). We validate scalability up to 2.5B activated parameters with 120B total parameters, and establish that activation density (top-m values) has greater impact on performance than total sparse parameter count.
In summary, our work makes three primary contributions: (I) Architectural advancement: We present the first memory-layer architecture competitive with state-of-the-ar 8-expert MoE models, closing significant performance gap in sparse model research. (2) Comprehensive analysis: We provide detailed ablation studies and comparative analysis revealing when and why UlrrMemmV2 outperforms MoE, particularly on memory-intensive tasks, while identifying trade-offs in different training phases. (3) Scalability validation: We demonstrate UltraMem V2's effectiveness at scale and establish design principles for activation density versus parameter count trade-offs, providing guidance for future memory-laye architectures.
2 RELATED WORK
MoE Architecture The concept of MoE was first introduced by (Shazeer et al. (2117). Since then, numerous studies (Fedus et al., 2022 Dai et al. 20224 Rajbhandari et al. 2022 Jiang et al. 2024) have been conducted improve its performance and efficiency. During this period, the general perception is that appropriately using smaller experts but activating a greater number can enhance the performance of MoE, typically activating two experts. Krajewski et al. (2024) systematically studied the influence of expert size and the number of activations, which is called "granularity"" They found that when the granularity was 8, MoE achieved the best performance and was significantly better than 2. The same conclusion was also discovered by (LMooEMMuennighoi) et al., 2024). Resent MOEs in the industrial sector (DeepSeek- (3 et al. 2024), Qwen3( Yang et al. 2025).
dots.Ilm (Huo et al. )025)) have all adopted this structure. However, they still face challenges in inference, such as high memory access costs and long inference latency, especially when dealing with large-scale models.
Memory Layer Architecture The idea of a memory layer was first explored by] Lample et al. (2019) with the introduction of the Product Key Memory (PKM). By activating embeddings instead of ex- pert, PKM aimed to expand the model's parameters while maintaining similar computation and less memory access. Subsequently, several improvements have been made PKM. For example, Kim & Jung (2020) introduced a concept similar to shared experts in MoE, allowing PKM and MLP to op- erate in parallel. Csordás al, (2023) made slight modification to PKM by removing the Softmax operation. He (2024) proposed PEER, which improved the activation of values in PKM by using an FFN with one inner dimension. Memory+(Berges et al.) 2024) also made some improvements the memory layer architecture. However, most of these memory layer architectures have only managed to match the performance of MoE models with one activated experts. The UltraMem (Huang et al.) 2024) was an attempt to address the limitations of existing memory layer architectures. It incorpo- rated Tucker Decomposed Query Key retrieval (TDQKR) and Implicit Value Expansion (IVE) to improve model performance while maintaining inference latency. However, it still can only matchs the performance of MoE with two activated experts. UlraMMmmV V2 builds on the previous work and aims overcome these limitations. By introducing several innovative improvements, UltraMem22
2

---

Published as a conference paper at ICLR 2026
achieve comparable performance MoE models with eight activated experts, filling the gap in the current research on memory layer architectures.
APPROACH
3.1 PRELIMILARY
Memory layers are structures designed to expand model capacity without a proportional increase in computational cost. We briefly review three common architectures: MoE, PKM and UltraMem.
MoE Layer, as shown in Figure ((), utilizes gating mechanism to selectively activate subset of parameters. Given hidden state x E RD,,, gate with parameters K Nx> computes routing scores for N experts.
S - X
(1)
The Top-M function selects the indices I of the m experts with the highest scores. These indices are used to retrieve the corresponding parameters from the expert pool, which consist of Pre-values U E Nx Dinner and Values V E RNxx Dinner, The final output o is the combination of the activated parameters, often weighted by the gating scores 8, for i E I. A common formulation is:
o = Σ SeY (SiLU(XU,) x V
(2)
iET
where Ug and V, are the parameters for the i-th expert.
PKM Layer, illustrated in Figure bb), employs key factorization to create a large memory from smaller, more efficient key-value sets. The input hidden state x is first projected to query q via linear layer grow and 9col- This query is then used to compute scores against multiple, smaller sets of factorized keys, row keys Keo N 1ss and column keys K E Nxxx,, retrieve values from N2 value pool by Product Quantizaion (PQ)(eegu et al. 2010). The scores from these factorized components aggregated identify the most relevant memory value in the full memory space. A Top-M function selects the indices I and scores Sgrid for the best-matching memory entries.
Syow Toopms (K Trww (x));
Senl = TTopN K cotsfooot (x)),
Syrid = Topx (Srow + s cot)-
(4)
The final output a weighted sum of the corresponding values V; retrieved from the value memory.
Σ
St grid V
(5)
IEI
Commonly PKM use multi-head trick, which similar to Multi-head attention Vaswani al. 2017). we omit this operation in the above formula description for brevity. This factorization allows for much larger memory capacity than could be addressed by single, monolithic key matrix, while keeping the number of parameters manageable.
UltraMem layer is illustrated in Figure ()c. It typically consists of keys Krow, Keof E R...D.r, tucker core C Rn and values V g= Given hidden state E D.,,, scores for activated values are computed by Tucker Decomposed Query-Key Retrieval (TDQKR):
Srac K Krow (x),
Scel K (x))
Syrid Trepx (STow x C Scot);
where Trow, Tcot : R:D: + TD. linearly convert the dimension of the input to DR. Cropu()) selects the top-m scores and set the remaining scores negative infinity. Then weighted pooling with Implicit Value Expansion (IVE) is conducted to generate the final output of memory layer:
$ --huffleevees grii)))
V $ Σ V x Σ W (VT 8p)).
(9)
where operation Shuffle aims eliminate row and column scoring, &, represents the scores corresponding to p-th virtual memory block, and W, E R D.,D: is linear projector for p-th virtual memory block. We have overlooked cumbersome operations and only present the key ones here.
unnecessary index topology prior introduced by
3

---

Published as a conference paper at ICLR 2026
Output
Keol
K Krow
value Inner
TDQCE
Activated Vallues
Keui
Ke
Index
TDQKER
scores
K
Top
TDQKR
Activated Pre-volues
Pre-values
ceol
Pre-value linear
querry linar
TDQKR
Hidden
Figure 1: Overall structure of UltraMem V2 layer.
3.2 OVERALL STRUCTURE
We present an improved structure called UltraMem shown in Figure |
Compared to Ultra-
Mem (Huang al. 2024), we highlight the improvements the following:
1) Every transformer block contains an FFN layer and an UltraMem layer.
2) The multiple linear layers in Implicit Value Expansion (IVE) removed and use only single linear layer. Meanwhile, we use separate queries for each tucker rank.
3) PEER (He 2024) is adopted, by which the embedding value changed to FFN with one inner dimension.
4) We improve the initialization of the parameters in the new structure.
5) We adjust the proportion of the calculation for memory layer.
3.3 DIFFERENT VIEW IN IMPLICIT VALUE EXPANSION
Given a set of row and column keys Ki Krow K inj [1,2, ......],, h). the Sprid in Equation g is
obtained by
S Stmm K you Brow St K (oo)).))
(10)
S Sorid = Topx (S IT ! bot SIT C cols SET CHiSi co.. ShT Ch
(11)
row
For a standard multi-head structure that requires h2 heads, generally h? row and column keys are needed. IVE introduces shared key pairs, so only h row and column keys needed to achieve h heads. Therefore, IVE has actually implemented a multi-head mechanism, there no need to explicitly define multi-head mechanism like that in Product Key Memory (PKM) (Lample al., 2019).
Considering that IVE adds linear layer for each head remap, but mapping different embeddings to corresponding linear layers will increase additional non-computational operations, affecting the inference speed. Meanwhile, we find that if the linear layer shared and the saved parameters are added to the FFN, better results can be achieved. This indicates that the parameter efficiency of nonshared linear layers is actually not high. As result, Equation is modified as:
V T $ T (V 8)
(12)
3.4 MILLION OF I-INNER-DIM EXPERTS INSTEAD EMBEDDINGS
PEER He 2024) uses FEN with one inner dimension replacing the value. Given val-v weight
matrix P "".,p, the final output is
W (VT (o(Px) 8 88)),
(13)
4

---

Published as a conference paper at ICLR 2026
where is the input, is the activate function. Consider standard SwiGLU FFN, given W1, W2 RM,, W3 E R...
o W3 (W2x (W.x))
(14)
We find that PEER is very similar to SwiGLU FFN 2020), where V, (PP) and & corre- sponds to Wa. W2x and o ( 1x), respectively. It should be noted that $ comes from Equations and 8] where TopM can be seen activation function to two parallel results simultancously is lossy. Therefore, we decide to remove the activation function in FFN based on PEER, leads the final output as
activate function. We argue that empirically, applying the
W (V ((xx) @ 6)).
(15)
For the sake of simplicity, this change will be uniformly abbreviated as PEER in this paper.
3.5 IMPROVED INITIALIZATION
When using a normal distribution to initialize the UltraMem layer, the standard deviation must be carefully designed; otherwise, the training process is very prone to divergence. The selection criteria for the initialization standard deviation are: (1) After initialization, the variance the output activation of the memory layer does not diverge with the increase in the number of layers; (2) The variance the output activation should not be too large. Considering that memory can be regarded an enhancement of FFN the memory layer consistent with that of FFN. The specific derivation process of the initialization variance can be found in the Appendix
certain extent, we make the initialization activation variance of
EXPERIMENTS
This section provides comprehensive experimental validation of the UltraMem V2 architecture. Our evaluation encompasses three primary objectives:
Performance Parity Validation: We demonstrate that UltraMem V2 achieves comparable performance state-of-the-art 8-expert MoE models, thereby bridging the substantial per- formance gap that has historically limited memory-layer architectures.
2. Architectural Advantage Analysis: We validate UltraMem superior performance on veetory-ntensive tasks, with particular emphasis on long-context memorization, multi- round memorization, and in-context learning capabilities. However, we also identify cer- tain limitations, including reduced effectiveness during early training phases and potential performance trade-offs in specific reasoning tasks compared to MoE.
3. Component Effectiveness Assessment: Verifying the effectiveness of core improvements through ablation studies, while reducing hyperparameter configuration requirements and simplifying the training pipeline.
Training Data Our experiments utilize both proprietary and open-source datasets. The proprietary training corpus comprises 3.9T tokens for Pre'raiining (PT) and 500B high-quality long context tokens for continued training (CT). For open-source comparisons, we employ the IT token dataset from OLMoE to ensure fair evaluation against existing baselines.
Evaluation Benchmarks We conduct evaluation across diverse benchmark suites encompassing both proprietary and open-source assessments. These include comprehensive evaluations of math, code, reasoning and knowledge capabilities. Additionally, we evaluate long-context performance through specialized benchmarks measuring long-context memorization, long-context reasoning, needle in haystack, and long-document retrieval capabilities. Detailed specifications of all evalua- tion datasets are provided in the Appendix
4.1 COMPARE TO MOE
We evaluate UltraMem V2 against both proprietary and open-source baselines across multiple train- ing stages and benchmarks. For proprietary models, compare against SeedMoF. variants with
5

---

Published as a conference paper at ICLR 2026
different parameter configurations. For open- models, we benchmark against OLMoE, Mem- ory+, and UltraMem architectures.
Proprictary models undergo two-stage training process: (1) pretraining (PT) on 1.6T tokens, followed by (2) continued training (CT) on 250B high-quality tokens. Selected models are further pretrained to 3.9T tokens with additional 32K context CT using 500B tokens. Open-source models undergo 500B or 1T tokens.
Training Protocol:
Model Configurations: For the proprietary model, UltraMem w-2..5//1200--topp2222 activates 2.5B parameters from 120B sparse parameters with 256 activated values per UltraMemV V2 layer. uses 768 activated values from 60B sparse parameters. We constrain row/column TopM to 128 to avoid quadratic intermediate variable explosion. For the open-source model, we make sure the same computation and parameters. The OLMoE has 64 experts, and each token activates 8 experts. Memory+ and UltraMem contains 4 memory layers and each memory layer has 2 heads with TopM 80. UltraMemV V2 contains 20 memory layers, which has head with TopM- 32. All three memory-layer--asse model activate same amount of value embeddings. Detail model hyperparameters shown Appendix D
UltraMem 22-..5//600B-oo/7777
Proprietary Model Comparison Table 1 presents comprehensive evaluation results across Open- Bench and HardBench benchmarks different training stages. Taaleeddemonstrate UltraMem V2's capabilities 32K long-context tasks. We observe several key findings:
Training Stage Dependency: UltraMemV2 exhibits distinct performance characteristics across PT and CT phases. After 1.6T PT. BatraMem/2-2.55/00---ooo underperforms on mathematical reasoning, coding, and reasoning tasks. However, following 250B CT, U1traMemV2 achieves competitive superior performance across all metrics, suggesting enhanced sensitivity to high-quality data and learning rate decay schedules.
2. Scaling behavior: After extended training (3.9T PT + 500B CT), 2222uMMm22-2..//6BBB top768 shows marginal improvements over SeedMoE-2.55/3000 on OpenBench but demon- strates clear advantages on HardBench. The diminishing returns may reflect general scaling limitations rather than architecture-speee specific issues, though this warrants further investiga- tion with larger parameter budgets.
3. Architecture trade-offs: Comparing
versus UltraMem V2- 2.5B/120B: 20B-16 we find that increasing activated values per layer (top768 vs top256) yields better performance than increasing total sparse parameters (60B 120B). This sug- gests activation density is more important than sparse parameter count, though higher acti- vation increases inference latency.
4. Long-Context Performance: UltraMem shows substantial improvements in memory- intensive tasks (long-context memorization: 23.5 vs 21.9, multi-round memorizing: 31.2 1.2 vs 25.0, in-context learning: 29.5 vs 21.6) compared SeedMoF. Performance variations in Kay-value retrieval and Multi-hop reasoning appear attributable to architectural differences rather than parameter count disparities.
Table I: Performance comparison of different models across various benchmarks. Models are grouped by training tokens with consistent background colors.
Training
Training
Openbench
tokens
Joss
knowledge
reasming
Math
code
All
knowledge
ressniing
code
1.6T PT
1.812
1.900
70.6
71.8
62.8
43.1
60.5
24.8
49.8
15:4
10.1
23.3
SeadMoFF
+250R CT
1.165
1,846
72.0
77.9
71.3
50.7
67.6
34.2
53.9
18.9
15.2
27.4
2
3.9T
1.791
1.891
73.3
74.0
65.6
45.6
63.1
25.9
51.8
16.5
10.9
23.5
1,837
79.5
ScedMoE:
1.6T
1.793
1.869
73.1
72.5
64.4
43.4
61.6
27.4
52.4
15.5
10.0
23.1
2,88808
2508 CT
1.123
1.796
79.1
76.9
71.2
54.4
68.1
35.6
56.
21.5
17.3
79.2
UlraaMemV.
L.67 PT
1.52
1.335
73.6
70.8
61.3
39.5
60.4
26.7
44.7
14.5
9.2
21.2
2.5B/120B 40g256
250B CT
1,066
1.803
80.7
76.4
71.8
52.4
68.3
35.5
54.4
20.4
16.7
27.9
1.855
UltraMemV.
250RR
1,770
80.1
72.5
2.88/00B-oop7688
3.9
1.748
1.847
76.4
74.2
62.6
45.6
62.8
27.8
53.4
16.5
11.8
24.5
+500 CT
0.975
1.784
81.7
79.0
73.9
56.9
70.7
38.9
57.5
23.8
19.4
31.7

---

Published as a conference paper at ICLR 2026
Table 2: Performance comparison on long-context tasks.
Long-contest memorizing
Multi-romnd memorizing
In-context learning
Find Needle
Key-val
Multi-hop reasoning
Reasoning
ScedMMEE-2.58/000
21.9 23.5
25.0 312
21.6 29.5
6.2 7.7
96.5
41.3
34.8
35.4
UlraMemV2-22500 top768
97.0
57.1
17.7
37.7
Open-Source Model Comparison Table evaluates AttraMemV2 against open-source alterna- tives under controlled parameter budgets. UltraMemV2 significantly outperforms memory- architectures (Memory+, UltraMem) while achieving competitive performance with OLMoE across 227M/1 .2B and IB/BB configurations. This validates UltraMem V effectiveness relative current state-of-the-art expert-routing MoE approaches.
Table 3: Comprehensive performance comparison across different open sourced models and bench- marks.
Model
Training
Eval
Evalnation Benchmarks
Tokens
Loss
Loss
ARC-C
ARC-E
Cummon- senerQA
Hellaas
MML.U-
Open- beokQ4
PIQA
SCIQ
Wine- grande
All
wag
Y3r
OU 12766-127M01
2.482
2.845
43.2
Mcmory--27MM1.222
3.566
2.900
10-2
5008
2.885
66.7
2,500
2.853
66.7
43.0
2.631
92.7
UlraMemV2 1877B
2,366
2,628
74.7
39.7
93.1
64.1
These results demonstrate UltraMem viability as a competitive alternative to current MoE paradigms.
4.2
STRUCTURE ABLATION
4.2.1 PEER
We investigate the impact of using PEER (He 2024) who first suggested replacing value embedding with an FFN having a I-dimensional inner layer. Our ablation study on compares two settings. The first, "Baseline", is configured with N 789 keys, value dimension 558 keys, pre-value
D. 288, and TopM 48. The second, "PEER", differs by using dimension D, = 288 (while D, remains 288), and TopM 24. Under this configuration, the calculation and memory access of the memory layer are completely consistent. As illustrated in Figure
2 which presents the training loss and downstream accuracy, PEER demonstrates signifi- cant advantage over the standard value embedding approach, noticed that the parameters involved in processing the activated top-m values is kept constant. This highlights the parameter efficiency and effectiveness of the FFN-based value processing proposed by PEER.
Training (Smnathed)
Open Benchmark
PEBR
PEER Buselin
Boaline
3.0
Acourais
0.26
2.8
0.24
2.8
2.6
0.22
0.20
15
20
30
30
70
1.00
150
200
150
200
300
Consumed tolens (B)
Coasumed tokens (B)
Figure 2: PEER ablation study on V2-MMMemV2-434M/5BBB Training loss (left) and Open-Oppenheeeee Benchmark accuracy (right) comparing PEER with Baseline using value embedding.
4.2.2 How LARGE DO D. AND D, NEED TO BE?
First, we set value dimention D, equal to value dimention D,, to explore the overall dimension configuration. In Table 9 the memory computational proportion 19.5% has D, D. H/8.

---

Published as a conference paper at ICLR 2026
Based on this configuration, we also tried D₂ De H/4, while ensuring the same number of parameters and the same computational load.
Training Loss Difference Curve (Smoothed)
Training Loss Difference Curve (Smoothed)
8) Du
0.025
Rp 1:1
Rg-
1110.-13 By D.
Loss Didference
0.015
Loss DiTerrnce
0.020
0.015
0.010
0.010
0.005
0.000
-0.005
200 300
600
20
70"
150 200
360 d01500
Consumed
Consumed tokens
Figure 3: Left: ablation on the sum of Du and Dp. Right: ablation on the proportion of D, and Dp-
Figure Ileft) compares the loss performance. The results indicate that the smaller the sum of De and 7- the more refined the division of memory. To maintain the same sparsity, larger key number N and TopM are required, which increases the number of combinations in the UltraMem part and thus leads to better performance. However, larger N and TopM will slow down the training and inference process, therefore, we did not further reduce the D, and Dp-
Next, we explored how to set D, and D, respectively. Based on the configuration with MCP- 17.0% in Table 9 we kept the sum of D. and Dp unchanged. and further increased D, and re- duced Dp. We tried the following two configurations in Table g Figure 3 illustrates the loss differences. Based on the loss result, we set D, D, = :
4.2.3 ULTRAMEMV2 LAYER NUMBER
Under normal circumstances, MoE is present in each layer. This will increase the participation of sparse parameters in the model pipeline. UlrraMem2 is distributed in the model at fixed-layer intervals. We conduct experiment in which the number of UltraMemV2 layers Lm. is gradually increased. Meanwhile, we make sure that the TopM x Lm and the total computation is fixed. The UltraMem V2 in this section is based on Seed-MoE and replaces MoE with SwiGLU FFN, inserting UltraMem V2 layers at regular intervals (fixed intervals of when each layer is available).
Figure shows the validation loss and open benchmark changes for UltraMem V2-430M/5B when inserting layers 2, 5, 10, and 20. We observe that validation loss quickly become indistinguishable when increasing the number of insertions of the UltraMem V2 layer, but sustained gains are observed on the Open Benchmark (obtain +2.3, +0.9, 1). Memory+ experiment, but with different conclusions. We speculate that the reason is that keep FFN in each layer, while they directly replace FFN with the Memory+ layer.
et al., 2024) also did a similar
Validttonn Loss Curve
Open- Henchmark
20 layers lapers
2.5
layees layers layers
0.26
topers
Accursey
2.3
2.2
0.20
6.18
15
Consumed tokens (8)
Consumed tokens
Figure 4: Effect of UltraMem layer number on training dynamics and performance. Validation loss (left) and open-benchmark accuracy (right) for UltraMem V-4330//5B with varying numbers of UI- traMem V2 layers (2, 5, 10, 20) under fixed computational budget. Although there is no obvious gain in the validation set loss after adding a certain number of layers, the model with more UltraMemV2 layers performs better in downstream tasks.
8

---

Published as a conference paper at ICLR 2026
We then perform continued training (CT) on UltraMem and MoE, but observed that in the Ul- V2.MemV2, the CT gains are smaller. Specifically, we pretrain the model with 1.6T tokens under constant learning rate, then continue training with 250B tokens under cosine decay learning rate. As shown in Table 4 the average score of UtraMem-430M/5B with UltraMem layers improves by 5.7 points after CT, but MoE improves by 7.8 points. In general, the larger the model, the less points it improves after CT. so UlraMem-440MM/BB should have at least gain of more than 7.8 points to be normal, which affects the effect of UltraMem post training. We find that the key to solving this problem lies in the number of UltraMem layers. When each transformer block has UltraMem layer, the gains of UltraMem V2 and Seed-Mor MoE become consistent on 1.25B/12.5B and 2.5B/25B models.
Table 4: Performance improvements after CT. "*" represents model structure based older version of Seed-MoE. Evaludation sets not included in the category reasoning, mathematics, code and knowledge are listed separately, including DROP et al. 2019), AGlEval(Zhong et al. 2023
Model
Reasoning
Math
Code
Knowledge
DROP
AGIEval
Average
Seed-MoE -2.5B/25B
+4.6
+10.2
+11.5
+5.2
+7.6
+7.8
UltraMe UlraMemm2214LL433MMBB -LL-4300/BB
+3.9
+13.1
+4.1
+5.1
+2.6
+5.7
Sead-MoE-1125BB122..1
+4.9
+11.3
+11.5
+6.7
+8.4
+7.5
UltaMMemV2-1.55B/12.555
+4.8
+15.6
+10.2
+5.7
+9.9
+8.3
Sced-Mol-2.5BB/55B
+6.1
+8.6
+9.0
+5.8
+8.2
+7.1
UtraMMmmV222.5B/25B
+5.7
+12.3
+11.2
+5.7
+8.5
+8.0
4.2.4
OTHER ABLATIONS
Appendix E shows more ablations, which reveal several key insights for UltraMem V2: Architectural simplifications (single value projector, separate queries per Tucker rank, single head) consistently improve performance. An optimal 17% memory computational proportion identified. Auxiliary losses, including Tucker core penalty and balance loss, are found to be ineffective. Layer-wise mem- ory sharing (S9-Ring, 96-Block) significantly boosts performance, offering scalable options. Lastly, while decaying learning rates for value parameters provide early gains, constant rate achieves better final performance over extended training.
4.3 TRAINING AND INFERENCE TIME
We trained UltraMem V2-2. 22-2.5///0B--t and SeedMoF-2.5B/600 on 1008 NVIDIA Hopper GPUs with batch size of 16M tokens per step. Both models achieve comparable training through- put: UltraMemV2 reaches 265B tokens/day while Seed-MoE reaches 262B tokens/day. However, the two architectures exhibit different resource constraints. UltraMem V2 is memory-bound con- suming more bandwidth due to atomic memory operations, while MoE compute-bout bound. Despite these different bottlenecks, they achieve similar overall training efficiency on current GPUs. This suggests that hardware with improved memory efficiency, particularly optimizations for atomic op- erations, would disproportionately benefit UltraMem enabling further training speedup.
We evaluated the inference time and memory access of MoE-I0B/200B across different batch sizes Figure The results show that UltraMem consis- tently achieves significantly lower memory access than MoE at commonly used batch sizes, while delivering up to 2x faster inference speed compared to Seed-MoE. We observe that smaller batch sizes yield even greater speedups; however, for practical considerations, we limit our evaluation to batch sizes of 32 and above. These results demonstrate the inference advantages enabled by Ultra- MemV2 extreme sparsity.
and Seed-
4.4 DISCUSSION
Integrating UltraMem V2 During Mid-Training While our experiments focus on training Ultra- Me from scratch, the architecture is amenable to integration at later training stages. Memory layers can be straightforwardly added during mid- training or post-training phases. A practical ap- proach is insert UltraMem V2 layers into existing model and initialization for either

---

Published as a conference paper at ICLR 2026
Memory access
ferennce latency (ms)
400
UltaMemV2
UltraMemV2
350
60
MoE
300
Memory (CB)
Latency (ms))
250
200
150
30
100
20
0
10
10
107
10
128 Batch size
256
512
Batch size
Figure 5: Changes in memory access and inference latency as the batch size varies. UltraMemVa and MoE both has 10B activated parameters and 200B total paremeters. Inference latency is evalu- ated on NVIDIA Hopper GPUs.
the value projection gating scalar. This ensures smooth integration without disrupting the pre- trained weights, allowing the memory layers to be gradually incorporated as training progresses. Re- cent work has demonstrated promising applications of this approach. For instance, [in et al. (2025) shows that integrating memory layers during fine-tuning can mitigate catastrophic forgetting. We believe that incorporating UltraMem V2 during continued training (CT) or fine-tuning phases repre- sents valuable research direction, particularly for scenarios requiring efficient model adaptation or knowledge retention.
Combining UltraMem V2 with MoE Architectures UltraMemV2 and MoE are not mutually ex- clusive, they can potentially be combined within single architecture. One straightforward approach is to use UltraMem V2 layers alongside traditional MoE layers within the same model, leveraging the complementary strengths of both architectures. However, converting a fully pretrained MoE model to UltraMem is more challenging due to the fundamental differences in their architec- tural designs. The most practical path forward is to design hybrid architectures from the outset or introduce UltraMem components during mid-training, discussed above. We leave detailed exploration of such hybrid architectures as promising future work.
5 CONCLUSION
In this paper, we developed UltraMem V2, new type of architecture that, for the first time, performs as well the top-tier 8-expert MoE models. By placing our new memory layers in every part of the model, we show that this approach is solid new option for building large and efficient Al.
Our work has few key takeaways. First, UltraMem matches the performance powerful MoE models on standard tests. Second, it's particularly good at tasks that require a great memory, like long-document understanding and multi-turn conversations, where significantly outperforms MoE. Third, we found that we could simplify the training process, removing the need for extra complex settings. Finally, we learn an important design lesson: better to activate more values than to simply increase the number of sparse parameters.
However, there are some limitations to be aware of. UltraMem gets off to a slow start; it doesn't perform as well as MoE models early in its training and needs lot of high-quality data to catch up. Its best performance also depends on putting memory layer in every single block of the model.
In summary, V1traMemV2 validates the potential of memory layer architectures for building effi- cient and powerful large-scale models. Future work should focus on improving its carly-stage train- ing dynamics and further exploring architectural trade-offss for diverse downstream applications.
THE USE OF LARGE LANGUAGE MODEL
The Large Language Model is only used to polish the content of this paper, correct grammatical errors, and make the paper more coherent and easy to understand.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, et al. Program synthesis with large language models. arXiv preprint arXiv:2108.0772 077332 2021.
Vincent-Pierre Berges, Barlas Oguz, Daniel Haziza, Wen-tau Yih, Luke Zettlemoyer, and Gargi Ghosh. Memory layers scale. arXiv preprint arXiv:2412.09764 2024.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language. Thirty-Fourth AAAI Conference on Artificial Intelligence, 2020.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger. Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fo- tios Chantzis, Elizabeth Barnes, Ariel Herbert-Voss William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin. Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob Mc- Grew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojciech Zaremba. Evaluating large language models trained code. 2021.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv:/803.054555 2018.
Karl Cobbe, Vincet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint 10. 2021.
Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber. Approximating two-layer feedforward net- works for efficient transformers. arXiv preprini 230.23300.0837. 2023.
Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Yu Wu, et al. Deepseekmoe: Towards ultimate expert specialization in mixture- of-experts language models. arXiv preprint S0Kv::401106066 2024.
Dheeru Dua, Yizhong Wang, Pradeep Dasigi, Gabriel Stanovsky, Sameer Singh, and Matt Gardner. DROP: reading comprehension benchmark requiring discrete reasoning over paragraphs. In Proc. of NAACI, 2019.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-39, 2022.
Jiaao He, Jiezhong Qiu, Aohan Zeng, Zhilin Yang, Jidong Zhai and Jie Tang. Fastmoe: fast mixture-of-expert training system. arXiv preprint 1rXiv:2103.13262, 2021.
Xu Owen He. Mixture of a million experts. arXiv preprint arXiv:2407.0413. 2024.
Dan Hendrycks, Collin Burns, Steven Basart, Andrew Critch, Jerry Li, Dawn Song, and Jacob Steinhardt. Aligning ai with shared human values. Proceedings of the International Conference on Learning Representations (ICLR), 2021a.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. Proceedings of the Interna- fional Conference on Learning Representations (ICLR), b021b.
11

---

Published as a conference paper at ICLR 2026
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt, Measuring mathematical problem solving with the math dataset. NeurIPS, 2021c.
Hongzhi Huang, Defa Zhu, Banggu Wu, Yutao Zeng, Ya Wang. Qiyang Min, and Xun Zhou. Over- tokenized transformer: Vocabulary generally worth scaling, arXiv preprint arXiv:2501.16975, 2025.
Yuzhen Huang, Yuzhuo Bai, Zhihao Zhu, Junlei Zhang, Jinghan Zhang, Tangjun Su, Junteng Liu, Chuancheng Lv, Yikai Zhang, Jiayi Lei, Yao Fu, Maosong Sun, and Junxian He. C-eval: A multi- level multi-discipline chinese evaluation suite for foundation models. In Advances in Neural Information Processing Systems, 2023.
Zihao Huang, Qiyang Min, Hongzhi Huang, Defa Zhu, Yutao Zeng, Ran Guo, and Xun Zhou. Ultra- sparse memory network. arXiv preprint arXiv:2411. 123664, 2024.
Bi Huo, Bin Tu, Cheng Qin, Da Zheng, Debing Zhang, Dongjie Zhang, En Li, Fu Guo, Jian Yao, Jie Lou, et al. dots. Ilml technical report. arXiv preprint 27X6v:2506.05767, 2025.
Herve Jegou. Matthijs Douze, and Cordelia Schmid. Product quantization for nearest neighbor search. IEEE transactions on pattern analysis and machine intelligence, 33(1):117-1288 2010.
Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bam- ford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral experts. arXiv preprint arXiv:240l.040880 2024.
Matt Gardner Johannes Welbl, Nelson F. Liu. Crowdsourcing multiple choice science questions. 2017.
Mandar Joshi, Eunsol Choi, Daniel Weld, and Luke Zettlemoyer. triviaqa: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension. arXiv e-prints, art. 1rXiv:1705.03553 2017.
Gyuwan Kim and Twe-Hwan Jung. Large product key memory for pretrained language models. arXiv preprint arXi::2010.0388/ 2020.
Jakub Krajewski, Jan Ludziejewski, Kamil Adamczewski, Maciej Pióro, Michat Krutul, Szymon Antoniak, Kamil Ciebiera, Krystian Król, Tomasz Odrzygóžżž, Piotr Sankowski, et al. Scaling laws for fine-gratt ome-grained mixture experts. arXiv preprint 2rXiv:2402.07871, 2024.
Guillaume Lample, Alexandre Sablayrolles, Mare' Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. Large memory layers with product keys. Advances in Neural Information Processing Systems, 32, 2019.
Jessy Lin, Luke Zettlemoyer, Gargi Ghosh, Wen-Tau Yih, Aram Markosyan, Vincent-Pierre Berges, and Barlas Oguz. Continual learning via sparse memory finetuning. arXiv preprint arXiv:2510.15500 151003 2025.
Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437, 2024.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? new dataset for open book question answering. In EMNLP, 2018.
Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Jacob Morrison, Sewon Min, Wei- jia Shi, Pete Walsh, Oyvind Tafjord, Nathan Lambert, et al. Olmoe: Open mixture-of-exppets language models. arXiv preprint arXiv:2409.02060, 2024
Toan Q Nguyen and Julian Salazar. Transformers without tears: Improving the normalization of self-attention. arXiv preprint arXiv: 1rXiv:19/00058895 2019.
12

---

Published as a conference paper at ICLR 2026
Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi. Am- mar Ahmad Awan, Jeff Rasley, and Yuxiong He. Deepspeed-me moe: Advancing inference and training to power next-generation ai scale. In International conference on machine learning, pp. 18332-18346 PMLR, 2022.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adver- sarial winograd schema challenge at scale. arXiv preprint arXiv:1907.10411 2019.
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202, 2020.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gate mixture-of-experts layer. arXiv preprint arXiv: / 011//7011..006 006538, 2017.
Mirac Suzgun, Nathan Scales, Nathanacl Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoc Le, Ed H Chi, Denny Zhou, and Jason Wei. Challenging big- bench tasks and whether chain-of-thouges can solve them. arXiv preprint arXiv:2210.0921 2022.
Alon Talmor, Jonathan Herzig, Nicholas Lourie, and Jonathan Berant. CommonsenseQQA: ques- tion answering challenge targeting commonsense knowledge. In Proceedings of the 2019 Con- ference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pp. nesota, Jume 2019. Association for Computational Linguistics. doi: 10.18653/11/N19- 1421. URL
Minneapolis, Min-
https::/acllattholll org/N19 1421
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Yuan Xie, Shaohan Huang, Tianyu Chen, and Furu Wei. Moec: Mixture of expert clusters. In Proceedings of the AAAI Conference Artificial Intelligence, volume 37, pp. 13807-13815, 2023.
An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao. Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint arXiv:2505.09388, 2025.
Yike Yuan, Ziyu Wang, Zihao Huang, Defa Zhu, Xun Zhou, Jingyi Yu, and Qiyang Min. Expert race: A flexible routing strategy for scaling diffusion transformer with mixture of experts. arXiv preprint arXiv:2503.16057 2025.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma- chine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
Wei Zhao, Mingyue Shang, Yang Liu, Liang Wang, and Jingming Liu. Ape210k: A large-scale and template-rich dataset of math word problems. arXiv preprint 20X9v:2009.118 11506, 2020.
Wanjun Zhong, Ruixiang Cui, Yiduo Guo, Yaobo Liang, Shuai Lu, Yanlin Wang, Amin Saied, Weizhu Chen, and Nan Duan. Agieval: human-centric benchmark for evaluating foundation models, 2023.
13

---

Published as a conference paper at ICLR 2026
A PRELIMINARY
Output
Actipted Vol
Fressalues
SCOTES
Act sated
Inflox
Fre-cllues
indes
scorfis
PO
10
Topm
Mest
Ricet
Pool
DAAIN
Hldden state
(a) MoE layer
(b) PKM layer
Outbut
Kiet
Acthvated
Arthored
Index
meetory
Artlated Walues
Weighted
Weighted
Inear
Weighted
TOQKER
Articated
Weighted
X
s0ores
Actistted Activsted scorss
Hidden
(c) UltraMem Inyer
Figure 6: Overall structure of (a) MoE layer, (b) PKM layer, (c) UltraMem layer.
B OPTIMIZED INITIALIZATION
First, we list the important theorems required for the derivation, starting with the Central Limit ThLorem(CLT))
If 21, 22, Tay... are random independent samples of of a random variable X, then:
Ex)) - ~(0
-
(16)
where oXX) is the standard deviation of the random variable X.
Notations:
σv is the standard deviation of the value in the memory layer.
à is the hidden size of the network.
dprew is the dimension of the "Pre-values".
d is the dimension of the "Values"
k is the top-k parameter.
nhead the head number.
L is the number of layers.
kinner is the multiple obtained by dividing the FFN inner dimension by the hidden size.
14

---

Published as a conference paper at ICLR 2026
To control the final output magnitude, adjust the initialization of the query and key norm such that the mean of the score obtained from "Top-k" is around 1, og is its standard deviation.
The initialization standard deviation of the linears attention and FFN Salazar 2019). We also set the same standard deviation for the linear layers in UlrraMem for convenience.
Nguyen & 5h
Now, we can derive the output variance according to Figure |
B. THE VARIANCE OF ULTRMEM- V2 LAYER OUTPUT
The imput of the memory layer is the output of layer normalization before FFN. So the input to memory has a mean of 0 and a variance I. According to CLT, the output of "pre value proj" has of 0 and variance of 0.4. "Activated Pre values" have mean of 0 and variance of oy. Thus, the matrix multiplication result of these two tensors has a mean of 0 and a variance of 0.4 of dprev. This matrix multiplication result is then multiplied by the top-k score, this final score has mean of 0 and variance of (0.4 + 0.4 o2) oy dprev-
Thus, the input of the "value proj" layer has mean of O, and a variance of (0.4 + 0.4 * of) of dprea # Thead- Finally, we can get the output variance using CLT:
T mem = (0.16 + 0.16 o².) * of * dgren * & Whead dzp/
(17)
B.2 THE VARIANCE OF TOP-K SCORE
There are many complicated operations in the "TDQKR", such tucker and SVD, which makes it difficult for us to estimate the variance accurately. Therefore, we generate a large amount of random data whose distribution aligns with the output distributions of the query and key normalization pro- cesses. These data are then fed into the TDQKR module. Through statistical methods and tuning the initialization of the query and key norm, we ensure that the average score the top-k results approximately 1. Subsequently, we compute the variance of these scores under this condition and incorporate it into formula 17
B.3 CALCULATE THE STANDARD DEVIATION FOR INITIALIZATION
To control the variance of the final output, we take the initialized output variance of the FFN reference and directly set the output variance of each UltraMem layer Tnem equal to that of the FFN ffn thereby deriving the initialized variance of the "Values" and "Pre-values" σv Let's start by calculating the output variance of the FFN.
The input to the FFN is the output of layernorm RMSNorm, and the moment of initialization, its mean is 0 and variance 1. Therefore, according to the Central Limit Theorem, the input to the Swish activation function has mean of 0 and variance of 0.4. Since the curve of the Swish activation function is quite similar to that of ReLU, truncated normal distribution is used subsequently estimate the distribution after Swish activation.
We use Mawi and Jswi to denote the mean and std of the input to the Swish activation func- tion; a and here. & n-Maaui convert the original normal distribution to the standard normal distribution. dawi
the boundary points of the truncate range, which are 0 and too respectively
(p(E) = 1 exp 15²): a a- faws b-pawi, ((x) ]1 (2f/2/2(22 the cumulative V2T
distribution function of the standard normal distribution; Let Z { {(a); Then, the mean of the activated gate is:
Pgate Mawii + (0(a) ((3)
7 Jani
(18)
The variance of the activated gate is:
BP(() - asp(a)
(p((a) v(3))
d gate = 2
(19)
Z
15

---

Published as a conference paper at ICLR 2026
Then the activated gate is multiplied by another linear output whose mean is 0 and variance is 0.4 by using CLT. We substitute the specific values and obtain that the inner activation of FFN has a variance of 0.16, and a mean of O.
To prevent the output variance of the final network from diverging as the number of network layers increases, the initialization standard deviation of the last linear layer in the FFN is usually further multiplied by factor of / 21. 2T. Thus, by using CLT, the variance of the FFN output is:
0.064 kinner
jfn
(20)
2L
We let mem offn to get the initialzation variance:
0.2 kinner # ov k Theaad (1 + o?) dprev de L
(21)
C EVALUATION BENCHMARK
For the proprictary model comparison, the evaluation benchmark including Open Benchmark, Hard Benchmark and Long-context Benchmark. Table [5] shows the components in Open Benchmark. Hard Benchmark contains more difficult tasks. Details of Long-context Benchmark is shown in Table For the menchmark in spen-source experiments, we evaluate models on Arc-C(Clark et al.,
2018
al. et al. 2018). PIQA(Bisk 2020), SCIQ(Johnnnes 2017), and Winogrande et al. 2019).
Are-B(Clark et al. 2018), CommonSenseeAA (Talmor al. 2019), Hellaswag
2019
MMLU-var Muennighohf et al., 2024), OpenbookQA
ct al.
Table 5: Open benchmarks across different domains
Code
Math
Knowledge
Reasoning
MBPP Austin al. 2021
Ape2 7han Zhao GSMSK(Cabne el a 201
2020
MMI Hendryycks 2021h C-Eaal Tipang el a. 2023
Arc.C Clark al.
2018
HumanEyal Chen at. 2021
BDIiSazzum et 2022
MATH Headrycks at 2021e
Trvva
DROP Dua DROP
WinoGrinde
Hellaswag Zellers 2019
Table Long-context evaluation tasks and their descriptions
Task
Description
Long-context memo-
Evaluate the model's ability to understand and recall information when there is long context
rizing
Multi-round
Evaluate the model's ability understand and recall information in the presence of multiple rounds of dialogue
rizing
In-context learning
Evaluate the model's capabilities under the given longer task demon- strations
Reasoning
Long context reasoning ability
Find needle
Evaluate the ability to quickly locate specific information in large amount of information
Key-val retrieval
Given large number of key-value pairs, evaluate the retrieval ability of the model when given keys
Multi-hop reasoning
Evaluate the model's ability to establish logical connections among dif- ferent context segments, thereby arriving at the answers to questions or making decisions
16

---

Published as a conference paper at ICLR 2026
D OPEN-SOURCE MODEL HYPERPARAMETERS
Table details the configurations of our open-source models. We denote the number and dimension of keys as knum and kdim, respectively, while vdim and pre-vdim represent the dimensions of the value and pre-value. Notably, the Memory+-227M/I ..B share values across memory layers, result- ing in larger knum compared to (22MMe although the number of values remains the same.
Table Model Configurations
Hiddeen size
Atin
Mem layer
Key number
activated param(M)
total param
Model
Layer
head
TopM
D. D>
OLMoE: OLMME-227MM11.
768
12
227
1.18
Memory-- 22-22277//222
Heemory-227M/L...2
768
12
1138
384
2x80
384
228
1.19
UhraMem-227M//1n
20
768
12
808
192
2x80
384
225
1.18
UhraMem
768
OLMoE-IB//B
16
2048
16
/
/
1070
6.71
MlnaMemM2-1MM77
2048
16
528
512
128
768
384
1079
Table 8: Value dimension and pre-value dimension Ablation Configurations
Hidden size
Mem layer
Key number
Model
D2
head TopM
D, Dp
Dp : De = 1:3 1152
24
964
524
94
216
72
D₃ De 1:8 1152
24
964
524
94
256
32
E OTHER ABLATIONS
E. SINGLE PROJECTOR AND MULTI QUERIES
Training Loss Difference Curve (Smoothed)
Open Benchmark
0.002
Baselite
0.52
Baseline
Single projector quectes
Single projector quectee
Loss Difiereac
0.48
-0.00 0.002
Artaracy
0.46
0.000
0.08
0.42
-0 -0.008
0.38
300
Consumed
(B)
Figure 7: Single projector and multi-query modifications on UltraMem M2-MMemM2-43550 V2-430M/5B. Training loss difference (left) and Open-Benchman accuracy (right) comparing baseline with proposed ap- proach. The modifications achieve 0.0026 loss reduction and 0.7-point accuracy improvement after 1.5T tokens.
Under the condition of maintaining the same parameter count and computational load, allocating more activations to the final value projector requires reducing activations in the FFN. Our research reveals that value projectors for the final derived value. Additionally, we employ separate queries for each tucker rank, which en- hances the accuracy of query-key operation results. After training 1.5T tokens, these two modifica- tions yield a loss reduction 0.0026 and a 0.7-point improvement in Open- ppen-Benchmark performance shown in Figure
not highly activation-efficient: thus, we use only single projector
E.2 HEAD
We also performed an ablation study on the number of heads in UltraMem V2-430M/SB. For head configuration, we double the Dy and TopM maintain the same computation and value retrieval. As
17

---

Published as a conference paper at ICLR 2026
Trsining (Smoothed)
hearhaark
heedw2 bead=1
heede2
hand-1
0.002
Loss Ditiereace
0.000
Aouurey
0.002
0.20
0.26
-0.005
0.24
-0.010
15 20
50
70
100
150
200
300
500
130
200
300
500
700
1000
Consumed takens (D)
Consumed tokens (0)
Figure 8: Ablation study on the number of heads in UltraMem22-430000 (left) and Open- Benchmark accuracy (right) comparing single head vs. two heads. Using a single head achieves 8e-4 loss reduction and 0.2-point accuracy improvement after IT tokens.
Training loss difference
illustrated in Figure after training with IT tokens, using a single head yields a marginal improve- ment over using two heads: the loss decreases by 8e-4 and the Open Benchmark score increases by 0.2 points. This phenomenon can be attributed the fact that while employing a single head in- creases the number of values to be retrieved, the dimensionality of each query and key also increases proportionally. Ultimately, the gain in retrieval accuracy from the increased dimensionality of the retrieval vector marginally outweighs the negative impact of the larger candidate pool for retrieval.
E.3 THE COMPUTATIONAL PROPORTION OF ULTRAMEMV2
The larger the keys dimension Dx. the greater the computational proportion of memory. To keep the total amount of computation unchanged, the computation of FFN will decrease. Therefore, D2 is too large, the computing power allocated to FFN will be too small, which will lead to poor performance; D& is too small, the query result in the memory layer will be very inaccurate, which will also result in poor performance.
To determine how configure DR. we conduct detailed ablation studies on UltraMem V2-500M/6B. We only adjusted Da and correspondingly tweaked the inner dimension of the FFN to ensure that the total computational cost and the total number of parameters remained consistent. The UltraMemV V2 related configurations are as follows: (MCP is short for Memory Computational Proportion in Table
Table 9: Key dimension Ablation Configurations
Hidden
Mem layer
Key number
Model
De
head TopM
De
Dp
size
MCP=12.0%
1152
964
144
MCP=14.5%
1152
24
964
432
94
144
144
MCP=17.0%
1152
24
964
524
94
144
144
MCP=19.5%
1152
24
964
610
94
144
144
MCP=23.0%
1152
24
964
730
144
144
MCP=25.09%
1152
24
964
796
94
144
144
MCP=27.5%
1152
24
964
880
94
144
144
As shown in Figure 9 the loss performance shows that 17% Memory Computational Proportion is the best configuration. When scaling up the model, if the hidden size increases by factor of a, both vdim and pre-vdim will also increase by a factor of a. This causes knum to increase only by factor of Va, and the proportion of computational load for memory will decrease as the model grows larger. Therefore, it is not suitable as a basis for scaling up the model. Finally, when we scale up the model, adopt Dr h/2 as reasonable configuration.
18

---

Published as a conference paper at ICLR 2026
Training (Smootbed)
MCP-17..0 17.0%
MC2-12.0%
0.000
KC.-1455%
Loss Ditierea.c
MCP-19.9%
MCP 23.0%
0.00
MCP-25 25.0%
XCP-27.55
0.002
20
30
50
70
100
150
200
300
5.00
700
Consumed tokens (3)
Figure 9: Training loss difference under different memory computational proportion.
E.4 AUXILARY LOSSES
In this section, we introduce two auxiliary losses which are NOT used in UltraMem Our expere- ments show no improvement under these auxiliary losses.
Tucker core penalty loss In TDQKR, when doing the first two TopM, Huang al. (Huang al. 2024) first do the Singular Value Decomposition of the tucker core and aggregate the row and column keys with the eigenvectors of the largest eigenvalues. To constrain this approximation error, they place constraints non-max ma-mmmxmmmm singular values
C = UAT (by SVD)
(22)
Laux
(max (0, A - -))
(23)
where, A denotes the singular values for C in descending order, with 7 serving as margin to prevent C' from degenerating into a rank-1 1 matrix, and a is the coefficient for the loss.
Balance loss in MoE solve the problem of dead experts and alleviate unbalanced computation of Expert Pa et al.) 2021). Due to the fact that UltraMem's parallelism segmented in the embedding dimension, there is problem of computational imbalance. We curious whether more balanced activation embedding will also improve the performance. Recall that the balance loss of MoB et al., 2023) constraint the result of the router, while UltraMem of as a MoE with row/column routers, so we propose the balance loss of UltraMem following MoE. Let u, t 1grx1 be the eigenvectors corresponding the largest eigenvalues, the probability to activate each row/column key is
be thought
P Prow = Softmaz(u Srow), Pcol = Softmaz(v Seat);
(24)
here omit the head index i for brevity. Ones get probabilities, can calculate the balance loss
N Lbalance = BN Σ fn Pus
(25)
7=1
where is the coefficient for the loss, N is the number of row/column keys, fu represents the frequency at which row/column keys is activated. Given batch with T tokens, Count,, is the number of times the n-th row/column key is activated, then
fn Countn/T. TT, Pr.
1 p T
Σ
(26)
row/col:
where Pn represents the average probability of the nith row/column key being selected.
E.4.1 EXPERIMENTS
In this subsection, conduct experiments UltraMemV2-43000 training 800B tokens.
19

---

Published as a conference paper at ICLR 2026
Tucker core penalty loss In contrast the approach in UltraMem (Huang al. 2024), which emphasizes the necessity of constraining non-man maximum eigenvalues of the Tucker core to mitigate approximation eπrоrs,
empirical findings suggest that such a constraint may be superfluous dur- ing training. We observe that the eigenspectrum of the Tucker core exhibits a sharp decay, where the principal eigenvalue, A1 , is naturally an order of magnitude larger than the subsequent eigenvalues. As illustrated in Figure [o]b, A is consistently more than four times larger than the second-largen eigenvalue, Ag. Given this substantial gap, the resulting approximation error from the non-maximum eigenvalues is negligible.
To further validate this, we conducte an ablation study by removing the Tucker core penalty loss proposed in UltraMem. Figure I shows the downstream task performance throughout training. The model trained without the penalty loss maintains comparable accuracy to the baseline in the early training stages. Notably, after training on 800B tokens, our model without the penalty exhibits a 0.3 point improvement in accuracy, indicating that forgoing the explicit eigenvalue constraint does not compromise, and may even slightly benefit, model performance.
Open-Benchmark
Tucker Core Eigenvalue dynamic
0.40
Baseline
2.5
Penalty runkO
No Penalty
Bassline racko
Penalty rank]
2.0
Buseline rank!
Aoarray
0.38
0.37
1.0
0.5
0.36
3x
4x100'
500
6x 10
700
0.0
15
20
30
50
70
100
150.20
300
500
700
Consumed lokeess (B)
Consumed toloons (B)
Figure 10: Tucker core penalty loss ablation study. Left: Open-Benchmark accuracy during training with and without Tucker core penalty loss. Right: Evolution of Tucker core eigenvalue magnitudes, showing the dominant eigenvalue λ (rank0) and second- second-largest eigenvalue A2 (rankl). The penalty loss maintains eigenvalue scparation without compromising downstream performance, validating that explicit eigenvalue constraints are unnecessary when the natural eigenvalue gap is sufficient.
Balance loss is a common technique in MoE to promote load balancing across experts. As demon- strated in OLMoE (Muennighoff al. 2024). it can lead improved training loss and downstream task accuracy. Typically, balanced expert utilization is also critical for training efficiency. However, in the UltraMem parallelism is implemented along the value dimension, which renders training speed independent of the value selection balance. We are nonetheless interested in the impact of the balance loss on model performance. To this end, we conduct experiments on UltraMem V2- 430M//B, incorporating the balance loss from Equation 25 with coefficient of 3 = 0.001.
As illustrated Figure I our findings reveal a dependency on the number of activated values (TopM). When smaller number of values activated (TopM 47 out of value number N 465124), the inclusion of the balance loss correlated with a clear improvement in both training loss and downstream accuracy. accuracy.Converse Conversely, when larger number of values activated (TopM 94), the application of the balance loss has a detrimental effect on performance, with both training loss and downstream accuracy degrading. These conflicting results indicate that the effectiveness of the balance loss is linked to the number of activated values. The loss provides a regularizing benefit when this number small, but this benefit is lost when the number of activated values becomes too large.
E.5 SHARED MEMORY
Inspired by shared memory models like PKM (Lample et al., 2019) and Memory+ (Berges et al., 2024), which utilize shared value table across multiple layers, we explore similar sharing paradigms. Unlike MoE where inter-layer expert sharing significantly inflates inference memory access for small batches, memory access in memory layer remains constant. However, TDQKR inherently increases index computation: it doubles with rank r = 2, and again doubles for each layer if, for instance, one value table is shared across four layers. To maintain computational parity,
20

---

Published as a conference paper at ICLR 2026
Training (Smoothes))
Open Benchmark
2.14
Balanta, togh-544
Anlence, tagh-94
Jalanc, to
Dalance, to
topM-47
0.30
togM-47
2.12
No Bolance, togh-47
No Bulance, toph=4?
Accuracy
0.36
1.08
2.10
0.34
0.32
0.30
2.06
6000
700
2010
000
Consumed tolonns (B)
Consumod tolcens (B)
Figure 11: Effect of balance loss on training dynamics with different numbers of activated val- ues. Training loss (left) and downstream accuracy on Open-Benchman (right) for V2 430M/5B with and without balance loss (3 0.001). Balance loss improves performance when fewer values are activated (TopM = 47) but degrades performance with more activated values (TopM = 94), indicating that the regularization benefit depends on the sparsity level.
this necessitates shrinking the FFN inner dimension, potentially degrading model performance due to reduced capacity. To address this, our experiments evaluate the efficacy of partially sharing the value table.
Experimental Setup Our experiments utilize the UltraMemV2-5000M configured with L 24 transformer layers, TopM 94 activated values, row/column key number n 964, and N = 9292296 values per table. We define three primary sharing strategies:
model as a baseline,
1. No-Nonn Each memory layer accesses its value table plus those of its g - nearest neighboring layers (e.g., (g- 1)/2 preceding and (g- 1)/2 succeeding, adjusted for bound- aries where g is odd; g even, it might be g/2 preceding and (g/2) succeeding, or similar asymmetric but consistent distribution).
2. Sg-Ring: Similar to Sg-NoRing, but with wrap-around, allowing layers near the begin- ning/end of the network access tables from the end/beginning. respectively, ensuring each layer can access g distinct tables.
3. Sg-Ble Sg-Block: Layers are grouped into contiguous blocks of gi layers within a block share all g value tables belonging that block, but do not access tables outside their block.
Table 10 provides an illustrative example of these schemes. As the number of shared tables increases, the effective key dimension for lookup also increases (e.g., for sharing 4 value tables, the composite key number n becomes 1928). We correspondingly reduce the FFN inner dimension ensure iso-computation across all configurations.
Table 10: Example of different shared value partern
Layer number
L=1
L = L=3
L
L=5
L - 24
S4-NoRing
1,2.3,4
1,2,3,4
2,3,4,5
3,4.5,6
4,5,6,7
21,22,23,24
S4-Ring
24,1,2,3
1,2,3,4
2,3,4,5
3,4.5,6
4,5,6,7
23,24.1,2
S4-Block
1,2,3,4
1,2,3,4
1,2,3,4
1,2,3,4
5,6,7,8
21,22,23,24
Ring vs. NoRing Topology We first ablate the S9-NoRing and S9-Ring configurations. As de- picted Figure [2] 59-NoRing achieves a lower training loss but exhibits a 1-point degradation in downstream accuracy compared to S9-Ring. This suggests S9-NoRing may be more susceptible to overfitting. We hypothesize this could be due to the values in tables at the network's extremities being more frequently shared and thus potentially over-specializing.
Optimal Number of Shared Layers We then investigate the impact of the number of shared lay- g within the Ring topology, specifically comparing S4-Ring, S9-Ring, and S16-Ring. These correspond to effective key number n 1928, 2896, and 3856, respectively. Figure [3 shows that increasing 9 provides continuous benefit training loss relative to the baseline. However, this
21

---

Published as a conference paper at ICLR 2026
Training Loss (Smoothed)
Open Benchmark
0.002
Zing NoRing
0.41
Rang
0.001
0.40
NoRing
0i0-0.01
0.000
MWw
1o05 Dif
Acouracy
0.38
0.000
0.37
-0.002
0.003
0.35
00.000
0.36
70 Consumed tolooms
6x100
Coasumed tokeas (B)
Figure 12: Comparison of Ring vs. NoRing topologies for shared memory configurations. Left: Training loss difference (smoothed) showing NoRing topology achieves lower training loss. Right: Open benchmark accuracy demonstrating the advantages of Ring, which achieving better final ac- curacy (Diff=0.014).
benefit appears to saturate at g 9, with S16-Ring showing negligible improvement over S9-Ring training loss. In terms of downstream accuracy, S9-Ring surpasses the baseline by point, while R16-Ring does not improve upon S9-Ring. This indicates trade-off: while access to larger pool of shared values can enhance representational capacity, the requisite FFN shrinkage to maintain computational parity eventually negates further gains.
Training Loss Curve Diff (Smoothed)
Open-Benchmark
0.010
0.008
Baseline
0.40
54-Ring
Baseline
S0 Ning
0.39
St-Ring
S16 Ring
SL6 Ring
1oos Dif
0.004
Accuracy
0.38
0.002
0.000
0.38
-0.000
0.35
0.004
-0 000
0.33
60
Consumed tokons
Consumed tokens (B)
Figure 13: Impact the number of shared layers in Ring topology contigurations. Left: Training loss difference (smoothed) relative to baseline showing S4-Ring, S9-Ring, and S16-Ring all achieve lower training loss, with benefits saturating around 9 shared layers. Right: Open benchmark accu- racy demonstrating S9-Ring achieves the best downstream performance with I 1-point improvement over baseline, while 516-Ring shows no improvement over S9-Ring.
Block-wise Sharing for Scalability Finally, considering the practicalities of large-scale model train- ing, particularly with pipeline parallelism, Ring-based sharing can introduce significant inter-stage communication overhead. Block-wise sharing (Sg-Block presents a more ongineering-friendy al- ternative. We evaluate S6-Block, with I, = 24 layers, this creates 4 blocks of 6 layers, offering a comparable degree of value table accessibility to S9-Ring (where each layer accesses 9 tables, but tables are re-used across layers). Figure [ comparess the baseline, S4-Ring, S9-Ring, and S6- Block. Results indicate that S6-Block, while marginally inferior to S9-Ring, still offers substantial improvements over the baseline. It is expected that within larger share range, the effect Sg- Block can be further enhanced. This validates Sy-Block promising strategy for large model training, balancing performance with practical implementation constraints.
E.6 VALUE LEARNING RATE SCHEDULE
The UltraMem Huang et al. 2024 demonstrates efficacy by employing a relatively high initial learning rate for its value parameters, which subsequently decayed. This strategy, while effective, introduces two additional hyperparameters: the initial learning rate multiplier and the decay dura- tion. To investigate the necessity of this decaying schedule and potentially simplify the training
22

---

Published as a conference paper at ICLR 2026
Training (Smonthed)
Open Benchmark
0.008
Baselin
Baseline
54
0.40
54-Ring
St-Jiing
0.39
S0-Ring
0.004
Sc Blocs
S6 Block
1,0s bir
Acouracy
0.002
0.37
0.000
0.36
0.35
-0.004
-0.006
0.33
60
Consumed tokons
Consumed takens (B)
Figure 14: Comparison of block-wise sharing topology (S6-Block) against ring-based topology. Left: Training loss difference (smoothed) showing S6-Block achieves comparable training loss im- provements to S9-Ring while being marginally inferior. Right: Open benchmark accuracy demon- strating S6-Block offers substantial improvements over baseline with performance Slightly worse than S9-Ring, validating block-wise sharing as a practical alternative for large-scale model training.
regimen, conduct ablation study M2traMemM2-S00MM We compare three settings for the pre-value and value learning rates, all trained for 1.4T tokens:
Baseline: An initial rate of 4x the main model's learning rate, linearly decaying to 1x by 350B tokens, mirroring the UltraMem approach.
2. Constant 1x: A constant rate of Ix the main model's learning rate.
3. Constant 1.5x: A constant rate of 1.5x the main model's learning rate.
Figure 15 depicts the training loss and downstream task accuracy for these configurations. The baseline initially achieves the lowest training loss, peaking in its advantage around 430B tokens. However, this gap progressively narrows, and by 1.4T tokens, the loss differences across all settings become negligible. A consonant trend is observed in downstream accuracy: the baseline shows its most significant lead around 400B tokens. Notably, upon completion of 1.4T tokens of training, the constant Ix learning rate setting surpasses the baseline by 0.4 points in downstream accuracy. The constant 1.5x setting also demonstrates superior performance to the Ix setting prior to 300B tokens, after which its relative performance deteriorates.
These findings suggest that while higher, decaying learning rate for the pre-value and value param- eters provides an early advantage training, particularly for shorter training budgets (fewer tokens), this benefit diminishes with extended training. In fact, for sufficiently long training horizons, main- taining a constant, moderate learning rate (e.g., Ix the main model's learning rate) can yield superior final performance, potentially obviating the need for a decaying schedule and its associated hyper- parameter tuning.
23

---

Published as a conference paper at ICLR 2026
Training Loss (Smoothed)
Open-Benchmar's Benchmark
Buseline
0.44
Baseline
0,020
Constact 1.5x
Constast 1.5x
Constant
0.42
Constant
1oos Ditt
0.015
0.40
0.010
0.38
0.000
2x 102
Consumed tokons (B)
Figure 15: Impact of value learning rate schedules on training dynamics and downstream perfor- mance. Left: Training loss difference (smoothed) showing the baseline (4x-1xx decay) initially achieves lowest loss but converges with constant rate settings by 1.4T tokens. Right: Open-domain benchmark accuracy demonstrating constant Ix learning rate achieves best final performance (0.4 point improvement), suggesting decaying schedules provide early training benefits that diminish over extended training horizons.
24


Reviews:
[{"review_id": "ACIjCcUSYJ", "reviewer": "Reviewer_3bZy", "summary": "This paper provides a memory‑layer architecture intended as an alternative to MoEs, which is an extension of UltraMem. The novelty with respect to the origianl UltraMem are, adding a memory layer to every Transformer block, simplifying the implicit value expansion (IVE) to a single linear projector, replacing value embeddings with an FFN with 1-dimensional inner layer, introducing a variance‑matching initialization, and rebalancing compute between memory and FFN. They claim parity with 8‑expert MoEs at similar active parameters/compute and advantages on long‑context tasks.", "questions": "How are KV‑cache footprint and router/TDQKR indexing costs accounted for in the iso‑compute comparisons in Tables 1–3?\n\nCould you report token/s and latency vs. batch size, plus HBM read/write estimates, for representative model sizes?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper clearly describes its position relative to MoE, PKM/UltraMem, and PEER.\n\nExperiments show that increasing the number of UltraMemV2 layers improves downstream accuracy even when validation loss plateaus.\n\nThe proprietary long‑context suite shows non‑trivial gains of 6.2 on multi‑round memorizing and 7.9 on in‑context learning.\n\nThe paper is explicit that UltraMemV2 underperforms early in training and benefits from continued training, and also notes dependence on per‑block placement.", "weaknesses": "The paper asserts matching compute and parameters, but does not report KV‑cache costs, routing FLOPs, or memory traffic for both MoE and UltraMemV2.\n\nThe claim that this work is the first memory layer to match 8‑expert MoE is not accurate in light of the Memory Layers at Scale paper [https://arxiv.org/abs/2412.09764].", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Nyv9gpaC42", "reviewer": "Reviewer_hfVp", "summary": "The paper presents UltraMemV2, a memory-layer architecture intended to close the performance gap between memory-layer sparse approaches and Mixture-of-Experts models. The authors introduce five main changes: placing memory layers in every transformer block, simplifying implicit value expansion to a single linear projection, using FFN-based value processing inspired by PEER, a new initialization for the memory layer, and rebalancing memory and FFN computation proportions. They evaluate UltraMemV2 across proprietary and open benchmarks and report that UltraMemV2 reaches parity with 8-expert MoE under matched compute/parameters while requiring much lower memory access. They highlight particularly strong gains on memory-heavy tasks and validate scaling up 120B parameter models.", "questions": "The method lags behind early in training as mentioned by the authors. Do you have more details and possible intuitions why UltraMemV2 is slower to train initially? And what does \"early\" mean in general - how much data does it need to catch up?\n\nIs there a reason for the drop in multi hop reasoning? Is this reflective of the model's overall reasoning capabilities on downstream tasks?\n\nCould you quantify the inference cost improvements that you alluded to eg the reduction in memory access?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "Strong architectural contributions: the five design changes proposed by the authors are all justified through ablations and contribute to improved model performance\n\nStrong empirical evaluation: multiple model scales up to 120B and a diverse selection of benchmarks make the authors claims very convincing.\n\nInitialization analysis: the paper contributes a new initialization scheme to stabilize training of the memory layer, which addresses a common failure mode for large sparse modules.\n\nPractical use: UltraMemV2 is a compelling architecture for deploying models under memory bandwidth constraints because of the relatively lower memory accesses.", "weaknesses": "The paper motivates UltraMemV2 with lower memory access and inference cost, but it would be more convincing to see latency and bandwidth comparisons vs. traditional MoE models\n\nProprietary data: this might be unavoidable but the proprietary nature of the benchmarks and data limits the reproducibility of the methods in this paper\n\nThe UltraMemV2 model has significantly worse benchmark performance on multi-hop reasoning. The paper would be improved if the authors investigated this further and demonstrated through other benchmarks whether UltraMemV2 has worse overall reasoning abilities or if it is specific to this benchmark.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0xoXqwy9cX", "reviewer": "Reviewer_f3Vh", "summary": "The paper proposed a sparse memory architecture by tweaking memory layer architecture to match the performance with MOEs with 8 experts. The paper does a good job of providing a Comprehensive analysis and detailed ablation studies. The paper also discusses scalability, though it’s unclear if the same scaling law of LLMs holds with trainable memory parameters. \nIt improves inference efficiency over prior work, by  simplifying value expansion with single linear projections, demonstrating that parameter efficiency of non-shared linear layers is actually not high. This paper also overcomes the limitation of number of memory layers, where previous work has shown degradation of performance if the number of memory layers are too high.", "questions": "1. A more rigorous scaling law analysis and discussion around scaling trends are required if this is proposed as an alternative architecture to MoEs\n2. Can this approach be introduce during mid-training or post training and achieve good performance. Any discussion around the performance gap with CT would be great. \n3. There are other parametric memory work, such as memory Layers at scale. It would be good to compare the results with such alternative approaches.  \n4. It would be good to share results on how high quality RAG combines with this approach and how does this approach compares with RAG for long context tasks\n\nMinor comments:\n1. Discussion about how can this be combined with MOEs", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "1. Good ablation for number of layers and  overcomes the limitation of number of memory layers\n2. Matched performence with MOEs with 8 experts\n3. Uses strong benchmarks and evaluation\n4. Simplifies the  value expansion, making inference more efficient", "weaknesses": "Overall the contribution is light, the paper aims to bridge the gap(in performance) between MOEs and Memory layer architectures. In terms of scientific novelty, some of the approaches seem incremental and this approach seems to combine multiple incremental tweaks to achieve performance improvements over baseline. For example, “Memory Layer at Scale”(Berges, 2024) paper demonstrated that  was multiple memory layers increase performance significantly over having a single layer(In their case, performance degraded going beyond 3 layers).  Another example is adoption of simplified value expansion with a small tweak of single linear projection.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EA1f9aQQMh", "reviewer": "Reviewer_rcay", "summary": "UltraMemV2 refines the Memory Layer architecture as a new sparse alternative to MoE models.\nWhile the original UltraMem only reached the performance level of a 2 active expert MoE, UltraMemV2 achieves performance comparable to an 8 active expert MoE, marking a significant advancement.\nThe authors introduce several key improvements, including a memory layer in each Transformer block, a simplified implicit value expansion (IVE), a PEER-based feedforward mechanism, improved initialization for stable training, and optimized compute ratios between memory to  FFN.\nWith these enhancements, UltraMemV2 shows notable gains (+6 to +8 points) on long-context, multi-step reasoning, and memory-intensive tasks, while efficient memory access comparable to existing MoE models.", "questions": "1. It is mentioned that UltraMemV2 cannot easily achieve higher sparsity. Did the authors analyze whether its sparsity scaling behavior differs from that of MoE models?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Considering that MoE architectures have become a de facto standard component in LLM training, the direction of this work is highly meaningful. Improving the cost efficiency and performance of such sparse architectures is an important and timely.\n\n2. The reported performance is very promising, matching the top-k=8 MoE configuration that is commonly used in recent LLM models.\n\n3. The experiments are conducted at a large scale up to 120B parameters with 2.5B active and trained on 4.4T tokens which makes the results convincing and demonstrates the method’s scalability.", "weaknesses": "1. It would be valuable to include more discussion or quantitative results about GPU hours, training/inference latency, and throughput, especially compared to existing MoE models.\nThese metrics would strengthen the claim that UltraMemV2 offers practical efficiency gains.\n\n2. While hyperparameter search is generally required for LLM training, the proposed model appears to be more sensitive to hyperparameter choices, such as initialization and learning-rate scheduling.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)