Research map / Architectures and efficiency / RNNs and normalization

Task-Induced Riemannian Metrics for Vision Transformer Feature Spaces

NeurIPS 2026 · Architectures and efficiency › RNNs and normalization · approach: Symmetry and equivariance in loss landscapes · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes a matrix-free diagnostic (kappa_cap, CV) for whether low-rank pullback metrics of ViT features under task decoders are learnable, plus the Spectral Pullback Network, a 310K-parameter importance head, and geometric token pruning evaluated on DPT depth, DINOv2 CLS, CLIP, and VGGT.

Problem

Euclidean distance or cosine similarity on ViT features ignores which feature directions a downstream task decoder is sensitive to; the task-induced pullback metric J^T J is too large to store or form for dense outputs.

Approach

Use a matrix-free captured-energy diagnostic kappa_cap(r) and a coefficient-of-variation measure to decide tractability, learn a low-rank pullback metric (SPN) via randomized power iteration with JVP/VJP supervision, distill token importance into a small head, and use a VAE bottleneck to make intractable (rich-spectrum) pairs tractable.

Key results

• Importance head reaches Spearman rho = 0.998 against the Jacobian-derived target on DINOv2 CLS L10 with about 310K parameters.
• Geometric token pruning reduces the additional depth error of ToMe-based selection by 25% on DPT depth at prune ratio 0.50.
• On VAE-compressed VGGT depth, RandNLA with q=1 reaches subspace alignment 0.713 versus a supervision ceiling of 0.710 and 0.005 for a random network.

Surprising findings

• On delocalized DPT depth, per-token importance reaches only Spearman rho of about 0.20, and plain Jacobian block norms reach only about 0.09.
• Importance-guided pruning loses to ToMe at the lowest prune ratios on DPT depth.

Limitations

• Per-token scoring is weak for delocalized dense decoders such as DPT depth (rho about 0.20).
• Wall-clock speedup from pruning appears only at 448x448 resolution.
• Importance-guided pruning underperforms ToMe at the lowest prune ratios on DPT depth.

Experimental setup

Datasets: ImageNet-val ImageNet-Sketch NYU-Depth-V2

Benchmarks: ImageNet-val ImageNet-Sketch NYU-Depth-V2

Metrics: captured energy fraction kappa_cap(20) coefficient of variation Spearman rho subspace alignment SILog degradation cosine-distance degradation top-1 accuracy FLOPs reduction

Baselines: ToMe Random pruning EViT DynamicViT A-ViT Token Cropr

Builds on: ToMe Last-Layer Fusion Hutchinson trace estimator Stochastic Lanczos Quadrature randomized SVD power iteration

Closest papers by meaning

• Hide & Seek: Transformer Symmetries Obscure Sharpness & Riemannian Geometry Finds It (ICML 2025 · similarity 0.87)
• Geometry-Aware Similarity Metrics for Neural Representations on Riemannian and Statistical Manifolds (NeurIPS 2026 · similarity 0.87)
• The Geometric Wall: Manifold Structure Predicts Layerwise Sparse Autoencoder Scaling Laws (NeurIPS 2026 · similarity 0.87)
• Follow the Energy, Find the Path: Riemannian Metrics from Energy-Based Models (NeurIPS 2025 · similarity 0.86)
• Learning Theory of Transformers: Local-to-Global Approximation via Softmax Partition of Unity (NeurIPS 2026 · similarity 0.85)
• Building Transformation Layers for Riemannian Neural Networks (NeurIPS 2026 · similarity 0.85)