Task-Induced Riemannian Metrics for Vision Transformer Feature Spaces¶
会议: NeurIPS 2026
arXiv: 2609.27988
OpenReview: https://openreview.net/forum?id=wQbmMf6PdD
代码: https://cyberiada.github.io/TaskInducedViTs/
领域: 其他
关键词: 黎曼度量, 视觉Transformer, 特征空间几何, Token剪枝, 谱拉回网络
一句话总结¶
本文针对视觉Transformer中欧氏距离与余弦相似度忽略下游任务敏感方向的缺陷，基于解码器雅可比矩阵构建了任务诱导的黎曼拉回度量，提出了矩阵自由的谱可处理性诊断器与谱拉回网络（SPN），并成功蒸馏出仅310K参数的重要性预测头，在冻结ViT主干下实现了高质量几何Token剪枝与图像检索。
研究背景与动机¶
在现代计算机视觉中，针对视觉Transformer（ViT）特征表征的分析、检索、Token合并与剪枝，几乎普遍依赖于欧氏距离或余弦相似度。这种几何假设隐含了一个极强的先验：特征空间中的所有扰动方向具有同等的重要程度。然而在实际部署中，下游任务解码器（如用于深度估计的DPT解码器、用于分类的CLS嵌入或用于相机位姿估计的VGGT头部）对不同方向的微小扰动表现出极大的敏感度各向异性。对某个ViT特征施加微小偏移，在某些特定特征方向上可能引起输出深度图或相机位姿的剧烈突变，而在另一些欧氏模长完全相同的方向上，输出结果却几乎保持静止。
造成这一现象的核心矛盾在于：表征空间天然存在于一个由下游任务目标诱导的局部非欧几何流形上，而传统的特征距离无法反映下游解码器的真实敏感性；但直接度量这种几何极其困难。若将探针层特征记为 \(F \in \mathbb{R}^{N \times D}\)，解码器输出相对于特征的雅可比矩阵为 \(J(F)\)，则衡量任务输出变化的拉回黎曼度量为高斯-牛顿型张量 \(g(F) = J(F)^\top J(F)\)。对于标准的ViT-B/14模型，特征维度 \(ND \approx 2 \times 10^5\)，\(g(F)\) 的参数量高达 \(\mathcal{O}(N^2 D^2) \approx 4 \times 10^{10}\)；而对于密集预测任务（如 \(M \approx 5 \times 10^4\) 的深度图），雅可比矩阵 \(J \in \mathbb{R}^{M \times ND}\) 包含上百亿个元素，内存与计算上根本无法显式构建或持久化存储。
面对显式构建不可行、任务敏感性可能弥散在庞大特征空间的挑战，本文转换了切入视角：无需在整个空间显式求出完整黎曼度量，而是探究任务敏感性是否高度集中在极少数低秩主导方向上，并利用矩阵无关（matrix-free）的雅可比-向量积在线诊断并学习局部几何。核心 idea：通过雅可比-向量积构建矩阵自由的捕获能量比率与空间聚集度诊断器，在可处理体制下训练低秩谱拉回网络（SPN）并将其蒸馏为310K参数的轻量Token重要性预测头，对不可行的高秩稠密任务则引入VAE瓶颈恢复谱集中性，最终在完全冻结主干与解码器前提下实现任务几何诱导的Token剪枝与检索。
方法详解¶
整体框架¶
本文构建了一套完整的任务诱导局部黎曼几何分析与应用管线。针对ViT在探针层 \(\ell\) 的中间特征 \(F^\ell \in \mathbb{R}^{N \times D}\)，将其与后续ViT层及任务解码器 \(\phi\) 复合为探针映射 \(\psi_\ell(F^\ell) = \phi(F^{\ell+1}, \dots, F^L) \in \mathbb{R}^M\)。该映射诱导的拉回半正定度量为 \(g(F^\ell) = J_{\psi_\ell}(F^\ell)^\top J_{\psi_\ell}(F^\ell)\)。整个方案分为四个紧密相连的阶段：首先运行基于随机化线性代数的矩阵自由诊断器，评估任务敏感性的谱捕获能量比率 \(\kappa_{cap}(r)\) 和Token空间聚集度（CV）；针对有效秩过高的稠密解码器，引入预训练的特征级VAE瓶颈压制输入维度以恢复低秩可处理性；随后在可处理体制下，利用无显式雅可比构建的RandNLA随机幂迭代监督训练谱拉回网络（SPN），预测正交主模态与特征能量；最后，将该几何度量蒸馏为单个仅约310K参数的Token重要性预测头 \(h_\theta\)，并在推理期指导无微调的几何Token软合并（CLS解码器）或硬剪枝配合末层融合LLF（稠密解码器）。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入特征与探针映射<br/>F in R^(N x D) 到输出 psi(F) in R^M"] --> B["矩阵自由度量诊断器<br/>评估谱捕获能量 kappa_cap 与空间聚集度 CV"]
B -->|判断体制: 高有效秩稠密任务| C["VAE特征瓶颈压缩<br/>低维隐空间约束雅可比上界"]
B -->|判断体制: 低秩可处理任务| D["谱拉回网络与RandNLA训练<br/>无显式雅可比学习前r个特征模态与能量"]
C --> D
D --> E["任务几何蒸馏重要性头<br/>310K轻量MHA预测单Token任务敏感分"]
E --> F["几何Token剪枝与融合<br/>CLS任务几何软合并 / 稠密任务硬剪枝+LLF"]关键设计¶
1. 矩阵自由度量诊断器：低开销解耦谱截断可行性与空间局域性
要对庞大的全特征黎曼度量 \(g(F) \in \mathbb{R}^{ND \times ND}\) 进行低秩截断，首要前提是确认前 \(r\) 个奇异方向是否能够覆盖大部分任务敏感性，以及敏感性是否在空间Token间弥散。本文设计了两个计算极度轻量的标量指标。第一个是捕获能量比率（Captured Energy Fraction）： $\(\kappa_{cap}(r) = \frac{\sum_{j=1}^r \sigma_j^2}{\|J\|_F^2} \in [0, 1]\)$ 其中分母 \(\|J\|_F^2 = \mathrm{Tr}(J^\top J)\) 通过 \(m=100\) 次Rademacher探针的Hutchinson迹估计器高效逼近，分子通过 \(q=2\) 轮随机幂迭代提取前 \(r=20\) 个奇异值。定理1证明了 \(\kappa_{cap}(r)\) 是任何将度量约束在 \(r\) 维子空间的紧致灵敏度上界。第二个指标是奇异向量在Token维度的变异系数（Coefficient of Variation, CV）： $\(\mathrm{CV} = \frac{1}{r} \sum_{j=1}^r \frac{\mathrm{std}_t(\|v_j^{(t)}\|_2)}{\mathrm{mean}_t(\|v_j^{(t)}\|_2)}\)$ 其中 \(v_j^{(t)} \in \mathbb{R}^D\) 表示第 \(j\) 个右奇异向量对应第 \(t\) 个Token的分块。若 \(\mathrm{CV} \approx 0\)（极度离域），定理2（因子化度量不可行性定理）证明了任何独立的单Token评分机制对敏感方向的辨别力与随机猜测完全无异；只有当 \(\mathrm{CV}\) 显著大于0时，单Token重要性评分才具备几何有效性。二者将主干-解码器划分为四大体制（高维谱、低秩离域、边际低秩高聚集、结构性低秩），为后续架构设计提供了精确指导。
2. 谱拉回网络与RandNLA训练：无显式雅可比矩阵的低秩黎曼度量学习
当诊断确认低秩逼近可行时，测试期若直接求雅可比反向传播则会带来严重延迟并丧失加速意义。为此，本文设计了谱拉回网络（Spectral Pullback Network, SPN）\(g_\theta(F)\)，直接以输入特征 \(F\) 预测谱分解形式的度量： $\(v^\top g_\theta(F) v = \sum_{j=1}^r \lambda_j(F;\theta) \big(u_j(F;\theta)^\top v\big)^2 + \varepsilon \|v\|^2, \quad U(F;\theta)^\top U(F;\theta) = I_r\)$ 其中正交模态矩阵 \(U \in \mathbb{R}^{ND \times r}\) 通过将特征投影后利用QR分解正交化得到，能量预测向量 \(\lambda \in \mathbb{R}_{>0}^r\) 通过全局池化、MLP与Softplus生成，\(\varepsilon > 0\) 为各向同性底数正则化项。为了规避构建雅可比矩阵的维度灾难，SPN使用基于随机化数值线性代数（RandNLA）的幂迭代监督算法训练：初始化高斯随机矩阵 \(\Omega \in \mathbb{R}^{ND \times r_0}\)，交替调用向量-雅可比积（VJP）与雅可比-向量积（JVP）完成 \(q\) 轮迭代 \(\Omega \leftarrow J^\top (J\Omega)\)，并进行QR正交化得到监督目标 \(\hat{V}_r\)。损失函数由不变于模态排列的Grassmannian子空间对齐损失 \(\mathcal{L}_{\text{sub}} = 1 - \frac{1}{r} \|U(F;\theta)^\top \hat{V}_r\|_F^2\) 与能量对数均方误差 \(\mathcal{L}_{\text{en}}\) 联合组成。
3. 任务几何蒸馏重要性头与几何Token剪枝：轻量推理与特征漂移补偿
在Token剪枝等仅需离散排名的下游应用中，完整运行SPN仍然存在冗余。本文将SPN进一步蒸馏为参数量仅约310K的轻量重要性预测头 \(h_\theta: \mathbb{R}^{N \times D} \to \mathbb{R}^N\)（包含1个4头自注意力块、逐Token MLP与Softplus激活）。该预测头直接拟合加权雅可比分块能量真实目标： $\(\mathrm{imp}^*(F)_t = \sqrt{\sum_{j=1}^r \sigma_j^2 \|v_j^{(t)}\|_2^2}\)$ 在训练时通过对数回归损失与成对排序损失联合优化。剪枝决策由定理4（一阶泰勒展开误差界）保证： $\(\|\phi(F^L) - \phi(\hat{F})\|_2 \le \sum_{t \in \mathcal{T}} \|J_{:,t}(\hat{F})\|_F \|\rho_t\|_2 + \frac{\beta}{2} \sum_{t \in \mathcal{T}} \|\rho_t\|_2^2\)$ 其中第一项对应Token任务敏感度，第二项为剪枝带来的特征漂移残差 \(\rho_t\)。对于只依赖全局CLS Token的分类解码器，采用基于局部任务马氏距离 \(d_{\text{task}}(a,b) = \sqrt{(F_a - F_b)^\top Q_a (F_a - F_b)}\) 的重要性加权软合并（Soft Merge）；对于重构2D网格的密集DPT解码器，采用重要性硬剪枝并在最终ViT Block前引入免调优的末层融合（Last-Layer Fusion, LLF），令冻结的被剪Token充分关注未剪Token，有效消除累积特征漂移。
4. VAE特征瓶颈压缩：突破高阶密集预测任务的不可行壁垒
对于VGGT深度估计与点云重建等复杂多视角大尺度稠密任务，其有效秩高达200~280，远远超过任何低秩度量的表达预算（体制一）。由于输出维度 \(M\) 由任务规格固定，压缩雅可比秩的唯一数学路径是压缩有效输入维度 \(ND\)。本文在读取的四层拼接特征 \(x_{\text{vis}}\) 与下游任务头之间插入预训练的拓扑对齐 \(S^2\)-VAE：编码器将庞大的Token序列投射为由32个128维Register Token组成的低维瓶颈隐空间 \(z \in \mathbb{R}^{4096}\)，解码器重构特征后再送入冻结的任务预测头。由于复合映射的所有信息流必须穿透维度为4096的隐层 \(z\)，根据矩阵乘积秩不等式，雅可比矩阵的秩被物理严格限制在 \(\mathrm{rank}(J_{\tilde{\psi}}) \le 4096\)。这一压缩将VGGT深度的有效秩从282陡降至30~53，使原本处于不可行体制的高维密集任务平滑过渡至低秩可处理的体制二，令后续SPN与重要性头能够在低维隐空间中顺利训练收敛。
损失函数 / 训练策略¶
SPN网络的训练损失函数由子空间投影损失与能量回归损失加权构成： $\(\mathcal{L}_{\text{SPN}} = \mathcal{L}_{\text{sub}} + \alpha \mathcal{L}_{\text{en}} = \left(1 - \frac{1}{r} \|U(F;\theta)^\top \hat{V}_r\|_F^2\right) + \alpha \cdot \mathrm{MSE}\big(\log \lambda(F;\theta), \log \hat{S}^2\big)\)$ 当相邻奇异值间距较小（\(\sigma_j / \sigma_{j+1} < 1.05\)）时，为防止模态对换引起的振荡，该模态对应的能量损失被屏蔽。重要性头 \(h_\theta\) 采用目标蒸馏与排序联合损失优化： $\(\mathcal{L}(\theta) = \mathbb{E}_F\big[\|\log h_\theta(F) - \log \mathrm{imp}^*(F)\|^2\big] + w_{\text{rank}} \mathcal{L}_{\text{rank}}\big(h_\theta(F), \mathrm{imp}^*(F)\big), \quad w_{\text{rank}} = 0.5\)$ 所有模块均在完全冻结ViT主干与解码器权重的前提下进行轻量训练，训练样本量仅需约2000张ImageNet验证集图像即可完全收敛。
实验关键数据¶
主实验¶
主实验在ImageNet-val（224分辨率）上评测了DINOv2-B/14分类CLS嵌入以及Depth-Anything V2单阶段DPT深度估计在不同剪枝比例 \(\eta\) 下的性能表现。在CLS任务中衡量与未剪枝特征的余弦距离衰减（\(\times 100\)），在DPT深度估计任务中衡量相对于未剪枝结果的附加尺度不变对数误差（SILog \(\times 100\)）。
| 任务 / 评估指标 | 剪枝方法 | \(\eta=0.05\) | \(\eta=0.10\) | \(\eta=0.20\) | \(\eta=0.30\) | \(\eta=0.40\) | \(\eta=0.50\) |
|---|---|---|---|---|---|---|---|
| DINOv2 CLS 余弦距离衰减 (\(\times 100\)) |
Random | \(0.17 \pm 0.03\) | \(0.36 \pm 0.05\) | \(0.86 \pm 0.09\) | \(1.49 \pm 0.14\) | \(2.31 \pm 0.22\) | \(3.44 \pm 0.31\) |
| ToMe score | \(0.67 \pm 0.11\) | \(0.78 \pm 0.07\) | \(1.12 \pm 0.11\) | \(1.75 \pm 0.13\) | \(2.56 \pm 0.25\) | \(3.82 \pm 0.42\) | |
| Importance (本文) | \(\mathbf{0.001 \pm 0.000}\) | \(\mathbf{0.007 \pm 0.001}\) | \(\mathbf{0.060 \pm 0.005}\) | \(\mathbf{0.25 \pm 0.03}\) | \(\mathbf{0.80 \pm 0.09}\) | \(\mathbf{2.09 \pm 0.13}\) | |
| DPT Depth 附加SILog误差 (\(\times 100\)) |
Random | \(4.59 \pm 0.09\) | \(3.70 \pm 0.10\) | \(\mathbf{3.19 \pm 0.19}\) | \(3.16 \pm 0.17\) | \(3.20 \pm 0.15\) | \(3.29 \pm 0.18\) |
| ToMe score | \(\mathbf{4.29 \pm 0.07}\) | \(\mathbf{3.66 \pm 0.09}\) | \(3.50 \pm 0.10\) | \(3.52 \pm 0.13\) | \(3.67 \pm 0.13\) | \(4.09 \pm 0.13\) | |
| Importance (本文) | \(4.88 \pm 0.06\) | \(4.86 \pm 0.07\) | \(3.28 \pm 0.15\) | \(\mathbf{3.02 \pm 0.16}\) | \(\mathbf{3.01 \pm 0.16}\) | \(\mathbf{3.08 \pm 0.15}\) |
在NYU-Depth-V2数据集的多阶段剪枝实验中，2阶段 [4, 8] 剪枝策略全面胜过ToMe：在 \(\eta=0.50\) 时，ToMe的附加SILog为 \(5.58\)，而本文Importance方法仅为 \(5.12\)；在 \(\eta=0.05\) 时，本文仅为 \(2.42\)（ToMe为 \(3.52\)），全分辨率下实现25%~35%的相对误差下降。
消融实验¶
表1记录了主干-解码器-探针层组合的谱可处理性诊断，表2记录了SPN在VAE压缩的VGGT深度任务上的子空间对齐度 \(V_1\)。
| 模型与探针层配置 | 捕获能量 \(\kappa_{cap}(20)\) | 空间聚集度 \(\mathrm{CV}\) | 全谱有效秩 \(r_{\text{eff}}^{\text{SLQ}}\) | 归属体制 | 核心特性说明 |
|---|---|---|---|---|---|
| DA V2 depth, L02 | \(95.6\%\) | \(0.026\) | \(34\) | 体制 (ii) | 低秩但全局离域，单Token评分不可行需全局注意力 |
| DA V2 depth, L10 (1-layer) | \(96.6\%\) | \(0.026\) | \(20\) | 体制 (ii) | 极高捕获能量，适合直接训练SPN |
| DINOv2 CLS, L10 | \(33.4\%\) | \(4.060\) | \(270\) | 体制 (iii) | 谱较分散但高空间聚集，重要性头可学到极高相关性 |
| VGGT camera, L16 | \(\approx 100\%\) | N/A | \(9\) | 体制 (iv) | 输出维度 \(M=9\)，结构性严格有界 |
| VGGT pointcloud, L08 | \(96.6\%\) | \(0.590\) | \(212\) | 体制 (i) | 全谱有效秩高，未压缩前低秩SPN逼近不足 |
| VGGT depth, L16 (full) | \(76.2\%\) | \(0.630\) | \(282\) | 体制 (i) | 高度不可行，需引入VAE瓶颈压制至 \(r_{\text{eff}} \in [30, 53]\) |
| SPN 训练监督配置 (VGGT Depth + VAE) | 子空间对齐度 \(V_1\) | 说明 |
|---|---|---|
| Random init (无监督基线) | \(0.005\) | 随机参数无几何表征能力 |
| RandNLA \(q=1\) (本文极简配置) | \(\mathbf{0.713}\) | 单步幂迭代即逼近理论确定性监督上限 |
| Spectral (精确SVD全量监督) | \(0.710\) | 理论上限（受单网络泛化均摊误差限制） |
关键发现¶
- 离域性决定因子化剪枝的有效边界：在DINOv2 CLS高聚集体制（\(\mathrm{CV}=4.06\)）下，重要性头预测与真实雅可比目标的Spearman秩相关系数高达 \(\rho = 0.998\)；而在高度离域的DPT深度任务（\(\mathrm{CV}=0.026\)）中，若直接拟合朴素分块范数，相关性仅为 \(\rho \approx 0.09\)，采用 \(\sigma^2\) 加权的谱目标才能恢复至 \(\rho \approx 0.20\)。
- 激进剪枝下任务敏感性支配特征冗余性：在DPT深度估计高剪枝率（\(\eta \ge 0.30\)）下，移除低任务敏感性Token带来的收益远超过特征级余弦相似度匹配，在 \(\eta=0.50\) 时比ToMe减少了25%的附加误差。
- 多层特征的雅可比正交性是不可行性的真正根源：消融表明，DPT读取4层特征时有效秩仅为36（层间相关性高），而VGGT跨越24层中的4层导致雅可比几近正交，有效秩叠加飙升至282，证明层间几何正交性而非单纯的密集解码造成了高维谱爆炸。
亮点与洞察¶
- 雅可比拉回几何的矩阵自由在线诊断：首次将黎曼拉回度量与随机化迹估计及幂迭代结合，利用仅仅 \(\mathcal{O}(m + rq)\) 次无显式雅可比构建的JVP/VJP运算，在模型训练前即可精准判定低秩几何学习的可行性。
- VAE信息瓶颈作为有效秩降维器：创造性地将预训练特征VAE作为任务雅可比算子的物理秩截断器，成功将原本谱弥散、不可处理的高维3D视觉解码器拉回低秩可学习区间。
- 可迁移的度量检索范式：提出的蒸馏黎曼度量（DRM）在免调解码器的情况下，在深度敏感图像检索任务中取得了 \(\rho = 0.783\) 的高排序相关度，远超余弦相似度的 \(0.206\)。
局限与展望¶
- 作者承认的局限：目前的几何诊断仅在送入解码器的探针层展开，尚未覆盖Transformer内部各层注意力和MLP的微观几何；同时重要性头与特定（解码器, 探针层）绑定，跨任务的泛化迁移能力仍需单独训练验证。
- 实测性能开销瓶颈：在小分辨率（\(224 \times 224\)）下，序列长度仅为257，剪枝降低的理论FLOPs（约12.1%）被重要性预测和Token拼接的固定工程开销抵消，壁钟耗时并未加速（\(0.96\times\)），实际加速比仅在大分辨率（\(448 \times 448\) 达 \(1.12\times\)）时显现。
- 未来改进方向：可进一步探索将局部拉回度量扩展至大语言模型的Residual Stream分析与归因验证，或利用该度量驱动黎曼流匹配生成更加平滑的轨迹。
相关工作与启发¶
- vs ToMe (Bolya et al., ICLR 2023)：ToMe使用特征余弦相似度通过二分图匹配合并冗余Token，完全忽略了下游解码器的敏感方向；本文证明了在CLS任务中引入局部任务度量、在稠密任务中结合LLF硬剪枝，能在高剪枝率下大幅超越ToMe。
- vs DynamicViT / A-ViT / Token Cropr：此类方法需对ViT主干进行复杂的端到端微调；而本文方法严格保持主干与解码器全冻结，仅凭外置轻量头实现几何对齐剪枝。
- vs 经典黎曼度量学习：传统神经黎曼方法直接参数化SPD矩阵并受限于 \(O(d^2)\) 存储，无法承受ViT数万维特征；本文通过低秩谱分解与RandNLA随机化投影成功打破了维度壁垒。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次系统性将任务诱导黎曼拉回几何引入ViT特征空间分析，理论证明与诊断框架完备。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖DINOv2、CLIP、Depth-Anything、VGGT等多种单尺度/多尺度与高维3D预测任务，实验详实严谨。
- 写作质量: ⭐⭐⭐⭐⭐ 概念清晰，定理与实际算法映射精准，逻辑自洽且剖析透彻。
- 价值: ⭐⭐⭐⭐☆ 为ViT特征解释、免微调Token剪枝及任务感知检索开辟了坚实的微分几何新视角。