Memory Mosaics

The leaking average in expression (6) is far too simple to effectively encode long range dependencies.

“The leaking average in expression (6) is far too simple to effectively encode long range dependencies as demonstrated in(Voelker et al., 2019;Peng et al., 2023;Gu & Dao, 2023).”

• cites Voelker et al., 2019Not matched

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

• cites Peng et al., 2023 sourceDisputed

  The Additional information section says RWKV’s time-decay mechanism preserves and propagates some information throughout the model’s temporal context. Section J.2 reports strong performance on several long-context tasks, though it also notes weaker results on others. The paper does not assess the citing document’s expression (6) specifically, but its findings run contrary to the blanket claim that a decaying average is too simple to encode long-range dependencies effectively. (LLM, conf=0.76)

  Checked by a reasoning model — may be wrong.

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

  Section 3.1 says time-invariant models struggle with selective copying because they cannot choose what to retain based on input content. But Figure 2 says they can solve the standard copying task, and the paper does not assess the citing document’s specific leaky average in expression (6). The source supports a narrower limitation, not the blanket claim that this average is too simple to encode long-range dependencies. (LLM, conf=0.82)

  Experimental — accuracy not yet measured

Implementing associative memories with kernel smoothing provides a direct illustration of the connection between self-attention and associative memories.

“Implementing associative memories with kernel smoothing therefore provides a particularly direct illustration of the connection between self-attention and associative memories (e.g., (Ramsauer et al., 2020)).”

• cites Ramsauer et al., 2020 sourceSupported

Unrolling equation (3) along an input sequence yields an expression that very much resembles masked self-attention.

“Unsurprisingly, unrolling equation (3) along an input sequence (x 1 . . .x D ) of duration D yields an expression that very much resembles masked self-attention (Vaswani et al., 2017).”

• cites Vaswani et al., 2017 sourceSupported

Memory Mosaics outperform all previously tested architectures on the REGBENCH in-context learning benchmark.

“Memory mosaics outperform all previously tested architectures in this benchmark.”

• cites Akyürek et al., 2024 sourceCould not verify

  Section 4.2 lists the architectures evaluated, but does not include Memory Mosaics. Section 4.5 reports that Transformers outperformed the other model classes tested on REGBENCH; it does not compare them with Memory Mosaics. Searches for “Memory Mosaics” and a REGBENCH performance comparison found no passage establishing the claimed result. (LLM, conf=0.76)

  Experimental — accuracy not yet measured

Because contextual memory contents are updated at inference time, recruiting some of them into the routing circuitry provides means to create new routes based on the first observations of a new sequence, suggesting an efficient alternative to capsule networks.

“Because the contents of contextual memory units are updated at inference time, recruiting some of them into the routing circuitry provides the means to create new routes on the basis of the first observations of a new sequence, suggesting an efficient alternative to capsule networks (Sabour et al., 2017).”

• cites Sabour et al., 2017 sourceSupported

The Memory Mosaic architecture outperforms decoding transformers on out-of-distribution tasks such as in-context learning.

“and outperforms them on o.o.d.tasks such as in-context learning.”

The comparatively transparent Memory Mosaic architecture matches the i.i.d. performance of decoding transformers on a language modeling task.

“showing that this comparatively transparent architecture matches the i.i.d.performance of decoding transformers on a language modeling task”

The gradient training algorithm splits the overall prediction task into disentangled prediction sub-tasks assigned to each memory unit (predictive disentanglement).

“This meta-learning interpretation reveals a remarkable phenomenon that we call predictive disentanglement : the gradient training algorithm splits the overall prediction task (e.g., predicting the next token in a natural language sentence) into disentangled prediction sub-tasks assigned to each memory unit.”

Memory Mosaics possess some disentanglement and compositional capabilities that have long eluded machine learning systems.

“Like transformers, Memory Mosaics possesses some of the disentanglement and compositional capabilities that have long eluded machine learning systems (Lake & Baroni, 2018).”

• cites Lake & Baroni, 2018Not matched

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

Memory Mosaics perform as well or better than transformers on medium-scale language modeling tasks.

“We illustrate these capabilities on a toy example and also show that memory mosaics perform as well or better than transformers on medium-scale language modeling tasks.”

All baseline methods perform well in the IID test set (even with only 100 training environments).

“All baseline methods perform well in the IID test set (even with only 100 training environments).”

The ROPE attention patterns do not extend nicely beyond the training context length.

“The ROPE attention patterns do not extend nicely beyond the training context length.”

The Memory Mosaic attention pattern is mostly flat, save for higher scores for the most recent tokens.

“In contrast, the Memory Mosaic attention pattern is mostly flat, save for higher scores for the most recent tokens.”

Single head network prediction error shows a sharp transition after lcm(p1, p2, p3) observations, when the network switches from predicting the future moon position by repeating the last observation to predicting by finding a matching memorized configuration.

“The prediction error shows a sharp transition after lcm(p 1 , p 2 , p 3 ) observations (red vertical line), when the network switches from predicting the future moon position by repeating the last observation to predicting by find a matching memorized configuration.”

Memory Mosaics perform very much like decoding transformers on the sort of language modeling task that made them famous.

“Such learning machines not only resemble decoding transformers (Section 3) but also perform very much like decoding transformers on the sort of language modeling task that made them famous (Section 7).”

Both models perform very similarly on the qualitative evaluation task of twenty-four prompts.

“Both models perform very similarly on this task.”

• cites Eldan & LiNot matched

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

The Memory Mosaic attention patterns remain mostly flat on extended contexts.

“In contrast the Memory Mosaic attention patterns remain mostly flat.”

The training process yields parameter matrices Wφ, Wψ, and Wz that dedicate one memory unit to each moon.

“Figure 10 shows how the training process yields parameter matrices W φ , W ψ , and W z , that dedicate one memory unit to each moon.”

In the interval between max period and lcm, the three-heads network returns accurate predictions for moon configurations very different from previously observed ones by combining individual moon predictions.

“In this interval, accurate predictions are returned for moon configurations that can be very different from the previously observed ones.Instead the network combines individual moon predictions, each well supported by the past observations.”

For the single-head network, prediction error shows a sharp transition after lcm(p1,p2,p3) observations, when the memory contains a full set of moon configurations.

“For the single head network (Figure 4), the plots show a sharp transition after lcm(p 1 , p 2 , p 3 ) observations, that is, when the memory contains a full set of moon configurations (red vertical line).”

The similarity of the Transformer and Memory Mosaic curves is especially striking when one recalls that the Memory Mosaic does not use position encoding.

“The similarity of the Transformer and Memory Mosaic curves is especially striking when one recalls that the Memory Mosaic does not use position encoding.”

The Memory Mosaic has a small advantage for very small depths (N b = 1 and N b = 4).

“The Memory Mosaic has a small advantage for very small depths (N b = 1 and N b = 4) but this advantage does not persist when the number of blocks increases.”

The Memory Mosaic's advantage does not persist when the number of blocks increases.

“The Memory Mosaic has a small advantage for very small depths (N b = 1 and N b = 4) but this advantage does not persist when the number of blocks increases.”

Baseline methods learned the training environments (good IID).

“This table, together with Figure 9, show that baseline methods learned the training environments (good IID) but not the meta-learning ability (poor OOD).”

For the three-heads network, the prediction error curve drops after seeing exactly p1, p2, and p3 observations, whenever the orbit of an additional moon has been memorized.

“For the three-heads network (Figure 5), the prediction error curve drops after seeing exactly p 1 , p 2 , and p 3 observations (black lines), that is whenever the orbit of an additional moon has been memorized.”

Memory Mosaics dominate the REGBENCH benchmark, substantially outperforming transformers, recurrent neural networks, and state-space models for training set sizes covering three orders of magnitude.

“Memory Mosaics dominate this benchmark, substantially outperforming transformers, recurrent neural networks, and state-space models for training set sizes covering three orders of magnitude.”

• cites Akyürek et al.Not matched

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

Baseline methods did not learn the meta-learning ability (poor OOD).

“This table, together with Figure 9, show that baseline methods learned the training environments (good IID) but not the meta-learning ability (poor OOD).”

The ALIBI attention patterns show the vanishing contribution of distant tokens.

“The ALIBI attention patterns show the vanishing contribution of distant tokens.”

The three-heads network produces accurate predictions after a time equal to max(p1,p2,p3), long before the combined period lcm(p1,p2,p3).

“Observe how the network produces accurate predictions after a time equal to the period max(p 1 , p 2 , p 3 ) of the slowest moon (last black line), long before the combined period lcm(p 1 , p 2 , p 3 ) (red line) of the moon system.”

The transformer performance plateaus after 100 to 150 tokens on Simple English Wikipedia, which is a bit shorter than a typical tiny story.

“The transformer performance plateaus after 100 to 150 tokens, which is a bit shorter than a typical tiny story.”

This effect is connected to the leaky average coefficient λφ.

“This effect is connected to the leaky average coefficient λφ, as shown in Figure16.”

Because the associative memory retrieval function (2) converges to stationary conditional expectations E(V|K), each memory unit is driven to produce a good conditional expectation estimate as soon as possible.

“Because the associative memory retrieval function ( 2) is known to converge to stationary conditional expectations E(V |K), each memory unit is driven to produce a good conditional expectation estimate as soon as possible.”

• AssumptionRests on an assumption: The associative memory retrieval function converges to stationary conditional expectations E(V|K).

When multiple memory units are available, training can distribute the overall prediction task among them by optimizing value extraction functions psi to produce values v_T that are more efficiently modeled by their respective memory units.

“When multiple memory units are available, the training procedure can also distribute the overall prediction task among the available memory units.As long as the memory units outputs can still be combined to address the overall task, the training algorithm can optimize the parameters of the value extraction functions ψ to produce values v T that more efficiently modeled by their respective memory”

An associative memory can be viewed as a device that estimates a conditional probability distribution P(V|K) on the basis of a sample of key-value pairs.

“This exchangeability property suggests that we can also view an associative memory as a device that estimates a conditional probability distribution P (V |K) on the basis of the sample (k 1 , v 1 ) . . .(k n , v n ) of key-value pairs.”

• AssumptionRests on an assumption: An associative memory stores key-value pairs without consideration for their temporal ordering, except perhaps when duplicate keys are involved.

Each memory unit operates as a little predictor whose outputs y_T can be interpreted as a conditional expectation that estimates features of the near future (v_T) based on past observations (k_T).

“Each memory unit operates as a little predictor whose outputs y T can be interpreted as a conditional expectation (1) that estimates features of the near future (v T ) of the input time series on the basis of its past observations (k T ).”

With a frozen value extraction function psi, training can make each memory unit statistically more efficient by tuning the key extraction function phi, i.e., learning a similarity metric (kernel).

“Let us first assume that each memory unit has a frozen value extraction function ψ.The training procedure can still make each memory unit statistically more efficient by tuning the parameters of the key extraction function φ, that is, by learning how to compare the current prediction context (x T , x T -1 , x T -2 . . . ) with past prediction contexts (x t , x t-1 , x t-2 . . . ) for t < T .”

• cites Bach et al., 2004 sourceNot checked — no full text of the source

Because each memory unit operates independently at inference time, predictive disentanglement works best when the overall prediction task is disentangled into smaller sub-tasks that can be modeled independently and efficiently.

“Because each memory unit operates independently at inference time, this works best when the overall prediction task is disentangled into smaller prediction sub-tasks that can be modeled independently and efficiently.”

The training procedure should be seen as a meta-learning process, distinct from the memory-based learning that occurs at inference time when new key/value pairs are added into the memories.

“The training procedure should therefore be seen as a meta-learning process, distinct from the memory-based learning that occurs at inference time when new key/value pairs are added into the memories.”

Predictive disentanglement is not a mysterious effect of scale but a property of the architecture.

“showing that this is not a mysterious effect of scale but a property of the architecture.”

The predictive disentanglement principle explains how training decomposes the overall task in interesting ways.

“identifying and illustrating the predictive disentanglement principle which explains how training decomposes the overall task in interesting ways”

When all key vectors share a same squared norm, Gaussian kernel smoothing becomes equivalent to attention.

“The close connection between this Gaussian kernel smoothing and attention (Bahdanau et al., 2015) is obvious when all key vectors k i share a same squared norm because expression (2) becomes”

• cites Bahdanau et al., 2015Not matched

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

There is no leak of future information in the elementary memory unit because the output y_T does not depend on v_T but only on previously stored key/value pairs.

“Although the value v T depends on the near future, the output y T does not depend on v T but merely leverages the previously stored key/value pairs to estimate v T .Therefore there is no leak of future information: each memory unit is a little machine that predicts a bit of future information (described by v T ) on the basis of recent information (described by k T ) and previously stored key/valu”

The retrieval function of an associative memory is a conditional expectation over the estimated distribution.

“The retrieval function is then a conditional expectation over this estimated distribution:”

Memory Mosaics achieve compositional and in-context learning capabilities in a comparatively transparent way, unlike transformers.

“Unlike transformers, memory mosaics achieve these capabilities in comparatively transparent way”

Such a conditional expectation can be constructed with Gaussian kernel regression.

“Such a conditional expectation can be constructed with Gaussian kernel regression,”

An n block deep Mosaic has the same number of parameters as an n block deep Transformer.

“This amounts to saying that a n block deep Mosaic has the same number of parameters than a n block deep Transformer, its performance is closer to that of a n + 1 block Transformer.”

Minimizing the sum of prediction losses favors memories that produce useful value estimates after fewer time steps.

“Minimizing their sum -the area under the curve-favors memories that produce useful value estimates after fewer time steps.”

Position encoding in Transformers is mostly useful to implement an initial induction head in the first two blocks.

“This suggests that position encoding in Transformers is mostly useful to implement an initial induction head in the first two blocks.”

These two mechanisms do not allow a head to selectively address a token by position.

“These two mechanisms are useful to implement bigram or n-gram induction heads in a single layer, but they do not allow a head to selectively address a token by position (we use a single scalar leaky average coefficient per head).”

Predicting Simple English Wikipedia articles using models trained on BABISTORIES is a challenging out-of-distribution task.

“Predicting Simple English Wikipedia articles using models trained on BABISTORIESis therefore a challenging out-of-distribution task.”

An n block deep Mosaic's performance is closer to that of an n+1 block Transformer.

“This amounts to saying that a n block deep Mosaic has the same number of parameters than a n block deep Transformer, its performance is closer to that of a n + 1 block Transformer.”

The absolute position encoding scheme cannot be extended to longer contexts.

“Because the absolute position encoding scheme cannot be extended to longer contexts, we provides a comparison with transformers using ROPE (Su et al., 2024) and ALIBI (Press et al., 2022).”

• cites Su et al., 2024Not matched

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

• cites Press et al., 2022Not matched

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

The chatbot creative writing phenomenon is fundamentally similar to the three-moons example, where unseen configurations are accurately predicted by combining individual moon predictions.

“This phenomenon is fundamentally similar to that illustrated in Figure 5, where moon configurations unlike any previously seen configurations are accurately predicted because the network has learned how to combine individual moon predictions.”

The outputs of disentangled sub-tasks can be recombined to provide predictions for inputs globally very different from training inputs, but whose disentangled components are individually predictable.

“Their outputs can then be recombined to provide predictions for inputs that are globally very different from the training inputs, but whose disentangled components are individually predictable, as illustrated in Section 5.”

The predictive disentanglement principle explains why training a Memory Mosaic breaks the overall prediction task into pieces that are more efficiently memorized when they are considered independently.

“Second, the predictive disentanglement principle explains why training a Memory Mosaic breaks the overall prediction task into pieces that are more efficiently memorized when they are considered independently (Section 5).”

The second astronomer only needs to log daily moon positions until each moon returns to a previously observed position, for a number of days equal to max(p1,p2,p3).

“In contrast, the second astronomer only needs to log daily moon positions until each of the moons returns to a previously observed position, for a number of days equal to the period max(p 1 , p 2 , p 3 ) of the slowest moon.”

These two mechanisms are useful to implement bigram or n-gram induction heads in a single layer.

“These two mechanisms are useful to implement bigram or n-gram induction heads in a single layer, but they do not allow a head to selectively address a token by position (we use a single scalar leaky average coefficient per head).”

The total number of persistent memory unit slots is close to the total number of FFN hidden units.

“The total number of persistent memory unit slots is therefore close to the total number of FFN hidden units.”

As λφ increases, the pick at the end of attention map becomes wider.

“As λ φ increases, k t in Eq 6 effectively takes a longer history into the account, and thus the pick at the end of attention map becomes wider.”

As λφ increases, kt in Eq 6 effectively takes a longer history into account.

“As λ φ increases, k t in Eq 6 effectively takes a longer history into the account, and thus the pick at the end of attention map becomes wider.”