Task aware geometry improves vision transformer feature analysis
Task-Induced Riemannian Metrics for Vision Transformer Feature Spaces
Computer Vision and Pattern RecognitionMachine Learning
Summary
Vision transformers create complex features to interpret images, but typical methods assume all directions in their feature space are equally important, which may not be true. The authors introduce a way to measure and learn the real geometry of these features related to specific tasks without storing huge amounts of data. They also develop tools to approximate and predict important parts of these features, which helps improve tasks like depth estimation from images. Their method works across different vision transformer models and helps reduce errors without needing to retrain the original models.
What this means in practice
- •For computer vision engineers: Improve feature analysis in vision transformers for better accuracy in image-related tasks like depth estimation without retraining models.
- •For machine learning practitioners: Optimize token selection in transformer architectures to reduce computational load while maintaining quality, using learned task-induced metrics.
Authors
Andrew Bond, Ege Erdem Özlü, Tuna Çimen, Ilkin Umut Melanlioglu, Tolga Birdal, Erkut Erdem, Aykut Erdem
Abstract
Methods operating on Vision Transformer (ViT) feature spaces typically rely on Euclidean distance or cosine similarity. This assumes that every direction is equally meaningful, but there is no reason to believe the true task geometry has this property. The task-sensitive geometry of the feature space is given by the pullback metric $g(F) = J(F)^\top J(F)$, where $J$ is the Jacobian of the decoder's output fed to a task-specific distance, with respect to the features. Storing the full $g$ is infeasible at modern scales, and for dense outputs such as depth maps even forming $J$ is impractical. We show that whether a low-rank approximation of this metric can be learned depends on the model-decoder pair, and we characterize this with a matrix-free diagnostic $κ_{cap}(r)$ computable with a low number of Jacobian-vector products. For tractable pairs, we develop the Spectral Pullback Network (SPN), which learns a low-rank version of the metric from randomized power iteration, and we distill it into a $310$K-parameter importance head that predicts token importance directly from the features. When the Jacobian spectrum is too spread out for a low-rank approximation, passing the decoder's input features through a VAE bottleneck can restore tractability. Across DPT, DINOv2, CLIP, and VGGT backbones, $κ_{cap}(r)$ predicts which learned-metric architectures are viable. The importance head reaches Spearman $ρ= 0.998$ on DINOv2 CLS, and our geometric token pruning reduces the additional depth error of ToMe-based token selection by $25\%$ on DPT depth at prune ratio $0.5$, without fine-tuning the ViT. Project page: https://cyberiada.github.io/TaskInducedViTs/