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

TL;DR

This paper explores replacing complex self-attention in transformers with simpler sequential modules by leveraging sparsity patterns, achieving efficient models with reduced parameters and latency.

Contribution

It introduces a layer-wise distillation framework that uses attention sparsity to effectively replace attention with simpler modules in pretrained vision transformers.

Findings

Sparser attention layers lead to smaller accuracy drops upon replacement.

Attention sparsity-guided distillation reduces the student-teacher gap.

Explicit attention sparsity improves efficiency in attention replacement.

Abstract

Self-attention serves as the core foundation of large-scale transformer pretraining, but its quadratic token interaction cost makes inference expensive. Replacing attention with simpler sequential modules is appealing, yet naive substitution is often lossy, especially at larger scales. This paper revisits attention replacement through the lens of sparsity. Based on the observation of diverse sparsity patterns across transformer layers, we posit that pretrained transformers decompose the complex token dependency across tokens into various sequence-to-sequence mappings of diverse complexities, where some layer functionalities can be approximated and replaced with much simpler sequential modules without loss. We evaluate this premise using a plug-and-play layer-wise distillation framework to approximate and replace attention functionalities in pretrained vision transformer models. Controlled group-wise replacements under a fixed training budget reveal a clear pattern: substituting layers with sparser attention incurs substantially smaller accuracy drops than replacing denser ones. We further impose explicit attention sparsity on the pretrained ViT via AViT-style token retention and perform sparsity-guided distillation for sequential replacing models, where we see increasing teacher sparsity consistently reduces the student-teacher gap. The proposed method achieves efficient attention replacement for reduced parameter size and latency through the guidance of attention sparsity.