기술 요약: 선택적 학습을 위한 세그먼트 수준의 어트리뷰션(Segment-Level Attribution)

문제 정의

대규모 추론 모델(LRM)은 확장된 사고 사슬(Chain of Thought, CoT)을 생성함으로써 높은 성능을 달나합니다. 그러나 이러한 추론 궤적(trace)은 종종 상당한 중복성을 겪으며, 여기서 오직 아주 적은 부분의 토큰만이 최종 답변에 의미 있게 기여하고, 나머지는 반복적인 내용, 잘린 생각, 또는 정보가 없는 장황한 부연 설명으로 구성됩니다. 이러한 중복성은 문제가 됩니다. 왜냐하면 이러한 장황한 궤적에 대해 지도 미세 조정(SFT)을 수행할 경우, 모델이 비효율적인 패턴을 모방하게 되어 사소한 지속 과정에 학습 능력을 낭비하고 잠재적으로 추론 성능을 저하시키기 때문입니다.

중요한 추론 부분을 식별하는 기존 방법들은 다음과 같은 한계에 직면해 있습니다:

• 토큰 수준 분석은 의미적 무결성을 간과하며, 중복성을 의미 있는 추론 단위 관점에서 해석하는 데 실패합니다.
• 간접적인 세그メント 수준 지표(예: 퍼플렉시티 또는 엔트로피)는 허수(false positive, 표면적인 스캐폴딩 텍스트를 과도하게 강조함)와 허오(false negative, 낮은 엔트로피를 가지지만 정확도에 결정적인 독립적 검증 단계를 필터링함)가 발생하기 쉽습니다.
• 가지치기(Pruning) 기반 접근 방식은 CoT 감독(supervision)을 압축할 때 추론 정확도를 저해할 수 있습니다.

방법론

본 논문은 최종 답변 예측에 대한 추론 세그먼트의 영향력을 정량화하기 위해 통합 그래디언트(Integrated Gradient, IG) 어트리뷰션을 활용하는 세그먼트 수준의 선택적 학습 프레임워크를 제안합니다.

1. 세그먼트 중요도 지표

저자들은 전환 키워드(예: "Wait," "Alternatively")를 사용하여 긴 CoT를 세그먼트로 분할합니다. 그런 다음 토큰 수준의 IG 어트리뷰션을 두 가지 세그먼트 수준 지표로 집계합니다:

• 어트리뷰션 강도(Attribution Strength): 모델의 예측에 대한 세그먼트의 전반적인 영향력 크기를 측정합니다. 이는 세그먼트 길이에 따른 편향을 방지하기 위해 세그먼트 길이의 제곱근으로 정규화된, 세그먼트 내 절대 IG 값의 합으로 계산됩니다.
• 어트리뷰션 방향 일관성(Attribution Direction Consistency): 세그먼트의 어트리뷰션이 균일하게 양수/음수인지, 아니면 혼합된 형태인지를 포착합니다. 이는 부호가 있는 IG 절대값의 합을 절대값 IG 합으로 나눈 비율로 정의됩니다.
  • 높은 일관성: 얕거나 편향된 추론(예: 표면적인 명료화 또는 심각한 왜곡)을 나타냅니다.
  • 중간 정도의 일관성: 모델이 가능성을 탐색하고, 이해를 정교화하며, 잘못된 내용을 제거하는 성찰적 추론을 나타냅니다.

2. 중요한 세그먼트 식별

프레임워크는 이중 기준 임계값을 기반으로 "중요한" 세그먼트를 식별합니다:

1. 높은 강도: 세그먼트들은 정규화된 어트리뷰션 강도에 따라 순위가 매겨집니다. 누적 강도가 임계값 $\tau$(예: 전체 어트리뷰션의 80%)를 초과할 때까지 상위 $k$개의 세그먼트가 선택됩니다.
2. 중간 정도의 일관성: 상위 순위 세그먼트 중, 방향 일관성이 임계값 $\beta$(예: 0.8) 미만인 것들을 유지합니다. 이는 강도는 높지만 극단적인 균일성을 보이는(얕은 추론을 나타내는) 세그먼트를 필터링합니다.

3. 선택적 SFT (Selective SFT)

CoT를 압축(이는 일관성을 해침)하거나 전체 궤적에 대해 학습(노이즈를 포함함)하는 대신, 프레임워크는 선택적 SFT를 적용합니다.

• 손실 함수는 식별된 중요한 세그먼트에 속하는 토큰에 대해서만 계산됩니다.
• 중요하지 않은 세그먼트의 토큰에 대한 손실은 마스킹됩니다.
• 이는 암묵적인 규제(regularization) 역할을 하여, 추론 과정의 일관성을 유지하면서도 모델이 정보가 없는 내용에 과적합되는 것을 방지합니다.

주요 결과

실험은 여러 모델(R1-Distill-Qwen-1.5B, 7B, Qwen2.5-7B-Instruct)과 데이터셋(MATH500, AMC23, AIME24, GPQA, Minerva, OlympiadBench)을 대상으로 수행되었습니다.

• 정확도 향상: 제안된 방법은 표준 Full-CoT SFT보다 일관되게 우수한 성능을 보였습니다. 그리디 디코딩(greedy decoding) 하에서, R1-Distill-Qwen-1.5B 등에서 최대 **4.7%**의 정확도 향상을 달성했습니다.
• 효율성 이득: 이 접근 방식은 성능을 유지하거나 개선하면서도, Full-CoT SFT 대비 토큰 감소량을 최대 **18.2%**까지 줄여 출력 길이를 크게 단축했습니다.
• 절제 연구(Ablation Studies):
  • 세그먼트 vs. 토큰 수준: 세그gment-level 선택이 토큰 수준 선택(IG 또는 무작위 샘플링 기반)보다 우수한 성능을 보였으며, 이는 추론 단위를 보존하는 것이 일관성을 더 잘 유지함을 입증합니다.
  • 지표 비교: IG 기반의 강도-일관성 기준이 "First-Correct Solution", "Confidence-Gain", "Segment Perplexity", "Segment Entropy"와 같은 대안적 중요도 측정법보다 우수했습니다.
  • 가지치기(Pruning) vs. 선택적 학습: 가지치기에 적용했을 때, IG 기반 방법은 다른 가지치기 전략보다 더 나은 성능을 유지했으나, 저자들은 Selective SFT(손실 마스킹)가 정확도 보존 측면에서 가지치기보다 우월하다고 언급했습니다.

의의 및 주장

본 논문은 다양한 형태의 중복성(반복, 절단, 표면적 명료화)으로부터 진정으로 중요한 추론 세그먼트를 구별하는 원칙적인 방법을 제공한다고 주장합니다. Integrated Gradients를 활용함으로써, 이 방법은 퍼플렉시티와 같은 간접적 지표의 한계를 해결하며 추론 단계의 직접적 및 간접적 영향을 모두 포착합니다.

주요 기여는 모델이 긴 추론 궤적으로부터 더 효과적으로 학습할 수 있게 하는 세그먼트 수준의 선택적 학습 프레임워크를 제공하는 것입니다. 저자들은 이 접근 방식이 다음을 가능하게 한다고 단언합니다:

1. 추론 효능과 출력 효율성을 동시에 개선합니다.
2. CoT 압축이나 가지치기에서 흔히 발생하는 정확도 저하를 피합니다.
3. 강화 학습에서 타겟 콘텐츠에 정책 그래디언트 업데이트를 강조하는 것과 같이, 다른 문맥으로의 광범위한 적용 가능성을 제시합니다.

이 연구는 모든 추론 단계가 똑같이 가치 있는 것은 아니라는 점을 시사합니다. 높은 어트리뷰션 강도와 중간 정도의 방향 일관성(성찰적 추론을 나타냄)을 가진 세그먼트에 학습 신호를 집중함으로써, 모델은 더 적은 토큰으로 더 나은 성능을 달성할 수 있습니다.