Step-Resolved Data Attribution for Looped Transformers

We study how individual training examples shape the internal computation of looped transformers, where a shared block is applied for $\tau$ recurrent iterations to enable latent reasoning. Existing training-data influence estimators such as TracIn yield a single scalar score that aggregates over all loop iterations, obscuring when during the recurrent computation a training example matters. We introduce Step-Decomposed Influence (SDI), which decomposes TracIn into a length-$\tau$ influence trajectory by unrolling the recurrent computation graph and attributing influence to specific loop iterations. To make SDI practical at transformer scale, we propose a TensorSketch implementation that never materialises per-example gradients. Experiments on looped GPT-style models and algorithmic reasoning tasks show that SDI scales excellently, matches full-gradient baselines with low error and supports a broad range of data attribution and interpretability tasks with per-step insights into the latent reasoning process.

Some AI models solve hard problems by repeatedly refining their internal representations instead of writing out every intermediate reasoning step in text. This paper introduces Step-Decomposed Influence, a method for understanding which training examples affect each internal refinement step of such looped transformers. The method decomposes a standard training-data influence score into a trajectory over recurrent steps and uses TensorSketch to make the computation practical without storing full per-example gradients. Experiments on algorithmic reasoning tasks and a looped chat model show that this step-by-step view can reveal when training examples matter, detect cancellation hidden by scalar influence scores, and help explain why harder problems benefit from more internal computation.