Data Efficient Any Transformer-to-Mamba Distillation via Attention Bridge

TL;DR

CAB presents a data-efficient cross-architecture distillation framework that transfers attention priors from Transformer teachers to state-space Mamba students using a lightweight Attention Bridge. By aligning token-level projections via learnable B/C mappings and employing flexible layer-wise alignment, CAB achieves strong performance on vision and language tasks under limited data and exhibits faster convergence and favorable efficiency. Key findings include significant accuracy and perplexity improvements over baselines, improved cross-layer attention similarity in mid-to-deep layers, and substantial memory/runtime gains compared with full attention alignment methods. The work demonstrates a practical pathway to leverage Transformer expertise for efficient, scalable SSMs, enabling broader adoption of attention-informed sequence modeling in data-constrained settings.

Abstract

State-space models (SSMs) have emerged as efficient alternatives to Transformers for sequence modeling, offering superior scalability through recurrent structures. However, their training remains costly and the ecosystem around them is far less mature than that of Transformers. Moreover, the structural heterogeneity between SSMs and Transformers makes it challenging to efficiently distill knowledge from pretrained attention models. In this work, we propose Cross-architecture distillation via Attention Bridge (CAB), a novel data-efficient distillation framework that efficiently transfers attention knowledge from Transformer teachers to state-space student models. Unlike conventional knowledge distillation that transfers knowledge only at the output level, CAB enables token-level supervision via a lightweight bridge and flexible layer-wise alignment, improving both efficiency and transferability. We further introduce flexible layer-wise alignment strategies to accommodate architectural discrepancies between teacher and student. Extensive experiments across vision and language domains demonstrate that our method consistently improves the performance of state-space models, even under limited training data, outperforming both standard and cross-architecture distillation methods. Our findings suggest that attention-based knowledge can be efficiently transferred to recurrent models, enabling rapid utilization of Transformer expertise for building a stronger SSM community.

Figures (13)

• Figure 1: Towards Effective Attention-to-SSM Distillation. We highlight the structural complementarity between attention-based and SSM-based models, and the limitations of direct attention transfer, motivating our proposed alignment-based distillation approach.
• Figure 1: Architectural specifications of vision models.
• Figure 2: Overview of representative strategies for transferring attention-based inductive biases from Transformers to State Space Models (SSMs).
• Figure 3: Convergence speed and test loss comparison on ImageNet.
• Figure 4: Cosine similarity between Vim and pretrained ViT, comparing results with and without attention alignment. Higher similarity indicates better alignment of attention representations.