Compress the Easy, Explore the Hard: Difficulty-Aware Entropy Regularization for Efficient LLM Reasoning

기술 요약: 쉬운 것은 압축하고, 어려운 것은 탐색하라 (CEEH)

문제 정의

Chain-of-Thought (CoT) 프롬프팅은 중간 단계를 외현화함으로써 대규모 언어 모델(LLM)의 추론 능력을 크게 향상시켰습니다. 그러나 이러한 명시적인 추론은 종종 지나치게 장황한 출력을 생성하여, 과도한 추론 지연 시간과 계산 비용을 초래합니다. 기존의 추론 압축 방법들은 일반적으로 길이 제한(예: 명시적 길이 패널티)을 포함하는 강화 학습(RL)을 사용하는데, 이는 치명적인 실패 모드에 직면합니다: 공격적인 간결성 최적화가 급격한 **엔트로피 붕괴(entropy collapse)**를 유발한다는 점입니다.

정책이 길이 제한을 충족하기 위해 지나치게 결정론적으로 변함에 따라, 탐색 공간이 조기에 축소됩니다. 이로 인해 광범위한 연역, 자기 성찰 또는 오류 수정이 필요한 까다로운 문제들에 대한 유효한 추론 경로를 발견하는 데 실패하게 됩니다. 결과적으로, 이러한 방법들은 간결성을 위해 추론 정확도를 희생하는 경우가 많습니다. 여기에는 핵심적인 긴장이 존재합니다: 균일한 엔트로피 정규화는 탐색을 유지할 수 있지만 긴 궤적을 유도하며(압축에 역행), 균일한 길이 패널티는 궤적을 압축하지만 엔트로피를 붕괴시킵니다(정확도 저하).

방법론: CEEH

저자들은 탐색과 압축의 균형을 동적으로 조절하는 난이도 인식 프레임워크인 CEEH(Compress the Easy, Explore the Hard)를 제안합니다. 이 방법은 그룹 상대 정책 최적화(GRPO) 환경 내에서 두 가지 시너지 효과가 있는 구성 요소를 통합합니다.

1. 난이도 인식 엔트로피 정규화 (Difficulty-Aware Entropy Regularization)
균일한 정규화를 적용하는 대신, CEEH는 "어려운" 질문에 대해서는 탐색을 선택적으로 유지하고 "쉬운" 질문에 대해서는 공격적인 압축을 허용합니다.

• 동적 난이도 추정: 시스템은 **비대칭 지수 이동 평균(Asymmetric Exponential Moving Average, EMA)**을 사용하여 각 질문에 대한 역사적 정확도 추정치를 유지합니다. 질문의 역사적 정확도가 전체 데이터셋 평균보다 낮을 경우 해당 질문을 "어려운"($D_h$) 것으로 분류합니다. 이 임계값은 모델이 개선됨에 따라 진화합니다.
• 선택적 메커니즘:
  • 어려운 질문($D_h$)의 경우: 더 강력한 엔트로피 정규화를 적용합니다. 이는 정규화 계수를 증폭(예: 코사인 어닐링 스케줄에서 $5\times$ 승수 적용)하고, 고엔트로피 토큰 생성을 장려하는 엔트로피 기반 어드밴티지(entropy-based advantage) 항($\psi(H)$)을 포함함으로써 달성됩니다. 이는 조기 수렴을 방지하기 위해 다양한 탐색 공간을 보존합니다.
  • 쉬운 질문의 경우: 정규화를 완화하여 정책이 짧고 잘 확립된 추론 경로를 자신 있게 활용할 수 있도록 합니다.

2. 동적 최적 길이 패널티 (Dynamic Optimal-Length Penalty)
어려운 질문에 대한 엔트로피 정규화로 인해 발생할 수 있는 잠재적인 길이 팽창을 상쇄하기 위해, CEEH는 질문 수준의 동적 길이 패널티를 도입합니다.

• 메커니즘: 시스템은 각 특정 질문에 대한 역사적으로 가장 짧은 정답 길이($L_x$)를 추적합니다.
• 패널티 함수: 이 진화하는 베이스라인을 초과하는 정답 응답에 대해서만 패널티가 적용됩니다. 패널티는 상대적 편차인 $R_{len} = (L_{x, y_i} - L_x) / L_x$로 계산됩니다.
• 적응형 신호: 이 설계는 모델의 역량이 향상되고 응답 길이가 자연스럽게 감소함에 따라 압축 신호가 안정적이고 변별력을 유지하도록 보장합니다. 이는 모델이 어려운 문제에 대해 불필요하게 긴 출력을 생성하기 위해 엔트로피 정규화를 악용하는 것을 방지합니다.

학습 역학 (Training Dynamics)
전체 보상은 검증 가능한 보상(정확도)과 길이 패널티를 결합합니다: $R_{total} = R - \beta R_{len}$. 이 방법은 탐색을 장려하기 위해 상단 클리핑 임계값을 높이는 것(Clip-Higher)과 어드밴티지 표준화를 피하여 길이 패널티 신호의 감쇠를 방지하는 것과 같은 구체적인 구현 전략을 채택합니다.

주요 기여

1. 엔트로피 붕괴 식별: 본 논문은 엔트로피 관점을 통해, 기존의 추론 압축 방법에서 발생하는 정확도-길이 트레이드오프 실패의 주요 요인이 급격한 엔트로피 감소임을 밝혀냈습니다.
2. CEEH 프레임워크: 난이도 인식 엔트로피 정규화(어려운 태스크에서의 탐색 유지를 위함)와 동적 최적 길이 패널티(학습 안정화 및 압축 강제를 위함)를 통합한 새로운 프레임워크를 제시합니다.
3. 실증적 검증: CEEH가 강력한 베이스라인들과 비교하여 정확도를 희생하지 않으면서도 추론을 효과적으로 압축하여 우수한 정확도-효율성 트레이드오프를 달取得함을 입증했습니다.

실험 결과

이 방법은 R1-Distill-Qwen2.5 모델을 사용하여 6개의 수학적 추론 벤치마크(GSM8K, MATH500, AIME24/25, AMC, OlymBench)에서 평가되었습니다.

• 정확도 vs. 길이: CEEH는 베이스 모델과 비교하여 정확도는 비슷하거나 약간 더 높은 수준을 유지하면서도 응답 토큰 수를 30% 이상 일관되게 줄였습니다.
• 베이스라인과의 비교:
  • 프롬프팅/오프라인 방법: 종종 정확도가 요동치거나 다양한 난이도 수준에 걸쳐 일반화에 실패했습니다.
  • 길이 전용 RL (예: Length-Penalty): 길이를 줄이는 데는 효과적이었으나, 엔트로피 붕괴로 인한 근본적인 추론 능력 저하를 나타내는 Pass@k 성능 저하가 빈번하게 관찰되었습니다.
  • CEEH: 제어된 탐색을 통해 올바른 솔루션을 발견하는 모델의 능력을 보존함으로써 Pass@k 점수를 유지하거나 향상시켰습니다.
• 강건성: 이 방법은 서로 다른 규모(1.5B 및 7B 파라미터)의 모델과 다양한 추론 태스크 전반에서 잘 일반화되었습니다.
• 학습 효율성: CEEH는 균일한 패널티 전략에 비해 높은 학습 정확도와 더 나은 데이터 효율성을 달성했는데, 이는 어려운 질문에 대한 고엔트로피 탐색 덕분입니다.

의의 및 주장

본 논문은 CEEH가 어디를 압축하고 어디를 탐색할지를 명시적으로 분리함으로써 추론 압축과 능력 보존 사이의 긴장을 해결한다고 주장합니다.

• 추론 능력 보존: 어려운 인스턴스에 대한 엔트로피 붕괴를 방지함으로써, CEEH는 모델이 단순히 최적이 아닌 지름길 솔루션에 확신을 갖는 것이 아니라 복잡한 연역과 자기 수정을 수행할 수 있는 능력을 유지하도록 보장합니다.
• 안정적인 압축: 동적 최적 길이 패널티는 모델의 향상되는 역량에 적응하는 강력한 신호를 제공하여, 엔트로피 정규화가 균일하게 적용될 때 흔히 나타나는 "길이 팽창" 현상을 방지합니다.
• 일반화 가능성: 이 접근 방식은 인스턴스의 난이도에 따라 탐색을 적응적으로 관리한다면, LLM 추론의 효율성 이득이 반드시 정확도와의 트레이드오프를 수반하지는 않는다는 것을 보여줍니다.

본 연구는 RL 기반 LLM 최적화 과정에서 진정한 추론 능력을 유지하기 위해서는 원칙적인 엔트로피 제어가 필수적이라는 결론을 내립니다.