SuPER-Mamba improves long sequence model accuracy with surprisal routing
SuPER-Mamba: Surprisal Prioritization for Extended-Context Routing in Mamba
Computation and LanguageArtificial Intelligence
Summary
Long sequences can be hard for existing models to handle well. The authors created SuPER-Mamba, which uses a built-in signal called Surprisal to help decide how models process different parts of a long sequence. This method improves the model's accuracy without needing extra training. They tested it on multiple models and long-sequence benchmarks and found significant performance improvements.
What this means in practice
- •For natural language processing teams: Enhance long document understanding by improving the accuracy of pretrained long-sequence models without additional training or architecture changes.
- •For machine learning platform engineers: Integrate surprisal-based routing to optimize inference efficiency on extended-context models in deployed systems.
Authors
Zhetong Zhang, Gayan K Kulatilleke
Abstract
State space models provide efficient linear-complexity processing for long sequences, but their performance frequently degrades as context length grows. We propose SuPER-Mamba, a training-free framework that uses Surprisal as a model-intrinsic signal for long-range state routing in pretrained Mamba models. SuPER-Mamba computes position-wise Surprisal from next-token predictions, aggregates neighboring responses through Local Pooling, and uses the resulting scores to prioritize Global-Channel updates without query-dependent scoring or additional training. We evaluate SuPER-Mamba on three state-space models across the LongBench-E benchmark. The results show that SuPER-Mamba improves the average score from 8.06 to 15.69, corresponding to a 94.7% relative gain over the Base Models and a 13.1% improvement over the state-of-the-art prior baseline. Ablation and length-wise analyses further support the effectiveness of the proposed routing framework.