From Sparsity to Simplicity: Enabling Simpler Sequential Replacements via Sparse Attention Distillation

Technical Summary: From Sparsity to Simplicity (S2S)

Problem Statement
Self-attention mechanisms are the foundation of large-scale transformer pretraining, enabling flexible token mixing and long-range dependency modeling. However, their quadratic complexity with respect to sequence length results in high inference latency and memory overhead. While replacing attention with simpler sequential modules (e.g., RNNs, State Space Models) is an attractive avenue for efficiency, naïve substitution often leads to significant accuracy degradation, particularly at larger scales. Existing efficiency approaches often modify the attention computation itself or prune tokens but retain attention as the core operator. This paper addresses the challenge of replacing attention layers entirely with sequential modules without substantial loss in performance.

Methodology
The paper proposes a framework called Sparsity to Simplicity (S2S), which posits that pretrained transformers decompose complex token dependencies into block-wise mappings of varying effective complexity. The authors hypothesize that layers exhibiting sparse attention patterns are inherently easier to approximate with simpler sequential modules.

The methodology consists of three core components:

1. Unified Layer-wise Attention Replacement:

   • The authors introduce a plug-and-play framework where specific attention layers in a pretrained Vision Transformer (ViT) are swapped with sequential token mixers (specifically BiMamba and BiLSTM) while preserving the original layer interface (LayerNorm, residual connections, and MLPs).
   • The sequential modules adopt a multihead bidirectional structure to match the non-causal, full-context nature of ViT attention.
   • Layer-aligned supervision is used during distillation: the output of the replaced sequential mixer is compared directly to the output of the original attention layer at the same depth using a mean squared error (MSE) loss.

2. Sparsity-Guided Distillation:

   • To control and study the impact of sparsity, the authors integrate A-ViT (Adaptive ViT), which uses halting scores to generate binary token retention masks.
   • Stage 1 (Distillation): The student model (with sequential replacements) is trained under the same layer-wise token masks generated by the teacher (A-ViT). This ensures that the student and teacher operate on identical active tokens, isolating the difficulty of approximating the token-mixing function from differences in token routing. The loss includes similarity loss, a halting alignment term, and classification loss.
   • Stage 2 (Finetuning): The student is finetuned using its own learned halting modules to generate masks, allowing it to operate autonomously while maintaining task performance.

3. Layer-Aware Sparse Replacement Strategy:

   • Instead of uniform replacement, the strategy targets specific layer groups based on their natural sparsity profiles. The authors replace only the deepest layers (which exhibit the highest natural sparsity) and apply the sparsity-guided distillation framework to these specific blocks.

Key Contributions

• The S2S Hypothesis: The paper establishes that attention sparsity is a signature of mapping complexity. Sparse blocks in pretrained transformers can be replaced by simpler sequential modules with minimal loss, whereas dense blocks are harder to approximate.
• Controlled Replacement Framework: A unified method for swapping attention with sequential modules (Mamba/LSTM) while maintaining interface compatibility and enabling layer-selective analysis.
• Explicit Sparsity Control: The introduction of a distillation pipeline where teacher-induced sparsity (via A-ViT masks) is explicitly aligned with the student, demonstrating that increasing teacher sparsity consistently reduces the student-teacher performance gap.
• Optimized Trade-off: A strategy that combines natural layer-wise sparsity with explicit sparsification to achieve better accuracy-efficiency trade-offs than sparse attention baselines.

Experimental Results
Experiments were conducted on ImageNet using DeiT backbones (Tiny, Small, Base) with BiMamba and BiLSTM as substitutes.

• Natural Sparsity and Replaceability: Visualizations of token interaction maps show that while shallow layers exhibit dense, global mixing, deeper layers become markedly sparser and more selective. Controlled experiments replacing different layer groups reveal that replacing deeper (sparser) layers incurs significantly smaller accuracy drops than replacing shallower (denser) layers.
• Effect of Explicit Sparsity: When the token retention ratio is reduced via A-ViT, the accuracy gap between the attention-based teacher and the sequential student consistently shrinks. This confirms that explicit sparsification simplifies the target mapping, making it easier for sequential modules to learn.
• Accuracy-Efficiency Trade-off: The proposed layer-aware sparse replacement (replacing only the deepest 3 layers with BiMamba under matched sparsity budgets) achieves:
  • Accuracy comparable to or slightly exceeding sparse attention baselines (A-ViT) at similar token retention rates.
  • Significant efficiency gains: Up to 1.71× estimated speedup at 384² resolution compared to the dense baseline, with higher token-mixing throughput than sparse attention models at larger token counts.

Significance and Claims
The paper claims that attention replacement should be treated as a structured, layer-dependent problem rather than a uniform architecture swap. The significance of the work lies in demonstrating that:

1. Sparsity is a signal for replaceability: Pretrained models naturally decompose complexity such that sparse layers are the most amenable to simplification.
2. Explicit sparsification aids distillation: Forcing sparsity on the teacher makes the learning task for the sequential student easier, reducing the performance gap.
3. Strategic replacement outperforms uniform approaches: By focusing replacement on naturally sparse layers and guiding the process with controlled sparsity, one can achieve a superior balance of accuracy and latency compared to simply pruning attention or replacing the entire model.

The authors conclude that this perspective encourages future research into hybrid models that leverage pretrained transformer knowledge while integrating simpler, more efficient sequence modules where the model's internal structure permits.