Research map / Architectures and efficiency / Attention and transformers
Mean-Shift Feature Transformer
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 MSF-transformer, which replaces the transformer's attention-plus-projection block with a mean-shift (Gaussian-kernel KDE gradient ascent) feature transformation, plus an interleaved grouped projection to cut parameters.
Problem
Understanding and improving the feature transformation performed by multi-head attention in transformers, including its feature-norm bias and parameter cost of multi-head projections, for vision models.
Approach
Views the transformer's residual MHA update as a gradient ascent on a kernel density of input tokens; replacing the ill-posed exponential kernel with a Gaussian kernel yields a mean-shift update with an added PROBE projection, and interleaved grouped projections reduce parameters while improving back-propagation.
Key results
- MSF-transformer ViT-S reaches 79.79% top-1 on ImageNet-1K at 100 epochs versus 78.98% for the original ViT-S, with 23.68M vs 21.91M parameters.
- The 2-group interleaved MSF ViT-S reaches 79.55% with 20.14M parameters, outperforming the original transformer with fewer parameters.
- On ADE20K segmentation with SETR, ViT-S MSF improves mIoU from 41.81 to 44.01.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace original transformer modules with MSF-transformer | ViT-S · ImageNet-1K val | top-1 accuracy (%): 78.98 -> 79.79 | Helped |
| Use interleaved 2-group projection instead of full projection (ablation) | MSF-transformer · ImageNet-1K val, ViT-S | top-1 accuracy (%): 79.79 -> 79.55 | No clear effect |
| Use block-wise 2-group projection instead of full projection (ablation) | MSF-transformer · ImageNet-1K val, ViT-S | top-1 accuracy (%): 79.79 -> 78.87 | Hurt |
| Use a residual-based update instead of the mean-shift-constrained update (ablation) | MSF-transformer · ImageNet-1K val, ViT-S | test accuracy: 79.79 -> 79.28 | Hurt |
| Use the Gaussian-kernel original transformer instead of MSF-transformer (ablation) | ViT-S · ImageNet-1K val | test accuracy: 79.79 -> 79.31 | Hurt |
| Share one projection matrix across all MSF projections (ablation) | MSF-transformer · ImageNet-1K val, ViT-S | top-1 accuracy (%): 79.79 -> 61.13 | Hurt |
| Use two projection matrices with mean-shift update (ablation) | MSF-transformer · ImageNet-1K val, ViT-S | top-1 accuracy (%): 71.34 -> 77.16 | Helped |
| Replace mean-shift update with the original transformer formulation (ablation) | ViT-S · ImageNet-1K val | top-1 accuracy (%): 77.16 -> 76.25 | Hurt |
| Use interleaved 2-group projection instead of full projection (ablation) | Original transformer · ImageNet-1K val, ViT-S | top-1 accuracy (%): 78.98 -> 78.57 | Hurt |
| Use block-wise 2-group projection instead of full projection (ablation) | Original transformer · ImageNet-1K val, ViT-S | top-1 accuracy (%): 78.98 -> 78.05 | Hurt |
| Use MSF-transformer in place of the original transformer | ViT-S backbone · iNat2018 | accuracy (%): 66.84 -> 68.19 | Helped |
| Use MSF-transformer in place of the original transformer | ViT-S backbone · iNat2019 | accuracy (%): 73.45 -> 74.60 | Helped |
Surprising findings
- Sharing a single projection matrix to form a two-matrix mean-shift update (18.37M params) outperforms ResNet-50 (25.50M params) at 75.97%.
- Tying QUERY and KEY in the mean-shift model degrades accuracy, because Gaussian-kernel self-attention collapses toward maximal self-weights.
Limitations
- Gains are modest (around 0.5 to 1 point top-1 on ImageNet) relative to the original transformer.
- Experiments are limited to image classification with ViT/Swin backbones plus one segmentation and two fine-grained tasks.
- Block-wise grouping degrades performance, and tying QUERY and KEY hurts accuracy.
Experimental setup
Datasets: ImageNet-1K iNaturalist 2018 iNaturalist 2019 ADE20K
Benchmarks: ImageNet-1K iNaturalist 2018 iNaturalist 2019 ADE20K
Metrics: top-1 accuracy mIoU parameter count FLOPs
Baselines: ViT-Ti ViT-SS ViT-S ViT-B Swin-T Swin-S ResNet-50 ResNet-101 ResNeXt-50 ResNeXt-101 EfficientNet-B0 SETR
Builds on: Transformer ViT Swin Transformer mean-shift interleaved group convolution
Closest papers by meaning
- Token Transformation Matters: Towards Faithful Post-hoc Explanation for Vision Transformer (CVPR 2024 · similarity 0.84)
- Multi-criteria Token Fusion with One-step-ahead Attention for Efficient Vision Transformers (CVPR 2024 · similarity 0.84)
- LookHere: Vision Transformers with Directed Attention Generalize and Extrapolate (NeurIPS 2024 · similarity 0.84)
- GTA: A Geometry-Aware Attention Mechanism for Multi-View Transformers (ICLR 2024 · similarity 0.83)
- Making Vision Transformers Truly Shift-Equivariant (CVPR 2024 · similarity 0.83)
- ImagePiece: Content-aware Re-tokenization for Efficient Image Recognition (AAAI 2025 · similarity 0.83)
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.