Self attention encoders such as Bidirectional Encoder Representations from Transformers(BERT) scale quadratically with sequence length, making long context modeling expensive. Linear time state space models, such as Mamba, are efficient; however, they show limitations in modeling global interactions and can suffer from padding induced state contamination. We propose MaBERT, a hybrid encoder that interleaves Transformer layers for global dependency modeling with Mamba layers for linear time state updates. This design alternates global contextual integration with fast state accumulation, enabling efficient training and inference on long inputs. To stabilize variable length batching, we introduce paddingsafe masking, which blocks state propagation through padded positions, and mask aware attention pooling, which aggregates information only from valid tokens. On GLUE, MaBERT achieves the best mean score on five of the eight tasks, with strong performance on the CoLA and sentence pair inference tasks. When extending the context from 512 to 4,096 tokens, MaBERT reduces training time and inference latency by 2.36x and 2.43x, respectively, relative to the average of encoder baselines, demonstrating a practical long context efficient encoder.
Figures & tables
Figure 1: Overall architecture of MaBERT.
Encoder pattern
CoLA
SST-2
MRPC
QQP
MNLI-m
MNLI-mm
QNLI
RTE
MMMMMMMMMMMM
0.401 ± 0.014
0.878 ± 0.007
0.805 ± 0.017
0.831 ± 0.003
0.745 ± 0.012
0.740 ± 0.017
0.796 ± 0.008
0.561 ± 0.032
TTTTTTTTTTTT
0.428 ± 0.020
0.891 ± 0.012
0.834 ± 0.015
0.843 ± 0.002
0.796 ± 0.015
0.802 ± 0.018
0.864 ± 0.007
0.575 ± 0.030
MTTMTTMTTMTT
0.555 ± 0.017
0.913 ± 0.008
0.823 ± 0.015
0.855 ± 0.003
0.804 ± 0.015
0.806 ± 0.019
0.868 ± 0.008
0.586 ± 0.030
TMTMTMTMTMTM
0.525 ± 0.017
0.896 ± 0.011
0.826 ± 0.017
0.862 ± 0.004
0.801 ± 0.014
0.802 ± 0.018
0.873 ± 0.006
0.587 ± 0.031
MTMTMTMTMTMT
0.573 ± 0.019
0.897 ± 0.010
0.797 ± 0.018
0.861 ± 0.003
0.803 ± 0.015
0.804 ± 0.016
0.859 ± 0.007
0.582 ± 0.031
TMMTMMTMMTMM
0.528 ± 0.014
0.903 ± 0.008
0.832 ± 0.016
0.863 ± 0.004
0.803 ± 0.014
0.802 ± 0.026
0.870 ± 0.009
0.595 ± 0.031
Table 1: GLUE benchmark scores across interleaved Transformer–Mamba encoder patterns. Models are pretrained with 10% of total steps. M and T denote Mamba and Transformer layers, respectively. CoLA uses Matthews correlation coefficient; SST-2, MNLI-m, MNLI-mm, QNLI, and RTE use accuracy; MRPC and QQP use F1.
Figure 2: Mask-aware attention pooling in MaBERT.
Model
CoLA
SST-2
MRPC
QQP
MNLI-m
MNLI-mm
QNLI
RTE
BERT
0.522 ± 0.017
0.912 ± 0.012
0.853 ± 0.017
0.856 ± 0.006
0.826 ± 0.014
0.829 ± 0.011
0.876 ± 0.012
0.618 ± 0.027
ALBERT
0.503 ± 0.018
0.920 ± 0.012
0.855 ± 0.017
0.857 ± 0.006
0.829 ± 0.012
0.832 ± 0.013
0.880 ± 0.013
0.618 ± 0.030
Longformer
0.534 ± 0.018
0.924 ± 0.012
0.863 ± 0.016
0.858 ± 0.006
0.830 ± 0.013
0.831 ± 0.014
0.882 ± 0.012
0.626 ± 0.031
BigBird
0.528 ± 0.016
0.926 ± 0.011
0.864 ± 0.015
0.857 ± 0.005
0.831 ± 0.014
0.832 ± 0.012
0.881 ± 0.013
0.624 ± 0.029
DeBERTa
0.617 ± 0.015
0.934 ± 0.013
0.862 ± 0.014
0.868 ± 0.004
0.838 ± 0.015
0.842 ± 0.013
0.886 ± 0.019
0.648 ± 0.034
MaBERT
0.676 ± 0.018
0.933 ± 0.010
0.869 ± 0.017
0.879 ± 0.005
0.835 ± 0.016
0.837 ± 0.017
0.893 ± 0.012
0.654 ± 0.033
Table 2: Performance comparison of baselines and the proposed MaBERT on the GLUE benchmark. Results are reported after full-budget pretraining (100% steps). The best result for each task is highlighted in bold.
Figure 3: Average GLUE score across pretraining budgets.
Figure 6
Model
CoLA
SST-2
MRPC
QQP
MNLI-m
MNLI-mm
QNLI
RTE
Full
0.676 ± 0.018
0.933 ± 0.010
0.869 ± 0.017
0.879 ± 0.005
0.835 ± 0.016
0.837 ± 0.017
0.893 ± 0.012
0.654 ± 0.033
PSM only
0.641 ± 0.022
0.918 ± 0.011
0.849 ± 0.019
0.863 ± 0.006
0.816 ± 0.017
0.818 ± 0.018
0.874 ± 0.013
0.638 ± 0.036
MAP only
0.661 ± 0.018
0.922 ± 0.031
0.841 ± 0.023
0.860 ± 0.004
0.819 ± 0.017
0.820 ± 0.027
0.878 ± 0.014
0.597 ± 0.078
None
0.596 ± 0.027
0.903 ± 0.013
0.841 ± 0.021
0.847 ± 0.007
0.803 ± 0.020
0.805 ± 0.021
0.855 ± 0.015
0.614 ± 0.041
Table 3: Integration ablation results of MaBERT on the GLUE benchmark. PSM denotes padding-safe masking and MAP denotes mask-aware attention pooling. In the PSM only and None variants, MAP is disabled and CLS pooling is used for prediction.
Figure 5: Efficiency and scalability across sequence lengths.
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Model
CoLA
SST-2
MRPC
QQP
MNLI-m
MNLI-mm
QNLI
RTE
BERT
0.419 ± 0.017
0.874 ± 0.013
0.821 ± 0.016
0.837 ± 0.006
0.780 ± 0.012
0.791 ± 0.013
0.846 ± 0.011
0.578 ± 0.028
ALBERT
0.388 ± 0.019
0.868 ± 0.013
0.814 ± 0.017
0.832 ± 0.007
0.772 ± 0.015
0.781 ± 0.014
0.842 ± 0.013
0.566 ± 0.032
Longformer
0.401 ± 0.020
0.892 ± 0.014
0.820 ± 0.016
0.834 ± 0.008
0.776 ± 0.012
0.787 ± 0.015
0.858 ± 0.011
0.573 ± 0.033
BigBird
0.406 ± 0.019
0.897 ± 0.012
0.822 ± 0.016
0.835 ± 0.005
0.781 ± 0.014
0.792 ± 0.014
0.862 ± 0.012
0.577 ± 0.031
DeBERTa
0.423 ± 0.018
0.881 ± 0.011
0.817 ± 0.013
0.829 ± 0.004
0.789 ± 0.015
0.801 ± 0.014
0.829 ± 0.012
0.569 ± 0.036
MaBERT
0.574 ± 0.016
0.904 ± 0.009
0.837 ± 0.016
0.868 ± 0.003
0.809 ± 0.014
0.814 ± 0.015
0.867 ± 0.007
0.602 ± 0.030
Appendix
Table B1: GLUE results after pretraining with 10% of the total steps (mean ± standard deviation over five seeds; best per task in bold).
Model
CoLA
SST-2
MRPC
QQP
MNLI-m
MNLI-mm
QNLI
RTE
BERT
0.452 ± 0.018
0.887 ± 0.012
0.833 ± 0.017
0.841 ± 0.005
0.796 ± 0.011
0.801 ± 0.021
0.851 ± 0.010
0.591 ± 0.027
ALBERT
0.438 ± 0.018
0.884 ± 0.012
0.828 ± 0.016
0.840 ± 0.006
0.792 ± 0.014
0.798 ± 0.020
0.848 ± 0.012
0.586 ± 0.029
Longformer
0.455 ± 0.019
0.906 ± 0.013
0.836 ± 0.015
0.844 ± 0.005
0.799 ± 0.013
0.804 ± 0.021
0.861 ± 0.011
0.592 ± 0.030
BigBird
0.462 ± 0.021
0.910 ± 0.011
0.839 ± 0.013
0.846 ± 0.006
0.803 ± 0.013
0.806 ± 0.024
0.863 ± 0.013
0.596 ± 0.032
DeBERTa
0.497 ± 0.017
0.908 ± 0.012
0.831 ± 0.014
0.852 ± 0.004
0.803 ± 0.014
0.808 ± 0.019
0.847 ± 0.011
0.603 ± 0.034
MaBERT
0.612 ± 0.015
0.917 ± 0.008
0.848 ± 0.015
0.873 ± 0.005
0.815 ± 0.013
0.816 ± 0.020
0.868 ± 0.008
0.612 ± 0.029
Appendix
Table B2: GLUE results after pretraining with 25% of the total steps (mean ± standard deviation over five seeds; best per task in bold).
Model
CoLA
SST-2
MRPC
QQP
MNLI-m
MNLI-mm
QNLI
RTE
BERT
0.515 ± 0.018
0.902 ± 0.010
0.845 ± 0.016
0.846 ± 0.005
0.807 ± 0.011
0.813 ± 0.012
0.859 ± 0.009
0.598 ± 0.030
ALBERT
0.490 ± 0.019
0.898 ± 0.012
0.840 ± 0.017
0.842 ± 0.006
0.803 ± 0.014
0.809 ± 0.013
0.856 ± 0.011
0.598 ± 0.031
Longformer
0.506 ± 0.020
0.918 ± 0.011
0.848 ± 0.016
0.844 ± 0.005
0.809 ± 0.013
0.815 ± 0.014
0.868 ± 0.010
0.606 ± 0.033
BigBird
0.514 ± 0.016
0.920 ± 0.012
0.852 ± 0.014
0.843 ± 0.006
0.812 ± 0.012
0.817 ± 0.018
0.869 ± 0.012
0.607 ± 0.030
DeBERTa
0.548 ± 0.018
0.918 ± 0.011
0.842 ± 0.015
0.862 ± 0.004
0.811 ± 0.013
0.820 ± 0.014
0.868 ± 0.010
0.619 ± 0.031
MaBERT
0.634 ± 0.020
0.924 ± 0.008
0.863 ± 0.018
0.877 ± 0.004
0.825 ± 0.011
0.828 ± 0.021
0.886 ± 0.007
0.624 ± 0.037
Appendix
Table B3: GLUE results after pretraining with 50% of the total steps (mean ± standard deviation over five seeds; best per task in bold).
Mode
CoLA
SST-2
MRPC
QQP
MNLI-m
MNLI-mm
QNLI
RTE
MAP
0.676 ± 0.018
0.933 ± 0.010
0.869 ± 0.017
0.879 ± 0.005
0.835 ± 0.016
0.837 ± 0.017
0.893 ± 0.012
0.654 ± 0.033
ATTN
0.648 ± 0.015
0.924 ± 0.025
0.857 ± 0.022
0.861 ± 0.011
0.823 ± 0.011
0.818 ± 0.023
0.888 ± 0.016
0.639 ± 0.036
CLS
0.661 ± 0.011
0.922 ± 0.031
0.841 ± 0.023
0.860 ± 0.004
0.819 ± 0.017
0.820 ± 0.027
0.878 ± 0.014
0.597 ± 0.078
MaskedMean
0.651 ± 0.018
0.924 ± 0.029
0.857 ± 0.021
0.866 ± 0.010
0.823 ± 0.018
0.817 ± 0.023
0.869 ± 0.013
0.631 ± 0.009
Appendix
Table B4: Effect of pooling choices on GLUE performance for MaBERT (mean ± standard deviation over five seeds). MAP: mask-aware attention pooling; ATTN: attention pooling; CLS: [CLS]-token pooling; MaskedMean: mean pooling over non-padding tokens.
Model
Params
Mem@128
Mem@512
Mem@1024
Mem@2048
Mem@4096
BERT
109.484
441.364
454.748
471.274
507.313
579.392
ALBERT
11.685
70.649
90.971
117.992
173.382
283.109
Longformer
148.661
602.629
652.634
709.375
828.660
1066.699
BigBird
128.061
507.797
551.183
662.196
1115.224
2885.278
DeBERTa
184.417
742.110
825.994
1047.506
1908.529
5262.076
MaBERT
205.322
808.787
841.846
923.875
1238.917
2445.003
Appendix
Table B5: Model complexity and peak GPU memory footprint (MB) measured during a forward pass at each sequence length (lowest per column in bold; applied to Mem@* columns).
Model
Train@128 × 32
Train@512 × 32
Train@1024 × 32
Train@2048 × 32
Train@4096 × 32
BERT
42.372
189.501
487.204
1451.200
4894.397
ALBERT
50.941
225.783
545.223
1491.872
4662.279
Longformer
834.994
838.428
1709.029
3466.178
6985.296
BigBird
60.710
263.221
658.337
1993.710
6430.660
DeBERTa
74.079
371.969
1151.794
4251.722
16136.670
MaBERT
64.117
240.334
516.242
1240.849
3319.429
Appendix
Table B6: Training runtime (ms) per optimizer step with effective batch size fixed (lowest per column in bold). Positional extension to 4,096 follows Table A .
Model
Infer@128 × 1
Infer@512 × 1
Infer@1024 × 1
Infer@2048 × 1
Infer@4096 × 1
BERT
5.220
5.331
5.702
15.122
43.424
ALBERT
6.140
6.217
6.692
17.994
48.749
Longformer
36.794
36.866
38.420
42.468
55.676
BigBird
7.921
7.948
8.059
23.378
67.887
DeBERTa
13.737
14.038
17.456
57.592
206.592
MaBERT
9.160
9.193
9.263
13.914
34.703
Appendix
Table B7: Inference latency (ms) per forward pass with batch size 1 (lowest per column in bold). Positional extension to 4,096 follows Table A .
Linear attention transformers have become a strong alternative to softmax attention due to their efficiency. However, linear attention tends to be less expressive and results in reduced accuracy compared to softmax attention. To bridge the accuracy gap between softmax attention and linear attention, we manipulate Mamba-2, a very strong linear attention variant. We first simplify Mamba-2 down to its most fundamental and important components, evaluating which specific choices make it most accurate. From this simplified Mamba variant (Mamba-2S), we improve the A-mask and increase the order of the hidden state, resulting in a method, which we call 2Mamba, that is nearly as accurate as softmax attention, yet much more memory efficient for long context lengths. We also investigate elements to Mamba-2 that help surpass softmax attention accuracy. Code is provided for all our experiments.
Gabriel Mongaras, Eric C. Larson
Lyle School of Engineering Southern Methodist University Dallas, TX 75205
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.
Zhetong Zhang, Gayan K Kulatilleke
School of Information Technology and Electrical Engineering The University of Queensland
Balancing fine-grained local modeling with long-range dependency capture under computational constraints remains a central challenge in sequence modeling. While Transformers provide strong token mixing, they suffer from quadratic complexity, whereas Mamba-style selective state-space models (SSMs) scale linearly but often struggle to capture high-rank and synchronous global interactions. We present a consistency boundary analysis that characterizes when diagonal short-memory SSMs can approximate causal attention and identifies structural gaps that remain. Motivated by this analysis, we propose InfoMamba, an attention-free hybrid architecture. InfoMamba replaces token-level self-attention with a concept bottleneck linear filtering layer that serves as a minimal-bandwidth global interface and integrates it with a selective recurrent stream through information-maximizing fusion (IMF). IMF dynamically injects global context into the SSM dynamics and encourages complementary information usage through a mutual-information-inspired objective. Extensive experiments on classification, dense prediction, and non-vision tasks show that InfoMamba consistently outperforms strong Transformer and SSM baselines, achieving competitive accuracy-efficiency trade-offs while maintaining near-linear scaling.
Youjin Wang, Jiaqiao Zhao, Rong Fu +5
Renmin University of China, Beijing, China · University of Macau, Macau SAR, China · Central South University, Changsha, China +1