Accelerating SAM2 with Efficient Memory Attention Module via Spatiotemporal Token Pruning

SAM, a foundation model for promptable image segmentation, has been widely implemented in various applications. SAM2 extends SAM to process video by integrating a memory attention module. The transformer-based module propagates temporal information via a cross-attention mechanism between the current frame's representations and memory features (stored outputs of previous frames). Given the extensive length of memory features, the quadratic complexity of the attention mechanism (with respect to spatial tokens) scales linearly with the number of stored frames. To accelerate inference time, we propose a spatiotemporal token pruning method that reduces the sequence length of memory tokens. Our motivation stems from two findings: high attention sparsity in 99% of memory tokens and high attention accumulation in the other 1% of these tokens. We formulate the pruning strategy based on two principal requisites of the video dense-prediction task: boundary delineation and temporal trackability. This principle guides the design of our method, where we leverage three scoring functions to form a reference matrix that evaluates the importance of memory tokens. We apply these scoring functions layer-wise, reflecting the transformer layer's tendency to propagate either spatial or temporal information. Our method is a non-retraining approach, yet it can effectively compress up to 87% of memory tokens, accelerating per-frame inference time by 1.7× while causing a negligible disturbance to the original model performance in several benchmark datasets.