Task-Induced Riemannian Metrics for
Vision Transformer Feature Spaces

Andrew Bond1, Ege Erdem Özlu1, Tuna Çimen1, Ilkin Umut Melanlioglu1, Tolga Birdal2, Erkut Erdem3, Aykut Erdem1
1Koç University  ·  2Imperial College London  ·  3Hacettepe University
Paper arXiv Code BibTeX
Overview of the task-induced pullback metric and how we access it.

The decoder's Jacobian defines a pullback metric on the ViT feature space. We never build the metric, and access it only through Jacobian-vector products.

Step 1 · Diagnose

κcap(r) measures how much of the decoder's sensitivity a rank-r metric can capture. It takes about 140 Jacobian-vector products and no training.

Step 2 · Learn

The Spectral Pullback Network learns the metric from randomized power iteration. A 310K-parameter head distills it into one score per token.

Step 3 · Use

The scores drive token merging and pruning in a frozen ViT. A first-order bound relates the error to the importance of the removed tokens.

Abstract

Methods that operate on ViT features compare them with Euclidean distance or cosine similarity. Both weight every feature direction equally. Task decoders do not.

The geometry a task induces is the pullback metric g = J⊤J, built from the Jacobian of the decoder's output with respect to the features. Storing g is infeasible at ViT scale, and for dense outputs such as depth maps even forming J is impractical. Whether a low-rank version of g can be learned depends on the model–decoder pair, and the matrix-free diagnostic κcap(r) measures it.

For tractable pairs, the Spectral Pullback Network learns a low-rank metric from randomized power iteration, and we distill it into a 310K-parameter importance head that reads only the features. When the spectrum is too spread out, a VAE bottleneck restores tractability. Across DPT, DINOv2, CLIP and VGGT, the diagnostic predicts which architecture is viable. The head reaches Spearman ρ = 0.998 on DINOv2 CLS, and geometric token pruning cuts the additional depth error of ToMe-based selection by 25% at prune ratio 0.5. The ViT stays frozen throughout.

Run the diagnostic before you train

The diagnostic sorts a backbone–decoder pair into one of four regimes. The regime determines which architecture to train.

Backbone / decoder κcap(20) CV reff Regime
DA V2 depth (DPT), L02 95.6% 0.026 34 Low rank, delocalized.
The metric fits, but its directions are spread over all tokens. Learning it needs cross-token attention.
DA V2 depth (DPT), L10 96.6% 0.026 20
DINOv2 CLS, L10 33.4% 4.06 261 Marginal κ, high CV.
No faithful low-rank metric, yet a few tokens carry the sensitivity, so per-token scoring works.
VGGT camera, L16 ≈ 100% N/A 9 Rank-bounded.
Nine outputs, so the metric is low rank by construction.
VGGT pointcloud, L08 96.6% 0.59 212 Rich spectrum.
Beyond any practical rank budget. A VAE bottleneck brings the effective rank down to 30–53.
VGGT depth, L16 76.2% 0.63 282

r = 20, m = 100 probes, q = 2, over 50 images. CV measures how unevenly the task-sensitive directions spread across tokens.

Token pruning results

Lower is better. On DINOv2 CLS, importance-guided merging stays closest to the unpruned embedding at every ratio, by more than an order of magnitude at low ratios.

On dense depth the ordering changes with the ratio. Importance beats ToMe scoring from η = 0.20 upward and is best of the three from η = 0.30. ToMe scoring wins at the lowest ratios.

Method η = 0.05 0.10 0.15 0.20 0.30 0.40 0.50
CLS cosine-distance degradation (× 100) · DINOv2-B/14
Random 0.170.36 0.600.86 1.492.31 3.44
ToMe score 0.670.78 0.951.12 1.752.56 3.82
Importance (ours) 0.0010.007 0.0240.060 0.250.80 2.09
DPT additional SILog (× 100) · Depth-Anything V2, single stage
Random 4.593.70 3.313.19 3.163.20 3.29
ToMe score 4.293.66 3.513.50 3.523.67 4.09
Importance (ours) 4.884.86 3.853.28 3.023.01 3.08

ImageNet-val at 224, mean over 3 seeds, best per column in bold. On NYU-Depth-V2, a 2-stage schedule beats ToMe scoring at every ratio and resolution by 25–35%.

What the importance head looks at

The head concentrates on depth discontinuities, such as object outlines and foreground–background transitions. ToMe's cosine scores highlight tokens that differ from their neighbours.

Token importance maps for DPT depth.

Rows: input, importance head, ToMe cosine scores, and the depth predicted by the frozen decoder. Each map is annotated with its Spearman ρ against the Jacobian-derived target. These four images are examples, and the held-out average on this task is about 0.20.

Findings

  • Tractability is structural. Whether a faithful low-rank task metric exists depends on the backbone, decoder and probe layer, and can be measured before any training.
  • Spatial concentration is not spectral concentration. DINOv2 CLS has a rich spectrum but concentrates its sensitivity on a few tokens, so per-token scoring works. Dense depth is the opposite case.
  • The geometric target makes the head trainable. Regressing plain Jacobian block norms instead drops the correlation from ρ ≈ 0.20 to ρ ≈ 0.09 on dense depth.

BibTeX

@article{bond2026taskinduced,
  author  = {Bond, Andrew and {\"O}zlu, Ege Erdem and {\c{C}}imen, Tuna and
             Melanlioglu, Ilkin Umut and Birdal, Tolga and Erdem, Erkut and Erdem, Aykut},
  title   = {Task-Induced Riemannian Metrics for Vision Transformer Feature Spaces},
  year    = {2026},
}