🤖 machine learning

Hidden States Know Where Reasoning Diverges: Credit Assignment via Span-Level Wasserstein Distance

이 논문은 결과값(정답 여부)만으로도 은닉 상태(hidden state)의 분포 차이를 이용해 추론 과정의 오류 지점을 찾아낼 수 있음을 증명하고, 이를 통해 GRPO의 토큰별 보상 할당을 정교화하는 SHEAR 방법론을 제안하여 수학 및 코드 생성 성능을 향상시켰습니다.

원저자: Xinzhu Chen, Wei He, Huichuan Fan, Wenzhe Niu, Zhongxiang Sun, Xuanru Wang, Jiuchong Gao, Jinghua Hao, Renqing He, Weijie Yu

게시일 2026-04-28

📖 2 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Xinzhu Chen, Wei He, Huichuan Fan, Wenzhe Niu, Zhongxiang Sun, Xuanru Wang, Jiuchong Gao, Jinghua Hao, Renqing He, Weijie Yu

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

1. 현재의 문제점: "결과만 보고 혼내는 선생님" (GRPO의 한계)

지금까지 AI를 학습시키는 방식(GRPO라고 부릅니다)은 마치 결과만 보고 점수를 매기는 아주 엄격한 선생님과 같습니다.

학생이 수학 문제를 풀 때, 처음에는 아주 논리적으로 잘 풀다가 마지막 계산에서 딱 한 번 실수해서 틀렸다고 해봅시다. 기존의 선생님은 학생이 쓴 풀이 전체를 보고 **"틀렸으니까 이 문제 전체는 0점이야!"**라고 말하며, 문제 앞부분의 완벽했던 논리까지도 "너 이거 다 틀렸어!"라며 혼냅니다.

이렇게 되면 학생은 "내가 어디서부터 잘못된 건지" 알 수 없고, 잘했던 부분까지 억지로 고치려다 오히려 실력이 엉망이 될 수 있습니다.

2. 이 논문의 아이디어: "학생의 '눈빛'을 읽는 관찰자" (SHEAR)

이 논문의 저자들은 아주 기발한 생각을 했습니다. "선생님이 일일이 풀이 과정을 검사하기 힘들다면, 학생이 문제를 풀 때 변하는 '눈빛(Hidden States)'을 관찰하면 어떨까?" 하는 것이죠.

여기서 **'눈빛(Hidden States)'**이란 AI 모델 내부에서 정보가 처리되는 상태를 말합니다.

잘 풀고 있을 때의 눈빛: 확신에 차 있고, 논리가 일관되며, 눈빛이 반짝반짝합니다. (데이터 분포가 안정적임)
실수하기 시작할 때의 눈빛: 당황하기 시작하고, 논리가 꼬이며, 눈빛이 흔들립니다. (데이터 분포가 갑자기 변함)

저자들은 AI의 내부 상태(눈빛)를 수학적인 방법(Wasserstein Distance, 두 집단 사이의 거리 측정법)으로 분석했더니, 실수가 발생하는 지점에서 눈빛이 급격하게 변한다는 사실을 발견했습니다.

3. 해결책: "눈빛이 흔들릴 때만 집중 교정하기" (SHEAR 방식)

이제 이 '눈빛 변화'를 이용해 학습 방식을 바꿉니다. 이것이 바로 SHEAR라는 방법입니다.

관찰: AI가 문제를 풀 때, 내부 상태(눈빛)가 정답을 맞힌 그룹과 틀린 그룹 사이에서 얼마나 차이가 나는지 실시간으로 체크합니다.
판단: "오, 이 부분부터는 정답을 맞힌 애들의 눈빛과 틀린 애들의 눈빛이 완전히 달라지는데? 여기서 사고가 터졌구나!"라고 판단합니다.
교정:
- 잘못된 부분: 눈빛이 확 변하며 틀린 길로 들어선 부분에는 **"여기서부터 잘못됐어! 집중해서 고쳐!"**라며 강한 피드백을 줍니다.
- 잘한 부분: 눈빛이 안정적이고 정답 그룹과 비슷하다면 **"이 부분은 아주 잘했어, 그대로 가!"**라며 가벼운 피드백만 줍니다.

4. 결과: "스스로 오답 노트를 쓰는 천재 학생"

이 방식을 적용했더니 놀라운 결과가 나왔습니다.

추가 비용 없음: 사람이 일일이 "몇 번째 줄이 틀렸어"라고 알려줄 필요가 없습니다. AI가 자기 자신의 '눈빛'을 보고 스스로 어디가 틀렸는지 찾아내기 때문입니다. (Self-supervised)
성적 향상: 수학 문제 풀이와 코딩 능력에서 기존 방식보다 훨씬 높은 점수를 기록했습니다. 특히 복잡하고 긴 풀이가 필요한 문제에서 실력이 압도적으로 좋아졌습니다.

요약하자면!

이 논문은 **"AI가 문제를 풀 때 발생하는 내부적인 '당황함(데이터 분포의 변화)'을 포착하여, 틀린 지점만 콕 집어 집중 학습시키는 기술"**에 관한 것입니다. 덕분에 AI는 마치 스스로 오답 노트를 정교하게 작성하는 학생처럼, 훨씬 효율적이고 똑똑하게 성장할 수 있게 되었습니다.

[기술 요약] Hidden States Know Where Reasoning Diverges: Span-Level Wasserstein Distance를 이용한 크레딧 할당

1. 문제 정의 (Problem Statement)

대규모 언어 모델(LLM)의 추론 능력을 향상시키기 위해 **검증 가능한 보상(Verifiable Rewards, RLVR)**을 사용하는 강화학습(예: GRPO 알고리즘)이 널리 사용됩니다. 하지만 여기서 핵심적인 문제는 **'크레딧 할당(Credit Assignment)'**입니다.

기존 방식의 한계: 현재의 GRPO와 같은 알고리즘은 결과(Outcome)가 맞았는지 틀렸는지만 확인합니다. 만약 추론 과정 중 중간에 한 번의 실수로 결과가 틀렸다면, 그 이전의 올바른 추론 단계(tokens)들까지 모두 부정적인 보상을 받게 됩니다. 즉, 어떤 토큰이 성공에 기여했고 어떤 토큰이 실패를 유발했는지 구분하지 못하는 '거친(coarse-grained)' 할당 문제를 안고 있습니다.
PRM(Process Reward Model)의 한계: 단계별 보상을 주는 PRM이 대안이지만, 이는 사람이 직접 단계별로 라벨링을 하거나 별도의 보상 모델을 학습시켜야 하므로 비용이 매우 높고, 모델이 진화함에 따라 보상 모델과 정책 모델 간의 괴리(mismatch)가 발생할 수 있습니다.

2. 핵심 아이디어 및 방법론 (Methodology)

본 논문은 **"모델의 숨겨진 상태(Hidden States) 분포 안에 추론의 질(Quality)에 대한 신호가 이미 들어있다"**는 가설에서 출발합니다. 별도의 외부 모델 없이, 모델 스스로의 표현(Representation)을 사용하여 스스로를 감독(Self-supervision)하는 방식을 제안합니다.

SHEAR (Span-level Hidden state Enabled Advantage Reweighting)

제안된 방법론인 SHEAR는 다음과 같은 단계로 작동합니다.

Span Decomposition (스팬 분해): 추론 경로(Rollout)를 슬라이딩 윈도우(Sliding Window) 방식을 사용하여 겹치는 구간(Span)들로 나눕니다.
Wasserstein Distance 계산: 각 스팬의 숨겨진 상태 분포와, 반대 그룹(정답 그룹 vs 오답 그룹)의 스팬 분포 사이의 **Wasserstein Distance(Wasserstein 거리)**를 계산합니다. 이 거리는 두 분포가 얼마나 다른지를 측정하며, 추론이 어긋나는 지점에서 급격히 증가하는 특성을 보입니다.
Token-level Aggregation (토큰 단위 집계): 계산된 스팬 단위의 거리를 Max-pooling을 통해 개별 토큰 단위의 가중치( $\omega_t$ )로 변환합니다.
Advantage Reweighting (이점 재가중치 부여): 기존 GRPO의 Advantage( $A_i$ $A_{i}$ )에 이 가중치를 곱하여 $\tilde{A}_t = A_i \cdot \omega_t$ $\tilde{A}_{t} = A_{i} \cdot ω_{t}$ 로 업데이트합니다.
- 오답 경로: 추론이 틀어지는(분포가 크게 변하는) 지점의 토큰에 더 큰 음의 업데이트를 부여합니다.
- 정답 경로: 다른 오답들과 확연히 구분되는(분포가 독특한) 올바른 추론 지점에 더 큰 양의 업데이트를 부여합니다.

3. 주요 기여 (Key Contributions)

경험적 발견 (Empirical Finding): 스팬 단위의 숨겨진 상태 분포 간 Wasserstein 거리가 별도의 라벨 없이도 추론의 품질(연속 성공률)과 강력하게 상관관계가 있음을 증명했습니다.
이론적 토대 (Theoretical Foundation): **Separation Theorem(분리 정리)**를 통해, 추론이 어긋난 이후의 스팬(post-divergence)이 어긋나기 전의 스팬(pre-divergence)보다 통계적으로 유의미하게 큰 Wasserstein 거리를 가짐을 수학적으로 입증했습니다.
새로운 알고리즘 제안 (SHEAR): 추가적인 모델 학습이나 데이터 라벨링 없이, 기존 GRPO 파이프라인에 최소한의 수정만으로 적용 가능한 효율적인 크레딧 할당 알고리즘을 제안했습니다.

4. 실험 결과 (Results)

수학적 추론 (Mathematical Reasoning): 5개의 벤치마크(AIME, AMC, MATH500 등)에서 Qwen2.5-Math, Llama3.1 등 다양한 백본 모델을 대상으로 실험한 결과, 표준 GRPO를 일관되게 능가했습니다. 특히 복잡한 다단계 추론이 필요한 AIME와 AMC에서 큰 성능 향상을 보였습니다.
코드 생성 (Code Generation): HumanEval, MBPP, LiveCodeBench 등 5개 벤치마크에서도 GRPO 대비 성능 향상을 확인했습니다.
PRM과의 비교: 별도로 학습된 고성능 PRM(Process Reward Model)과 비교했을 때도 대등하거나 더 우수한 성능을 보였습니다. 이는 외부 모델의 불일치 문제 없이 모델 스스로의 표현을 활용하는 것이 더 효과적일 수 있음을 시사합니다.
효율성: 학습 시간 오버헤드는 모델 크기에 따라 약 7%~15% 수준으로 매우 낮아 실용적입니다.

5. 의의 (Significance)

이 논문은 **"모델의 내부 표현(Internal Representation)이 모델의 논리적 오류를 스스로 인지하고 있다"**는 점을 수학적/경험적으로 증명했다는 점에서 매우 중요합니다.

기존의 강화학습이 '결과'라는 희소한 신호(Sparse Signal)에만 의존했다면, SHEAR는 모델의 '내부 상태'라는 풍부한 신호(Dense Signal)를 활용하여 비용 효율적이고 정교한(Fine-grained) 학습을 가능하게 했습니다. 이는 향후 LLM의 추론 능력을 고도화하는 데 있어 별도의 라벨링 비용 없이도 모델을 스스로 개선할 수 있는 강력한 프레임워크를 제공합니다.