Research map / Architectures and efficiency / Attention and transformers
Visual Transformer with Differentiable Channel Selection: An Information Bottleneck Inspired Approach
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes DCS-Transformer, a transformer block with differentiable Gumbel-Softmax channel selection on attention queries/keys and MLP features, plus a separable variational upper bound on the information bottleneck loss, to compress visual transformers like MobileViT and EfficientViT.
Problem
Visual transformers are computationally expensive; existing channel pruning fixes channels in linear layers and ignores how channel choice affects attention affinity, which hurts accuracy after compression.
Approach
Select channels of queries and keys (via learned Gumbel-Softmax masks) to compute more informative attention weights, prune MLP input/output channels via differentiable selection, and directly minimize a separable variational upper bound on the information bottleneck loss during training.
Key results
- DCS-MobileViT-S reaches 81.0% top-1 on ImageNet-1k with 1.2 GFLOPs versus 78.4% for MobileViT-S at 1.4 GFLOPs.
- DCS-EfficientViT-B1 [r224] reaches 80.8% top-1 at 0.46 GFLOPs versus 79.4% at 0.52 GFLOPs for EfficientViT-B1.
- On COCO object detection with SSDLite, DCS-EfficientViT gets 29.0 mAP at 1.4 GFLOPs versus 28.4 for EfficientViT at 1.5 GFLOPs.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace MobileViT-XS transformer blocks with DCS-Transformer blocks | MobileViT-XS · ImageNet-1k val | top-1 accuracy: 74.8 -> 76.8 | Helped |
| replace MobileViT-S transformer blocks with DCS-Transformer blocks | MobileViT-S · ImageNet-1k val | top-1 accuracy: 78.4 -> 81.0 | Helped |
| replace EfficientViT-B1 [r224] transformer blocks with DCS-Transformer blocks | EfficientViT-B1 [r224] · ImageNet-1k val | top-1 accuracy: 79.4 -> 80.8 | Helped |
| replace EfficientViT-B1 [r288] transformer blocks with DCS-Transformer blocks | EfficientViT-B1 [r288] · ImageNet-1k val | top-1 accuracy: 80.4 -> 81.6 | Helped |
| replace ViT-S/16 transformer blocks with DCS-Transformer blocks | ViT-S/16 · ImageNet-1k val | top-1 accuracy: 81.2 -> 82.0 | Helped |
| replace Swin-T transformer blocks with DCS-Transformer blocks | Swin-T · ImageNet-1k val | top-1 accuracy: 81.3 -> 82.0 | Helped |
| use a DCS-EfficientViT-B1 backbone | EfficientViT-B1 · COCO val | mAPbox: 33.5 -> 34.8 | Helped |
| use a DCS-EfficientViT-B1 backbone | EfficientViT-B1 · COCO val | mAPm: 31.9 -> 33.2 | Helped |
| replace MobileViT-S with DCS-MobileViT-S as the SSDLite backbone | SSDLite · MS-COCO val | mAP@IoU 0.50:0.05:0.95: 27.7 -> 28.7 | Helped |
| train DCS-MobileViT-S without the IB loss (ablation) | DCS-MobileViT-S · ImageNet-1k val | top-1 accuracy: 81.0 -> 79.9 | Hurt |
| train DCS-MobileViT-XS without the IB loss (ablation) | DCS-MobileViT-XS · ImageNet-1k val | top-1 accuracy: 76.8 -> 75.7 | Hurt |
| train DCS-EfficientViT (r224) without the IB loss (ablation) | DCS-EfficientViT (r224) · ImageNet-1k val | top-1 accuracy: 80.8 -> 80.0 | Hurt |
Surprising findings
- Pruning attention outputs alone with channel selection in attention weights improves accuracy over the vanilla model (DCS-MobileViT-S w/o IB loss reaches 79.9% vs 78.4%).
Limitations
- Hyperparameters (lambda, eta, warm-up epochs) tuned by cross-validation, adding complexity.
- Search performed on a 100-class ImageNet subset rather than full ImageNet.
- Gains on detection and segmentation are reported only on a limited set of backbones.
Experimental setup
Datasets: ImageNet-1k COCO
Benchmarks: ImageNet-1k COCO
Metrics: top-1 accuracy FLOPs mAP APb APm IB loss
Baselines: MobileViT-S MobileViT-XS EfficientViT-B1 EViT ToMe SPViT VTC-LFC DeIT ViT-S/16 Swin-T
Builds on: MobileViT EfficientViT Gumbel-Softmax DARTS-style differentiable search Mask R-CNN
Closest papers by meaning
- Linear Differential Vision Transformer: Learning Visual Contrasts via Pairwise Differentials (NeurIPS 2025 · similarity 0.86)
- Vision Transformers Need Registers (ICLR 2024 · similarity 0.86)
- What Does It Mean to Be a Transformer? Insights from a Theoretical Hessian Analysis (ICLR 2025 · similarity 0.86)
- You Only Need Less Attention at Each Stage in Vision Transformers (CVPR 2024 · similarity 0.85)
- Simplifying Transformer Blocks (ICLR 2024 · similarity 0.85)
- Skip-Attention: Improving Vision Transformers by Paying Less Attention (ICLR 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.