用任务诱导黎曼度量改进 ViT 特征空间的 token 剪枝
Task-Induced Riemannian Metrics for Vision Transformer Feature Spaces
论文教你怎么用解码器雅可比给 ViT 特征空间算真正的任务几何，剪掉一半 token 还能让深度误差少涨 25%，做高效推理的可以看看。
这篇论文指出 ViT 特征空间常用的欧氏距离和余弦相似度忽略了任务相关的真实几何结构，提出用解码器 Jacobian 的 pullback 度量 g(F)=JᵀJ 来刻画。作者设计了免矩阵诊断量 κ_cap(r)，通过少量 Jacobian-vector 积判断低秩近似是否可行，并开发 Spectral Pullback Network (SPN) 从随机幂迭代学习低秩度量，蒸馏成 310K 参数的 token 重要性头。在 DINOv2 CLS 上重要性头达到 Spearman ρ=0.998，基于该度量的几何 token 剪枝在 DPT 深度任务、剪枝比例 0.5 时，把 ToMe 方案的额外深度误差降低 25%，且无需微调 ViT。实验覆盖 DPT、DINOv2、CLIP、VGGT 四个骨干。
Task-Induced Riemannian Metrics for Vision Transformer Feature Spaces
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 $κ_{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, $κ_{cap}(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/