주의력 헤드 영향 기반 추론 데이터 선택

초록

본 논문은 LLM의 추론 능력을 효율적으로 전이하기 위해, 모델 내부의 ‘검색 헤드’가 담당하는 토큰 복사 메커니즘을 활용한다. 중요한 헤드를 식별하고 이를 비활성화한 약화 모델과 원본 모델 간의 손실 차이를 ‘Attention Influence Score’로 정의해, 단계별·샘플별 중요도를 정량화한다. 이를 기반으로 단계 가중 SFT와 샘플 선택을 수행해 기존 길이·엔트로피·손실 기반 방법보다 높은 추론 정확도를 달성한다.

상세 분석

AIR는 기존 데이터 선택 방법이 “표면적인” 메트릭(길이, 엔트로피, 전체 손실 등)에 의존해 추론 단계별 인과적 중요성을 포착하지 못한다는 문제점을 정확히 짚는다. 논문은 두 가지 핵심 가정을 제시한다. 첫째, 최신 메커니즘 해석 연구에서 특정 어텐션 헤드가 ‘검색(head)’ 역할을 수행하며, 토큰‑레벨 복사(copy‑paste) 메커니즘을 담당한다는 점이다. 둘째, 올바른 추론 단계가 실행될 때 이 검색 헤드가 활성화되므로, 헤드의 비활성화가 해당 단계의 손실을 크게 증가시킨다면 그 단계는 모델에게 필수적이라고 판단할 수 있다.

헤드 식별 과정은 토큰 출현(C1)과 최대 어텐션 가중치(C2)를 동시에 만족하는 경우를 ‘검색 토큰’으로 정의하고, 전체 컨텍스트 대비 검색 성공률(R_h)를 계산한다. R_h가 상위 δ(예: 5%)에 속하는 헤드들을 ‘핵심 검색 헤드’로 선정한다. 여기서 δ는 실험적으로 조정 가능한 하이퍼파라미터이며, 작은 값일수록 보다 보수적인 헤드 선택이 이루어진다.

핵심 헤드를 마스크하는 약화 모델(θ_ref)은 해당 헤드의 어텐션 가중치를 균등 분포(1/L)로 강제함으로써, 정보 전달 경로를 사실상 차단한다. 이때 파라미터 자체는 변하지 않으므로, 손실 차이는 순수히 어텐션 메커니즘의 기여도만을 반영한다. 손실 차이 Δℓ(x_t)=ℓ(θ_ref, x_t)−ℓ(θ_base, x_t)는 양수일 경우 해당 토큰이 검색 헤드에 크게 의존했음을 의미한다.

샘플 수준 스코어 S_sample은 토큰별 Δℓ을 토큰별 기본 손실 ℓ(θ_base)로 정규화한 비율의 합으로 정의한다. 이는 샘플 간 난이도 차이를 보정하면서, 전체 추론 흐름에서 검색 메커니즘 의존도가 높은 샘플을 정량적으로 추출한다. 반면 단계 수준 스코어 S_step(k)는 한 단계에 속한 토큰들의 평균 Δℓ을 사용해, 단계별 중요도를 측정한다. 상위 P% 단계는 가중치 α(≥1)만큼 부스트되어 SFT 손실에 반영된다.

실험에서는 59K 규모의 CoT 데이터셋을 대상으로, AIR 기반 샘플 상위 10%를 선택하거나, 단계 가중 SFT를 적용한 모델이 기존 길이·엔트로피·전체 손실 기반 선택보다 일관되게 높은 정확도를 기록한다. 특히, “핵심 단계만 강조” 전략이 전체 학습 효율을 크게 높이며, 동일한 파라미터 수에서도 RL 기반 방법에 근접하거나 능가한다는 점이 주목할 만하다.

이 접근법의 강점은 (1) 사전 학습된 모델을 그대로 활용해 추가 학습 없이 메커니즘을 해석한다는 점, (2) 손실 차이 기반 스코어가 직접적인 인과관계를 제공해 데이터 선택의 신뢰성을 높인다는 점, (3) 단계‑샘플 이중 스코어링이 데이터 효율성을 극대화한다는 점이다. 한계로는 (a) 검색 헤드가 모델마다 다르게 분포할 수 있어 사전 헤드 탐색 비용이 발생한다는 점, (b) 현재는 토큰 복사에 초점을 맞추어 사실적 추론(예: 수학 연산)에는 추가적인 메커니즘이 필요할 수 있다는 점을 들 수 있다. 향후 연구는 다양한 모델 아키텍처와 복합 추론 유형에 대한 헤드 식별 자동화, 그리고 헤드 마스킹 외에 가중치 조정 방식 등을 탐색할 여지가 있다.