루프형 트랜스포머의 단계별 데이터 영향 분석

초록

**
본 논문은 공유 블록을 τ번 반복 적용하는 루프형 트랜스포머에서, 개별 학습 샘플이 어느 반복 단계에 영향을 미치는지를 정량화하는 Step‑Decomposed Influence (SDI)를 제안한다. 기존 TracIn이 전체 루프에 대한 단일 스칼라 점수만 제공하는 반면, SDI는 각 단계별 영향 궤적을 제공한다. 메모리 효율을 위해 TensorSketch 기반의 스케치‑during‑backprop 기법을 도입해 대규모 모델에서도 실시간으로 계산 가능하게 하였으며, 알고리즘ic reasoning, Sudoku, 대형 LLM 등 다양한 실험을 통해 정확도와 해석 가능성을 입증한다.

**

상세 분석

**
루프형 트랜스포머는 동일한 파라미터 블록을 τ번 순환 적용함으로써, 파라미터 수와 계산 깊이를 분리한다. 이 구조는 테스트 시 compute budget을 동적으로 조절할 수 있다는 장점이 있지만, 내부 상태가 단계별로 어떻게 변하는지, 그리고 어떤 학습 데이터가 어느 단계에 기여하는지는 기존 영향 추정 방법으로는 파악하기 어렵다. 기존 TracIn은 전체 파라미터에 대한 손실 그래디언트의 내적을 누적해 단일 스칼라를 산출한다. 이는 루프 내부에서 발생하는 미세한 시점 차이를 가려버린다.

논문은 먼저 손실에 대한 전체 파라미터 w_body의 미분을 단계별 외적 합으로 전개한다(Prop. 1). 구체적으로 각 토큰 j와 단계 t에 대해 ∂ℓ/∂h_{t,j}·∂h_{t,j}/∂w_body 를 구하고, 이를 토큰과 단계에 대해 합산하면 전체 그래디언트가 된다. 매트릭스 파라미터(W)와 편향(b) 각각에 대해 외적 형태(δ_{t,j}⊗a_{t,j})와 벡터 형태(δ_{t,j})로 표현할 수 있다.

이 전개를 기반으로 Step‑Decomposed Influence(I_t) 를 정의한다. I_t(z, z′)는 학습 샘플 z의 전체 그래디언트와 테스트 샘플 z′의 t‑단계 그래디언트 ϕ_t(z′) 사이의 내적을 학습률 η_k 로 가중합한 값이다. 모든 단계에 대해 합산하면 기존 TracIn과 정확히 동일함을 보이며(식 6), 따라서 정보 손실이 전혀 없다는 보존 정리를 제공한다.

실제 구현에서는 단계별 전체 그래디언트를 저장하는 것이 메모리·시간적으로 비현실적이다. 이를 해결하기 위해 두 가지 스케치 기법을 결합한다. 벡터 파라미터는 CountSketch(CS)를, 매트릭스 파라미터는 TensorSketch(TS)를 사용한다. TS는 외적 u⊗v 를 직접 스케치하는 방식으로, O(d_out + d_in + m log m) 시간에 스케치를 생성한다. 이 스케치는 선형성이 보장되므로, 백프로파게이션 중에 각 토큰·단계별 a_{t,j}와 δ_{t,j}를 바로 스케치하고 누적한다. 결과적으로 per‑example 그래디언트를 전혀 물리화하지 않고도 SDI를 추정할 수 있다.

이론적으로는 TS의 분산 상한을 기존 외적 스케치보다 엄격히 낮게 증명했으며, 실험에서는 m=1024 정도의 차원에서도 원본 그래디언트와 거의 동일한 상관관계를 보였다.

실험에서는 (1) 파리티 자동화 회로 복원, (2) Sudoku 풀이 시 추가 루프 단계가 어느 시점부터 효과가 감소하는지 “influence horizon” 탐지, (3) 330M 파라미터 LLM(NanoChat)에서 루프 진행에 따라 내부 표현이 점진적으로 성장하는 현상을 시각화했다. 모든 실험에서 SDI는 전체 TracIn과 평균 절대 오차 <0.02를 기록했으며, 단계별 영향 그래프를 통해 기존 스칼라 기반 분석으로는 놓쳤던 신호 상쇄 현상이나 특정 단계에 집중된 학습 데이터의 역할을 명확히 드러냈다.

결과적으로 SDI는 루프형 트랜스포머의 “생각 과정”을 정량화하는 도구로서, 모델 디버깅, 데이터 정제, 테스트‑타임 compute 조절 등 다양한 실용적 응용 가능성을 열어준다.

**