루프 트랜스포머의 단계별 데이터 기여도 분석

Step-resolved data attribution for looped transformers

본 연구에서는 공유 블록이 τ번 반복적으로 적용되어 잠재적 추론을 가능하게 하는 루프 트랜스포머의 내부 연산에 개별 학습 예제가 어떻게 영향을 미치는지 분석합니다. 기존의 학습 데이터 영향 추정 방법인 TracIn은 모든 반복 횟수를 통합하여 단일 스칼라 값을 제공하며, 이로 인해 학습 예제가 반복 연산 과정 중 언제 중요한 역할을 하는지 파악하기 어렵습니다. 본 연구에서는 TracIn을 반복 연산 그래프를 펼쳐서 분해하고, 각 반복 횟차에 대한 영향력을 할당하는 '단계 분해 영향 (Step-Decomposed Influence, SDI)' 기법을 제안합니다. SDI를 트랜스포머 규모에서 실용적으로 사용하기 위해, 각 예제에 대한 그래디언트를 명시적으로 계산하지 않고 텐서 스케치(TensorSketch) 구현을 제안합니다. 루프 GPT 스타일 모델 및 알고리즘 추론 작업에 대한 실험 결과, SDI는 뛰어난 확장성을 가지며, 낮은 오차로 전체 그래디언트 기반 방법과 유사한 성능을 보입니다. 또한, SDI는 다양한 데이터 기여도 분석 및 해석 작업에 활용될 수 있으며, 잠재적 추론 과정에 대한 단계별 통찰력을 제공합니다.

We study how individual training examples shape the internal computation of looped transformers, where a shared block is applied for $τ$ 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-$τ$ 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.