Learn To be Efficient: Build Structured Sparsity in Large Language Models

Summary

This paper proposes a method to increase the structured sparsity of models through training, called Learning-to-be-efficient (LTE). It introduces a new training loss that guides the model to activate fewer neurons while maintaining original performance. Simultaneously, LTE employs a threshold-based sigmoid routing strategy, allowing flexible expert selection instead of a predefined fixed number. To achieve acceleration, the authors further implement an efficient version based on Triton. Compared to previous sparse acceleration methods, LTE can be applied to activation functions other than ReLU, offering better generality.

Strengths

- The method demonstrates good generality and can be applied to various activation functions. - The designed separability loss is intuitive and effective. - Experiments on diverse types of tasks prove the effectiveness of the method. - Implementation of a Triton kernel achieves computational time acceleration.

Weaknesses

- Due to the mainstream dense models having an FFN to attention ratio of approximately 2:1, although good sparse acceleration is achieved in FFN, the overall speedup is not high. If experiments were conducted on models with a higher proportion of FFN, the acceleration effect would likely be better, potentially enhancing the impact of the work. - Some experimental and methodological details lack sufficient discussion. See the questions section for more details.

Questions

- Regarding the experimental analysis in Line 135: Since MoEfication approximates dense model computation, expert scores are only used for expert selection and do not scale the expert outputs. This ensures that when all experts are selected, the computation is entirely equivalent to the original dense model. Therefore, it should not be possible to have a situation where two experts are selected but only one contributes. Could the authors further elaborate on this phenomenon? - Concerning the expert grouping method, is there room for further improvement in clustering W_1? LLaMA uses GLU, where the intermediate representation is obtained by element-wise multiplication of two vectors, which is different from the vanilla FFN studied in the original MoEfication. - What principles guide the selection of thresholds, and is there any transferability to other LLMs?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes, the authors have adequately addressed the limitations.

Summary

The paper presents a new approach (LTE) aimed at improving the inference efficiency of large language models by developing structured activation sparsity. The method trains LLMs to activate fewer neurons in FFN layers while attempting to maintain task performance. The approach works by grouping neurons into experts and using a routing strategy (based on Sigmoid instead of Softmax) to select experts adaptively. The authors evaluate the method on RoBERTa, GPT2, and LLaMA models across various NLP tasks. They report that LTE outperforms existing baselines and provides FLOPs and a latency reduction by utilizing sparsity through a custom CUDA implementation.

Strengths

1. The paper presents a new method for inducing sparsity in LLMs. It is based on several MoE concepts but uses Sigmoid-based routing instead of the traditional Softmax. 2. The authors tested their method on multiple models, datasets, and task types. 3. The paper includes a custom CUDA kernel implementation to strengthen the applicability of the method.

Weaknesses

1. The two-step training increases the complexity of applying the approach 2. Dependency on multiple hyperparameters/thresholds. 3. The paper keeps claiming that existing methods focus on existing sparsity in pre-trained models, but inducing sparsity in training (even in LLMs) is not new and has been around for a while.

Questions

1. Since you are introducing a different training approach, how does the total training time compare to the baseline? 2. Can you provide more insights into the limitations of Softmax routing?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

1. Complexity overhead of the training. 2. Many of the presented insights are already in the MoE literature, except for the Sigmoid routing. 3. Limited number of baselines (using only two baselines, there is a lot of work about sparsity in LLMs). Imo, the paper would be much stronger if it compares against more (sparsity) baselines.

Summary

This article introduces a novel training algorithm, LTE, designed to train large language models (LLMs) to achieve more structured activation sparsity during inference. Thus, it enhances their efficiency without compromising performance until a very high sparsity.

Strengths

1. LTE performs excellently across all datasets: In multiple natural language understanding (NLU) tasks, LTE shows no significant performance degradation at high sparsity levels (80-95%) and maintains good performance even at sparsity levels exceeding 90% of FFN. 2. This article develops a CUDA kernel to speed up inference by reducing memory and computational overheads.

Weaknesses

1. Stage 1 of LTE will train all the model's parameters. How do you implement Dejavu? As far as I know, Dejavu is a post-training method that freezes the model's parameters. I'm not sure whether this is a fair comparison. 2. The LTE algorithm introduces a two-stage training process, which might be complex and computationally intensive to implement. This can be a barrier to practical adoption in resource-constrained environments.

Questions

1. Have you tried LTE also on the attention layers?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The article includes a limitation section.

Summary

This work aims to introduce structured sparsity to large language models (LLMs) to improve their execution efficiency. To achieve this, it enhances previous MoEfication methods by employing a sigmoid-based non-competitive routing function and a threshold-based expert selection, allowing for adaptive expert numbers. Experiments across various models and different language understanding and generation tasks validate the effectiveness of the proposed method.

Strengths

1. The paper is well-motivated and easy to follow. 2. The proposed method exhibits good soundness and can achieve real-device speed-up with the developed CUDA kernel. 3. The proposed method has been evaluated across both encoder and decoder language models, achieving a consistently improved accuracy-efficiency trade-off.

Weaknesses

1. The major concern is that the technical contribution and novelty of this work are somewhat limited. The use of MoE with sigmoid functions to avoid competition among experts has been adopted in previous works, such as [1][2], and an extension to MoEfication will intuitively work. The authors are expected to analyze the key differences that make the proposed method particularly suitable for MoEfication. [1] "Approximating Two-Layer Feedforward Networks for Efficient Transformers," R. Csordás et al., EMNLP'23. [2] "SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention," R. Csordás et al., arXiv'23. 2. One missing baseline is structured weight pruning [3][4]. Considering this work mainly targets a per-task fine-tuning setting, structured weight pruning can already achieve decent sparsity with non-trivial generation speed-up without performance dropping [3][4]. This comparison can inform the community whether weight sparsity or context sparsity is more cost-effective or if they can be applied together. [3] "A Simple and Effective Pruning Approach for Large Language Models," M. Sun et al., ICLR'24. [4] "Fluctuation-based Adaptive Structured Pruning for Large Language Models," Y. An et al., AAAI'24. 3. Task-specific fine-tuning for each task is too costly. The authors are highly encouraged to perform continued pretraining and evaluate across different tasks to validate the generalization capability of the proposed method. 4. I wonder how the task-specific fine-tuning is performed for the baseline methods. Specifically, is Dejavu/MoEfication performed on top of the fine-tuned model, or is the model fine-tuned after applying these techniques, or are there any smarter strategies? This question arises because the proposed method simultaneously updates both model weights and expert selection strategies, which may be a key reason why it outperforms the baselines.

Questions

My questions have been included in the weakness section. I'm willing to adjust my scores if my concerns are properly addressed.

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

This work does not suffer from notable negative societal impacts as it aims to improve the efficiency of LLMs.

Official Comment by Reviewer Jye3

Thanks for the response of the authors. Most of my questions have been addressed. I will raise my score.

Thanks for the Response

We thank the reviewer for the response. We hope these explanations and discussions have answered your questions. Thanks again for your insightful and constructive comments, which indeed help improve our paper. We are happy to answer further questions if you have any in the future.