cs.CVSep 23, 2026

Task-Induced Riemannian Metrics for Vision Transformer Feature Spaces

Authors: Andrew BondEge Erdem ÖzlüTuna ÇimenIlkin Umut MelanliogluTolga BirdalErkut ErdemAykut 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)J(F)g(F) = J(F)^\top J(F), where JJ is the Jacobian of the decoder's output fed to a task-specific distance, with respect to the features. Storing the full gg is infeasible at modern scales, and for dense outputs such as depth maps even forming JJ 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)κ_{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 310310K-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)κ_{cap}(r) predicts which learned-metric architectures are viable. The importance head reaches Spearman ρ=0.998ρ= 0.998 on DINOv2 CLS, and our geometric token pruning reduces the additional depth error of ToMe-based token selection by 25%25\% on DPT depth at prune ratio 0.50.5, without fine-tuning the ViT. Project page: https://cyberiada.github.io/TaskInducedViTs/

Explore similar work

Jul 2, 2026cs.CV

Understanding Geometric Representations in Self-Supervised Vision Transformers via Subspace Intervention

We introduce a controlled subspace intervention framework to investigate how self-supervised Vision Transformers (ViTs) encode dense geometric information. While linear probing is widely used to assess geometric representations, it treats features as a black box, failing to disentangle the underlying topology. To address this issue, we decompose the weights of converged linear probes to isolate the low-rank subspaces containing explicit geometric signals using Singular Value Decomposition (SVD). Our perspective yields three key insights: (1) Pre-training objectives determine how features are encoded. DINOv2 aligns spatial features for efficient linear extraction, while Masked Autoencoders (MAE) tend to disperse these signals, requiring a broader spatial context. (2) Explicit geometric representations are highly compressible, suggesting dense predictive heads could potentially be constrained to low-rank subspaces with minimal performance loss. (3) The layer-wise task affinity suggests that geometric precision peaks at intermediate layers before yielding to semantic abstraction in the final layers. By connecting internal encoding mechanics with downstream performance, these findings provide a basis for effective feature selection and lightweight decoder design. The source code is available at https://github.com/Zhou-Weichen/Geosubprobe.
Weichen Zhou, Yawen Zou, Chunzhi Gu +3
May 21, 2026cs.LG

ASAP: Attention Sink Anchored Pruning

Vision Transformers (ViTs) face severe computational bottlenecks due to the quadratic complexity of self-attention at high resolutions. Existing token reduction methods rely on local metrics - such as single-layer attention scores - that are inherently vulnerable to the attention sink phenomenon, where uninformative tokens are paradoxically preserved over salient foreground objects. We propose ASAP (Attention Sink Anchored Pruning), a training-free framework that recasts this sink as a feature. Modeling ViT information flow as a Lazy Random Walk, ASAP identifies the sink as a dominant accumulator of probability mass. By computing the diffusion distance to the sink within the cumulative transition matrix, ASAP partitions tokens via Radial Diffusion Clustering and compresses background redundancy through Transition Weight Pooling in a single shot. Extensive experiments across image, video, and vision-language tasks demonstrate ASAP outperforms state-of-the-art methods, accelerating throughput by up to 48% while maintaining - or even exceeding - baseline accuracy.
Jaehyuk Lee, Hanyoung Kim, Yanggee Kim +1
Aug 11, 2026cs.CV

Putting Registers to Work: Task Registers for Token Pruning in Vision Transformers

Token-pruning policies are usually designed for a single recognition pipeline, but pretrained Vision Transformers are reused across tasks with different spatial demands. We ask which parts of a pruning policy transfer across image classification, semantic segmentation, and object detection. For each pipeline, controlled probes freeze the no-pruning checkpoint and apply a series of parameter-free reduction criteria at one eligible layer at a time without retraining. The probes reveal three differences: segmentation and detection rank the criteria differently, classification is especially sensitive to attention-based pruning in the earliest layers, and the dense tasks prefer opposite recovery endpoints. These findings motivate Task-Adaptive Pruning (TAP). Existing register tokens serve as task-agnostic storage for feature artifacts. TAP instead introduces one task register per task and activates only the current one. Its evolving state ranks tokens, distributes an exact removal budget over depth, and sets the recovery scale for dense features. At a final keep rate of ρ=0.5ρ=0.5, our jointly adapted model, TAP-J, reaches 47.047.0 mIoU at 1.30×1.30\times encoder throughput on ADE20K and 53.753.7 box AP at 1.32×1.32\times encoder throughput on COCO while remaining competitive on ImageNet-1K.
Hongsen Cao, Mona Jaber, Shanxin Yuan +1