Memory Mosaics at scale

Transformers show an induction head mechanism, that is, predict b after sequence [..., a, b, ..., a].

“Transformers show an interesting induction head mechanism [Olsson et al., 2022], that is, predict b after sequence [ . . . , a, b, . . . , a ] .”

• cites Olsson et al., 2022 sourceSupported

Local-window memory approaches (Alibi position encoding, sliding-window attention) also struggle on the multi-unrelated-documents question-answering task for the same reason.

“Similarly, local-window memory approaches, such as Alibi position encoding Press et al. [2021] and sliding-window attention Beltagy et al. [2020], also struggle for the same reason.”

• cites Press et al., 2021 sourcePartly supported

  The Introduction (p. 1) identifies ALiBi as a method that penalizes attention to more distant tokens. Appendix B.2 (p. 24) suggests that ALiBi might not use contexts longer than those seen in training. The paper does not evaluate multi-unrelated-documents question answering or establish that ALiBi struggles on that task for the reason claimed. (LLM, conf=0.82)

  Experimental — accuracy not yet measured
• cites Beltagy et al., 2020 sourcePartly supported

  Section 3.1 (pp. 3–4) identifies sliding-window attention as a Longformer mechanism and says windowed attention alone is not flexible enough for task-specific representations, motivating global attention. The paper does not report that sliding-window attention struggles on the specific multi-unrelated-documents question-answering task or establish the cited sentence’s “same reason.” (LLM, conf=0.84)

  Experimental — accuracy not yet measured

Memory compression methods perform poorly on RULER long-context tasks as the required context length increases.

“It is clear that memory compression methods perform poorly as the required context length (i.e., required information storage space) increases.”

• Also stated in: § D Failures of memory compression baselines · p. 15 · plain

  “memory compression methods perform poorly as the required context length increases. Numbers are copied from Hsieh et al. [2024] Figure 4.”

• cites Hsieh et al., 2024Partly supported

  The Experiments & Results section (p. 6) reports that the evaluated models generally lose performance on RULER as sequence length increases. But it does not identify those models as memory-compression methods or report a result specifically for such methods; Related Work (p. 3) mentions compression as a separate approach. (LLM, conf=0.85)

  Experimental — accuracy not yet measured

Many memory compression algorithms (RNNs, xLSTM, rwkv, state-space models) fail on new-task storage and retrieval and in-context learning evaluation dimensions by construction.

“Many memory compression algorithms, such as RNNs, xLSTM [Beck et al., 2025], rwkv [Peng et al., 2023], and state-space models [Gu and Dao, 2023], fail on new-task storage and retrieval and in-context learning evaluation dimensions by construction.”

• cites Beck et al., 2025Not matched

  Couldn’t be matched to a paper in Crossref, so it can’t be checked.

• cites Peng et al., 2023 sourcePartly supported

  The Limitations section (page unavailable) says RWKV's recurrent architecture limits its ability to look back at previous tokens and may impair recall of details over very long contexts. The paper does not establish that RWKV fails new-task storage and retrieval or in-context learning evaluations by construction; its claim about impaired recall is qualified rather than categorical. (LLM, conf=0.87)

  Experimental — accuracy not yet measured
• cites Gu and Dao, 2023 sourcePartly supported

  Section 3.1 says time-invariant models struggle with selective copying and context-dependent retrieval because their fixed dynamics cannot select information from context. But the claim extends that limitation to state-space models generally: the paper introduces a selective state-space model to address it, and its Induction Heads results report that Mamba learned the task. The source therefore does not support saying state-space models fail these dimensions by construction. (LLM, conf=0.88)

  Experimental — accuracy not yet measured

Many memory compression algorithms (RNNs, xLSTM, rwkv, state-space models) fail on the multi-unrelated-documents question-answering task by construction because they cannot store all articles before reading the question.

“Many memory compression algorithms, such as RNNs, xLSTM [Beck et al., 2025], rwkv [Peng et al., 2023], and state-space models [Gu and Dao, 2023], fail on this task by construction because they cannot store all articles before reading the question.”

• cites Beck et al., 2025Not matched

  Couldn’t be matched to a paper in Crossref, so it can’t be checked.

• cites Peng et al., 2023 sourcePartly supported

  The paper’s “9 Limitations” section says RWKV’s recurrent architecture limits its ability to look back at previous tokens and may impair recall of details over very long contexts. Section L also says it cannot retrospectively readjust the weight of earlier information. Neither section establishes that RWKV fails the specified multi-document question-answering task by construction or cannot store all the articles before seeing the question. (LLM, conf=0.85)

  Experimental — accuracy not yet measured
• cites Gu and Dao, 2023 sourcePartly supported

  Section 3.1 describes recurrent models as compressing context into a finite state and says their effectiveness depends on how well that state preserves necessary information. It does not address the multi-unrelated-documents question-answering task or establish that state-space models necessarily fail it because they cannot store every article before the question arrives. (LLM, conf=0.86)

  Experimental — accuracy not yet measured

The incompetence of memory compression algorithms on the multi-unrelated-documents question-answering task has been experimentally demonstrated by Hsieh et al. [2024] and Li et al. [2024].

“This incompetent of memory compression algorithms has also been experimentally demonstrated by Hsieh et al. [2024] and Li et al. [2024].”

• cites Hsieh et al., 2024Partly supported

  Section 5, Task Error Analysis (p. 9), reports that Yi's question-answering performance degrades as distracting paragraphs are added, supporting the claim of difficulty with this kind of QA task. But Section 4, Experiments & Results (p. 6), describes experiments on long-context language models, not memory compression algorithms; the paper does not experimentally demonstrate those algorithms' incompetence. (LLM, conf=0.86)

  Experimental — accuracy not yet measured
• cites Li et al., 2024Not matched

  Couldn’t be matched to a paper in Crossref, so it can’t be checked.

A memory mosaics v2 trained on one trillion tokens still performs better on these tasks than a transformer trained on eight trillion tokens.

“A memory mosaics v2 trained on one trillion tokens still perform better on these tasks than a transformer trained on eight trillion tokens.”

Memory mosaics v2 demonstrate superior new-task learning capabilities with fewer examples and less prior knowledge from human designers.

“Our memory mosaics v2 demonstrate superior new-task learning capabilities (with fewer examples and less priori knowledge from human designers).”

Memory mosaics v2 significantly outperform transformers on carrying out new tasks at inference time (new-knowledge storage and in-context learning).

“and significantly outperforms transformers on carrying out new tasks at inference time (second and third dimensions).”

• Also stated in: § 5.3 In-context learning · p. 8 · boosted

  “In summary, the experiments demonstrate that memory mosaics v2 not only outperform transformer by a significant margin (more than 10%) on in-context learning, but also consistently improve performance as more demonstrations are provided.”

The improvements of memory mosaics v2 over transformers on new tasks cannot be easily replicated by simply increasing the training data for transformers.

“These improvements cannot be easily replicated by simply increasing the training data for transformers.”

Scaling memory mosaics to large language model sizes (llama-8B scale) and real-world datasets preserves their favorable compositional and in-context learning properties.

“This work shows that these favorable properties remain when we scale memory mosaics to large language model sizes (llama-8B scale) and real-world datasets.”

Memory mosaics v2 match transformers on learning training knowledge (persistent-knowledge storage).

“Throughout the evaluation, memory mosaics v2 match transformers on the learning of training knowledge (first dimension)”

• Also stated in: § 5.1 Persistent-knowledge storage and retrieval · p. 5 · hedged

  “The results (Table 1) show that memory mosaics v2 and transformers perform similarly in this evaluation dimension, suggesting that both models are capable of effectively storing and retrieving persistent knowledge.”

Removing long-term memory from memory mosaics v2 after training achieves comparable transformer performance on the 13 persistent-knowledge benchmarks while using fewer parameters and computations.

“Interestingly, removing long-term memory from memory mosaics v2 after training achieves a comparable transformer performance on the 13 persistent-knowledge benchmarks, while using fewer parameters and computations.”

Removing long-term memory from memory mosaics v2 after training degrades its performance on six language benchmarks (42.1% vs 34.9%).

“Table 9 shows six language benchmarks in which removing long-term memory from memory mosaics v2 after training degrades its performance.”

• Also stated in: § F Additional results on persistent-knowledge storage and retrieval · p. 16 · boosted

  “removing the 'longterm memory' after training dramatically hurt the performance (42.1% vs 34.9%).”

Memory mosaics v2 large outperforms public base models (similar scale) across all task lengths on RULER question-answer tasks.

“Memory mosaics v2 large outperforms these models across all task lengths.”

• Also stated in: § G Additional results on new-knowledge storage and retrieval · p. 17 · plain

  “Memory Mosaics v2 large outperforms these models across all task lengths, despite that Memory Mosaics v2 uses 1/4 generation lengths (32 tokens) of other public base models (128 tokens).”

Memory mosaics v2 and baseline transformers perform closely on 19 common language benchmarks.

“Table 1 evaluates both memory mosaics v2 and baseline transformers on 19 commonly used language benchmarks, showing that they perform closely on these benchmarks.”

Memory mosaics v2 consistently improve classification performance as they see more demonstration shots.

“memory mosaics v2 consistently improve classification performance as it sees more demonstration shots (blue curves).”

In memory mosaics, attention scores on near tokens heavily depend on token positions.

“That is, attention scores on near-tokens (positions) heavily depend on positions, while attention scores on far-tokens are almost invariant to token positions.”

• Also stated in: § 3.3 3-level memory · p. 4 · plain

  “The (averaged) attention scores in transformer heavily depends on token positions (curly curves), while the attention scores in memory mosaics at far tokens (e.g. position 0 to 450) are almost invariant to positions (flat curves).”

Removing long-term memory from memory mosaics v2 after training degrades performance on RULER question-answer tasks by 20% ∼ 30%.

“Table 10 shows that removing long-term memory from memory mosaics v2 after training degrades the performance on the RULER question-answer tasks by 20% ∼ 30%.”

Transformers with ROPE position encoding struggle to extrapolate context length.

“Unfortunately, transformers (with ROPE position encoding) struggle to extrapolate context length, as shown in Table 3.”

Memory compression methods fail on all cases in in-context learning tasks (Tacred few-shot classification).

“In this challenging in-context scenario, memory compression methods just don't work at all.”

• Also stated in: § D Failures of memory compression baselines · p. 15 · plain

  “Memory compression methods fail on all cases. Numbers are copied from Li et al. [2024] Table 4.”

• cites Li et al., 2024Not matched

  Couldn’t be matched to a paper in Crossref, so it can’t be checked.

In memory mosaics, attention scores on far tokens are almost invariant to token positions.

“That is, attention scores on near-tokens (positions) heavily depend on positions, while attention scores on far-tokens are almost invariant to token positions.”

• Also stated in: § 3.3 3-level memory · p. 4 · plain

  “The (averaged) attention scores in transformer heavily depends on token positions (curly curves), while the attention scores in memory mosaics at far tokens (e.g. position 0 to 450) are almost invariant to positions (flat curves).”

Memory mosaics v2 outperforms transformers on 4k task-length RULER question-answer tasks by 1.4% to 5.6%.

“Memory mosaics v2 outperforms transformers on 4k tasklength by 1.4% ∼ 5.6%.”

The stochastic long-term memory training setup enhances context-length extrapolation ability by more than 15%.

“The experimental results in Appendix G Table 12 show that this training setup enhances context-length extrapolation ability by more than 15%.”

When both models are fine-tuned at 32k context length, memory mosaics v2 significantly outperforms transformers by 12.3% to 14.8% on RULER question-answer tasks.

“At this increased difficulty level, memory mosaics v2 significantly outperforms transformers by 12.3% to 14.8%.”

For semantic label in-context learning tasks, transformers trained on 8T tokens match the performance of memory mosaics v2 trained on 1T tokens.

“For semantic label tasks (Figure 6), × 8 times more training data helps transformers (8T data) match the performance of memory mosaics v2 (1T data).”

• Also stated in: § 6.1 Comparison of two approaches · p. 9 · plain

  “Transformer with × 8 times more training data (8T, dash red line) starts to match the performance of Memory Mosaics v2 (1T, solid blue line).”

Memory mosaics v2 trained on 1T tokens outperforms transformers trained on 1T tokens on new-knowledge storage and retrieval (RULER question-answer tasks) by 12.3% (41.1% vs 53.4%).

“Training on the same number of tokens (1T), transformers lag behind memory mosaics v2 by 12.3% (41.1% vs 53.4%).”

Transformers struggle to maintain performance and exhibit degraded performance as more demonstrations are provided.

“In contrast, transformers struggle to maintain their performance and exhibit counterintuitively degraded performance as more demonstrations are provided (red curves).”

The 10% advantage of memory mosaics v2 holds even when comparing under the same number of parameters or the same computational budget.

“This 10% advantage holds even when comparing under the same number of parameters or the same computational budget.”

A single fine-tuning mini-batch (one optimization step) on memory mosaics v2 yields a 22% accuracy improvement.

“Surprisingly, a single fine-tuning mini-batch (one optimization step) on memory mosaics v2 yields a 22% accuracy improvement.”

Memory mosaics v2 outperforms transformer by more than 10% under either the same FLOPs or the parameters budgets.

“The results show that memory mosaics v2 outperforms transformer by more than 10% under either the same FLOPs or the parameters budgets.”

• Also stated in: § E Model efficiency comparison · p. 16 · plain

  “Memory mosaics v2 large outperforms transformer large by more than 10%.”

The other 6 common language benchmarks perform poorly when long-term memory is removed.

“In contrast, Appendix Table 9 indicates that the other 6 benchmarks perform poorly when long-term memory is removed.”

Two fine-tuning mini-batches on memory mosaics v2 are sufficient to reach optimal performance.

“Two fine-tuning mini-batches on memory mosaics v2 are sufficient to reach the optimal performance.”

Memory mosaics v2 and transformers perform similarly in persistent knowledge storage and retrieval.

“The results (Table 1) show that memory mosaics v2 and transformers perform similarly in this evaluation dimension, suggesting that both models are capable of effectively storing and retrieving persistent knowledge.”

Memory mosaics v2 trained on 4k context length successfully extrapolates context length ×4 to ×8 times without any fine-tuning.

“Memory mosaics v2 not only outperforms transformer on 4k task-length, but also successfully extrapolate the context length × 4 ∼ × 8 times without any fine-tuning.”

A transformer fine-tuned with 800 mini-batches still lags behind memory mosaics v2 fine-tuned with a single mini-batch.

“In contrast, a transformer fine-tuned with 800 mini-batches still lags behind memory mosaics v2 fine-tuned with a single mini-batch.”

Memory mosaics v2 significantly outperform transformers by more than 10% on in-context learning classification tasks.

“Memory mosaics v2 significantly outperform transformers by more than 10%.”

• Also stated in: § 5.3 In-context learning · p. 8 · boosted

  “In summary, the experiments demonstrate that memory mosaics v2 not only outperform transformer by a significant margin (more than 10%) on in-context learning, but also consistently improve performance as more demonstrations are provided.”

Memory mosaics v2 significantly outperform transformers on all classification tasks.

“Memory mosaics v2 significantly outperform transformers on all classification tasks.”

The stochastic long-term memory size setup boosts context length extrapolation ability by more than 15%.

“The stochastic long-term memory size setup boost context length extrapolation ability by more than 15%.”

Memory mosaics v2 significantly outperform transformers on in-context learning with a large margin (more than 10%).

“Memory mosaics v2 significantly outperform transformers on in-context learning with a large margin (more than 10%).”

Augmenting a transformer with long-short-term attention does not help it overcome the limitations of in-context learning.

“However, Figure 5 shows that augmenting a transformer with long-short-term attention does not help it overcome the limitations of in-context learning.”

Memory mosaics v2 significantly outperform transformers on all anonymous-label classification tasks.

“Again, memory mosaics v2 significantly outperforms transformers on all classification tasks.”

Transformers trained on 8T tokens still lag behind memory mosaics v2 trained on 1T tokens by 6.5% (46.9% vs 53.4%) on new-knowledge storage and retrieval (RULER question-answer tasks).

“However, the resulting transformer (trained on 8T tokens) still lags behind memory mosaics v2 (trained on 1T tokens) by 6.5% (46.9% vs 53.4%).”

• Also stated in: § 6.1 Comparison of two approaches · p. 10 · plain

  “In summary, × 8 more training data helps transformers in certain new task learning benchmarks. However, the resulting transformers (8T data) still lag behind memory mosaics v2 trained on 1T data.”

Removing long-term memory after training does not degrade performance of 13 common language benchmarks.

“Table 2 shows that removing long-term memory after training does not degrade the performance of 13 common benchmarks.”

Memory Mosaics [Zhang et al., 2025] have demonstrated appealing compositional and in-context learning capabilities on medium-scale networks (GPT-2 scale) and synthetic small datasets.

“Memory Mosaics [Zhang et al., 2025], networks of associative memories, have demonstrated appealing compositional and in-context learning capabilities on medium-scale networks (GPT-2 scale) and synthetic small datasets.”

• cites Zhang et al., 2025Not matched

  Couldn’t be matched to a paper in Crossref, so it can’t be checked.

New-knowledge storage and retrieval ability is a prerequisite for learning new tasks via memory-based methods because the data of new tasks must be adequately stored before learning.

“The new-knowledge storage and retrieval ability is a prerequisite for learning new tasks via memorybased methods (e.g., Gaussian kernel regression), because the data of new tasks must be adequately 'stored' before learning (Note that memory-based methods are lazy methods).”

• AssumptionRests on an assumption: Memory-based methods are lazy methods.

For a given distribution, the optimal bandwidth in kernel regression depends on the number of examples (key-value pairs).

“That is, for a given distribution, the optimal bandwidth depends on the number of examples (key-value pairs in associative memory).”

• cites Hastie et al., 2009 sourceNot checked — no full text of the source

The stochastic long-term memory training setup encourages allocation of position-invariant signals to long-term memory and position-dependent signals to short-term memory.

“This stochastic long-term memory training setup encourages the allocation of positioninvariant signals to long-term memory and position-dependent signals to short-term memory (as shown in Figure 1).”

The induction head mechanism contributes to transformers' in-context learning ability.

“This mechanism contributes to their in-context learning ability.”

Memory compression algorithms lack the ability to store large amounts of information before getting a command on how to process the information.

“The reason is that these memory compression algorithms lack the ability to store large amounts of information before getting a command on how to process the information.”

The 13 common language benchmarks are almost exclusively based on information stored in persistent memory and retrieved by short-term memory.

“This suggests that these 13 tasks are almost exclusively based on information stored in persistent memory and retrieved by short-term memory.”

The multi-unrelated-documents question-answering tasks are notably more challenging than typical needle-in-a-haystack benchmarks owing to their high information entropy.

“These tasks are notably more challenging than typical 'needle-in-a-haystack' benchmarks [Kamradt, 2023], owing to their high information entropy.”

• cites Kamradt, 2023Not matched

  Couldn’t be matched to a paper in Crossref, so it can’t be checked.

RULER question-answer tasks rely on long-term memory to perform well.

“This indicates that the ruler question-answer tasks rely on long-term memory to perform well.”

A model with limited new-knowledge storage ability will struggle to learn information that exceeds its storage capacity.

“Similarly, a model with limited new-knowledge storage ability will struggle to learn information that exceeds its storage (memory) capacity.”

Memory mosaics construct associative memories using keys to represent the recent past and values to represent the near future.

“Inspired by studies of induction head mechanism, Memory Mosaics [Zhang et al., 2025] construct associative memories using keys to represent the recent past and values to represent the near future (Figure 2 left):”

• cites Zhang et al., 2025Not matched

  Couldn’t be matched to a paper in Crossref, so it can’t be checked.

Memory mosaics reveal appealing in-context learning capabilities on small synthetic datasets.

“The resulting memory mosaics also reveal appealing in-context learning capabilities on small synthetic datasets.”

• cites Zhang et al., 2025Not matched

  Couldn’t be matched to a paper in Crossref, so it can’t be checked.

Memory mosaics v2 is not simply a transformer variation but represents a different architecture.

“These phenomena imply that memory mosaics (v2) is not simply a transformer variation but represents a different architecture.”

In memory mosaics, semantically similar cases may receive different key features because the averaging weights are fixed and independent of the semantic input x.

“The averaging weights in Equation 7 are fixed and independent of the semantic input x . As a result, semantically similar cases, such as 'tom-and-jerry' and 'tom- - -and- - -jerry', may receive different key features.”