How Does Reasoning Flow? Tracing Attention-Induced Information Flow for Targeted RL in LLMs

기술 요약: FlowTracer

문제 정의

강화 학습(RL)은 특히 검증 가능한 보상이 존재하는 수학적 문제 해결 및 코드 생성과 같은 도메인(RLVR)에서, 대규모 언어 모델(LLM)의 단계별 추론 능력을 향상시키기 위한 핵심적인 도구로 자리 잡았습니다. 그러나 LLM에 RL을 적용할 때 발생하는 지속적인 병목 현상은 **토큰 수준의 크레딧 할당(token-level credit assignment)**입니다.

GRPO(Group Relative Policy Optimization)와 같은 현재의 RL 레시피는 일반적으로 생성된 모든 토큰을 동일하게 취급하며, 최종 결과에 따라 균등한 보상이나 패널티를 부여합니다. 이러한 접근 방식은 결정적인 추론 단계와 **일상적인 포맷팅 또는 유창한 미사여구(filler)**를 구분하지 못한다는 단점이 있습니다. 최근 연구들이 더 세밀한 크레딧 할당을 위해 내부 모델 신호(예: 엔트로피, 어텐션 통계)를 활용하려는 시도를 하고 있으나, 이러한 방법들은 정보 전파의 글로벌 구조를 무시하는 점 단위의 휴리스틱에 의존하는 경우가 많습니다. 결과적으로, 초기 단계의 중요한 전제 조건들은 긴 경로를 거치며 신호가 희석되어 과소 평가될 수 있으며, 정답 근처의 토큰들은 단순히 위치가 가깝다는 이유만으로 과도하게 가중치가 부여될 수 있습니다.

방법론: FlowTrace

저자들은 FlowTracer를 제안합니다. 이는 답변을 목표로 하는 추론 흐름을 어텐션 유도 유향 비순환 그래프(Attention-Induced Directed Acyclic Graph, DAG) 상에서 추적하여 토큰 수준의 크레딧을 도출하는 프레임워크입니다. 이 방법은 추론을 프롬프트에서 최종 답변으로 이어지는 '영향력의 흐름(flow of influence)'으로 모델링한다는 전제하에 작동합니다.

1. 어텐션 유도 DAG 구축

토큰 시퀀스는 노드가 토큰을 나타내고 유향 에지가 의존성을 나타내는 시간 순서 기반의 DAG로 모델링됩니다. 에지 용량($W_{ik}$)은 레이어와 헤드 전체에 걸쳐 집계된 어텐션 가중치(평균값)로부터 도출되며, 이는 토کن $i$에서 토큰 $k$로의 잠재적 정보 전달 강도로 해석됩니다. 확률적 커널과 달리, 이 가중치들은 반드시 1의 합을 가질 필요가 없으므로 신호의 증폭이나 감쇄를 허용합니다.

2. 답변 타겟형 흐름 보존 (Answer-Targeted Flow Conservation)

가공되지 않은 어텐션 그래프에는 크레딧 신호를 희석시키는 광범한 "막다른 길(dead-end)" 분기(예: 미사여구 토큰, 폐기된 가설)가 포함되어 있습니다. 이를 해결하기 위해 FlowTracer는 Doob-h 유사 재가중치(Doob-h-like reweighting) 메커니즘을 도입합니다.

• 도달 가능성 잠재력 ($h$): 각 토큰에 대해 잠재 함수 $h(i)$를 계산합니다. 이는 토큰 $i$에서 답변 영역에 연결된 가상 싱크(virtual sink)까지의 총 영향력 가중 경로 합을 나타냅니다.
• 재가중치 적용: 에지 용량은 $W'_{ik} = W_{ik} \frac{h(k)}{h(i)}$로 재가중됩니다.
• 이론적 보장: 이 변환은 비제로 도달 가능성을 가진 노드들에 대해 국소 흐름 보존(local flow conservation) ($\sum_{k>i} W'_{ik} = 1$)을 강제합니다. 이는 답변으로 이어지지 않는 경로를 효과적으로 필터링하고, 중간 토크들이 경로 길이 혹은 무관한 분기로 인해 유효 질량을 잃거나 얻지 않도록 보장합니다.

3. 흐름 처리량 및 크레딧 할당

그래프가 답변에 맞춰 조건화되면, FlowTracer는 슈퍼 소스(프롬프트와 연결됨)로부터 단위 흐름을 주입하고 슈퍼 싱크(답변과 연결됨)에서 이를 수집합니다.

• 처리량 계산: 각 토کن $k$를 통과하는 흐름 $f(k)$를 계산합니다. 총 처리량이 높은 토큰은 장거리 의존성을 매개하는 추론 허브(reasoning hubs) 또는 **집계 체크포인트(aggregation checkpoints)**로 식별됩니다.
• 보상 셰이핑 (Reward Shaping): 이러한 처리량 값은 RL 손실 함수(특히 GRPO) 내에서 비균등 스케일링 인자 $\gamma_{i,t}$를 생성하는 데 사용됩니다. 높은 흐려를 보이는 토큰은 더 높은 가중치(예: $\gamma_{flow} = 1.5$)를 할당받아 결정적인 추론 단계에 대한 학습 신호를 증폭시키는 반면, 영향력이 낮은 미사여구에 대한 업데이트는 억제합니다.

주요 기여

1. 글로벌 구조적 크레딧 할당: 점 단위의 휴리스틱(엔트로피, 그래디언트 크기)과 달리, FlowTracer는 정보 전파의 글로벌 토폴로지를 모델링하여 프롬프트와 답변을 연결하는 "추론 백본(reasoning backbone)"을 식별합니다.
2. 흐름 보존형 재가중치: Doob-h 유사 변환의 도입은 흐름 보존을 만족하면서 답변을 타겟팅하는 깨끗한 서브그래프를 추출할 수 있게 하여, 가공되지 않은 어텐션 가중치에 내재된 구조적 편향을 제거합니다.
3. 고유량(High-flow) 토큰의 인과적 검증: 고유량 토큰과 저유량 토큰을 마스킹하는 인과적 개입 실험을 통해, 저자들은 고유량 토큰이 추론의 핵심 동력임을 입증했습니다. 이 토큰들을 섭동(perturbing)하면 답변의 유의미한 변화와 정답 역전이 발생하는 반면, 저유량 토큰을 섭동할 때는 영향이 미미했습니다.
4. 효율적인 통합: 이 방법은 미들 레이어에서 어텐션 맵을 추출하기 위해 단 한 번의 추가 순전파(forward pass)만을 요구하며, 2.2%~4.5%의 미미한 계산 오버헤드만으로 일관된 성능 향상을 제공합니다.

실험 결과

FlowTracer는 Qwen3-4B/8B 및 Llama-3.1/3.2 모델을 사용하여 수학적 추론 벤치마크(AIME24/25, AMC23, MATH500, OlympiadBench, GSM8K) 및 기타 추론 작업(Countdown, CrossThinkQA)에서 평가되었습니다.

• 성능 향상: FlowTracer는 표준 GRPO 베이스라인 및 기타 크레딧 할당 전략(Random, High-entropy, Gradient, Correlation, Attention)보다 일관되게 우수한 성능을 보였습니다.
  • Qwen3-8B에서 수학 벤치마크 평균 정확도 **43.4%**를 달여, GRPO 베이스라인(39.4%)보다 4.0% 높았습니다.
  • Qwen3-4B에서는 8K 컨텍스트 설정에서 GRPO 대비 **3.8%**의 이득을 보였으며, 특히 AIME25에서는 **+5.8%**라는 놀라운 개선을 보였습니다.
• 롱 컨텍스트 확장성: 컨텍스트 길이가 1K에서 8K로 증가함에 따라 FlowTracer와 베이스라인 간의 성능 격차가 벌어졌으며, 이는 노이즈를 필터링하고 긴 추론 체인에서 핵심 토큰을 찾아내는 탁월한 능력을 나타냅니다.
• 일반화: 이 방법은 Llama 제품군에도 효과적으로 전이되어, Llama-3.1-8B에서 평균 정확도 +1.4%, Llama-3.2-3B에서 **+1.1%**의 향상을 가져왔습니다.
• 절제 연구 (Ablation Studies):
  • 토큰 선택: 흐름 점수 기준 상위 40%의 토큰을 선택했을 때 최적의 결과를 얻었으며, 하위 k개 토큰을 선택하면 성능이 저하되었습니다.
  • 레이어 선택: 미들 레이어(L/3 ~ 2L/3)의 어텐션 맵이 흐름 분석을 위한 가장 돌출된 신호를 제공했습니다.
  • 연속적 vs 하드 마스킹: 하드 Top-40% 마스크가 흐름 값의 헤비 테일(heavy-tailed) 분포으로 인해 어려움을 겪는 연속적 재가중치 방식보다 더 안정적이고 효과적이었습니다.

의의 및 주장

본 논문은 FlowTrace가 균등한 보상과 고립된 토큰 휴리스틱을 넘어, RL4LLM의 핵심 과제인 크레딧 할당 문제를 해결한다고 주장합니다. 어텐션의 내부 기하학을 활용하여 정보 전파의 글로벌 구조를 재구성함으로써, 이 프레임워크는 결정적인 추론 단계와 일상적인 미사여구를 효과적으로 구분해 냅니다.

저자들은 이 접근 방식이 학습 신호를 고영향 토큰에 정확히 집중시켜 복잡한 추론 작업에서 일관된 성능 향상을 제공한다고 단언합니다. 나아가, 본 연구는 LLM의 내부 어텐션 구조가 더 효율적이고 해석 가능한 모델 정렬을 위한 강력한 구조적 신호를 제공하며, 생성 모델 내의 "추론 백본"을 식별하고 강화하는 원칙적인 방법을 제시함을 시사합니다. 이 방법은 다양한 모델 아키텍처와 태스크 도메인에 걸쳐 견고하게 작동하는, 추론 지향적 RL 학습을 위한 범용 플러그인으로 제시됩니다.