# Task-Induced Riemannian Metrics for Vision Transformer Feature Spaces

arXiv: https://arxiv.org/abs/2609.27988 · HTML: https://mlsanity.com/paper/2609.27988/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 33 of its 52 references that are in the map.
- First posted on arXiv 23 Sep 2026.

## Builds on (most cited first, 30 of 33)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [The Unreasonable Effectiveness of Deep Features as a Perceptual Metric](https://mlsanity.com/paper/1801.03924/) (2018)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Axiomatic Attribution for Deep Networks](https://mlsanity.com/paper/1703.01365/) (2017)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Vision Transformers for Dense Prediction](https://mlsanity.com/paper/2103.13413/) (2021)
- [VGGT: Visual Geometry Grounded Transformer](https://mlsanity.com/paper/2503.11651/) (2025)
- [A Simple Unified Framework for Detecting Out-of-Distribution Samples and Adversarial Attacks](https://mlsanity.com/paper/1807.03888/) (2018)
- [Depth Anything V2](https://mlsanity.com/paper/2406.09414/) (2024)
- [Energy-based Out-of-distribution Detection](https://mlsanity.com/paper/2010.03759/) (2020)
- [SmoothGrad: removing noise by adding noise](https://mlsanity.com/paper/1706.03825/) (2017)
- [Token Merging: Your ViT But Faster](https://mlsanity.com/paper/2210.09461/) (2022)
- [Learning Robust Global Representations by Penalizing Local Predictive Power](https://mlsanity.com/paper/1905.13549/) (2019)
- [DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification](https://mlsanity.com/paper/2106.02034/) (2021)
- [Quantifying Attention Flow in Transformers](https://mlsanity.com/paper/2005.00928/) (2020)
- [Not All Patches are What You Need: Expediting Vision Transformers via Token Reorganizations](https://mlsanity.com/paper/2202.07800/) (2022)
- [AdaViT: Adaptive Tokens for Efficient Vision Transformer](https://mlsanity.com/paper/2112.07658/) (2021)
- [Flow Matching on General Geometries](https://mlsanity.com/paper/2302.03660/) (2023)
- [Analyzing the Structure of Attention in a Transformer Language Model](https://mlsanity.com/paper/1906.04284/) (2019)
- [Latent Space Oddity: on the Curvature of Deep Generative Models](https://mlsanity.com/paper/1710.11379/) (2017)
- [Understanding the Latent Space of Diffusion Models through the Lens of Riemannian Geometry](https://mlsanity.com/paper/2307.12868/) (2023)
- [The Power Spherical distribution](https://mlsanity.com/paper/2006.04437/) (2020)
- [Data Augmentation in High Dimensional Low Sample Size Setting Using a Geometry-Based Variational Autoencoder](https://mlsanity.com/paper/2105.00026/) (2021)
- [Riemannian Metric Learning via Optimal Transport](https://mlsanity.com/paper/2205.09244/) (2022)
- [Token Cropr: Faster ViTs for Quite a Few Tasks](https://mlsanity.com/paper/2412.00965/) (2024)
- [A Review on Riemannian Metric Learning: Closer to You than You Imagine](https://mlsanity.com/paper/2503.05321/) (2025)

## Abstract (arXiv)

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 kappacap(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 310K-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, kappacap(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/