Yanzhi Wang
2024
Rethinking Token Reduction for State Space Models
Zheng Zhan
|
Yushu Wu
|
Zhenglun Kong
|
Changdi Yang
|
Yifan Gong
|
Xuan Shen
|
Xue Lin
|
Pu Zhao
|
Yanzhi Wang
Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing
Recent advancements in State Space Models (SSMs) have attracted significant interest, particularly in models optimized for parallel training and handling long-range dependencies. Architectures like Mamba have scaled to billions of parameters with selective SSM. To facilitate broader applications using Mamba, exploring its efficiency is crucial. While token reduction techniques offer a straightforward post-training strategy, we find that applying existing methods directly to SSMs leads to substantial performance drops. Through insightful analysis, we identify the reasons for this failure and the limitations of current techniques. In response, we propose a tailored, unified post-training token reduction method for SSMs. Our approach integrates token importance and similarity, thus taking advantage of both pruning and merging, to devise a fine-grained intra-layer token reduction strategy. Extensive experiments show that our method improves the average accuracy by 5.7% to 13.1% on six benchmarks with Mamba-2 compared to existing methods, while significantly reducing computational demands and memory requirements.
Pruning Foundation Models for High Accuracy without Retraining
Pu Zhao
|
Fei Sun
|
Xuan Shen
|
Pinrui Yu
|
Zhenglun Kong
|
Yanzhi Wang
|
Xue Lin
Findings of the Association for Computational Linguistics: EMNLP 2024
Despite the superior performance, it is challenging to deploy large language models (LLMs) due to their massive parameters and computations. While pruning is a promising technique to reduce model size and accelerate the inference, the traditional pruning techniques can hardly be applied for LLMs as they need to finetune the model on the full dataset with multiple epochs consuming massive data and hardware resources. To deal with this problem, post-training pruning methods are proposed to prune LLMs in one-shot without retraining. However, their accuracy after pruning may suffer from certain performance degradation due to the lack of retraining with massive data. To address this issue, in this paper, we first formulate the post-training problem for layer-wise LLM compression to simultaneously prune multiple weights in LLMs. Next, we provide an optimal solution for this problem and design our post-training pruning algorithm for both unstructured and semi-structured sparsity. Our extensive experiments demonstrate the superior performance of the proposed methods in comparison to SOTA baselines across various LLM families including transformer-based LLMs and Mamba-based LLMs.
Search
Co-authors
- Zhenglun Kong 2
- Xuan Shen 2
- Xue Lin 2
- Pu Zhao 2
- Zheng Zhan 1
- show all...