Learning Dynamic Evidence Routes for Vision Transformer Probing
TL;DR
SSMProbe is introduced, an explicitly inspectable probe that replaces invariant pooling with a Sinkhorn-learned evidence route followed by a diagonal S4 decoder and uses the mathematical transparency of state-space models to turn a frozen ViT readout into an auditable evidence-routing analysis.
Abstract
Probing frozen vision transformers typically uses permutation-invariant aggregation (GAP or $\texttt{[CLS]}$), treating patch tokens as an unstructured set. Content-dependent probes such as self-attention are useful accuracy controls, but they do not expose a fixed token schedule or fixed position weights for auditing. We introduce $\textbf{SSMProbe}$, an explicitly inspectable probe that replaces invariant pooling with a Sinkhorn-learned evidence route followed by a diagonal S4 decoder. The S4 decoder is a linear time-invariant (LTI) system whose final state has fixed, position-dependent coefficients, so the probe-induced routed sequence can be audited as a concrete object rather than inferred only from accuracy. Our central measurement is the geometry of routed evidence: which patch tokens are moved to influential positions by this diagnostic, whether those tokens form spatially organized regions or random-like dispersed sets, and how the fixed S4 kernel weights them. Across MAE, BEiT, DINOv2, and supervised ViT, this route geometry separates MAE's dispersed, nearly random-like routes from the more spatially organized routes of BEiT, ViT, and DINOv2, with DINOv2 retaining a distinct strong $\texttt{[CLS]}$ profile. SSMProbe uses the mathematical transparency of state-space models to turn a frozen ViT readout into an auditable evidence-routing analysis.
