Ultra-Sparse Memory Network

Authors: Zihao Huang, Qiyang Min, Hongzhi Huang, Yutao Zeng, Defa Zhu, Ran Guo, zhou Xun
Primary area: foundation or frontier models, including LLMs
Keywords: Large language model, sparse model, scaling law

Abstract:
It is widely acknowledged that the performance of Transformer models is logarithmically related to their number of parameters and computational complexity. While approaches like Mixture of Experts (MoE) decouple parameter count from computational complexity, they still face challenges in inference due to high memory access costs. This work introduces UltraMem, incorporating large-scale, ultra-sparse memory layer to address these limitations. Our approach significantly reduces inference latency while maintaining model performance. We also investigate the scaling laws of this new architecture, demonstrating that it not only exhibits favorable scaling properties but outperforms MoE. In experiments, the largest UltraMem we train has \textbf{20 million} memory slots. The results show that our method achieves state-of-the-art inference speed and model performance within a given computational budget, paving the way for billions of slots or experts.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730151942914

thread_type:
official_review

<thread_object>
creator: reviewer_bpzv

<content>
summary:
This paper introduces a novel layer designed to increase model size efficiently. Specifically, the approach refines the existing PKM method and incorporates TDQKR as a sparse method to enhance performance. Additionally, it proposes implicit value expansion to address memory access challenges in large models. A comparative analysis of memory access between the proposed method and MOE is also provided. Experimental results demonstrate significant performance improvements.

strengths:
1. The techniques proposed in this paper are novel. On one hand, the paper thoroughly studies the techniques in PKM, offering valuable empirical insights for future research. On the other hand, it integrates algorithms with memory management to address the memory access issue.   
2. The paper includes an efficiency analysis focused on memory access.
3. Experimental results show that the method significantly outperforms existing approaches on large models, being up to six times faster than MOE at the same scale.

weaknesses:
1. The writing could be improved, particularly since it introduces concepts beyond the domain of algorithms. These should be better explained as background information. For example, the implicit value expansion technique is somewhat confusing when the paper discusses virtual and physical memory, and could benefit from further clarification.

2. The experimental results in the main paper do not fully support the claims made in the methods section. First, while the paper asserts that the method improves PKM, there is no direct comparison except for validation loss. Second, although the paper claims that its memory access method outperforms MOE, this is only demonstrated in the appendix. The main table merely shows FLOPs and model parameters, which do not sufficiently illustrate the method's advantages (including the running time).

### Minor:
1. The dimensions in Equation (6) are inconsistent.

questions:
1. The proposed IVE aims to reduce memory access. Given that it handles the general form of matrix multiplication, could it also be applied to MoE or other methods that involve extensive memory access?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730219257224

thread_type:
official_review

<thread_object>
creator: reviewer_ch2j

<content>
summary:
In this paper, the authors propose UltraMem, a memory-efficient method designed to replace MLP layers or Mixture of Experts (MoE) in Transformer architectures, particularly for large language models. UltraMem builds on a refined Product Key Memory (PKM), utilizing a 2D grid structure with row-column key decomposition to reduce retrieval costs. Additionally, it employs Tucker Decomposition to further minimize memory usage and computational overhead. Experiments across various language tasks demonstrate UltraMem’s promising effectiveness and efficiency.

strengths:
This paper addresses an interesting and impactful problem with practical applications in large language model (LLM) research.

- The authors propose an alternative, memory-efficient approach to achieving the performance of Mixture of Experts (MoE) models.
- They introduce a sparse memory access mechanism and a 2D Product Key Memory structure, which restricts memory access to only the most relevant slots, enhancing efficiency.
- **Scalability**: The use of Tucker Decomposition improves the scalability of the method, allowing it to handle larger models effectively.
- **Experiments**: The experiments demonstrate promising results, achieving comparable performance to MoE models of similar scale while maintaining greater memory efficiency.

weaknesses:
Although the idea presented in this paper is novel, the experimental performance assessment shows some limitations:

1. **Lack of a Proper Baseline**: The authors do not use a well-established MoE baseline to evaluate the performance of their method. Since UltraMem is proposed as a memory-efficient alternative to MoE, it would be beneficial to compare it against real-world MoE implementations from existing open-source models, providing a more comprehensive analysis.

2. **Limited Experiments with Popular Models**: The experiments primarily use custom models, limiting the generalizability of the results. It would be valuable to assess UltraMem’s effectiveness on popular open-source models like Llama, Mistral, or others.

3. **Limited Analysis of Sparsity Levels**: The paper could benefit from a deeper investigation into the effects of different sparsity levels on memory efficiency and performance, as this is a key component of UltraMem’s approach.

questions:
here are some minor writing issues, such as in Line 094, where it states 'introduce the origin large memory.' These should be easy to address.

My main questions and observations are:

1. **Comparison with Fine-Grained Mixture of Experts**: How does this method compare to fine-grained Mixture of Experts as in [1], where the hidden dimension is split between experts?
2. **Training Speed**: How does UltraMem’s training speed compare to that of MoE?
3. **Placement of UltraMem Blocks**: On what basis were the positions of the UltraMem blocks within the architecture chosen?
4. **Figure Quality**: Improving the quality of the figures and plots would enhance clarity.

**Suggestions**:

1. It would be beneficial to conduct experiments with popular open-source models like Llama-3.2-1B, Mistral, or others.
2. Since Section 3, "Why UltraMem Instead of MoE," discusses UltraMem’s advantages over MoE, it would be valuable to include experiments demonstrating that UltraMem consistently outperforms MoE in real-world scenarios, using a comprehensive baseline model for comparison.
3. Additionally, an ablation study comparing UltraMem's memory efficiency and performance against MoE based on [1] would further strengthen the analysis if applicable.

[1] *Scaling Laws for Fine-Grained Mixture of Experts*.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730628221369

thread_type:
official_review

<thread_object>
creator: reviewer_jfrz

<content>
summary:
This paper proposes a novel deep learning architecture called UltraMem, designed to reduce memory access costs during the inference process of large language models, thereby improving inference efficiency. The core innovation of UltraMem lies in its introduction of an ultra-sparse memory layer, which allows the model to activate only a small number of necessary memory units when processing tasks. This approach reduces the number of memory accesses, effectively lowering inference latency. Specifically, , UltraMem surpasses MoE with the same parameters and computation as model capacity increases.

strengths:
1. This paper introduces the UltraMem architecture, which demonstrates innovation by significantly reducing inference latency while maintaining computational efficiency. 
2. The paper presents extensive experiments comparing the performance of UltraMem with traditional models (such as MoE and dense models), verifying UltraMem’s advantages in inference speed, memory access costs, and scalability. 
3. The experiments show that UltraMem’s memory access volume grows much more slowly with batch size compared to MoE, enhancing its practical utility.

weaknesses:
1. The paper lacks references and descriptions of the architecture diagrams within the main text. For example, each step in Figure 4 is not referenced in the text. Additionally, certain terms in the architecture diagrams, such as “fetch values,” are not explained in the text, making the paper difficult to follow.
2. The paper claims that the proposed UltraMem method has stronger scalability. However, UltraMem was only tested on models with 151M, 680M, and 1.6B parameters, without experiments on larger models, such as the 6.5B model or beyond. This parameter range is insufficient to fully verify UltraMem’s scalability in large-scale models. It is recommended to extend the experimental scale to explore UltraMem’s performance and scalability at extremely large parameter sizes (e.g., 10B and above).
3. Several modules proposed in the paper, such as Implicit Value Expansion (IVE) and Multi-Core Scoring (MCS), collectively enhance UltraMem’s performance. However, the independent effects of each module are not adequately evaluated. For instance, Table 2 provides ablation experiments but lacks detailed analysis of the independent effects of key modules like IVE and MCS across different tasks or sparsity settings. It is recommended to expand on Table 2 by adding experiments that assess the independent impact of each optimization module on model accuracy.
4. Figure 4 illustrates the process of Tucker decomposition but lacks analysis of how different decomposition parameters, such as rank r, affect model accuracy. To clarify the impact of different Tucker decomposition configurations on model performance, it is recommended to include more comprehensive ablation experiments to quantify Tucker decomposition’s specific role in UltraMem. Additionally, experiments should be added to analyze the effects of different values of E in the IVE method on experimental results.

questions:
see weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730693349558

thread_type:
official_review

<thread_object>
creator: reviewer_dbfb

<content>
summary:
Motivated by the high memory access cost bottleneck of Mixtures of Experts (MoEs) for Large Language Model (LLM) inference, the authors propose an alternative methodology, UltraMem, based on Product Ley Memory (PKM) which improves memory access above MoEs, but suffers in performance compared to them. The authors make several observations to improve the vanilla PKM methodology, basing their method on this baseline. The authors then not three problems with the PKM architecture, notably that it does not scale to large value sizes, product key decomposition biases retrieval, and unbalanced GPU computation/communication for large models. The authors propose to decompose the large memory layer of PKM into many smaller memory layers distributed across layers, allowing execution of memory and transformer layers to overlap. The authors propose to use a Tucker decomposition instead of the product key decomposition of PKM. Finally, rather than explicitly maintaining a large memory table for values, the authors propose a virtual memory approach. The authors evaluate UltraMem across several tasks compared to dense and an MoE equivalent model, and compare the performance across these tasks, along with the inference and memory accesses for each, demonstrating significantly faster inference with similar performance to an MoE model.

strengths:
* Memory access is perhaps the main bottleneck in inference for contemporary hardware, and so this is a well-motivated and impactful direction to be exploring.
* The results appear to show impressive improvements over Mixtures of Experts (MoEs) in inference time and memory access while maintaining validation loss/perplexity. 
* The authors methodology is detailed and thoughtful, for example deriving the correct initialization for their method (which would appear to also apply somewhat to PKM).
* Many of the figures/graphics themselves are good in design and would be helpful to understanding the methodology if it wasn't for the other serious issues with readability and captions (see below).
* Changes to PKM (before the UltraMem specific changes) are listed in 2.2 as the bag of tricks that improve PKM's performance, and are separate from the author's proposed UltraMem structure. It would be beneficial if more papers took such an approach to be clear about the differences between improved training methodology for baseline methods and a newly proposed method. Ablation in particular is nice about being clear how these changes improve baseline PKM performance.
* Real-world inference and memory access results on GPU hardware.
* Large-scale experiments, in particular results are evaluated on various model sizes, from 151M param up to 1.6B. FLOPS for the sparse models.

weaknesses:
* Overall the paper reads poorly, more as a collection of experimental details and results rather than a cohesive story. I know that's frustrating and perhaps vague feedback to get as an author but I think it's really important to point out as it makes the paper quite hard to read as it is, and reduces the impact of the author's work. To be clear, I do appreciate the experimental and methodological details themselves. However, I believe the authors could do the work much more justice by taking a step back and framing their research story better. In particular, I would recommend focusing first on section 3 as higher-level motivation of the method before diving into building on top of PKM. 
* All the figures containing the majority of results, and some of the methodology figures (Fig 4) are **far** too small. They are so small that the figures are completely unreadable on paper. I measured the font sizes in the figures out of curiousity, and most of the fonts are 2-3pt!
* Given that the proposed method is largely based on PKM, PKM should be a baseline in the results. The authors mention PKM's "..effectiveness is significantly inferior to that of MoE.", without explaining on what measures this is true, or giving the reader to make that judgement in the results. I don't actually doubt the authors on this, but it must be demonstrated in the experiments.
* While as noted above it's great that the "bag of tricks" for PKM is listed, along with the ablation, it also would seem important that the performance of the improved PKM should be evaluated as a baseline over just vanilla PKM, and compared with the improvements from the author's proposed method.
* The related work is **far** too short given all the work done in this field, much of which is referred to by the authors and built upon in the proposed method. To be honest I'm never a fan of related-work at the end of the paper, but I feel like this paper in particular would have greatly benefited by the related work coming before the methodology rather than at the end of the paper, as instead of having to bring up related work throughout the paper as it is built upon, it could have just been earlier summarized and then referenced throughout, allowing the story to focus on the methodology more clearly.
* Relies too much on validation loss/perplexity for the evaluation of whether generalization performance is maintained in the main paper/figures. The authors do have six other metrics for tasks in Table 1 and figures in the appendix for these tasks (which they should explain better in the paper/background rather than quickly citing all of them in a list). Given that the authors are in a sense proposing a form of sparsity, and given the results of Jaiswal et al. for pruning (Compressing LLMs: The Truth is Rarely Pure and Never Simple, ICLR 2024) I believe it's more important to focus on the performance in task-specific measures.

questions:
Suggestions:
* Focus on the high-level motivation before diving into details of methodology and experiments.
* Move the related work up front in the paper, and put in background required for your method. Include MoEs, PKM, Tucker decompositions, Megatron, any background relevant to IVE and MCS, etc.
* Fix readability of all the text in the figures.
* Figure 2 (and to some extent Figure 3) are poorly labelled/captioned, and it's largely left to the reader to figure out what each side represents and how it's related to the method. At a minimum label the subfigures or describe in caption what left/right are. Would also suggest positioning floats at top of page.
* When including a large table of metrics (i.e. in Table 1), it's helpful to the reader to explain whether lower or higher is better for each, e.g. with an arrow in the header.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732016919948

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We extend our deepest gratitude to all the reviewers for their time and valuable feedback. We have carefully considered your comments and have made the following revisions to the manuscript, highlighted in red for your convenience:
1. We have rewritten the introduction to provide a higher-level analysis of the causes behind UltraMem, facilitating a clearer understanding of our motivation.
2. The related work has been enriched and moved to Section 2. In this section, we have highlighted points pertinent to our study, making it easier for the readers to grasp the prior work and the key challenges in this field.
3. All the reviewers pointed out issues with the images. We have redrawn all the images to ensure text clarity, added captions, and updated the corresponding citations in the manuscript to aid reader comprehension.
4. Minor typos in the equations have been corrected, and we have provided explanations for previously undefined concepts, such as physical memory and virtual memory.
5. We have moved some improvements related to Megatron to the appendix due to the extensive amount of content in the main text and its lower relevance to the core contributions of this paper.
6. Clarifications on the specific benchmarks used have been added. Our benchmarks encompass knowledge, reasoning, reading comprehension, and comprehensive ability.
7. A new comparative experiment between UltraMem and MoE under varying sparsity parameters has been added. The results indicate that, with continuous increases in sparse parameters while keeping activation parameters constant, UltraMem's inference time does not experience significant growth, whereas MoE's inference time grows exponentially.
8. We have added three sets of ablation experiments to assess the impact of IVE, TDQKR, and MCS under different hyperparameters, demonstrating the rationale behind our choices of E=4, r=2, and h=2.
In addition to submitting the revised manuscript, we will also provide individualized responses to each reviewer to address their specific concerns and questions. We hope that these revisions and responses will clarify any ambiguities, address all the reviewers' concerns, and lead to reconsidering increasing the rating for our manuscript.


**References**

[1] Jegou H, Douze M, Schmid C. Product quantization for nearest neighbor search. IEEE transactions on pattern analysis and machine intelligence, 2010.

[2] Roller S, Sukhbaatar S, Weston J. Hash layers for large sparse models. NeurIPS, 2021.

[3] Geva M, Schuster R, Berant J, et al. Transformer feed-forward layers are key-value memories. EMNLP, 2021.

[4] Chi Z, Dong L, Huang S, et al. On the representation collapse of sparse mixture of experts. NeurIPS, 2022.

[5] Zhou Y, Lei T, Liu H, et al. Mixture-of-experts with expert choice routing. NeurIPS, 2022.

[6] Krajewski J, Ludziejewski J, Adamczewski K, et al. Scaling laws for fine-grained mixture of experts. ICLR, 2024.

[7] Bershatsky D, Cherniuk D, Daulbaev T, et al. LoTR: Low tensor rank weight adaptation. arXiv, 2024.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732248072628

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Many thanks for your comments and acknowledgment of our work. If you have any additional comments or suggestions, we sincerely look forward to the opportunity to engage in further discussions with you and are more than happy to revise our manuscript further.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732779414698

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response 2

comment:
Thank you to all the reviewers for your invaluable comments and suggestions. Based on your feedback, we have made the following revisions to our paper:

1. We have conducted experiments for PKM-151M, 680M, and 1.6B, which serve as baselines for UltraMem. Currently, PKM-151M-x12 has completed its evaluation. PKM-680M-x12 is 30% through its training, and based on available results, we have predicted its performance at the end of training, which is now included in Table 1. PKM-1.6B-x12 is still in the early stages of training, and thus reliable results are not yet available; these will be updated in the final version of the paper. Table 6 has been supplemented with the hyperparameter configurations of the PKM models, and Figure 11 has been updated to include the current evaluation curves for PKM.

2. We have conducted additional ablation studies, specifically implementing various modifications independently on the PKM structure to assess the benefits of each change. We are currently unable to compile the results, but they will be added in the coming days and reflected in the final version of the paper.

Once again, we thank the reviewers for your valuable suggestions. These iterations have made the conclusions of Ultramem more robust, and we hope the revisions will prompt reconsideration of our score.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733298018610

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Final Author General Response

comment:
Dear Reviewers and Area Chairs,

We sincerely appreciate the attention to detail and valuable suggestions offered by the reviewers throughout the review process, as well as the guidance provided by the Area Chairs. As we approach the end of the rebuttal phase, we would like to summarize the core contributions of our manuscript, "Ultra-Sparse Memory Network," and reflect on the interactions we've had concerning the feedback provided.

### **Main Contributions**

Our paper, "Ultra-Sparse Memory Network," introduces an ultra-sparse memory layer to reduce memory access costs during LLM inference, stands as a significant advancement. This not only preserves the performance levels of conventional models, such as Mixtures of Experts (MoE), but also surpasses them in terms of inference speed and **scalability**, while maintaining similar training speed compared to MoE. Our submission highlights the following key advancements:
1. **Efficient Memory Access Reduction**: We effectively demonstrate how UltraMem decreases the required memory access for model operations, leading to significant improvements in inference speeds.
2. **Maintenance of Performance While Scaling**: Similar to traditional dense models and existing sparse models like MoE, our method maintains competitive performance metrics even as the model scales.
3. **Superior Inference Speed and Scalability**: During ongoing scaling, UltraMem's performance matches or exceeds those of MoE, while maintaining a significant speed advantage in inference, ranging from 2 to 6 times faster.

These contributions were positively acknowledged as pioneering by the reviewers:
- "…an impressive approach to reducing computational load while maintaining strong model performance…" (Reviewer DbFB)
- "...real-world inference and memory access results on GPU hardware..." (Reviewer DbFB)
- "…a significant step forward in addressing LLM inference bottlenecks…" (Reviewer jFrZ)
- "...achieving comparable performance to MoE models of similar scale while maintaining greater memory efficiency..." (Reviewer ch2J)
- "...significantly outperforms existing approaches on large models, being up to six times faster than MOE at the same scale...." (Reviewer bPzv)
</content>
</thread_object>

</thread>


Metareview:
metareview:
The Ultra-Sparse Memory Network (UltraMem) paper presents a significant advancement in the efficiency and scalability of large language models (LLMs) by introducing an ultra-sparse memory layer. This approach addresses the memory access cost bottleneck in LLMs, particularly during inference, where it outperforms the Mixture of Experts (MoE) models in terms of speed, ranging from 2 to 6 times faster. The reviewers have noted several key strengths and weaknesses in the submission.

Strengths:
* Efficiency and Scalability: UltraMem demonstrates impressive improvements in inference speed and memory access efficiency while maintaining model performance. This is a critical advancement for practical applications of LLMs.
* Methodological Approach: The authors' methodology is detailed and thoughtful, with significant improvements to the baseline Product Key Memory (PKM) model. The paper's approach to clearly delineating changes to the baseline and the newly proposed method is commendable.
* Experimental Rigor: The submission includes extensive experiments, including comparisons with dense models and MoE, validating UltraMem's advantages across various tasks and model sizes.

Weaknesses:
* Presentation and Readability: The initial version of the paper was criticized for poor readability and a lack of cohesion. However, the authors have made substantial revisions to address these issues.
* Baseline Comparisons: The lack of direct comparisons with PKM in larger model sizes and the absence of comparisons with popular open-source models were noted as limitations. However, the authors have since initiated experiments to address these concerns.
* Ablation Studies: Some reviewers called for more comprehensive ablation studies to independently validate the effectiveness of each proposed modification. The authors have responded positively, committing to include these studies in the final version of the paper.

Based on the reviewers' feedback and the authors' responses, I recommend accepting the "Ultra-Sparse Memory Network" paper for presentation at the ICLR 2025 conference. The paper introduces a novel and impactful approach to improving LLM efficiency, backed by rigorous experimentation and a clear commitment to addressing reviewer feedback. The revisions have significantly improved the paper's presentation and addressed the primary concerns raised by the reviewers, positioning UltraMem as a valuable contribution to the field of large language model research.