Kernelized Parameter Sensitivity Atlases for Amortized Conditional Interpretability
Organizations: Technische Universität Darmstadt · Daffodil International School · University of Dhaka
Abstract
Gradient-based attribution provides input-specific information about model parameters, but computing it via backpropagation for many samples is expensive, while static importance scores discard input dependence. We derive KPSA (Kernelized Parameter Sensitivity Atlas) from a block decomposition of the empirical neural tangent kernel. It measures normalized parameter-group sensitivity, embeds those sources in an explicit query space, and uses a lightweight kernel to estimate sensitivities for unseen inputs. We derive bounds that separate finite-source from irreducible representation error and give margin-based top- recovery. This paper studies two interpretability applications. Firstly, KPSA identifies parameter groups with high susceptibility to perturbations across vision, forecasting, and protein-modeling setups. Across a fivefold range of parameter budgets (-), its selected groups attain - of the perturbation susceptibility of a first-order backpropagation oracle. In vision and forecasting, these parameter groups form sparse sensitivity circuits that extend across depth and vary systematically with the input. Secondly, we use KPSA on an activation-attribution problem for MLP features in transformers. It retains of the gap between a backpropagation oracle and a constant-sensitivity baseline on vision while increasing batched query-time throughput from 551 to 2,041 images/s and reducing peak GPU memory from 4.55 to 0.44GB. On the protein-modeling task, it recovers of this gap. An ablation of KPSA's kernel aggregation shows that nearest-neighbor retrieval in KPSA space can remain competitive. The empirics show that compact kernel atlases can store useful input-dependent parameter sensitivity while amortizing repeated gradient attribution, a step toward scalable conditional interpretability. Code can be found at https://github.com/Axym-Labs/kpsa .
Figures & tables
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
| Use | Setting | Model | Functional | Eligible groups and query kernel | ||
|---|---|---|---|---|---|---|
| Influence | ImageNet-1k | ViT-B/16 (86.6M) | True-class logit minus strongest alternative | 36,864 coupled MLP features; DINOv2-S CLS, 1,000-prototype RBF map, scale .10 | 8,000 | 169 |
| Influence | ETTh1 | Chronos-Bolt-small (47.7M) | Mean 24-step median forecast minus last observation | 2,048 final-decoder FF features; final encoder state, exact RBF, scale .25 | 1,600 | 128 |
| Influence | Weather | Chronos-Bolt-small (47.7M) | Same forecasting functional | Same frozen ETTh1 setting | 1,600 | 128 |
| Influence | Swiss-Prot | ESM2-35M (33.5M) | True masked-residue logit minus strongest alternative | 1,920 layer-3 outgoing FF features; sequence mean, exact RBF, scale .10 | 1,600 | 48 |
| Circuit | ImageNet-1k | ViT-B/16 | ImageNet logit margin | Full 12-block MLP stack; DINOv2-S CLS, 1,000-prototype RBF map, scale .10 | 8,000 | 110 |
| Circuit | ETTh1 | Chronos-Bolt-small | Forecasting functional above | Full 12-block FF stack; final encoder state, exact RBF, scale .25 | 1,600 | 128 |
| Setting | Method | Gap recovered (%) | Gradient energy (%) | |
|---|---|---|---|---|
| Vision | Semantic KPSA | 34.7 [23.7, 47.1] | 71.8 [68.9, 74.7] | 169 |
| Vision | Categorical KPSA | 39.4 [28.8, 50.9] | 75.8 [73.1, 78.4] | 169 |
| Vision | Semantic nearest | 28.1 [16.4, 40.5] | 69.5 [66.2, 72.6] | 169 |
| Vision | Constant sensitivity | 0.0 [0.0, 0.0] | 46.8 [45.1, 48.4] | 169 |
| Vision | Direct gradient | 100.0 [100.0, 100.0] | 100.0 [100.0, 100.0] | 169 |
| ETTh1 | Semantic KPSA | 46.5 [34.7, 56.2] | 72.9 [69.3, 76.2] | 128 |
| Setting | Method | Attribution gap recovered (%) | |
|---|---|---|---|
| Vision | Semantic KPSA | 66.5 [51.4, 84.1] | 76 |
| Vision | Categorical KPSA | 76.5 [62.8, 92.5] | 76 |
| Vision | Semantic nearest | 74.8 [61.4, 90.1] | 76 |
| Vision | Constant sensitivity | 0.0 [0.0, 0.0] | 76 |
| Vision | Activation attribution | 100.0 [100.0, 100.0] | 76 |
| Protein | Semantic KPSA | 68.3 [46.3, 98.5] | 128 |
| Method and budget | Gap recovered (%) | Images/s | Peak GiB | Atlas MiB |
|---|---|---|---|---|
| Activation attribution | 100.0 | 551 | 4.24 | 0.0 |
| Semantic KPSA, 0.02% | 64.1 | 2,041 | 0.41 | 112.5 |
| Semantic KPSA, 0.1% | 80.9 | 2,041 | 0.41 | 112.5 |
| Semantic KPSA, 32-feature groups, 1/12 MLP parameters | 59.3 | 2,088 | 0.41 | 3.5 |
Neural Network Interpretability
Neural network interpretability studies how models represent information and make predictions, and how to explain, test, or control their internal computations.
See all 330 papers