Learning Spatial Decay for Vision Transformers

TL;DR

This paper introduces Spatial Decay Transformer (SDT), a novel vision transformer model that employs a content-aware gating mechanism to generate dynamic, data-dependent spatial decay, improving attention mechanisms for spatially-structured vision tasks.

Contribution

It adapts content-aware gating from language models to vision transformers, introducing a unified framework for data-dependent spatial decay in 2D attention mechanisms.

Findings

Consistent performance improvements on ImageNet-1K classification.

Effective integration of spatial priors with learned content representations.

Establishment of data-dependent spatial decay as a new paradigm in vision transformers.

Abstract

Vision Transformers (ViTs) have revolutionized computer vision, yet their self-attention mechanism lacks explicit spatial inductive biases, leading to suboptimal performance on spatially-structured tasks. Existing approaches introduce data-independent spatial decay based on fixed distance metrics, applying uniform attention weighting regardless of image content and limiting adaptability to diverse visual scenarios. Inspired by recent advances in large language models where content-aware gating mechanisms (e.g., GLA, HGRN2, FOX) significantly outperform static alternatives, we present the first successful adaptation of data-dependent spatial decay to 2D vision transformers. We introduce \textbf{Spatial Decay Transformer (SDT)}, featuring a novel Context-Aware Gating (CAG) mechanism that generates dynamic, data-dependent decay for patch interactions. Our approach learns to modulate spatial attention based on both content relevance and spatial proximity. We address the fundamental challenge of 1D-to-2D adaptation through a unified spatial-content fusion framework that integrates manhattan distance-based spatial priors with learned content representations. Extensive experiments on ImageNet-1K classification and generation tasks demonstrate consistent improvements over strong baselines. Our work establishes data-dependent spatial decay as a new paradigm for enhancing spatial attention in vision transformers.