Rethinking Token Reduction for State Space Models (2024)

Abstract: 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.

## introduction there are growing research interests and efforts in ssms in recent years. building on the foundation laid by the kalman filter model @xcite , ssms have evolved to address long-range dependencies and are optimized for parallel training. several works @xcite @xcite have proposed ssmbased models capable of processing sequence data across a variety of tasks and modalities. a notable recent contribution, mamba @xcite , integrates time-varying parameters into ssms, allowing the model to selectively propagate or forget information. additionally, mamba introduces a hardware-aware parallel algorithm that accelerates both training and inference. unlike quadratic attention mechanisms, which become prohibitively expensive with longer sequence lengths, mamba's subquadratic-time architecture is more efficient and better suited for handling long sequences. the exceptional scaling performance of mamba underscores its potential as an effective alternative to the transformer model @xcite for generative language modeling tasks. in line with existing research efforts aimed at enhancing the efficiency of transformer models @xcite , exploring the efficiency of ssms is crucial for facilitating real-time applications. while weight pruning and quantization are prevalent techniques for optimizing transformer models @xcite @xcite , token reduction @xcite @xcite has proven effective in improving transformer efficiency due to the token length dimension or number of token is independent of the model architecture. given that ssm blocks also process input tokens similarly to transformer models, applying existing state-of-the-art (sota) token reduction techniques @xcite @xcite to ssms appears to be a straightforward post-training approach to enhance their efficiency, especially when scaling to billions of model parameters. this can achieve faster serving and lower peak memory usage, facilitating the wider deployment of large-scale ssms like mamba. however, as illustrated in figure 1 , this application of token reduction to ssms, while offering some benefits of faster inference with fewer tokens, results in significant performance drops. in this paper, after applying existing transformer token reduction techniques to ssms and observing their failures, we conduct an insightful analysis to understand the patterns and reasons for their failures on ssms. based on our analysis, we propose a unified post-training token reduction method for ssms to preserve performance and improve efficiency. we first employ a decoupling strategy that computes the importance of each token and classifies them into two sets: less important tokens and more important tokens. following this, we devise a fine-grained intra-layer token reduction strategy for the hidden states and residual connections of mamba. our approach uses a hybrid token reduction strategy (combining and taking advantages of pruning and merging) on hidden state tokens, meticulously designed to balance preserving essential information and eliminating redundancy. our unified strategy can be generalized to other model architectures like transformers. in summary, the main contributions of our work are as follows: token reduction. token reduction is an effective strategy to enhance computational efficiency by reducing the number of processed tokens or patches @xcite @xcite @xcite . it enables significant acceleration without requiring additional weights or specialized hardware, aiming to selectively retain the most informative tokens. several innovative approaches have been developed for transformers. for example, evit @xcite uses the attentiveness of the [cls] token with respect to other tokens to identify the most important tokens. dynamicvit @xcite and spvit @xcite add layers that employ the gumbel-softmax trick to selectively prune less informative tokens. agile-quant @xcite leverage the activation-aware token pruning technique to reduce the outliers for llms. tome @xcite measures dot product similarity between token keys to determine redundancy and merge accordingly. pumer @xcite proposed a token reduction framework for large-scale vlms with text-informed pruning and modalityaware merging strategies to progressively reduce the tokens of input image and text. however, the dynamics of information flow between tokens and the learning mechanisms in models like mamba @xcite ) remain largely unexplored. the absence of attention layers in mamba makes current token reduction methods ineffective. furthermore, the inclusion of the ssm module prevents the effective use of existing token reduction methods. ## methodology to tackle the problem, we first rethink the token importance metric for ssms. we then introduce a novel approach for unified token reduction by token importance classification that combines the advantages of both token pruning and token merging to facilitate faster and memory-efficient computation across ssm layers. ## conclusion in this paper, we introduced a unified post-training token reduction method for ssm architectures like mamba. we addressed the limitations of existing token reduction techniques by combining token importance and similarity to create a fine-grained reduction strategy. our method includes multiple design choices for effective intra-layer optimizations. experiments show significant reductions in computational demands and peak memory usage, while maintaining competitive accuracy, outperforming baseline methods on benchmarks.