SuPER-Mamba: Surprisal Prioritization for Extended-Context Routing in Mamba

Zhetong Zhang; Gayan K Kulatilleke

cs.CLOct 12, 2026

SuPER-Mamba: Surprisal Prioritization for Extended-Context Routing in Mamba

Authors: Zhetong Zhang, Gayan K Kulatilleke

Organizations: School of Information Technology and Electrical Engineering The University of Queensland

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.

Figures & tables

Figure 1: Motivation of SuPER-Mamba. A pretrained Mamba produces position-wise Surprisal from next-token prediction probabilities while processing the input context. Relatively predictable local content typically produces lower and more stable Surprisal responses, whereas less predictable local changes may produce larger responses. SuPER-Mamba treats these responses as model-intrinsic predictive feedback and further investigates the relationship between such local predictive feedback and signals for long-range state allocation.

Figure 2: Overview of the SuPER-Mamba framework. Phase 1 performs offline Context Scoring. (a) Surprisal Calculation computes position-wise Raw Surprisal Scores from the next-token prediction probabilities of the pretrained Mamba. (b) Local Pooling aggregates neighboring predictive responses using category-specific Mean or Max Pooling to obtain Pooled Surprisal Scores. Phase 2 uses these scores for long-range state routing. The historical Top- $K$ Content positions selected by the Pooled Surprisal Scores are combined with the most recent $R=32$ Content positions to form the final routing set. For long-range state updates, Instruction and Query positions do not participate in Content Selection and retain their original state updates, while among Content positions, only the selected positions update the predefined Global Channels. All Content positions continue to participate in the original recurrent computation, so the complete input is neither removed nor truncated.

Category	Dataset	Base Mamba2	LongMamba	LAMB	SuPER-Mamba (Ours)
Single-doc QA	Qasper	1.76	4.75	5.63	5.00
Single-doc QA	MultiFieldQA-en	5.03	11.12	14.76	13.81
Multi-doc QA	HotpotQA	2.89	5.45	7.71	6.43
Multi-doc QA	2WikiMQA	2.31	5.19	8.01	10.09
Summarization	GovReport	7.66	14.67	16.38	13.70
Summarization	MultiNews	11.31	14.72	17.49	14.45

Table 1: Main LongBench-E results on Mamba2-1.3B. Results are grouped by task category. Different datasets are evaluated using their official LongBench-E task-specific metrics, including QA F1 for Qasper, MultiFieldQA-en, HotpotQA, 2WikiMQA, and TriviaQA, ROUGE-L F1 for GovReport, MultiNews, and SAMSum, classification score for TREC, count score for Passage Count, retrieval score for Passage Retrieval-en, and code similarity for LCC and RepoBench-P. Each dataset score is the arithmetic mean over the 0–4k, 4–8k, and 8k+ context-length ranges. The Average row reports the arithmetic mean over all 13 dataset-level scores. Best results are shown in bold , and second-best results are underlined .

Category	$K$	Kernel	Pooling
Single-doc QA	1024	18	Max
Multi-doc QA	1024	10	Mean
Summarization	2048	18	Mean
Few-shot	1024	18	Max
Synthetic	256	18	Mean
Code	1024	12	Mean

Table 2: Category routing configurations used in the full LongBench-E evaluation.

K

denotes the number of historical Content positions selected based on the Pooled Surprisal Scores and allowed to update the Global Channels. Kernel denotes the Local Pooling window size

k

, and Pooling denotes Mean or Max aggregation. All task categories additionally retain the most recent

R=32

Content positions.

Architecture	Base Model	LAMB	SuPER-Mamba
Mamba1-1.4B	8.52	12.98	17.93
Mamba2-1.3B	8.09	21.30	21.56
Zamba2-1.2B	7.57	7.32	7.59
Average	8.06	13.87	15.69

Table 3: Overall LongBench-E results of SuPER-Mamba on Mamba1-1.4B, Mamba2-1.3B, and Zamba2-1.2B. Each dataset is evaluated using its official LongBench-E task-specific metric, and the results are equally averaged across the three context-length ranges. The Average row reports the arithmetic mean of the overall scores of the three models.

Variant	Average Score
Base Model	8.09
Random Control	19.33
Raw Surprisal	16.33
+ Local Pooling	20.85
Full SuPER-Mamba	21.56

Table 4: SuPER-Mamba component ablation results on Mamba2-1.3B. All results are evaluated on the complete set of 13 LongBench-E datasets using their official task-specific metrics. Each dataset score is obtained by equally averaging the results over the three context-length ranges.

Panel of Figure 3 — Figure 3: Hyperparameter sensitivity analysis on a LongBench-E subset using Mamba2-1.3B. The horizontal axis represents the six task categories, and the vertical axis reports the corresponding category-level LongBench-E score obtained by equally averaging the results of the datasets within each category. Each configuration is evaluated on 50 samples from each of the 13 datasets. (a) analyzes the routing budget $K$ with $k=18$ and Mean Pooling fixed, testing $K\in\{256,512,768,1024,1536,2048,3072\}$ . (b) analyzes the Local Pooling kernel $k$ with $K=1024$ and Mean Pooling fixed, testing $k\in\{3,5,7,9,12,15,18,21\}$ . Boxes summarize the result distributions across all tested settings for each category, gray diamonds indicate the common reference configuration with $K=1024$ , $k=18$ , and Mean Pooling, and red circles indicate the best observed setting for each category in the corresponding parameter analysis.

Method	0 to 4k	4 to 8k	8k+
Base Model	14.98	5.45	3.85
LAMB	22.41	22.32	19.15
SuPER-Mamba	23.64	21.18	19.85

Table 5: LongBench-E results of Mamba2-1.3B across different context-length ranges. The 0–4k, 4–8k, and 8k+ columns correspond to the three LongBench-E context-length ranges. Each column uses the official task-specific metric for each dataset and reports the equal average across all 13 datasets.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Category	$K$	Kernel	Pooling
Qasper	Single-doc QA	1024	18	Max
MultiFieldQA-en	Single-doc QA	1024	18	Max
HotpotQA	Multi-doc QA	1024	10	Mean
2WikiMQA	Multi-doc QA	1024	10	Mean
GovReport	Summarization	2048	18	Mean
MultiNews	Summarization	2048	18	Mean

Appendix

Table 6: Complete mapping from LongBench-E datasets to task categories and their corresponding category routing configurations. Different categories use different routing parameter configurations while following the same routing mechanism during inference.

Category	$K$	Kernel $k$	Pooling
Single-doc QA	1024	18	Max
Multi-doc QA	1024	10	Mean
Summarization	2048	18	Mean
Few-shot	1024	18	Max
Synthetic	256	18	Mean
Code	1024	12	Mean

Appendix

Table 7: Final category routing configurations used for the full LongBench-E evaluation.

K

denotes the selection budget for historical Content positions, Kernel

k

denotes the Local Pooling window size, and Pooling denotes Mean or Max aggregation. These configurations are determined from the parameter analysis on a LongBench-E subset and remain fixed throughout the subsequent full evaluation.

Dataset	Random Control	Raw Surprisal	+ Local Pooling	Full SuPER-Mamba
Qasper	4.12	3.36	4.85	5.00
MultiFieldQA-en	13.00	11.70	14.40	13.81
HotpotQA	6.54	5.71	6.82	6.43
2WikiMQA	8.83	8.39	10.00	10.09
GovReport	12.58	10.65	13.80	13.70
MultiNews	12.89	11.85	14.31	14.45

Appendix

Table 8: Dataset-level component ablation results on LongBench-E using Mamba2-1.3B. Each dataset is evaluated using its official LongBench-E task-specific metric. Samples are divided into the 0–4k, 4–8k, and 8k+ context-length ranges, and the three range-level scores are averaged arithmetically to obtain the dataset-level score. The Average row reports the arithmetic mean over all 13 dataset-level scores. Bold values indicate the best result among the four component settings.

Variant	0–4k	4–8k	8k+	Average
Random Control	21.45	18.51	18.03	19.33
Raw Surprisal	19.17	15.87	13.94	16.33
+ Local Pooling	23.05	20.45	19.06	20.85
Full SuPER-Mamba	23.64	21.18	19.85	21.56

Appendix

Table 9: Length-wise component ablation results on LongBench-E using Mamba2-1.3B. The 0–4k, 4–8k, and 8k+ columns correspond to the three context-length ranges. Each range uses the official LongBench-E task-specific metric for each dataset and reports the equal average across all 13 datasets. Bold values indicate the best result in each column.

Long-Context Language Model Inference

Long-context language model inference makes processing and generating over lengthy prompts faster and more memory-efficient.

See all 224 papers

Mamba-3: Improved Sequence Modeling using State Space Principles
Mar 16, 2026Aakash Lahoti, Kevin Y. Li, Berlin Chen +5Sequence Modeling Efficient Language Model Inference
Carnegie Mellon University · Princeton University · Together AI +1
MaBERT:A Padding Safe Interleaved Transformer Mamba Hybrid Encoder for Efficient Extended Context Masked Language Modeling
Mar 3, 2026Jinwoong Kim, Sangjin ParkHybrid Transformer-State Space Model Architectures Long-Context Language Modeling
Graduate School of Industrial Data Engineering, Hanyang University, Seoul, Republic of Korea
Sessa: Selective State Space Attention
Apr 20, 2026Liubomyr HorbatkoState Space Models Long-Context Language Modeling

Mar 16, 2026J/K move · Enter open · S save

Mamba-3: Improved Sequence Modeling using State Space Principles

Aakash Lahoti, Kevin Y. Li, Berlin Chen, Caitlin Wang+4

Carnegie Mellon University · Princeton University · Together AI+1

Sequence Modeling Efficient Language Model Inference State Space Models State Tracking

Open paper AskPDF arXiv

Scaling inference-time compute has emerged as an important driver of LLM performance, making inference efficiency a central focus of model design alongside model quality. While the current Transformer-based models deliver strong model quality, their quadratic compute and linear memory make inference expensive. This has spurred the development of sub-quadratic models with reduced linear compute and constant memory requirements. However, many recent linear models trade off model quality and capability for algorithmic efficiency, failing on tasks such as state tracking. Moreover, their theoretically linear inference remains hardware-inefficient in practice. Guided by an inference-first perspective, we introduce three core methodological improvements inspired by the state space model (SSM) viewpoint of linear models. We combine: (1) a more expressive recurrence derived from SSM discretization, (2) a complex-valued state update rule that enables richer state tracking, and (3) a multi-input, multi-output (MIMO) formulation for better model performance without increasing decode latency. Together with architectural refinements, our Mamba-3 model achieves significant gains across retrieval, state-tracking, and downstream language modeling tasks. At the 1.5B scale, Mamba-3 improves average downstream accuracy by 0.6 percentage points compared to the next best model (Gated DeltaNet), with Mamba-3's MIMO variant further improving accuracy by another 1.2 points for a total 1.8 point gain. Across state-size experiments, Mamba-3 achieves comparable perplexity to Mamba-2 despite using half of its predecessor's state size. Our evaluations demonstrate Mamba-3's ability to advance the performance-efficiency Pareto frontier.

Abstract

Long-Context Language Model Inference

Explore similar work

Mamba-3: Improved Sequence Modeling using State Space Principles

MaBERT:A Padding Safe Interleaved Transformer Mamba Hybrid Encoder for Efficient Extended Context Masked Language Modeling

Sessa: Selective State Space Attention