기술 요약: IG-Lens

문제 정의

디코더 전용 트랜스포머(Decoder-only transformers)는 종종 잔차 연결(residual stream)이 다음 토큰에 대한 실행 중인 추정치를 정교화하는 반복적 추론 과정으로 개념화됩니다. 여기서 핵심적인 해석 가능성 질문은 다음과 같습니다: 정확히 어느 레이어에서 예측된 토큰의 확률이 결정되는가?

기존의 레이어별 판독(layer-wise readout) 도구들은 확률 공간에서 이를 정밀하게 답변하는 데 실패합니다:

• Logit Lens & Tuned Lens: 이들은 각 레이어에서의 확률 수준을 보고하지만, 가산적 분해(additive decomposition)를 제공하지 않습니다. 이들의 추정치는 편향되어 있으며 깊이에 따라 단조성을 띠지 않으므로, 레이어별 확률의 합은 수학적 의미를 갖지 못합니다.
• Direct Logit Attribution (DLA) & Embedding Methods: 이들은 가산적이지만 **로짓 공간(logit space)**에서 작동합니다. 소프트맥스(softmax) 함수는 매우 비선형적이기 때문에, 로짓의 가산적 분해가 실질적인 관심 대상인 확률의 가산적 분해로 직결되지 않습니다.
• Layer Conductance & Layer Integrated Gradients (IG): 이 방법들은 그래디언트를 레이어별로 적분하지만, 각 레이어를 독립적인 베이스라인에 귀속시킵니다. 결과적으로 각 레이어 항들이 전체 예측 확률 변화량으로 "망원경처럼 접히며 합쳐지는(telescope)" 성질을 갖지 못합니다.

방법론: IG-Lens

저자들은 IG-Lens라는 방법을 도입하여, 트랜스포머 레이어 전반에 걸친 확률의 정확한 가산적 분해를 제공합니다.

핵심 개념

IG-Lens는 베이스라인으로부터 최종 레이어까지의 은닉 상태(hidden states)를 통과하는 단일 경로를 따라 Integrated Gradients(IG)를 망원경식(telescoping)으로 적용합니다.

1. 경로 정의: 방법론은 $0 = L_0 < L_1 < \dots < L_k = n$인 레이어 시퀀스를 선택하고, $h_{L_0} \to h_{L_1} \to \dots \to h_{final}$로 이어지는 은닉 상태의 경로를 정의합니다.
2. Head-Once Readout: 결정적으로, 판독 함수 $f(h)$는 단일 은닉 벡터에 단 한 번 적용되는 최종 언임베딩(unembedding)으로 정의됩니다: $f(h) = \text{softmax}(\text{head}(\text{norm}(h)))_{y_t}$. 이 함수에는 어텐션이나 MLP가 포함되지 않습니다. 모든 토큰 간의 혼합(cross-token mixing)은 분리된 은닉 상태를 생성하는 순전파 과정에서 이미 발생했다고 가정합니다.
3. Telescoping Attribution: 레이어 $L_j$의 기여도는 $h_{L_{j-1}}$에서 $h_{L_j}$로 이어지는 직선 경로를 따른 $f$의 세그먼트 IG로 정의됩니다. 그래디언트 정리(gradient theorem)에 의해, 세그먼트에 대한 그래디언트의 적분은 양 끝점 값의 차이와 같습니다. 이러한 세그먼트들을 합치면 다음과 같은 망원경식 합(telescoping sum)이 됩니다:
   $$\sum_{j=1}^k IGL_j = f(h_{final}) - f(h_{base}) = \Delta p$$
   이를 통해 소프트맥스와 LayerNorm이 선형화되거나 고정되지 않고 적분 경로 내에 포함됨으로써, 합산된 기여도가 정확히 타겟 확률의 변화량($\Delta p$)과 일치하도록 보장합니다.

예측 인지 추정기 (Prediction-Aware Estimator)

적분은 정확하지만, 유한 단계 추정기는 이산화 오차(discretization error)를 유발할 수 있습니다. IG-Lens는 기본값으로 예측 인지 추정기(IDGI에서 영감을 받음)를 채택합니다:

• 원시 그래디언트에 기반하여 단계를 할당하는 대신, 관측된 타겟 확률의 변화($f(h^{(s)}) - f(h^{(s-1)})$)에 따라 각 단계를 할당합니다.
• 스칼라 확률 판독의 경우, 이는 각 세그먼트를 양 끝점 값의 망원경식 합으로 축약시킵니다.
• 결과: 어떤 단계 수($m \ge 1$)에서도 부동 소수점 정밀도 수준에서 완결성(completeness)이 유지됩니다. 모델이 그래디언트 민감도는 보이지만 출력 변화는 없는 단계들은 자연스럽게 필터링되어 리만 이산화 오차를 제거합니다.

구현

이 방법은 역전파(backward pass/autograd)를 필요로 하지 않는 **단일 패스 배치 알고리즘(single-pass batched algorithm)**을 사용합니다.

• 모든 보간 지점(interpolation points)을 하나의 텐서로 쌓아 포인트와이즈(pointwise) 함수 $f$를 한 번의 순전파로 평가함으로써, 그래디언트 계산의 오버헤드를 피합니다.
• 이러한 효율성은 특히 "head-once" 설계 덕분에 가능하며, 이는 토큰 지점들 사이의 독립성을 보장합니다.

주요 기여

1. 정확한 가산적 분해: IG-Lens는 레이어 전반에 걸쳐 (로짓이 아닌) 확률의 정확한 가산적 분해를 제공하는 최초의 방법이며, 여기서 기여도의 합은 $\Delta p$와 일치합니다.
2. Telescoping Identity: 본 논문은 세그먼트 적분이 전체 예측 변화량으로 접히며, 소프트맥스를 적분 경로 내에 유지한다는 수학적 증명(Proposition 1)을 제공합니다.
3. 예측 인지 추정: 기본 추정기는 단계 수에 관계없이 기계 정밀도의 완결성을 보장하며, "민감하지만 반응은 없는" 단계를 필터링합니다.
4. 효율적인 구현: 역전파 호출 없이 전체 어트리뷰션 맵을 계산하는 단일 패스 순전파 평가를 통해 계산 효율성을 확보했습니다.
5. 위치 선정: 저자들은 기존 방법들이 왜 확률 공간에서 실패하는지(로짓 공간의 가산성, 편향, 또는 망원경식 구조의 부재)를 명확히 규명하고, IG-Lens를 로짓 공간 방법론에 대응하는 확률 공간의 대응물로 포지셔닝합니다.

결과 및 평가 계획

• 수치적 검증: 통제된 테스트와 실제 모델(Llama-3.2-1B-Instruct)을 통해, 기본 추정기가 레이어별 엔드포인트 차이를 약 $10^{-7}$(float32 정밀도) 수준으로 재현함을 확인했습니다. 망원경식 합은 단계 수 $m$에 관계없이 부동 소수점 정밀도 내에서 정확히 성립합니다.
• 예시 출력: "What is the capital of Vietnam?"이라는 프롬프트 실험에서, 이 방법은 로그 렌즈(log-lens) 값이 낮음에서 높음으로 전환되는 특정 레이어 세그먼트에 어트리뷰션 질량을 성공적으로 집중시켰습니다. 각 토큰의 어트리뷰션 합은 $\Delta p$와 부동 소수점 정밀도까지 일치합니다.
• 평가 전략: 저자들은 "온셋(onset, 시작점)"에는 내재적인 정답(ground truth)이 없다는 점을 인정합니다. 따라서 이들은 **인과적 절제(causal ablation)**를 통해 방법을 검증할 것을 제안합니다. 만약 IG-Lens가 레이어 $L$을 온셋으로 식별했다면, $L$을 절제했을 때 낮은 온셋 점수를 가진 레이어들을 절제했을 때보다 타겟 확률을 더 크게 감소시켜야 합니다. 이 테스트는 실행 비용이 저렴하며 주장을 반증할 수 있는 능력을 갖추고 있습니다.

의의 및 주장

본 논문은 IG-Lens를 반복적 추론 관점에 대한 정량적 정교화로 포지셔닝합니다. "이 레이어에서의 실행 중인 예측은 무엇인가?"를 묻는 Logit Lens와 달리, IG-Lens는 "최종 확률 중 얼마만큼이 각 레이어에 책임이 있는가?"를 묻고, 그 책임의 합이 전체와 같음을 보장합니다.

제한적인 주장 및 한계:

• 조건부 한계(Conditional Marginal): 저자들은 IG-Lens가 조건부 한계를 측정한다고 명시합니다. 즉, $h_{L-1}$이 주어졌을 때 $h_L$로부터 norm+head가 읽어낼 수 있는 추가적인 확률을 측정합니다. 이는 $h_L$이 상위 어텐션 및 MLP 블록을 통해 라우팅되는 효과를 배제합니다.
• 트레이드오프: 이 방법은 "총 인과적 효과(total causal effect via the full forward pass)" 대신 "합 = $\Delta p$"를 선택했습니다. 총 인과적 효과를 위해서는 구성 요소 분해(DLA와 같은)가 여전히 적절한 도구입니다.
• 비단조성(Non-Monotonicity): 이 방법은 확률 궤적이 비단조적인 경우(확률이 구축되었다가 부분적으로 취소되는 경우)를 처리하기 위해 부호가 있는 어트리뷰션(signed attributions)을 사용하여, 절대값 기반 방법보다 더 충실한 레이어 기여도를 제공합니다.

중심적인 미결 과제는 여전히 경험적인 영역에 있습니다: IG-Lens의 온셋이 (특히 비단조적 궤적을 가진 토큰에 대해) 인과적 절제가 예측을 가장 많이 감소시키는 레이어를 예측하는지 여부입니다.