Ultra-Sparse Memory Network

We greatly appreciate the time and effort you have taken to review our manuscript. In response to your insightful comments, we address each point individually.

For the weakness:

1. Improved Clarity and Background Information: We acknowledge the need for a clearer presentation of complex concepts. To this end, we have revised the introduction and related work sections, redrawn figures, and included additional annotations to enhance understanding, particularly regarding the implicit value expansion technique and its relation to virtual and physical memory. Additionally, Figure 5 has been modified for better comprehension.

2. Support for Claims in the Methods Section:

   • Comparison with PKM: In our ablation studies, UltraMem demonstrates a 0.1 improvement in loss over PKM, a gain typically associated with doubling the dense computational power. Evidence from PEER (https://arxiv.org/abs/2407.04153) showing MoE's superiority over PKM further supports our assertion that UltraMem significantly outperforms PKM even without direct comparison. The reasons for not using PKM as a baseline in larger models (151M, 680M, 1.6B) include the high training costs and inherent structural issues of PKM. PKM itself is not well-suited for scaling up training, as a very large PKM is challenging to parallelize and results in significantly longer training times.

   • Performance Over MoE: Figures 1(b) and 1(c) in the main text illustrate UltraMem's inference speed and memory usage compared to MoE across varying batch sizes, clearly showcasing UltraMem's advantages. Tables 1 and 2 list the models' FLOPs and parameters, with UltraMem showing superior performance in benchmarks under the same activation and total parameters compared to MoE. Furthermore, we have added a new figure (Figure 6.c) comparing the inference speed of UltraMem and MoE as sparse parameters increase. This data shows MoE's inference speed growing exponentially, while UltraMem's remains almost unchanged, highlighting a significant advantage.

For the minor error:

1. We thank for the detailed reading of our paper. It is indeed a typo in Eq6 where the matrix dimensions are not properly matched. We have revised the Eq6 in our latest submitted version.

For the questions:

1. It's a good idea to apply IVE in MoE. Generally, memory networks and MoEs are both sparse approximations of a large inner FFN. It is possible to expand an expert to a series of virtual experts, who share the same FFN parameter while own different projectors to make them function differently. Also, the pre-pooling technique in Eq15, which saves quantities of computation, is still applicable for IVE in MoE.

We hope these revisions and expansions adequately address your concerns and strengthen the paper's claims. If there are any further questions or suggestions, please feel free to let us know. We are looking forward to providing additional clarification and discussion.

For the questions:

1. The MoE model in our paper is configured with Granularity (G=2) and Expansion (E=12), which represents a relatively strong baseline. Additionally, the Mixtral MoE employed a similar configuration.

2. UltraMem-1.6B-x12 consumed 103B tokens per day, while MoE-1.6B-x12 consumed 113B tokens per day, both on 128 A100 GPUs. Please note that while the MoE training has been extensively optimized internally, UltraMem still has room for further optimization.

3. In the original PKM paper, experimental observations suggest placing PKM in the middle of the model for optimal performance. Our early experiments confirmed that PKM achieves the best results when placed in the middle, while positioning it in the shallower layers tends to yield poorer outcomes. Consequently, for inserting UltraMem, we generally avoid the initial few layers and distribute it evenly across the subsequent layers. Regarding how many layers to skip, we primarily consider engineering aspects for training and inference acceleration. Based on our experience, we recommend skipping five layers for each UltraMem layer to approximately ensure adequate overlap for value retrieval.

4. We have redrawn all the figures to ensure the text is clear and have added annotations and references to facilitate reader understanding.

For the suggestions:

1. As previously mentioned, we have aligned our studies with the open-source 6.5B dense model, thereby conducting comparisons with established open-source models. This alignment ensures that our evaluations reflect competitive benchmarks within the field.
2. Validating performance in real-world scenarios is indeed crucial, and we fully agree with your perspective. We have plans to train larger models and validate them in real-world scenarios after performing SFT and RLHF. In our pre-training evaluations, our benchmarks encompass a range of scenarios, including knowledge, reasoning, and reading comprehension. In these tests, UltraMem consistently matches or exceeds the performance of MoE, demonstrating its robust capabilities. This success gives us greater confidence to expand to larger models in the future.
3. The MoE model used in our studies already corresponds to the configuration from [1] with (G=2) and (E=12) settings. This ensures that our analysis and comparison of memory efficiency and performance between UltraMem and MoE are grounded on a well-established benchmark.

We sincerely appreciate your guidance, which has helped us enhance the clarity and depth of our paper. If there are any further questions or suggestions, please feel free to let us know. We are looking forward to providing additional clarification and discussion.

Thank you for your valuable and insightful feedback on our manuscript. In response to your insightful comments, we address each point individually.

For the weakness:

1. Establishing a Proper Baseline with MoE: We acknowledge the importance of establishing a relevant MoE baseline for a fair comparison. It is essential to note the challenges in using existing open-source MoE models, as they often do not provide access to the training data used, or their architectural parameters differ significantly from our 1.6B UltraMem model. Therefore, we implemented a comparable MoE model with identical parameters (1.6B MoE 2in34) using the same training data to ensure a fair and meaningful comparison.

2. Experiments with Popular Models: In our paper, the 6.5B dense model has been aligned with the open-source RedPajama 7B dense model, and our evaluation metrics match those reported in their blog (https://www.together.ai/blog/redpajama-7b). This RedPajama 7B dense model replicates the results of the LLaMA 7B model. Thus, we are effectively comparing our models against established open-source models. The reason we replicated the RedPajama 7B dense model was to ensure the reliability of our training framework, as detailed in the table below. | model | MMLU-EM | BoolQ-EM | NaturalQuestions (closed book) - F1 | NaturalQuestions (open book) - F1 | QuAC - F1 | HellaSwag - EM | TruthfulQA - EM | avg | |---------------------------------------|---------|----------|-------------------------------------|----------------------------------|-----------|----------------|-----------------|-------| | Redpajama-incite- 7b | 0.323 | 0.694 | 0.258 | 0.6 | 0.323 | 0.702 | 0.205 | 0.444 | | Reproduced 7b (6.5b in paper) | 0.246 | 0.707 | 0.272 | 0.612 | 0.346 | 0.691 | 0.298 | 0.453 |

3. Analysis of Sparsity Levels: We agree with your suggestion and have thus explored different sparsity levels more thoroughly. In response, we have now included additional experiments in the manuscript that compare UltraMem and MoE across various sparsity settings (Figure 6.c), enriching our analysis of how sparsity influences both memory efficiency and performance. The experiments show that, for UltraMem, the inference time remains largely stable even as sparse parameters exponentially grow, provided that the activation parameters (top-m) are constant. In contrast, MoE's inference time escalates significantly under analogous conditions.

Thank you for your valuable and insightful feedback. Below, we outline the actions we’ve taken in response to your comments, addressing each point individually.

For the weakness:

1. We sincerely apologize for any inconvenience caused by our writing and figures. We have updated the manuscript to ensure that each step in Figure 4 is clearly referenced within the text. Additionally, we have enhanced the figure captions with more detailed annotations to help readers better understand the methodologies depicted.

2. We acknowledge the importance of scaling up to 10B models and certainly plan to pursue this as part of our future work. Currently, only a few studies have successfully validated models on such a large scale. In our experiments, training a 1.6B x12 UltraMem and an equivalent MoE requires 128 A100 GPUs for 10 days to process 1T tokens, whereas replicating a 6.5B dense model demands 512 A100 GPUs over approximately 14 days. While we recognize the necessity for further research in this direction, the associated computational costs remain a substantial consideration.

3. Independently ablating each modification is crucial to accurately assess the effectiveness of each component of the model, thanks very much for pointing this out. However, it is possible that applying modifications 'a' and 'b' separately on the baseline might each yield a loss improvement of -0.02, but when applied simultaneously, the combined benefit may only result in total loss improvement of -0.03 rather than -0.04. This is because individual modifications can interact and interfere with each other's effects. Therefore, we believe that sequential integration of modifications offers a more accurate understanding of their effectiveness.

4. Following your recommendation, we have included additional experiments to explore different Tucker decomposition settings (rank r = 2, 3, 4), IVE settings (E = 4, 9, 16), and MCS settings (h=2, 4, 8), which are detailed in Table 3 to better elucidate their effects on model performance. The experiments show that, for IVE, as E increases, there is a consistent improvement in model performance alongside a notable increase in computational cost. However, as E rises, marginal gains decrease, leading us to recommend E=4. For TDQKR and MCS, increasing r and h does not significantly change the computational load, but the effectiveness no longer shows marked improvement, hence we suggest using r=2 and h=2.

We appreciate your guidance that has significantly enhanced the depth and clarity of our paper. If there are any further questions or suggestions, please feel free to let us know. We are looking forward to providing additional clarification and discussion.

For the suggestions: Thank you for your detailed suggestions. We have taken your feedback into careful consideration and have made the following revisions to our manuscript:

1. We have rewritten the introduction to focus on high-level motivation, clearly outlining the current issues, what our methodology aims to address, and the extent of its capabilities.

2. The section on related work has been moved to the front of the paper as the second section. We've enriched this part with detailed background information on MoEs, PKM, Tucker Decomposition to provide readers with a thorough understanding of the current research landscape. Megatron is not central to our research direction; therefore, it is omitted from the related work section.

3. All figures have been redrawn to ensure all text is clearly visible. Additional annotations have been included to aid in understanding.

4. We have improved the captions and labeling in Figures 2 and 3, providing detailed descriptions of what each side represents and how it relates to the method. Subfigures are now clearly marked, and we have adjusted the positioning of floats to appear at the top of the pages.

5. We have included arrows in the headers to indicate whether higher or lower values signify improvement, aiding in clearer interpretation of the metrics.

We hope these adjustments meet your expectations and contribute to the improved readability and comprehension of our paper. If there are any further questions or suggestions, please feel free to let us know. We are looking forward to providing additional clarification and discussion.

We greatly appreciate the time and effort you have taken to review our manuscript. In response to your insightful comments, we address each point individually.

For the weakness:

1. We have revised the introduction to bolster clarity and cohesion by integrating key points from the original Section 3. This adjustment aims to elucidate our method’s motivation at an early stage, thereby enhancing overall readability.

2. Recognizing the importance of figure legibility, we have redrawn all figures to ensure that all text is clear and easily readable. Additionally, we have improved the figure captions for better comprehension. We hope these updates effectively address your concerns.

3. We acknowledge your concern and clarify that PKM was indeed used as a baseline in our ablation studies, where UltraMem showed a significant improvement in loss by 0.1. This is a considerable enhancement, usually observed with a doubling of computations in dense models. We refrained from including PKM as a baseline in our larger model trials (151M, 680M, 1.6B) primarily due to the high training costs of PKM. PKM suffers from significant memory access issues in the middle of the model, which severely reduces training efficiency. This makes it impractical to scale up the model effectively, which is one jumping-off point of our research. Specifically, training a 1.6B x12 UltraMem required 128 A100s over 10 days for 1T tokens (similar to MoE). Further tuning of PKM requires a certain amount of training cost.

4. For similar cost reasons, we had to limit the scope of certain experiments.

5. To improve logical flow, we have repositioned the related work section to precede the methodology. This restructuring allows for a more coherent narrative where the discussion of state-of-the-art techniques and their limitations in MoE and PKM provides necessary context. Furthermore, we enriched this section with a detailed discussion on Tucker decomposition and its relevance to our work, thereby facilitating better comprehension of the techniques involved.

6. Thank you for your valuable suggestions. In the evaluation section, we have categorized the downstream metrics in the paper as follows

   • Knowledge: MMLU, TriviaQA, GPQA, ARC
   • Reasoning: BBH, BoolQ, HellaSwag, WinoGrande
   • Reading Comprehension: DROP
   • Overall Performance: AGIEval

   These downstream metrics cover the categories of benchmark tasks discussed in Jaiswal's work, including both zero-shot and few-shot (in-context learning) settings. Furthermore, since we are focusing on the pre-trained models (without alignment), the instruction-following setting in Jaiswal's paper may not be the most suitable evaluation setting for our work.

   Regarding your point about using loss and perplexity to measure a model's general capabilities, it is indeed an interesting question! In general, as mentioned in Jaiswal's paper, during inference, "...all compressed models are derived from the same dense counterpart with high similarity, and aforementioned differences(scale, training strategy, architecture) don’t exist, making their evaluation more challenging". However, in our case, we are focusing on pre-training (training models from scratch), and the models with different sparsity levels are non-homogeneous (e.g., the cosine similarities of the same modules' weights among them are less than 0.2). In such cases, the training/validation loss will be an efficient and reliable indicator of general performance, as the breadth of the training/validation corpus will far exceed that of specific task datasets. In fact, when exploring the scaling laws of pre-training, the validation/test loss is commonly used as a metric for measuring a model's general capabilities, as seen in the paper of OpenAI's scaling law (https://arxiv.org/abs/2001.08361) and the technical report of LLaMA-3 (https://arxiv.org/abs/2407.21783).

First, I would like to thank the authors for their revised paper and rebuttal explaining the changes/what they addressed, it's clear they have taken the feedback seriously and made a significant effort to effect those changes within the rebuttal period.

Quickly reviewing the revised paper, and going through the weaknesses I identified in the initial version:

• The figures are much improved, now legible, and in some cases with better captions, although some of the later figures could also do with better captions still.
• I like the new text for the introduction, it much better motivates the method - especially in explaining the inference issues with MoEs!
• Moved background to before method, which I think works much better for explaining the author's method
• Expanded background section with tensor decomposition serves the authors' work much better.
• Improved results table with arrows indicating if metric is desired to be small/large

Further changes I think would help the author's work:

• Expand the caption in Figure 4/5, they are complex figures, and now I can read it properly, I found the caption very brief and not helpful. Ideally you want your figures to be standalone (as much as that's ever possible).

Authors rebuttal for PKM baseline/limitations of experiments due to compute (Weakness 3/5): I think this is, to some degree, a fair point: that PKM is a costly baseline to run for the larger models and that's what motivated the research, and I'm with the authors on that motivation completely. Also in general, as much as possible as a reviewer I avoid demanding prohibitively expensive experiments to keep research accessible. However, in the context of MoE research, all experiments are going to be very computationally expensive, and specifically in your case you are trying to motivate claims are you are doing better than the relatively expensive PKM so this seems to me to be a minimum requirement to have those complete baselines to me. In summary, while I do understand this is all very computationally expensive, if we let papers make claims with MoEs specifically go unverified because experiments are expensive, it would be bad for the research community in the long run. Also want to be clear I wouldn't expect all such results to be feasible in the rebuttal period, but was hoping perhaps for some commitment to have them by final paper.

Summary: Overall I am much happier with the revised paper and I do think it addresses a lot of the low-hanging issues with the initial submission on the layout of the paper, figures, and much (but not all) of the writing, and will be re-evaluating my rating of the work in that context. However, while I sympathize with the computational issues, I believe the lack of some key baselines is a major issue still that restrains me from giving a much stronger rating.

Summary of Revisions

In response to the constructive feedback from the reviewers, we have made several significant adjustments and additions:

For common response:

1. Improved Logical Flow: Following Reviewer DbFB’s insightful suggestions, we restructured our paper to enhance coherence. We moved the related work to the beginning to detail MoE and PKM limitations and expanded the discussion on Tucker decomposition. Additionally, we refined the introduction to incorporate essential insights earlier, enhancing clarity and readability.
2. Enhanced Clarity and Detailing: In response to all reviewers, we significantly improved the manuscript's clarity, especially around the technical descriptions of our architecture and the experimental setup, ensuring that our methods are understandable and replicable.

For individual response:

1. Reviewer DbFB expressed concern about the lack of a fundamental baseline, specifically the performance of the PKM model at the scales of 151M, 680M, and 1.6B. Upon reevaluating our experiments, we agreed with the reviewer's perspective. During the rebuttal period, we conducted these three experiments, making the comparison with UltraMem more solid.
2. Reviewer jFrZ noted that our ablation studies were missing certain verificatory experiments. These include evaluating the performance of IVE, TDQKR, and MCS under various configurations and assessing the benefits of ablating each component separately. We fully agree with this observation. Consequently, we conducted four additional sets of ablation experiments. These experiments were designed to validate the optimal configuration for IVE, TDQKR, and MCS, as well as to independently ablate each component. The outcomes of these ablations robustly support our model structure and hyper-parameter configurations.
3. Reviewer ch2J was concerned about the lack of comparisons with public models and the absence of inference speed evaluations at different sparsity levels. We addressed the first concern by comparing our dense baseline (Dense-6.5B) with the open-source model Redpajama-incite-7b in our paper, showing similar performance, which confirms our model's validity against public models. Furthermore, we conducted extra experiments comparing the inference speeds of UltraMem and MoE when increasing sparsity parameters with fixed activation parameters. These ablations consistently demonstrated UltraMem's superiority, showing minimal impact on its speed with increased sparsity, unlike MoE, which significantly slows down.
4. Reviewer bPzv felt that comparing UltraMem and PKM solely based on validation loss was not sufficiently convincing and that comparisons with MoE lacked detailed metrics such as memory access, FLOPs, and parameter counts. Upon re-evaluating our experiments, we agreed with the need for more extensive PKM comparisons. Consequently, we trained PKM models at scales of 151M, 680M, and 1.6B. Evaluating these against UltraMem, the results further validated UltraMem's effectiveness. Regarding the missed comparisons, we had actually detailed these metrics in the main text of the original manuscript, which might have been overlooked due to clarity issues. The revised version of the paper now clearly highlights these comparisons to ensure better understanding.

We believe that the majority of the reviewers' concerns have been addressed through additional ablation experiments, improved clarity of expression, and detailed explanations. These enhancements, driven by reviewers' feedback, have significantly fortified the technical solidity and clarity of our paper.

We are grateful once more for the engaging and constructive review process, which has undeniably enriched our work.

Best regards,

Authors of UltraMem