ExPLAIND: Unifying Model, Data, and Training Attribution to Study Model Behavior

TL;DR

ExPLAIND provides a theoretically grounded framework that unifies attribution to model components, data, and training dynamics by extending the Exact Path Kernel to AdamW. It introduces a tensor-based accumulation of influence scores that can be aggregated to deliver parameter-, data-, and step-level explanations, and validates these scores through parameter pruning experiments. The framework is applied to Grokking in a modulo-addition Transformer, uncovering a sequence of learning phases and a cyclic data-geometry alignment among embeddings and decoder layers. Overall, ExPLAIND offers a principled, holistic view of modern training dynamics and model behavior with practical implications for interpretability and model auditing.

Abstract

Post-hoc interpretability methods typically attribute a model's behavior to its components, data, or training trajectory in isolation. This leads to explanations that lack a unified view and may miss key interactions. While combining existing methods or applying them at different training stages offers broader insights, such approaches usually lack theoretical support. In this work, we present ExPLAIND, a unified framework that integrates all these perspectives. First, we generalize recent work on gradient path kernels, which reformulate models trained by gradient descent as a kernel machine, to realistic settings like AdamW. We empirically validate that a CNN and a Transformer are accurately replicated by this reformulation. Second, we derive novel parameter- and step-wise influence scores from the kernel feature maps. Their effectiveness for parameter pruning is comparable to existing methods, demonstrating their value for model component attribution. Finally, jointly interpreting model components and data over the training process, we leverage ExPLAIND to analyze a Transformer that exhibits Grokking. Our findings support previously proposed stages of Grokking, while refining the final phase as one of alignment of input embeddings and final layers around a representation pipeline learned after the memorization phase. Overall, ExPLAIND provides a theoretically grounded, unified framework to interpret model behavior and training dynamics.

Figures (10)

• Figure 1: The ExPLAIND framework is based on a decomposition of the model along its components, data, and training steps. Explanations are obtained by accumulating the resulting influence scores.
• Figure 2: We prune the CNN weights and achieve test accuracy comparable to liPruningFiltersEfficient2017. Across all sparsity levels, the KL divergence of our model outputs is consistently lower. We report the means and standard deviations over 5 runs.
• Figure 3: Training statistics, influences, and validation of our representation pipeline hypothesis.
• Figure 4: At different training stages, we show the influences $\Psi_S(\Theta_{dec}, x, x')$ of the training examples $x'$ on the test samples $x$ of the decoder layer $\Theta_{dec}$, summed over the output dimensions, for predictions on the test set and the training set, accumulated over the preceding $50$ steps. The plots are labeled and ordered by the sum of inputs $a+b$ and the corresponding result ($a+b \bmod 113$). Corresponding figures for the other layers are provided in the Appendix in \ref{['app:kernel_matrices']}.
• Figure 5: Similarity $Sim_{\Theta_{layer}}(x, x')$ of predictions of the test set of the Transformer model accumulated over different training stages. All layers and other training stages shown in Appendix \ref{['fig:app-epk-val-sim']}.

Theorems & Definitions (12)

• Theorem 3.1: Extension of bellExactKernelEquivalence2023 to AdamW
• proof : Sketch of Proof
• Corollary 3.2: Gradient Descent with Momentum
• proof
• Corollary 3.3: Loss and intermediate activations
• proof
• Definition 4.1: Tensor of Influences
• Definition 4.2: Accumulated influence
• Definition 4.3: Influence of the regularization
• Definition 4.4: Similarity from the model's perspective