Kernelized Parameter Sensitivity Atlases for Amortized Conditional Interpretability

Davide Wiest; Ahanaf Hasan Ariq; Abdul Monaf Chowdhury

cs.LGOct 8, 2026

Kernelized Parameter Sensitivity Atlases for Amortized Conditional Interpretability

Authors: Davide Wiest, Ahanaf Hasan Ariq, Abdul Monaf Chowdhury

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- $k$ 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 ( $0.02$ - $0.1\%$ ), its selected groups attain $39$ - $85\%$ 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 $81\%$ 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 $68\%$ 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 $F$	Eligible groups and query kernel	$N$	$n$
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

Appendix

Table A1: Frozen experimental instances used in the primary causal and circuit figures. ImageNet-200 is a balanced 200-class subset of ImageNet validation. The circuit-structure rows use a wider parameter scope than the corresponding primary causal rows; the separate efficiency variants are specified in Section A.5 .

Setting	Method	Gap recovered (%)	Gradient energy (%)	$n$
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

Appendix

Table A2: Parameter-influence results at a 0.1% eligible-parameter budget. Entries are query means with paired 95% bootstrap intervals. Gap recovery uses constant sensitivity and the direct query gradient as the 0% and 100% anchors. Gradient energy is measured relative to the direct query gradient.

Setting	Method	Attribution gap recovered (%)	$n$
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

Appendix

Table A3: Neuron-localization results at a 0.1% eligible-parameter budget. Entries are query means with paired 95% bootstrap intervals. Constant sensitivity and action-matched activation attribution define the 0% and 100% anchors.

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

Appendix

Table A4: Measured vision accuracy–cost points from Figure 5 . Throughput and peak memory use batch size 32 on one RTX 5090. Atlas storage is persistent float32 storage and excludes model weights. GiB and MiB use powers of two. The two feature-atlas rows share the same query cost because only the final selection budget changes.

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

Sep 29, 2026

Weights Read and Write Features: Scalable Parameter Decomposition Grounded in Activation Space

Activation space and parameter space provide complementary views of model computation. Activations represent information, while weights read, transform, and write that information. Yet existing interpretability methods largely study the two spaces separately, leaving the connection between represented information and parameter-level computation underexplored. We introduce Activation-Supported Parameter Decomposition (ASPD), which jointly decomposes activation and parameter spaces and grounds each learned weight component in the activation features it reads or writes. This grounding constrains otherwise non-unique parameter decompositions using the model's internal activations, while an internal reconstruction objective provides a local learning signal at the weight matrix being analyzed. Together, these properties enable scalable, interpretable, and causally editable parameter decomposition in pretrained large language models, demonstrated on Qwen-3-8B. The learned read--write components can also be composed into parameter-level mechanism circuits. We use ASPD to recover mechanisms underlying the classic IOI circuit and trace semantic transformations through model weights.

Tue M. Cao, Lisiane Pruinelli, My T. Thai

Department of Computer and Information Science and Engineering, University of Florida, Gainesville, FL 32611, USA · College of Nursing, University of Florida, Gainesville, FL 32611, USA

Mechanistic Interpretability Circuits in Language Models

Ask PDF

Jun 5, 2026

When Attribution Patching Lies: Diagnosis and a Second-Order Correction

A central goal of mechanistic interpretability is to identify which internal components causally drive a language model's behavior. Because these importance estimates serve as the evidence for identifying circuits, systematic errors can lead to the misidentification of the underlying mechanisms. While activation patching provides a gold-standard causal metric, its computational cost is prohibitive at scale. Practitioners instead rely on attribution patching, a gradient-based, first-order approximation whose reliability remains poorly understood. In this work, we characterize the source of this unreliability, demonstrating that the dominant error stems from the non-linearities in the downstream network rather than local curvature at the patched component. This insight yields three practical tools: (i) a reliability score to detect untrustworthy estimates, (ii) error bounds quantifying potential attribution mis-specifications, and (iii) a Hessian-vector-product (HVP) correction that eliminates the leading-order error with only one additional backward pass. In evaluations across five model families (124M-9B parameters) and both random-token and naturalistic (name-swap) perturbations, HVP is the only second-order correction feasible at larger scale, where standard baselines like Integrated Gradients become computationally prohibitive. In comparative experiments, a multi-step HVP variant matches or exceeds the accuracy of Integrated Gradients at significantly lower compute, outperforming prior second-order baselines. These improvements lead to higher-fidelity circuit recovery on standard benchmarks and support a Screen-Flag-Fix workflow that targets computational effort only toward the components flagged as unreliable.

Luyang Zhang, Jialu Wang

1Carnegie Mellon University · University of California, Santa Cruz

Gradient-Based Attribution LLM Interpretability Causal Attribution Mechanistic Interpretability Circuits in Language Models Activation Patching

Ask PDF

Apr 15, 2026

Path-Sampled Integrated Gradients

We introduce path-sampled integrated gradients (PS-IG), a framework that generalizes feature attribution by computing the expected value over baselines sampled along the linear interpolation path. We prove that PS-IG is mathematically equivalent to path-weighted integrated gradients, provided the weighting function matches the cumulative distribution function of the sampling density. This equivalence allows the stochastic expectation to be evaluated via a deterministic Riemann sum, improving the error convergence rate from

O(m^{-1/2})

O(m^{-1})

for smooth models. Furthermore, we demonstrate analytically that PS-IG functions as a variance-reducing filter against gradient noise - strictly lowering attribution variance by a factor of 1/3 under uniform sampling - while preserving key axiomatic properties such as linearity and implementation invariance.

Firuz Kamalov, Fadi Thabtah, R. Sivaraj +1

Department of Computational Sciences, Canadian University Dubai, UAE · Abu Dhabi School of Management, Abu Dhabi, UAE · Department of Mathematics, Dr. B.R. Ambedkar National Institute of Technology, India

Gradient-Based Attribution Integrated Gradients Feature Attribution

Ask PDF

Abstract

Neural Network Interpretability

Explore similar work

Weights Read and Write Features: Scalable Parameter Decomposition Grounded in Activation Space

When Attribution Patching Lies: Diagnosis and a Second-Order Correction

Path-Sampled Integrated Gradients