Memory Layers at Scale

Thank you for the kind review! We will try to address the comments in the order they were raised.

MoE Results

While we cannot fully explain the poor performance of MOE models (especially for the 1.3b setting), we can provide the following additional information:

• Our MOE implementation reuses code from a recent SoTA MOE architecture paper (reference not given to protect anonymity), which we believe to be well optimized for MOE.
• We have independent evidence from other internal teams that match our MOE results It is quite surprising that MoE performs slightly worse than the dense model for NQ and TQA on Figure 5. MoE does perform better than dense at the 1.3b scale on all other tasks of Table 2, PIQA, ObQA and HotPot.

Memory+ layout

Are there any ablation studies, showing that whether layout plays a role on the performance?

Table 3 (top) provides some insights into how the memory layers layout (number of layers, centering and spacing) affects performance. In the table, adding more Memory+ layers beyond 3 decreases performance. We believe this is because for every memory layer we add, we remove the FFN layer, resulting in effectively less computation. Also, we share the value parameters across all the Memory+ layers of the model, thus the total parameter count does not increase. As a result, there is a sweet spot after which adding more memory layers doesn’t help. It is possible to add memory layers in addition to or parallel to an FFN layer, however we did not consider this as it would break our compute-controlled comparisons.

Decoding Efficiency of Memory+

I'm not fully convinced that the decoding efficiency of Memory+ layer is consistent with dense baseline across all the batch size. When the batch size is increasing, the activated memory ratio will increase, where the memory access demand also increases.

You are indeed correct that Memory+ layers memory access grows with batch size. However, for the top-k value that we pick (4 heads, 32 top-k, 128 total), the memory demand remains comparable to the dense model for small batch sizes (note that the analysis for one forward pass through the entire model, not just a single memory layer). For large batch sizes, all models are compute bound, thus the curve follows the dense models curve closely for the full range of batch sizes. However, MOE models quickly need a lot more memory bandwidth even at small batch sizes, and are thus severely memory bound in this regime.

Model Design Philosophy

Our design is built on top of the work of Large Memory Layers with Product Keys (https://arxiv.org/abs/1907.05242) in which they use two independent sets of sub keys that they then combine into a complete key. This is done for efficiency reasons, since otherwise it would be infeasible to compute top-k over millions of keys. This likely causes some degradation in key lookup (similar to product quantization). On the other hand, the value parameters are kept in a flat table and are independent of each other.

Relation between performance and sparsity ratio

Is there any experiment about the relation between performance and sparsity ratio? In other words, when the sparsity ratio is further increasing, can the performance also improve consistently?

In Table 2 at the 1.3b scale, we study the relation between increasing the Memory+ parameter count (which increases sparsity) and performance. At the scales studied, it seems increasing the parameter count increases performance. However, we were unable to scale further in this research work and could not verify if further increase always results in better performance.

Another way to increase sparsity would be to decrease k. We find that doing this (e.g. go from top-128 to 64) hurts performance. We also did ablations on these values, but largely came to the same conclusions as the original product-keys (Lample et al) paper, so decided not to repeat those results here. Overall, varying heads or k while keeping their product the same effects results minimally. Increasing the total keys improves little beyond 128, while having substantial memory lookup and gpu memory costs, while decreasing to 64 has non-negligible accuracy degradation. In the end we stick to 128, which was also the default for the original paper.

We thank the reviewer for the supportive feedback! For benchmarks beyond factuality, we tried to include a variety of standard benchmarks on table 2 for the 8B models. In addition to this, here are results for the GSM8K math benchmark that we recently ran (8B, 1 trillion tokens):

I am curious how this may affect out-of-distribution tasks or reasoning in new domains unseen during pretraining

While this is a valid and interesting question, at this time we do not have any data to help us speculate on this issue. However, we are continuing research with memory models, including in new domains of knowledge, and hope to contribute in this direction in future work.

We thank the reviewer for their thoughtful review. Please find our answers and clarifications below:

Evidence for claims

it would be helpful to include a similar plot for non-factual tasks (e.g. HumanEval or MMLU)

We opted to not provide numbers for these benchmarks for the small model sizes, since they are often not meaningful (e.g. MMLU is ~25% which is chance level). However, here are the nll values for the two mentioned benchmarks at 1.3b scale for reference (lower is better):

I can’t find a compute-matched comparison anywhere in the paper

All of our experiments are compute (FLOPs) matched within the same scale category. E.g. for the 1.3b scale, the dense, MOE and memory layer models all have (almost) identical flop cost. For the MOE-memory layer comparisons, we also attempted to match parameter count. We will make it more clear in the paper that the comparisons are FLOP-controlled, as this is one of the main claims of the paper.

How would varying the number of experts and size of experts affect the performance relative to memory layers

For the main experiments, we set the number of memory values at 2^20 for the memory models, and tried to set the MOE experts to match the parameter count. This choice results in varying number of MOE experts for different scales (e.g. 16 experts for the 134m model, 4 experts for the 1.3b model). While we agree that having more experiments at various parameter / expert counts would be ideal, we did not have the needed resources to run them. We would guess that having more but smaller experts would behave more similarly to a memory layer. In fact, the PEER work which we compare against, investigates the limit of this, where they have up to a million rank-1 experts. Our work performs comparably or better, while being simpler and more compute efficient.

Other questions and comments

How does the efficiency of a dense model compare to the memory model at training and inference?

We provide such an analysis in figure 4, plotting latency against batch size for dense, MOE and memory models. Memory models have similar efficiency with dense models in both the small batch-size (typical inference) and large batch-size (typical training) regimes, while MOE inference latency is much higher. This is a roofline analysis. In practice, we observe ~10% lower training throughput compared to a dense model during training at 8B scale, due to communication overheads.

What is the value of k chosen and how does it affect performance? Thank you for catching this oversight! We used 4 heads and k=32 (total keys 128) for all experiments, we will add this to the paper. We also did ablations on these values, but largely came to the same conclusions as the original product-keys (Lample et al) paper, so decided not to repeat those results here. Overall, varying heads or k while keeping their product the same effects results minimally. Increasing the total keys improves little beyond 128, while having substantial memory lookup and gpu memory costs, while decreasing to 64 has non-negligible accuracy degradation. In the end we stick to 128, which was also the default for the original paper.

Typos and nits

Thank you for pointing these out, we will edit this in the final version of the paper.

Thank you so much for the positive feedback! We will do our best to improve based on the reviews. Here is how we plan on addressing the comments:

Experimental details

We will add an appendix with the details about the experimental setup including model dimensions and training hyper-parameters. Here are some of these details. For 134m base models, we use dim of 768, 12 layers and attention with 12 heads. For 373m base models, we use dim of 1024, 24 layers and attention with 16 heads. For 720m base models, we use dim of 1536, 22 layers and attention with 12 heads. For 1.3b base models, we use dim of 2048, 22 layers and attention with 16 heads. For 8b base models, we use dim of 4096, 32 layers and attention with 32 heads. In all Memory and Memory+ experiments, we use 4 heads and 32 top-k values. We used a learning rate of 3e-4 for the 134m to 1.3b models and 1e-4 for the 8B models.

Code release for EmbeddingBag

Our code is already open source (withholding reference due to anonymity), however here is the relevant file with the kernel implementation for your review: https://justpaste.it/cb6xz .

Unbalanced key usage

Regarding unbalanced usage of keys, we notice that there can be a few steps at the very beginning of training during learning rate warmup where the keys become quite unbalanced. This issue goes away on its own and the distribution naturally smooths out during training. In general, we did not need to do any regularization to enforce balance. We will add a plot showing how this distribution evolves during the first few steps of training in the appendix.

Thank you for your detailed review and thoughtful comments! We will try to address the feedback in order of priority, starting with the discussion of performance as this was deemed to be the “biggest weakness” of the paper.

Discussion of performance

The biggest weakness of this paper is that the discussion of performance with respect to implementation is lacking, especially regarding parallelism and batching

The reviewer’s intuition and observations in this section about the performance of memory layers vs. MOE or dense models is correct. Memory layer models have a large advantage compared to MOE at inference time, where decoding batch sizes are small. At large batch sizes, the three models are in principle equivalent. This is the typical setting for training. In practice, memory layers do have some additional communication overhead, however we have been able to realize training throughput that is within 10% of the dense model training speed using our optimized implementation.

I would have liked to see a chart that compares the lookup speed of Dense vs. MoE vs. Top-K for different sized batches of query vectors

We provide such a chart in figure 4, where we plot the decoding time (dominated by memory lookup for small batch sizes) of the 3 architectures against the batch size.

MOE results

I am surprised by the poor performance of the MoE models

While we cannot fully explain the poor performance of MOE models (especially for the 1.3b setting), we can provide the following additional information:

• Our MOE implementation reuses code from a recent SoTA MOE architecture paper (reference not given to protect anonymity), which we believe to be well optimized for MOE.
• We have independent evidence from other internal teams that match our MOE results

That being said, the 1.3b MOE results do not fit the scaling trends. We don’t know the cause of this, but it might be because the MOE configuration which is parameter matched to memory layer at this scale only has 4 experts, which is quite small.

If there are 2^20 values, then I'm a bit confused about why there are only 984m total parameters

The value dimension in the memory layer is equal to the model dimension, which is different for each dense model scale. For the 134m model, this is 768. This is indeed confusing, and we will provide a table of model parameters for each scale in the appendix.

Theoretical claims

The reviewer is mostly correct, however the way product keys are implemented, we have 2 sets of half-keys, and the full set of keys is defined as the product of these sets. As a result, if A is in the first set and B is in the second set, then AB is by definition part of the key set. For the example given by the reviewer, if A’B’, AX and BY are in the keys, then so are A’B, AB’, which both have higher scores than A’B’.

Additional references

Thank you for providing additional references, we will include them in the final version.