Source-linked AI summary

Task-Induced Riemannian Metrics for Vision Transformer Feature Spaces

Andrew Bond, Ege Erdem Özlü, Tuna Çimen, Ilkin Umut Melanlioglu, Tolga Birdal, Erkut Erdem, Aykut Erdem

arXiv:2609.27988v1cs.CVcs.LG

TL;DR

Euclidean or cosine comparisons assume feature directions are equally meaningful, although task decoders respond differently to feature changes and dense-output Jacobians are impractical to form. The paper uses matrix-free diagnostics to select learned-metric architectures and reports token-importance and pruning results, while noting inference-time speed limits at standard resolution.

  • Problem

    Feature-space Euclidean or cosine comparisons treat directions as equally relevant, while storing the pullback metric is infeasible and dense-output Jacobians may be impractical to form.

    Source: §1 · p. 1
  • Method

    The paper estimates task sensitivity with a matrix-free κcap(r) diagnostic, learns low-rank metrics with the SPN, distills token importance into a small head, and uses a VAE bottleneck when needed.

    Source: §1 · p. 2, §6 · p. 8
  • Results

    Across evaluated backbone–decoder pairs, the diagnostic sorts cases into regimes that predict viable architectures; geometric pruning reduces ToMe-based DPT depth selection’s additional error by 25% at prune ratio 0.50.

    Source: §1 · p. 2, §4.3 · p. 5, §8 · p. 9
  • Takeaways & Limitations

    The diagnostic can guide architecture choice, and VAE compression can move otherwise intractable feature spaces into a tractable regime.

    Source: §8 · p. 9

    At 224 × 224, pruning is slightly slower in wall-clock time than the unpruned model despite removing 12.1% of theoretical FLOPs; speedup appears only at higher resolution and is modest.

    Source: M Limitations · p. 31

Abstract

from arXiv · show

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 $\kappa_{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, $\kappa_{cap}(r)$ predicts which learned-metric architectures are viable. The importance head reaches Spearman $\rho = 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/

1 Introduction

The paper replaces feature-space similarity assumptions with decoder-induced geometry, then asks when that geometry can be approximated and deployed efficiently. It proposes diagnostics and learned approximations, and reports gains for importance prediction and token pruning.

1 Introduction

  • The pullback metric measures how feature perturbations change task outputs, capturing directions that Euclidean distance or cosine similarity treat as equally relevant.The metric is the decoder-output Jacobian product JᵀJ.Source: §1 · pp. 1–2
  • The matrix-free κcap(r) diagnostic measures the fraction of decoder sensitivity captured by the top-r singular directions, while CV indicates whether that sensitivity is concentrated across tokens.Both measurements use Jacobian-vector products; the paper also gives theoretical results on rank-r sensitivity ceilings, per-token scoring, and power-iteration convergence.Source: §1 · p. 2
  • The Spectral Pullback Network learns a low-rank metric using randomized power iteration, and a roughly 310K-parameter importance head predicts per-token scores from features alone.A VAE bottleneck can make a low-rank approximation possible when the original feature-space spectrum is too spread out.Source: §1 · p. 2
  • For CLS decoders, geometric merging replaces cosine-based merging; for dense decoders, hard pruning is paired with Last-Layer Fusion to re-attach pruned tokens before decoding.A first-order error bound guides both approaches.Source: §1 · p. 2
  • Geometric pruning reduces ToMe-based token selection’s additional DPT depth error by 25% at prune ratio 0.50.ToMe remains competitive at the lowest ratios.Source: §1 · p. 2

2 Related Work

The paper positions its approach against metric-learning, pullback-geometry, token-importance, and pruning methods. Its distinction is to study decoder-induced geometry on frozen ViT features and use that geometry for token selection.

Metric learning and pullback geometry

  • Unlike methods that fit a Mahalanobis metric or output an SPD matrix, this work estimates a decoder-output pullback metric on frozen ViT features and represents it through its leading eigenspace.The paper identifies JᵀJ as the Gauss–Newton form of the decoder-output Fisher metric.Source: §2 · p. 2

Token importance in ViTs

  • Prior token-importance analyses use gradient saliency or attention proxies, whereas this work studies the full pullback JᵀJ because large attention or gradient magnitude alone does not distinguish task-sensitive directions..Source: §2 · p. 2

Token pruning and merging

  • Unlike cosine-based ToMe and learned or attention-derived pruning scores, this method measures token similarity with decoder-induced geometry; several prior pruning methods fine-tune the backbone, unlike this frozen-backbone setting.The paper also introduces theorems about when per-token scoring can succeed.Source: §2 · p. 3

3 Background and Problem Formulation

The paper defines a local feature-space geometry by composing the remaining ViT layers with a task decoder, then measuring how perturbations affect decoder outputs. It treats this geometry as a possibly singular sensitivity form and specifies the output metric as a modeling choice.

Problem formulation

  • A probe map composes the remaining ViT layers with a task decoder to map features at layer ℓ to a vectorized task output; its Jacobian defines the pullback metric.The examples include DPT depth, a CLS embedding, and VGGT camera pose.Source: §3 · p. 3
  • For a feature perturbation v, vᵀgv = ∥Jv∥² is the squared decoder-output change, so large values identify task-sensitive directions.A task-specific output distance can be composed with the probe map.Source: §3 · p. 3

What kind of metric is g?

  • Because the output dimension is smaller than the feature dimension, g is positive semidefinite and singular, with a kernel of feature perturbations ignored by the decoder; the deployed metric adds an isotropic floor ε.The paper makes no manifold or submanifold assumptions about the features.Source: §3 · p. 3

Choice of output metric

  • The Euclidean output metric is a modeling choice: any positive-definite output metric H can be incorporated by replacing J with LJ, and the VGGT camera test found unchanged tractability under an SO(3)-aware H.The two output metrics also produced token-importance rankings with Spearman ρ = 0.995.Source: §3 · p. 4

4 When is Metric Learning Tractable?

The section develops matrix-free measures of whether task sensitivity is low-rank and spatially concentrated, then uses them to characterize which metric-learning architectures suit different backbone–decoder pairs. The resulting regimes distinguish when cross-token modeling, per-token scoring, or feature compression is appropriate.

4 When is Metric Learning Tractable?

  • The diagnostic asks whether a fixed-rank metric can capture a pair’s task sensitivity and whether that sensitivity is concentrated on a few tokens or diffuse.Its measurements are designed to answer these questions before approximation or training.Source: §4 · p. 4

4.1 Captured Energy, Effective Rank, and Spatial Concentration

  • κcap(r) measures the fraction of task sensitivity captured by the top-r singular directions, while effective rank summarizes spectral spread and token-block CV measures spatial concentration.The top-r vectors reveal how task-sensitive energy is distributed across tokens.Source: §4.1 · p. 4
  • Low token CV marks delocalized directions, where per-token scoring is uninformative; higher CV indicates concentration on fewer tokens and makes per-token scorers viable.Theorem 1 characterizes the rank-r sensitivity ceiling, and Theorem 2 bounds the weakness of factored metrics on delocalized directions.Source: §4.1 · p. 4, §4.2 · p. 4

4.2 Structural Characterization of Tractability

  • Block power iteration has a geometric subspace-alignment convergence rate governed by (σr+1/σr)^2q, but the theorem’s constant can be large and flat spectra may require about 15 iterations.The authors increase q until the recovered subspace stabilizes in practice.Source: §4.2 · p. 5

4.3 Regime Taxonomy

  • The six measured cases fall into four regimes, each calling for a different architecture; two power-iteration steps suffice with a large spectral gap, while flatter spectra need more.Using several nearby feature layers does not itself make a decoder intractable.Source: §4.3 · p. 5
  • VGGT depth and pointcloud heads have effective ranks exceeding any practical rank budget by an order of magnitude; the VAE pathway of §6 brings VGGT depth effective rank down to [30, 53].Source: §4.3 · p. 5
  • Depth Anything V2 cases are low-rank but delocalized, so per-token metrics cannot represent their sensitivity and full metric learning needs cross-token attention.Per-token scores remain usable but weak for pruning, which also uses Last-Layer Fusion and multi-stage schedules.Source: §4.3 · p. 5
  • DINOv2 CLS has a rich spectrum that a rank-20 metric captures only partly, but high spatial concentration makes per-token scoring viable; its importance head reaches ρ = 0.998.CLIP CLS behaves similarly.Source: §4.3 · p. 5

5 Learning the Task-Induced Metric

The section develops a way to predict task-sensitive feature geometry without requiring the decoder’s Jacobian at inference. It combines a learned low-rank metric, randomized power-iteration supervision, and a distilled token-importance head.

5 Learning the Task-Induced Metric

  • The section seeks feature-only prediction of a task-induced metric so inference does not require producing decoder outputs and backpropagating to obtain the Jacobian.The paper notes that this Jacobian-based inference would lose the possible gains.Source: §5 · p. 6

5.1 Why Naive Supervision Fails

  • Distance matching leaves local geometry underdetermined, while random-direction regression poorly overlaps the rank-r subspace; power iteration amplifies alignment with task-sensitive directions.For a uniform direction, the stated overlap probability is r/ND ≈10−5.Source: §5.1 · p. 6

5.2 The Spectral Pullback Network (SPN)

  • The SPN predicts a rank-r-plus-isotropic metric from features, using orthonormal mode estimates and positive energy predictions for the top singular directions.Its inference-time metric-vector product costs O(NDr), rather than O((ND)2).Source: §5.2 · p. 6

5.3 RandNLA Training

  • Randomized power iteration supplies SPN supervision through Jacobian-vector and vector-Jacobian products without explicitly forming J.The method alternates these products and re-orthonormalizes between rounds; the reported convergence result is exponential in the number of rounds.Source: §5.3 · pp. 6–7

5.4 Importance Head and Inference-Time Distillation

  • The paper distills per-token importance into a roughly 310K-parameter head that predicts scores from features alone, with no Jacobian access or SPN call at inference.Its target uses the σ2-weighted top-r projection, which preserves spatial signal even when plain block norms are uniform across tokens.Source: §5.4 · p. 7

5.5 Empirical Validation

  • Spearman ρ = 0.998 ± 0.001 on DINOv2 CLS L10 measures the importance head’s rank correlation with its Jacobian-derived target; on DPT, the spectral target reaches ρ ≈0.20 versus ≈0.09 for plain block norms.The SPN also reaches its deterministic supervision ceiling on the VAE-compressed VGGT depth pipeline.Source: §5.5 · p. 7

6 Extending to Intractable Settings via VAE Compression

For VGGT depth and pointcloud, sensitivity is too spread out for a practical low-rank metric in the original feature space. The paper inserts a VAE bottleneck to reduce the effective spectrum and evaluates SPN training in the compressed pipeline.

Intractable settings

  • VGGT depth and pointcloud have effective ranks 282 and 212 at the decoder-head inputs, exceeding practical rank budgets and making a faithful original-space rank-r SPN infeasible.For dense outputs, the output dimension can be in the tens of thousands.Source: §6 · p. 8

VAE bottleneck

  • A pretrained VAE encodes VGGT’s concatenated decoder-input features into 32 register tokens of dimension 128, and its decoder reconstructs features for the frozen task head.Every path to the output passes through the 4096-dimensional latent, bounding the compressed Jacobian rank by 4096.Source: §6 · p. 8

Empirical validation

  • With VAE compression, VGGT depth’s effective rank falls from 282 to 30–53 and r0.90 is about 40, an approximately 6× reduction into the tractable regime.The VAE-compressed VGGT depth pipeline also reaches δ1 ≈0.99 reconstruction quality.Source: §6 · p. 8

Experimental validation

  • On VAE-compressed VGGT depth, one power-iteration step yields subspace alignment V1 = 0.713, close to the deterministic-supervision ceiling V1 = 0.710 and above the random model’s 0.005.The compressed pullback is an approximation through VAE-reconstructed features, expected to agree in task-sensitive directions when reconstruction preserves task-relevant content.Source: §6 · p. 8

7 Token Merging and Pruning via Geometric Importance

The section uses predicted task-sensitive importance to merge tokens for CLS decoders and hard-prune tokens for dense decoders. Its results compare these strategies with baselines while retaining the decoder’s output as the target of preservation.

7 Token Merging and Pruning via Geometric Importance

  • A first-order error bound motivates targeting low-sensitivity tokens, while the dense pipeline’s correction term depends on how much removed-token content is recovered.The bound guides both merging and pruning.Source: §7 · p. 8

7.1 CLS Decoders: Soft Merge

  • For CLS decoders, the method merges the ηN lowest-importance tokens into their nearest remaining tokens using importance-weighted averaging and replaces cosine matching with a single-token pullback distance.The experiments use the SPN prediction for the matching distance.Source: §7.1 · pp. 8–9

7.3 Results

  • 25% less additional error than ToMe scoring on DPT depth at η = 0.50 is reported for importance-guided pruning, while the two-stage schedule beats ToMe by 25–35% on NYU-Depth-V2 across ratios and resolutions.Importance-guided merging has the lowest degradation at every DINOv2 CLS ratio; the table caption states it outperforms ToMe at every CLS ratio and at η ≥0.2 for DPT.Source: Table 2 · p. 9, §7.3 · p. 9

7.2 Dense Decoders: Hard Prune + Last-Layer Fusion

  • For dense decoders, the method hard-prunes low-importance tokens, freezes their features, and uses Last-Layer Fusion to reinsert them before the final ViT block for updated attention.Last-Layer Fusion adds no parameters.Source: §7.2 · p. 9

7.3 Results

  • The reported pruning removes 12.1% of theoretical backbone FLOPs at η = 0.20, but a wall-clock speedup appears only at 448 × 448, where it is 1.12×.ImageNet top-1 accuracy remains between 0.66 and 0.68 for every method and ratio, versus 0.674 unpruned.Source: §7.3 · p. 9

8 Conclusion

The paper studies decoder-induced geometry on frozen ViT feature spaces and develops tools for diagnosing and learning low-rank approximations. Its conclusion reports that the diagnostic predicts viable architectures and that learned importance scores improve token pruning without fine-tuning the ViT.

8 Conclusion

  • ρ = 0.998 on DINOv2 CLS; geometric pruning reduces ToMe-based selection’s additional depth error by 25% at prune ratio 0.5.The backbone and decoder remain frozen.Source: §8 · p. 9
  • Across DPT, DINOv2, CLIP, and VGGT probes, κcap(r) and sensitivity spread sort backbone–decoder pairs into regimes that predict viable architectures; a VAE bottleneck can make intractable pairs tractable.Source: §8 · p. 9

A.3 Pullback Metrics

  • The pullback metric measures the squared output change caused by an infinitesimal feature perturbation, and directions in the Jacobian’s kernel leave the output unchanged to first order.Source: §A.3 · pp. 14–15

A.4 Specialization to ViT Probe Maps

  • The probe map composes the remaining ViT layers, task decoder, and optional task metric; its pullback on the feature space is positive semidefinite when the Jacobian has a nontrivial kernel.Source: §A.4 · p. 15
  • For VGGT camera pose, an SO(3)-aware block-scaled output metric leaves κcap at 0.99, yields importance rankings correlated at ρ = 0.995, and agrees on 97% of 50%-prune tokens.Source: §A.4 · p. 15

B.4 Algorithm Pseudocode

  • The matrix-free diagnostic uses a Jacobian operator and randomized SVD, with JVP/VJP products rather than explicitly forming the Jacobian.Source: §B.4 · p. 16

C.1 Theorem 1: Captured Sensitivity Ceiling

  • κcap(r) is both the quality of the rank-r metric and an upper bound on the sensitivity retained by any restriction to an r-dimensional subspace.Source: §C.1 · p. 17

C.2 Theorem 2: Impossibility of Factored Metrics

  • The proposed connection between attention entropy and singular-vector delocalization is motivational, not rigorous, because the needed regularity conditions are unstated.Source: §C.2 · p. 18

C.3 Theorem 3: Convergence of Power Iteration

Power-iteration convergence varies with the spectrum, while pruning outcomes depend on task sensitivity, residual drift, and decoder setting. The section connects these factors to pruning methods and reports where importance-guided choices outperform ToMe.

C.3 Theorem 3: Convergence of Power Iteration

  • Theorem 3 gives a geometric convergence rate for a fixed sketch, but at small q its constant can dominate the spectral-gap term, making the bound loose or vacuous.The paper chooses q by an offline check that increases it until the recovered subspace stabilizes.Source: §C.3 · p. 18
  • At q = 0, 2, 5, subspace-alignment errors are 0.198, 0.011, 0.001 for DPT depth, 0.617, 0.254, 0.083 for DINOv2 CLS, and 0.802, 0.438, 0.214 for CLIP CLS.Two steps suffice where the spectral gap is large, while flatter spectra converge more slowly; this is an empirical observation.Source: §C.3 · p. 18

C.4 Theorem 4: Token Pruning Error Bound

  • With roughly uniform residuals across pruned tokens, minimizing the first-order pruning-error term amounts to dropping tokens with the smallest importance; LLF reduces the residual term.The token importance score proxies the decoder-Jacobian block norm, while the pruning bound follows from a Taylor expansion around the modified final-layer tensor.Source: D Token Pruning: Details and Extended… · p. 19, §C.4 · p. 19

D Token Pruning: Details and Extended Discussion

  • For CLS decoders, task-induced distance weights token differences in decoder-sensitive directions; importance-weighted merging replaces cosine similarity in the ToMe-style rule.The per-token marginal metric is the rank-r pullback metric restricted to perturbations of that token alone.Source: D Token Pruning: Details and Extended… · p. 20
  • Dense decoders use hard pruning and reinsert pruned tokens; Last-Layer Fusion runs the final ViT block on the full sequence and reduces residual drift without adding parameters or training.The importance head is trained on single-layer probe maps and transfers to the full multi-layer DPT decoder.Source: D Token Pruning: Details and Extended… · p. 20
  • On CLIP-ViT-B/16 at η = 0.50, ToMe scoring degrades sharply to 33.21 against 0.76 for importance (Wilcoxon p < 10−50); on NYU, 2-stage pruning beats ToMe by 25–35% across three resolutions.Importance also retains a more than 100× advantage at low ratios on ImageNet-Sketch; ImageNet top-1 remains 0.66–0.68 across methods and ratios.Source: D Token Pruning: Details and Extended… · p. 20
  • Importance is best for CLS at all ratios; on ImageNet DPT, it beats ToMe from η = 0.20 and is best from η = 0.30, while 2-stage LLF beats ToMe across tested NYU ratios and resolutions.For single-stage NYU pruning, importance beats ToMe at η = 0.05–0.30 but is slightly worse at η = 0.50; 4-stage pruning fails at η ≥ 0.30.Source: D Token Pruning: Details and Extended… · p. 21
  • At 224 × 224 and η = 0.20, pruning removes 12.1% of theoretical FLOPs but runs at 0.96× wall-clock speed; at 448 × 448, 16.7% FLOPs saved yields 1.12× speedup.The importance head and LLF together cost about 2–3% more than ToMe scoring.Source: D Token Pruning: Details and Extended… · p. 21

E Full Sweep Results

The tractability sweep shows that low-rank feasibility and token localization vary across probes, while multi-layer behavior depends on how far apart the layers are. For dense depth pruning, the two-stage schedule is more reliable on NYU than single-stage pruning.

E Full Sweep Results

  • Depth-Anything V2 single-layer probes are low-rank and concentrated, while DINOv2 CLS is low-rank but delocalized, with κcap falling from 58% at L02 to 33% at L10.Depth-Anything has κcap ≥96%, effective rank ≈4, and tail-richness ratio ≈1.2; the DINOv2 decline may reflect more attention paths at later layers.Source: E Full Sweep Results · p. 22
  • VGGT depth is high-rank at the head input, but VAE compression reduces its effective rank to 30–53 and makes the SPN trainable.Before compression, κcap(20) is 76.2% and the effective rank is 282 at xvis.Source: E Full Sweep Results · p. 22

E.1 Multi-layer probes and layer spacing

  • The four-layer Depth Anything V2 probe has effective rank 36.0 versus 282.6 for VGGT’s widely spaced layers, a ≈7.8× gap attributed to greater layer orthogonality.Depth Anything V2’s layers are only three layers apart and have effective layer multiplicity ≈1.1×; VGGT’s widely spaced layers are nearly orthogonal, with multiplicity ≈4×.Source: §E.1 · p. 22
  • At the multi-layer Depth Anything V2 probe, CV rises from ≈0.026 at L02 to 2.381, preserving spatial locality; low CV at L02 follows ten later global-attention blocks.The measured multi-layer intractability comes from nearly orthogonal layers, not dense decoding or multi-layer reading alone.Source: §E.1 · p. 22
  • At η = 0.50, single-stage importance wins on ImageNet at 224 and 336 but loses to ToMe at 448; on NYU, strategies are within ∼1 point and random is nominally best.The single-stage result is not uniformly robust across dataset and resolution.Source: §E.1 · p. 22, Table 4 · p. 23
  • The 2-stage schedule beats ToMe scoring at every ratio on NYU-Depth-V2, with ≈25–35% relative reduction; single-stage pruning loses its advantage at η = 0.50.At NYU@336 and @448, the same schedule retains the 25–35% relative reductions.Source: §E.1 · pp. 23–24, Table 5 · p. 23

F Importance Head Ablations

The ablations identify a compact importance head that can match Jacobian-derived targets, while additional evaluations test its robustness across backbones, datasets, and pruning settings. Qualitative maps suggest that it selects depth-relevant boundaries, unlike ToMe’s cosine scores.

F Importance Head Ablations

  • The importance head reaches Spearman ρ ≈0.20 on DPT depth with the σ2-weighted target, versus ρ ≈0.09 with plain Jacobian block norms.DPT is the delocalized case, where the block-norm signal is approximately constant across tokens.Source: F Importance Head Ablations · p. 24

E.2 Ranking-loss weight wrank.

  • Pruning quality is nearly the same with ranking-loss weights 0, 0.5, and 1.0; DINOv2 quality saturates by 2,000 training images, while DPT shows no consistent training-set-size trend.On DPT, 4,000 images are best at the two lowest ratios and 500 at the others.Source: §E.2 · p. 24, §E.4 · p. 25

E.3 Architecture size.

  • Pruning quality plateaus at the default 310K-parameter head, and larger heads give no measurable improvement; two attention heads are slightly better than four except at η = 0.05, where they tie.The sweep varies hidden width over {128, 256, 512} and attention heads over {2, 4, 8}.Source: §E.3 · p. 24

G Additional Backbone: CLIP CLS Pruning

  • At η = 0.50 on CLIP CLS, the importance head’s degradation is 44× smaller than ToMe’s, whose degradation approaches the random baseline.The head and ToMe are evaluated on CLIP-ViT-B/16 CLS pruning at ImageNet-val @ 224.Source: Table 10 · p. 25, G Additional Backbone: CLIP CLS Pruning · p. 25
  • On ImageNet-Sketch, importance retains a more than 100× advantage at low ratios, but the gap closes at η = 0.50.For η = 0.05, importance gives 0.001 versus 0.519 for ToMe; at η = 0.50, the values are 2.169 and 2.171, respectively.Source: G Additional Backbone: CLIP CLS Pruning · pp. 25–26

H Additional Baselines and Classification Accuracy

  • The importance head is best at every ratio or tied with EViT, the strongest frozen-backbone baseline; both training-free baselines beat ToMe scoring.EViT uses CLS attention weights, whereas the importance head uses a single forward pass on features; the baseline comparison is a separate single-seed run.Source: H Additional Baselines and… · p. 26, Table 12 · p. 26
  • Across methods and prune ratios, ImageNet top-1 accuracy stays between 0.66 and 0.68, versus 0.674 unpruned; at η = 0.50, importance scores 0.670 and ToMe 0.668.The methods are not separated by this accuracy measure because the weighted k-NN classifier uses the CLS token, which is never pruned.Source: H Additional Baselines and… · p. 26

I Token Importance Visualizations

  • On four illustrative DPT images, importance scores concentrate on depth discontinuities and reach ρ ∈[0.87, 0.95], while ToMe scores are nearly uncorrelated at ρ ∈[−0.14, 0.27].The illustrative images exceed the held-out average agreement of ≈0.20; ToMe highlights differences from neighboring tokens rather than depth-relevant regions.Source: I Token Importance Visualizations · pp. 26–27, Figure 2 · p. 27

J SPN VAE Validation: Full Setup and Results

The paper tests whether randomized power-iteration supervision trains an SPN on VAE-compressed VGGT depth features. Random initialization gives little alignment, while one-step RandNLA reaches the deterministic-supervision ceiling, supporting the compressed pipeline in this setting.

J SPN VAE Validation: Full Setup and Results

  • The experiment probes VAE-compressed VGGT depth at latent z ∈R4096, where the decoder reconstructs the head-input features used by the frozen depth head.An exact SVD of the 4096-column Jacobian provides the rank-r target subspace for evaluating predicted modes.Source: J SPN VAE Validation: Full Setup and… · p. 27
  • RandNLA at q = 1 reaches subspace alignment 0.713 versus 0.710 for Spectral supervision, while random initialization scores 0.005.The random-to-structured-supervision improvement is ≈140×; the ceiling is ≈0.71 because the feature-conditioned SPN cannot reproduce each image’s SVD exactly.Source: J SPN VAE Validation: Full Setup and… · p. 28, Table 13 · p. 28
  • The results support both the VAE pathway for making an intractable dense decoder trainable and RandNLA’s ability to reach the supervision ceiling at its cheapest setting.The ≈0.29 gap from perfect alignment is attributed to amortized prediction, not a limitation of the SVD or RandNLA.Source: J SPN VAE Validation: Full Setup and… · p. 28

K Task-Induced Retrieval

Task-Induced Retrieval uses a learned rank-r task geometry to compare feature maps without rerunning the decoder, and its distance predicts depth-decoder similarity better than Euclidean or cosine baselines in both evaluated settings.

K Task-Induced Retrieval

  • DRM scores feature differences with a rank-r, feature-dependent pullback seminorm and uses its induced pseudometric as a task-sensitive alternative to Euclidean or cosine distance.The metric aggregates token differences using task-dependent importance weights and projects them onto a rank-k subspace.Source: K Task-Induced Retrieval · pp. 28–29
  • The path-based pseudometric is impractical to compute directly at feature-map scale, so the evaluation approximates path length along the straight segment between feature maps.The reported symmetrized distance is a first-order approximation and does not generally satisfy the triangle inequality.Source: K Task-Induced Retrieval · p. 29
  • Training combines regression on frozen depth-decoder distances with distillation toward Jacobian-derived token importance, calibrating overall scale and shaping per-token contributions.Without regression, the scale and projection are unsupervised; without distillation, token weights lack a geometric anchor.Source: K Task-Induced Retrieval · p. 29
  • On the in-distribution depth-retrieval evaluation (train ImageNet, test ImageNet, same-distribution split), DRM’s predicted ranking matches the reference with Spearman ρ = 0.783, versus 0.223 for L2 and 0.206 for cosine.Source: K Task-Induced Retrieval · p. 29, Table 14 · p. 30
  • Cross-dataset, DRM reaches Spearman ρ = 0.457 versus 0.338 for cosine and MAE 0.246 versus 0.627, retaining its lead on every reported metric.The methods were trained on ImageNet and evaluated on NYU-Depth-V2.Source: K Task-Induced Retrieval · p. 30, Table 15 · p. 30

L Glossary of Notation

  • The glossary defines κcap(r) as top-r captured Jacobian sensitivity, CV as per-token energy variation, and η as the fraction of tokens removed.Source: L Glossary of Notation · p. 30

M Limitations

The paper bounds its claims to measured probe layers and tested pruning configurations, while identifying limits in VAE approximation, head transfer, inference cost, and evaluation scope. It argues that decoder-sensitive directions complement feature quality and outlines extensions beyond pruning and ViTs.

M Limitations

  • The measured κcap regimes apply to probe layers feeding decoders; internal attention or MLP outputs remain untested, and VAE-reconstructed geometry is supported empirically without an analytic discrepancy bound.The VAE pathway is supported by reduced effective rank and an SPN result.Source: M Limitations · pp. 30–31
  • Each importance head is trained for one decoder–probe-layer pair; transfer to another decoder has not been tested, so each new task currently needs its own head.Source: M Limitations · p. 31
  • Pruning benefits depend on configuration and dataset: single-stage pruning loses at low ImageNet ratios and at η = 0.50 on NYU, while four-stage pruning fails at high ratios; two-stage wins in tested settings.Single-stage pruning also loses to ToMe scoring at 448 × 448 on ImageNet.Source: M Limitations · p. 31
  • Pruning targets output quality rather than throughput: at 224 × 224 scoring and token reinsertion outweigh attention savings, while speedup is modest even at 448 × 448.At 448 × 448, speedup is 1.12× versus 1.15× for ToMe scoring.Source: M Limitations · p. 31
  • The evaluation is diagnostic, using one A40 and 50–2,000 images; ToMe is the primary dense-pruning baseline, with no comparable dense-decoder baseline sweep.Additional frozen-backbone baselines cover DINOv2 CLS only.Source: M Limitations · p. 31

N Discussion

  • The discussion presents decoder-sensitive feature directions as complementary to learning better features, recoverable from the decoder Jacobian without retraining the backbone.Source: N Discussion · p. 31

O Future Work

  • Proposed extensions include task-induced flow matching and image interpolation, per-token κcap for dense-prediction reliability, and diagnostic applications to language-model attribution and OOD detection.Other directions include cross-task transfer, optimal-transport retrieval, and estimating sectional curvature.Source: O Future Work · p. 31
Loading 2609.27988v1…