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

TL;DR

ExPLAIND is a unified, theoretically grounded framework that combines model, data, and training attribution methods to better understand complex model behaviors and training dynamics.

Contribution

It generalizes gradient path kernels to realistic optimizers, introduces novel influence scores, and jointly interprets model components and data during training.

Findings

Accurately replicates CNN and Transformer models using kernel reformulation.

Influence scores are effective for parameter pruning.

Analyzes Grokking, confirming and refining its stages.

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.