跳转至

Learning Dynamic Evidence Routes for Vision Transformer Probing

会议: NeurIPS 2026
arXiv: 2605.00915
OpenReview: https://openreview.net/forum?id=1CCRjuvAbv
论文: https://arxiv.org/abs/2605.00915
领域: 可解释性
关键词: 视觉Transformer探测, 状态空间模型, Sinkhorn最优传输, 动态证据路由, 空间组织度

一句话总结

提出了结合可微Sinkhorn排序路由与对角S4解码器的探针架构SSMProbe,利用S4因果线性时不变递推产生的固定位置系数与空间可预测性度量,首次对冻结ViT的多样化预训练特征展开证据路由几何与时序加权解耦审计。

研究背景与动机

探针(Probing)是评估和理解冻结视觉Transformer(ViT)表征质量的核心工具。然而,传统线性探针大多依赖置换不变的全局平均池化(GAP)或分类标记([CLS]),将二维空间中的图块标记(patch tokens)退化为无序集合,无法揭示分类证据在空间维度上的调度顺序。尽管自注意力(Self-Attention)或深层MLP等内容自适应探针具备更强拟合能力,但其注意力权重高度依赖样本输入和动态查询交互,缺乏透明、静态的位置加权系数,导致研究人员无法清晰审计“图块的空间排列顺序”与“序列时序衰减权重”之间具体的数学交互作用。

这一矛盾在多源预训练范式(如MAE的自编码重建、BEiT的离散掩码恢复、DINOv2的自监督对比学习以及标准有监督ViT)日益分化的背景下愈发显著。不同预训练目标驱动网络在深层形成了迥异的标记表征形态:例如DINOv2倾向将语义信息高度压缩进全局[CLS],而MAE则促使各个patch保留密集的重建线索。常规固定扫描启发式规则(如单向Raster、4向VMamba交叉扫描、蛇形扫描或对角扫描)强加了刚性的几何假定,完全忽略了下游视觉任务对关键辨识特征的动态聚焦需求。

针对这些缺陷,本文提出了全新的切入视角:将图块排列视为一维最优传输排序问题,并接入具备严格线性时不变(LTI)因果递推性质的状态空间模型作为读取器。核心 idea:构建由可微Sinkhorn路由网络与对角S4层组成的探针SSMProbe,利用S4解析导出的固定位置衰减系数显式解耦证据排列与时序读取,并提出超出先验位置偏差的超额空间信息增益指标(ERSI)来量化不同ViT骨干网络路由的紧凑性与空间几何规律。

方法详解

整体框架

SSMProbe的整体流程接收冻结ViT输出的patch标记集合,经过轻量级线性打分器计算每个patch的相对重要性标量,通过可微Sinkhorn最优传输算法与等距一维坐标点求解最优匹配,生成平滑双随机置换矩阵。重排后的特征序列输入单通道对角S4层展开因果时序聚合,最终基于S4终端隐含状态输出类别预测分布。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入图像 x 与冻结ViT编码器"] --> B["Patch标记提取<br/>T = [h1, ..., hN] ∈ R^{N×d}"]
    B --> C["可微Sinkhorn图块路由<br/>线性打分标准化与一维最优传输"]
    C --> D["对角S4因果序列解码<br/>复数模态状态更新与LTI因果递推"]
    D --> E["固定位置系数解析解耦<br/>W·z_out 线性分类与证据几何分析"]
    E --> F["下游分类 logits 与空间规律度量 ERSI"]

关键设计

1. 可微Sinkhorn图块路由:将空间特征排序转化为一维熵正则最优传输 针对确定性启发式空间扫描(如Raster或蛇形扫描)无法随图像语义自适应聚焦前景这一缺陷,该设计采用单层无偏置线性打分网络 \(s_i = \mathbf{w}^\top h_i\),并在图像内部将标量得分标准化为 \(\tilde{s}_i = (s_i - \mu(s)) / (\sigma(s) + \epsilon)\)。将排序任务建模为这组归一化得分与目标等距一维网格点 \(p_j = j / (N-1)\) 之间的二分图匹配问题,代价值定义为二次欧氏距离 \(\mathbf{C}_{i,j} = (\tilde{s}_i - p_j)^2\)。为保证端到端梯度的平滑反向传播,引入熵正则化最优传输目标 \(\mathbf{P}_\tau = \operatorname{argmin}_{\mathbf{P} \in \mathbb{R}_+^{N\times N}} \langle \mathbf{P}, \mathbf{C} \rangle - \tau \mathcal{H}(\mathbf{P})\),利用标准Sinkhorn-Knopp矩阵平衡迭代算法(\(K=20, \tau=0.1\))快速求得双随机置换矩阵 \(\mathbf{P}^{(K)}\)。重排后的输入序列为 \(\tilde{T} = (\mathbf{P}^{(K)})^\top T\)。该设计不仅实现了任务驱动的平滑排序,而且其行最大值熵与覆盖率显著优于直接连续松弛的NeuralSort与SoftSort,避免了排序极化或边界塌缩。

2. 对角S4因果序列解码器:基于LTI系统构造解析固定位置权重 为了使探针具备明确的序贯时间衰减行为且不随输入token动态改变权重系数,模型选取单层对角化结构化状态空间模型(Diagonal S4)。其连续时间动态由系统矩阵对角参数化为 \(\lambda_i = -\operatorname{softplus}(a_i) + \mathrm{i}\omega_i\)(严格保证其实部小于零),并通过双线性双纽变换(Bilinear Discretization)离散化映射至单位圆盘内部: $\(\bar{\mathbf{A}} = (I - \Delta \mathbf{A}/2)^{-1}(I + \Delta \mathbf{A}/2), \quad \bar{\mathbf{B}} = (I - \Delta \mathbf{A}/2)^{-1}\Delta \mathbf{B}\)$ 在处理长度为 \(N\) 的重排token序列时,输出末尾时间步隐状态 \(z_N = \sum_{k=1}^N \mathbf{C}\bar{\mathbf{A}}^{N-k}\bar{\mathbf{B}}\tilde{t}_k + \mathbf{D}\tilde{t}_N\)。经线性分类头 \(W\) 映射后,整个前向过程可严格等价分解为对重排位置的定常线性加权 \(\sum_{k=1}^N \alpha_k \tilde{t}_k + b\),其中固定位置系数为: $\(\alpha_k = W\mathbf{C}\bar{\mathbf{A}}^{N-k}\bar{\mathbf{B}} + \mathbf{1}_{\{k=N\}}W\mathbf{D}\)$ 相比于Transformer探针中随输入样本内容相互作用的动态注意力权重 \(a_j(T, \pi)\),S4的 \(\alpha_k\) 在训练完成后完全固定。这使得研究者能够精确计算任意图块 \(i\) 在S4中的期望读取位置 \(\tau_i = \sum_{t=1}^N t \mathbf{P}^{(K)}_{i,t}\),并对模型谱半径、前馈残差跳跃路径(\(\mathbf{D}\)项)以及时间记忆跨度展开严格的解析消融与扰动分析。

3. 证据路由空间几何与超额空间信息增益(ERSI):剥离先验位置偏差的几何测度 以往分析图块空间分布常采用全局紧凑度比值 \(R_k = D_k / D_k^{\mathrm{rand}}\)(平均成对欧氏距离与随机选择集合的比值),但该类统计指标无法剔除数据集本身固有的边缘或中心位置偏置(Location Bias)。为此,本文提出了超额空间信息增益(Evidence-Route Spatial Information, ERSI)。首先从标定集中统计网格位置 \(i\) 被选入前 \(k\) 位的经验边缘先验频率 \(m(i)\),构建可交换零模型: $\(p_0(i_t \mid i_{<t}) \propto m(i_t) \mathbf{1}\{i_t \notin i_{<t}\}\)$ 随后引入包含单一连续空间尺度参数 \(\sigma\) 的几何替代模型,赋予靠近已选集合 \(i_{<t}\) 的临近候选点更高转移概率: $\(p_{\mathrm{geo}}(i_t \mid i_{<t}; \sigma) \propto m(i_t) \exp\left(-\frac{d(i_t, i_{<t})^2}{2\sigma^2}\right) \mathbf{1}\{i_t \notin i_{<t}\}\)$ 通过在标定子集上极大化似然估计确定最优尺度 \(\hat{\sigma}\),并在未见测试集上计算平均每个路由patch所节省的编码比特: $\(\mathrm{ERSI}_k = \frac{1}{k \log 2} \left[ \log p_{\mathrm{geo}}(\pi_{1:k}; \hat{\sigma}) - \log p_0(\pi_{1:k}) \right]\)$ 当 \(\mathrm{ERSI}_k > 0\) 时,数学上证明了探针发现的辨识证据具有超出数据集位置先验的真实空间聚集性;配合4邻域连通分量占比 \(L_k\)、凸包面积比 \(A_k\)、最小生成树长度 \(S_k\) 及空间粗粗网格信息熵 \(H_k\),构成了严密的几何拓扑审计矩阵。

损失函数 / 训练策略

训练采用独立优化器梯度隔离协议(Per-Head Gradient Isolation)。冻结骨干网络单次前向输出并在内存中通过 detach() 切断梯度,多组探针头在同一批次特征上并行独立前向,各自计算交叉熵损失并由独立的 AdamW 优化器(学习率 \(10^{-3}\),余弦退火,无权重衰减,batch size 256)更新参数。ImageNet-1K 探针统一训练 5 个 epoch,CUB-200-2011 和 Stanford Cars 迁移任务训练 100 个 epoch。

实验关键数据

主实验

在ImageNet-1K上对冻结MAE进行探针分类评测,系统对比了置换不变基线、确定性静态扫描、内容池化以及可微Sinkhorn探针的精度表现。

类别 探针方法 机制类型 Top-1 准确率 (%) 相比强基线增益
下界基线 GAP 全局平均池化 (置换不变) 58.10 ± 0.03 基线
下界基线 CLS 分类标记线性探测 56.65 ± 0.03 -1.45
启发式扫描 Raster scan 单向行优先 S4 扫描 63.64 ± 0.05 +5.54 (vs GAP)
启发式扫描 4-dir (VMamba) 四向正交SS2D S4 扫描 64.68 ± 0.07 +6.58 (vs GAP)
启发式扫描 4-dir (Snake) 四向蛇形S4 扫描 64.76 ± 0.06 +6.66 (vs GAP)
启发式扫描 4-dir (Diagonal) 四向对角S4 扫描 64.78 ± 0.07 +6.68 (vs GAP)
随机置换 Random-Dynamic Perm 动态随机序列 + S4 63.95 ± 0.06 +5.85 (vs GAP)
仅内容加权 Attention Pool 单Query交叉注意力池化 67.70 ± 0.04 +9.60 (vs GAP)
仅内容加权 Content-Weighted Pool 线性打分 Softmax 加权 67.68 ± 0.06 +9.58 (vs GAP)
完整方法 Sinkhorn (SSMProbe) 动态可微最优传输 + 对角S4 70.33 ± 0.05 +5.55 (vs Diagonal)
结构对比 Bi-GRU (Raster) 双向GRU序列扫描 69.31 ± 0.12 +1.02 (vs Bi-GRU)
结构对比 DeepSets 置换等变MLP池化 68.48 ± 0.05 +1.85 (vs DeepSets)
参考上限 Transformer ref. 宽通道自注意力 (3M参数) 71.61 ± 0.14 -1.28 (接近重型模型)

在跨预训练范式(MAE、BEiT、有监督ViT、DINOv2)的细粒度迁移数据集(CUB-200-2011与Stanford Cars)上,探针精度与空间几何度量呈现显著分化:

骨干网络 CUB CLS (%) CUB GAP (%) CUB Sinkhorn (%) Cars CLS (%) Cars GAP (%) Cars Sinkhorn (%) CUB \(\mathrm{ERSI}_{20}\) Cars \(\mathrm{ERSI}_{20}\) 路由紧凑度 \(R_{20}\) (CUB/Cars)
MAE 28.88 ± 0.10 19.40 ± 0.10 56.67 ± 7.54 16.64 ± 0.11 17.62 ± 0.05 54.86 ± 4.35 0.318 0.336 0.687 / 0.567
BEiT 80.80 ± 0.13 88.27 ± 0.13 88.12 ± 0.32 57.28 ± 0.28 76.41 ± 0.10 82.21 ± 0.64 0.175 0.276 0.609 / 0.599
ViT (Supervised) 85.45 ± 0.10 85.14 ± 0.13 85.22 ± 0.19 57.01 ± 0.21 56.89 ± 0.10 62.38 ± 0.15 0.113 0.296 0.669 / 0.561
DINOv2 89.49 ± 0.21 68.55 ± 0.23 82.28 ± 1.60 87.54 ± 0.11 75.13 ± 0.09 84.83 ± 0.60 0.484 0.328 0.697 / 0.680

消融实验与扰动分析

在MAE骨干网络上对排序机制、时序解码及后验扰动进行系统解耦测试:

测试维度 实验配置 评测数据集 / 样本规模 关键指标 结果说明与机制解读
排序算法对比 NeuralSort MAE / CUB (全量测试) Top-1: 29.51% 梯度饱和严重,目标位置覆盖率仅0.5%,排列极化
排序算法对比 SoftSort MAE / CUB (全量测试) Top-1: 54.07% 连续松弛平滑但质量中等,质量集中于两极外侧10%
排序算法对比 Sinkhorn (Ours) MAE / CUB (全量测试) Top-1: 60.32% 双随机归一化促使覆盖率达57.1%,信息熵高达0.868
序列顺序扰动 正常学习路由 MAE / ImageNet (512图) Top-1: 73.87 ± 0.70% 保持Sinkhorn学得的一维时序到达顺序
序列顺序扰动 序列完全翻转 MAE / ImageNet (512图) Top-1: 7.62 ± 0.49% 下降 66.25 pp;破坏了前紧后松的位置系数权重
序列顺序扰动 路由后完全打乱 MAE / ImageNet (512图) Top-1: 18.20 ± 0.65% 下降 55.66 pp;证明S4对图块输入顺序具有强敏感依赖
序列顺序扰动 相邻位置微调互换 MAE / ImageNet (512图) Top-1: 73.71 ± 0.62% 仅微幅下降 0.16 pp;全局大尺度骨架支配排序增益
前景对齐验证 Top-20 图块前景覆盖 CUB (5,794全量测试) 前景占比: 55.40% 远超随机挑选(12.60%)与位置先验(33.25%),超额+22.15%
Token剔除测试 移除Top-20 Sinkhorn标记 MAE / ImageNet (512图) 精度降幅: 19.9 ± 1.2 pp 显著高于剔除CLS余弦高分标记(3.1 pp)或随机标记(0.9 pp)

关键发现

  • 不同预训练骨干的读出偏好存在本质差异:DINOv2表现出极强的 [CLS] 统治力(CUB上高出GAP达20.94个百分点),而MAE在固定扫描下严重欠拟合(仅6.7%~10.5%),但引入可微路由后飙升至56.67%,证明MAE表征中富含高度离散但空间可组织的细粒度线索。
  • 全局宏观序列决定识别成败,局部抖动具备高鲁棒性:在S4解码器前将序列全局反转导致准确率暴跌66.25个百分点,而相邻图块的局部置换仅产生0.16个百分点的微弱波动,说明S4长程状态转移更依赖全局粗粒度的“背景→前景核心”推进路径。
  • 空间信息增益真实存在且独立于中心偏置:所有骨干网络在各数据集上的 \(\mathrm{ERSI}_{20}\) 均显著大于0(最高达0.484 bits/patch),且成对紧凑度比值 \(R_{20} < 0.70\),有力证实了学得的路由并非拟合相机先验中心,而是自适应聚集于物体语义部件。

亮点与洞察

  • 探针解码器属性与可解释性解耦的数学契合:巧妙选用离散LTI特性的对角S4状态空间模型作为探针读取器,其解析展开式产生的恒定位置系数 \(\alpha_k\) 摆脱了注意力探针中内容与位置纠缠的动态权重黑盒,使位置因果分析具备理论封闭解。
  • 打分标准化与一维连续最优传输的端到端桥接:将离散排序问题松弛为基于方差归一化得分与规则一维目标点间的Sinkhorn迭代,既避免了argmax不可导与硬排列梯度消失,又以仅增加768个参数的极低代价实现了全图块拓扑重排。
  • 建立超额空间信息增益标准:提出剥离边际位置偏差后的 \(\mathrm{ERSI}\) 统计量,为自监督表征评估提供了超越单纯分类Top-1精度的空间组织度度量标尺。

局限与展望

  • 下游解码器与打分权重的联合依赖:Sinkhorn打分向量与S4分类头联合端到端优化,所学路由受到特定下游任务目标和S4时序偏置的双重牵引,不能直接等同于预训练骨干本身的纯无监督本征注意力。
  • 评估场景目前集中于图像级细粒度分类:主要实验建立在ImageNet-1K、CUB-200-2011和Stanford Cars等分类基准上,未来有待拓展至密集预测任务(如目标检测、语义分割或开放词表定位)中的patch拓扑路由审计。
  • 对大尺度几何形变敏感:实验发现当输入图像遭受大幅度剪裁(保留90%维度)时,路由重合度降至69.30%,提示当前以全局相对打分为基础的路由机制对边缘截断具备一定的敏感性。

相关工作与启发

  • 对比传统探针 (Linear Probes / Alain & Bengio 2016):传统方法局限于静态GAP或[CLS]分类,忽视了空间图块序列的内部拓扑;SSMProbe通过学习排列矩阵,释放了密集patch维度的诊断能力。
  • 对比视觉状态空间主干 (VMamba / Liu et al. 2024; Vim / Zhu et al. 2024):VMamba等工作探索将SSM作为骨干架构,依赖人工设计的4向交叉硬编码扫描;SSMProbe则将SSM降维作为探针诊断工具,利用可微最优传输学习数据自适应路由。
  • 对比可微排序网络 (NeuralSort / Grover et al. 2019; SoftSort / Prillo et al. 2020):NeuralSort易陷入局部极值且位置覆盖率不足1%,SoftSort易在边界塌缩;SSMProbe的双随机Sinkhorn迭代覆盖率高达57.1%,与序列模型的长记忆衰减更为匹配。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次将Sinkhorn最优传输排序与S4的LTI时序解析特性结合用于ViT表示探测,视角独特且数学推导闭环。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖4大代表性预训练骨干、3大经典视觉基准,并设计了精确的数千万次排列置信度分析与扰动消融。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑严密,符号定义标准,图表数据与附录技术细节完备扎实。
  • 价值: ⭐⭐⭐⭐⭐ 为深入理解视觉Transformer深层特征组织、表征空间拓扑及探索轻量化自适应序列建模提供了极其重要的诊断基准。