Can kernelized parameter sensitivity atlases amortize conditional interpretability?

A stored atlas replaces per-query gradients

Gradient-based attribution answers a conditional question: which parameters matter for this input and this model functional? A fresh backward pass gives a high-resolution answer, but its latency and memory costs recur for every query, while a single global importance vector is cheap but erases input dependence. The paper reframes the repeated-query problem as statistical estimation: from a finite source set whose gradients were measured once, estimate the sensitivity profile of an unseen query using only a compact representation and a forward computation.

The target is defined through the empirical neural tangent kernel (NTK), the inner product of parameter gradients. Because the profiled parameters are split into disjoint groups, the NTK decomposes exactly into per-group blocks, and each diagonal block equals the expected squared first-order susceptibility to a random perturbation of that group. Normalizing these diagonal blocks gives a sensitivity profile s(x) that sums to one across groups, so it records where sensitivity lies while discarding the input-dependent total gradient scale.

KPSA measures these profiles on a fixed source set, embeds each source in a categorical or learned query space, and transfers profiles to new inputs with a nonnegative retrieval kernel. Vision uses a Nyström approximation of that kernel; forecasting and protein modeling evaluate it against all sources. The retrieval kernel is distinct from the model NTK: the NTK defines the target, the retrieval kernel moves it between inputs.

More sources cannot fix a lossy representation

Without a query gradient, the object actually available is not the realized profile s(x) but its conditional mean given the query representation Z: m(z) = E[s(X) | Z = z]. The paper separates the two error sources formally. A finite-source bound splits the gap into representation error, local target-instability error, and finite-source error, and only the last shrinks as sources are added.

This is the paper's sharpest conceptual claim. If the representation merges inputs with different sensitivity profiles, no number of source gradients recovers the discarded variation. The authors define an irreducible sensitivity risk for a query space and show that no predictor using only Z can beat it. A separate margin condition then converts uniform profile error into exact top-k recovery: the retrieved top-k groups match the oracle's whenever the sensitivity gap exceeds twice the error bound.

The practical reading is that the query space, not the atlas size, sets the ceiling. The paper's own diagnostics reflect this: the categorical query space is competitive in vision but contributes little on the two forecasting transfers and the protein task, and nearest-neighbor retrieval is slightly stronger on ETTh1 but slightly weaker on Weather and protein. The authors present this as evidence about the representation rather than support for a modality-independent best kernel.

Sparse circuits recover 39-85% of oracle susceptibility

The first application asks whether one atlas can pick sparse parameter-group sets whose exact perturbation effects stay input dependent. Across a fivefold range of parameter budgets (0.02-0.1%), the selected groups attain 39-85% of the perturbation susceptibility of a first-order backpropagation oracle, measured with exact antithetic parameter perturbations rather than gradient proxies.

At the 0.1% budget, semantic KPSA recovers 34.7%, 46.5%, 50.2%, and 33.2% of the constant-to-direct causal gap on ImageNet, ETTh1, Weather, and Swiss-Prot respectively, while its selected sets contain 71.8%, 72.9%, 74.6%, and 42.8% of the same-budget direct gradient energy. Those two numbers answer different questions: gradient energy captured is a first-order quantity, whereas the causal gap is measured by finite forward perturbations, which can retain nonlinear error.

The circuit claim is narrower than it may sound. At 0.1%, KPSA selections occupy an average of 6.0 ViT blocks and 5.8 Chronos blocks, and their overlap rises with representation similarity: from the lowest-similarity bin to the top 1%, the Jaccard coefficient increases by 0.177 in vision and 0.471 in forecasting, against 0.089 and 0.210 after shuffling the representation-profile pairing. The exact susceptibility gain over a layer-count-matched constant selection is 26.3% of the direct-gradient scale in vision and 15.8% in forecasting. The authors explicitly do not claim a signed directed computational graph; circuit here means an input-conditioned set of native parameter groups with measured local causal utility.

Neuron localization trades accuracy for throughput and memory

The second application localizes MLP features in transformers for exact zero-deactivation, evaluated against an action-matched gradient-times-activation reference. At the 6,400-source confirmation operating point, feature-level KPSA recovers 80.9% of the activation-attribution gap while processing 2,041 images/s, 3.7 times the reference throughput, and using 0.41 rather than 4.24 GiB of peak allocated memory. Its persistent atlas is 112.5 MiB.

Compressing the partition changes the tradeoff. A 32-feature partition cuts atlas storage to 3.5 MiB and retains 59.3% gap recovery, sharing the same query cost because only the final selection budget changes. On protein modeling, semantic KPSA is the strongest approximate method at 68.3%, compared with 56.6% for nearest retrieval and 4.4% for the categorical kernel.

The amortization is finite rather than free. The measured profile-only construction cost for the cost-balanced 1,600-source atlas breaks even after roughly 15,000 batched queries, and that figure excludes representation encoding and prototype fitting. Below that query volume, the one-time build has not paid for itself.

Fixed model, fixed partition, unsigned energy

The guarantees are scoped tightly, and the paper says so. KPSA estimates sensitivity for a fixed model, functional, parameter partition, representation, and source distribution. It cannot recover variation the representation discards; fixed bandwidths retain a smoothing-bias floor even as sources grow; normalized unsigned energy encodes neither sign nor directed computation; exact top-k recovery requires a margin; and the bridge from first-order sensitivity to finite perturbations is local in parameter space.

Two further boundaries matter for anyone considering deployment. Distribution shift is left as an open statistical problem, as are adaptive bandwidths and signed or interaction-aware atlases. And the i.i.d. finite-source rate is not invoked for the fixed, overlapping forecasting windows, although the deterministic ranking statements still apply there.

The ablation also cautions against assuming the kernel does the work. Nearest-neighbor retrieval in KPSA space can remain competitive with kernel aggregation, which suggests that much of the benefit comes from the query space and stored profiles rather than from a sophisticated weighting scheme.

Choose the atlas when queries repeat and the representation fits

The paper's own comparison points to a specific decision rule rather than a general endorsement. The atlas is worth building when the same frozen model and functional will be queried many times, when the representation preserves the sensitivity variation that matters, and when the query volume clears the measured break-even of roughly 15,000 batched queries for the 1,600-source vision atlas.

When those conditions hold, the payoff is a shift in cost structure: per-query reverse-mode differentiation is replaced by a forward representation plus kernel retrieval over a compact atlas, with 81% of the vision activation-attribution gap retained at 2,041 images/s and 0.44 GB peak memory. When the representation is weak, the evidence is equally direct: the categorical kernel reaches only 4.4% gap recovery on protein modeling, and no amount of additional source data repairs that.

The most useful distinction the paper draws is therefore between two failure modes that look alike from the outside. A disappointing atlas may be under-sourced, which more gradients can fix, or it may be representation-limited, which they cannot. The reported diagnostics, held-out profile retrieval and effective source counts, are the paper's way of telling those cases apart before trusting a retrieved sensitivity circuit.

Sources used in this answer