Memory Layers at Scale

Probably not a substantial claim: may not pass S3 (load-bearing).

Product-key networks were introduced to make memory lookup more efficient and scalable.

“Product-key networks (Lample et al., 2019) were introduced to make the memory lookup more efficient and scalable.”

• cites Lample et al., 2019 sourceSupported

May combine two claims.

The field focused on alternatives such as mixture-of-experts, which more closely resemble dense networks and are easier to scale.

“Instead, the field focused on alternatives such as mixture-of-experts (Shazeer et al., 2017), which more closely resemble dense networks and are thus easier to scale.”

• cites Shazeer et al., 2017 sourcePartly supported

  Section 1.3 (p. 3) describes substantial research on mixture-of-experts, and Section 2 (p. 3) explains how sparse gating lets the model use thousands of experts while evaluating only a handful per example. The paper supports MoE as a scalable approach, but it does not say MoEs more closely resemble dense networks or that this resemblance makes them easier to scale. (LLM, conf=0.84)

  Experimental — accuracy not yet measured

Probably not a substantial claim: may not pass S3 (load-bearing). May combine two claims.

Earlier work introduced end-to-end trainable memory layers and incorporated them into neural computational systems.

“Earlier works introduced end-to-end trainable memory layers (Sukhbaatar et al., 2015) and incorporated them as part of neural computational systems (Graves et al., 2014).”

• cites Sukhbaatar et al., 2015Not matched

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

• cites Graves et al., 2014 sourceSupported

Probably not a substantial claim: may not pass S3 (load-bearing). May combine two claims.

In mixture-of-experts, each FFN layer contains multiple experts, and only a subset of them participate in computation for each input.

“In MOE, each FFN layer is composed of multiple”

• cites Shazeer et al., 2017 sourcePartly supported

  Sections 1.2 (p. 2) and 2 (p. 3) describe an MoE layer containing multiple feed-forward expert networks, with sparse gating selecting only some experts for each input. They do not say that each FFN layer contains experts; the paper places its MoE layer between LSTM layers. (LLM, conf=0.85)

  Experimental — accuracy not yet measured

May combine two claims.

Using a shared pool of memory parameters across all memory layers keeps the parameter count the same and maximizes parameter sharing.

“In contrast to previous work (Lample et al., 2019), we use a shared pool of memory parameters across all memory layers, thus keeping parameter count the same and maximizing parameter sharing.”

• cites Lample et al., 2019 sourcePartly supported

  The Training details (p. 6) describe placing multiple memories at different transformer layers, so the citation identifies relevant previous work. The paper says memory heads share values within a memory (Multi-head memory attention, p. 5), but our copy does not establish whether parameters are shared across memory layers. The claimed contrast in cross-layer sharing is therefore not fully supported; keeping the parameter count the same is the citing authors’ description of their own method. (LLM, conf=0.78)

  Experimental — accuracy not yet measured

May combine two claims.

Memory layers also improve significantly on coding tasks (HumanEval, MBPP) and general knowledge tasks (Hellaswag, MMLU).

“while also improving significantly on coding (HumanEval, MBPP) and general knowledge (Hellaswag, MMLU).”

May combine two claims.

The paper demonstrates scaling laws with memory capacities up to 128 billion parameters, pretraining to 1 trillion tokens, and comparing against base models up to 8 billion parameters.

“demonstrating scaling laws with up to 128B memory parameters, pretrained to 1 trillion tokens, comparing to base models with up to 8B parameters.”

May combine two claims.

The PEER architecture performs similarly to Memory for the same number of parameters, but lags behind Memory+.

“The PEER architecture performs similarly to Memory for the same number of parameters, while lagging behind Memory+.”

May combine two claims.

Memory models improve drastically over dense baselines and generally match the performance of models with twice as many dense parameters on QA tasks.

“We can see from table 1 that Memory models improve drastically over the dense baselines, and generally match the performance of models with twice the number of dense parameters on QA tasks.”

May combine two claims.

In the shared-memory ablation, model performance first improves and then declines as more memory layers are added, peaking at about three memory layers.

“We see that as we add more memory layers, performance initially increases.However, as we're effectively removing dense parameters from the model for each added memory layer, eventually the model performance degrades, revealing a sweet spot at around 3 memory layers (table 3, left).”

May combine two claims.

Memory+ improves further over Memory, with performance generally falling between dense models with two to four times higher compute.

“Memory+ improves further over Memory, with performance falling generally between dense models with 2x-4x higher compute.”

May combine two claims.

Adding random key-value pairs or a fixed softmax-sink key to the top-k during pre-training gives minor improvements, but these have a negative impact on training speed and the gains are not consistent for larger model sizes.

“For key sampling improvements, including the random keys and the fixed (sink) key, we see minor improvements, however these have some negative impact on training speed in our implementation, and the gains were not consistent for larger model sizes, therefore we excluded them from our experiments, leaving this direction open for future exploration.”

May combine two claims.

The custom forward pass achieves 3 TB/s of memory bandwidth, close to the H100 specification of 3.35 TB/s, compared to less than 400 GB/s with PyTorch's implementation.

“Our forward pass optimizes memory accesses and achieves 3TB/s of memory bandwidth, which is close to our H100 specification of 3.35TB/s (compared to less than 400GB/s with PyTorch's implementation).”

May combine two claims.

The memory layer operation is expected to be limited solely by GPU memory bandwidth, but PyTorch's implementation has multiple inefficiencies in practice.

“As the number of floating-point operations is negligible, we expect this operation to be solely limited by the GPU memory bandwidth, but find multiple inefficiencies in PyTorch's implementation in practice.”

Probably not a substantial claim: may not pass S3 (load-bearing).

Like mixture-of-experts, memory-augmented models aim to increase parameter space without adding significant computational cost.

“Like MOE, memory augmented models also aim to augment the parameter space of the model without adding significant computational cost.”

May combine two claims.

Memory layers are rarely used in modern architectures and have not been optimized for hardware accelerators.

“Such components are rarely used in modern architectures and have not been optimised for hardware accelerators.”

May combine two claims.

Memory layers with sparse activation are almost entirely limited by memory bandwidth, whereas dense layers are predominantly limited by FLOPs.

“In contrast to dense layers which are predominantly FLOP-bound, memory layers with their sparse activation pattern are almost entirely memory bandwidth bound.”

Probably not a substantial claim: may not pass S3 (load-bearing), S4 (based).

Scaling laws for specific task families and capabilities such as factuality remain understudied.

“Nevertheless, scaling laws for specific task families and capabilities like factuality remain understudied.”

Probably not a substantial claim: may not pass S3 (load-bearing). May combine two claims.

Memory networks were first proposed by Weston et al. (2015) and later with end-to-end training by Sukhbaatar et al. (2015).

“Memory networks were proposed initially in (Weston et al., 2015), and with end-to-end training in (Sukhbaatar et al., 2015).”

• cites Weston et al., 2015Not matched

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

• cites Sukhbaatar et al., 2015Not matched

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

May combine two claims.

Little research has been done to improve memory layer performance, partly because they are rarely used and not optimized.

“In addition to, and partly as a result of this, little research was done to improve their performance.”

Probably not a substantial claim: may not pass S3 (load-bearing).

Memory layers can be implemented with a simple, low-cost key-value lookup where keys and values are embeddings.

“Such memory layers can be implemented with a simple and cheap key-value lookup mechanism where both keys and values are encoded as embeddings (Weston et al., 2015).”

• cites Weston et al., 2015Not matched

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

May combine two claims.

Using an associative memory for simple associations would be more efficient and more natural than using feed-forward networks.

“using an associative memory for this purpose would be both more efficient and more natural.”

• AssumptionRests on an assumption: Simple associations are an important subset of information that language models need to learn.

Probably not a substantial claim: may not pass S3 (load-bearing), S4 (based). May combine two claims.

The parallel memory implementation is independent of other model parallelism schemes such as tensor, context, or pipeline parallelism, and operates on its own process group.

“The implementation is independent of other model parallelism schemes such as tensor, context or pipeline parallelism, and operates on its own process group.”

May combine two claims.

The memory layer implementation here is orders of magnitude more scalable than previous works, yet a substantial engineering effort is still needed to make it efficient enough for large-scale production.

“While the memory layer implementation presented here is orders of magnitude more scalable than previous works, there still remains a substantial engineering task to make them efficient enough for large scale production uses.”

• AssumptionRests on an assumption: Modern GPU architectures and dense neural architectures have co-evolved and been optimized together for decades.

Probably not a substantial claim: may not pass S3 (load-bearing), S4 (based).

Deep networks encode information at different levels of abstraction across layers.

“Deep networks encode information at different levels of abstraction across different layers.”

Classified as probably not a substantial claim.

Adding memory to multiple layers may help the model use its memory in more versatile ways.

“Adding memory to multiple layers may help the model use its memory in more versatile ways.”

May combine two claims.

PEER works similarly to memory layers in practice, but requires twice as many parameters for the same number of keys.

“PEER works similarly to memory layers in practice, but requires twice the number of parameters for the same number of keys.”

• cites He, 2024Not matched

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

May combine two claims.

Memory layers can be scaled efficiently and give practitioners a new direction for trading off memory against compute.

“They can be efficiently scaled, and provide practitioners with an attractive new direction to trade-off memory with compute.”

Probably not a substantial claim: may not pass S3 (load-bearing), S4 (based).

Models with the same base model configuration have negligible differences in FLOPs.

“Models with the same base model configuration have negligible differences in FLOPs.”

• AssumptionRests on an assumption: Models with the same base model configuration have negligible differences in FLOPs.

May combine two claims.

A naive nearest-neighbour search compares every query-key pair, which becomes prohibitive for large memories.

“A naive nearest-neighbour search requires comparing each query-key pair, which quickly becomes prohibitive for large memories.”

May combine two claims.

Sparse and dense layers are both needed and likely complementary.

“showing sparse and dense layers are both needed and likely complementary (see section 5.4).”

May combine two claims.

Raising the key dimension adds dense parameters to the model, so it cannot be increased without limit while keeping comparisons fair.

“However, increasing the key dim does add more dense parameters to the model, and thus we cannot increase it indefinitely without breaking our fair comparisons.”

May combine two claims.

Memory layers typically have a much larger number of keys and values, which makes sparse lookup and updates necessary.

“Second, memory layers typically have larger scale in terms of the number of keys and values, making sparse lookup and updates a necessity.”

May combine two claims.

Memory layers share parameters, so they have the same memory footprint as a single memory layer.

“As noted earlier, memory layers share parameters, thus have identical memory footprint to a single memory layer.”

Probably not a substantial claim: may not pass S4 (based).

Memory layers are naturally memory-intensive, largely because of their many trainable parameters and optimizer states.

“Memory layers are naturally memory-intensive, mostly due to the large number of trainable parameters and associated optimizer states.”

Probably not a substantial claim: may not pass S4 (based).

Memory layers have distinct scaling challenges because they are light on compute and heavy on memory.

“Being light on compute, and heavy on memory, memory layers have distinct scaling challenges.”

Probably not a substantial claim: may not pass S3 (load-bearing). May combine two claims.

The PEER model is conceptually similar to a memory layer, but it retrieves a pair of embeddings that combine into a rank-1 matrix.

“The PEER model is conceptually similar to a memory layer, but instead of retrieving a single value embedding, it retrieves a pair of embeddings, which combine into a rank-1 matrix.”

• cites He, 2024Not matched

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

May combine two claims.

Like memory layers, mixture-of-experts and PEER increase parameter count without significantly increasing FLOPs.

“Like memory layers, these methods increase the number of parameters in the model without significantly increasing FLOPs.”

Probably not a substantial claim: may not pass S3 (load-bearing).

The evidence presented so far that memory layers improve factuality is only high-level empirical evidence.

“We have so far presented only high level empirical evidence that memory layers improve factuality of models.”

May combine two claims.

For some benchmarks, small models can perform very poorly and their accuracy numbers are noisy.

“For some bencmarks, the performance of small models can be very low, and accuracy numbers noisy.”

• AssumptionRests on an assumption: For some benchmarks, small model performance is very low and accuracy numbers are noisy.

Probably not a substantial claim: may not pass S4 (based). May combine two claims.

Memory layers with sparse activations complement dense networks by increasing capacity for knowledge acquisition while remaining computationally light.

“Memory layers with their sparse activations nicely complement dense networks, providing increased capacity for knowledge acquisition while being light on compute.”

May combine two claims.

In memory layers, keys and values are trainable parameters, whereas in attention layers they are activations.

“First, the keys and values in memory layers are trainable parameters, as opposed to activations.”

• cites Vaswani et al., 2023Not matched

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

Probably not a substantial claim: may not pass S3 (load-bearing). May combine two claims.

PEER builds on product-key networks by replacing vector values with rank-one matrices, bridging memory architectures and mixture-of-experts.

“The recent PEER (He, 2024) builds on this work, replacing vector values with rank-one matrices, forming a bridge between memory architectures and MOE.”

• cites He, 2024Not matched

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

May combine two claims.

Fast approximate vector similarity techniques could be used, but incorporating them is challenging when keys are continually trained and need re-indexing.

“While fast approximate vector similarity techniques (Johnson et al., 2019) could be used here, it's a challenge to incorporate them when the keys are being continually trained and need to be re-indexed.”

• cites Johnson et al., 2019Not matched

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