English

Step-resolved data attribution for looped transformers

Machine Learning 2026-02-11 v1 Artificial Intelligence

Abstract

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 \textit{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.

Cite

@article{arxiv.2602.10097,
  title  = {Step-resolved data attribution for looped transformers},
  author = {Georgios Kaissis and David Mildenberger and Juan Felipe Gomez and Martin J. Menten and Eleni Triantafillou},
  journal= {arXiv preprint arXiv:2602.10097},
  year   = {2026}
}
R2 v1 2026-07-01T10:30:15.048Z