기술 요약: TRE: 신뢰 영역에서의 탐색 장려

문제 정의

엔트로피 정규화(Entropy Regularization)는 조기 수렴을 방지하고 탐색을 장려하기 위해 강화 학습(RL)에서 표준적으로 사용되는 기법이지만, 본 논문은 대규모 언어 모델(LLM)에 이를 직접 적용할 경우 이점이 미미하거나 오히려 성능을 저하시키는 경우가 많다는 점을 지적한다. 저자들은 이러한 실패의 원인을 LLM 특유의 거대한 어휘 집합과 긴 생성 호라이즌(generation horizon)으로 인해 발생하는 누적된 꼬리 위험(cumulative tail risk) 때문이라고 분석한다.

고전적인 RL 환경(예: 로보틱스)에서는 유효한 행동(action)이 조밀하게 존재하며, 무작위 탐색을 하더라도 대개 유효한 상태 공간 내에 머무른다. 반면, LLM 생성 과정에서 유효한 토큰은 방대한 어휘 집합(예: 약 150,000개의 토큰) 내에서 **극도로 희소한 매니폴드(infinitesimally sparse manifold)**를 형성한다. 표준적인 전역 엔트로피 최대화(global entropy maximization)는 확률 분포를 무차별적으로 평탄하게 만들어, 그럴듯한 후보군으로부터의 확률 질량을 구문론적 또는 논리적으로 유효하지 않은 '헤비 테일(heavy tail)' 영역으로 이동시킨다. 단일 단계에서 발생하는 노이즈는 미미할 수 있으나, 이는 긴 생성 시퀀스에 걸쳐 기하급수적으로 누적된다. 이러한 누적은 사고 체인(Chain-of-Thought, CoT) 추론의 일관성을 해치며, 생성 호라이즌이 길어질수록 유효한 추론 경로를 유지하는 것을 통계적으로 불가능하게 만든다.

방법론: 신뢰 영역 엔트로피 (Trust Region Entropy, TRE)

차원의 불일치(전역 엔트로피 최대화와 언어 생성의 희소성 간의 차이)를 해결하기 위해, 저자들은 탐색 범위를 사전 학습된 모델이 이미 타당하다고 판단하는 '신뢰 영역'으로 엄격히 제한하는 **신뢰 영역 엔트로피(TRE)**를 제안한다.

핵심 메커니즘

TRE는 전체 어휘 집합 $A$에 대해 엔트로피를 최대화하는 대신, 신뢰 영역 $\mathcal{A}_{TR}(s_t) \subset A$ 내에서 국소적으로 엔트로피를 최대화한다. 이를 통해 탐색이 오직 높은 확신도를 가진 후보들 사이에서만 이루어지도록 보장하며, 유효하지 않은 토큰으로 확률이 새어나가는 것을 방지한다.

신뢰 영역은 다음 두 가지 주요 전략을 통해 정의된다:

1. TRE-K (크기 기반): 신뢰 영역을 단계 $t$에서 가장 높은 로짓(logit) 값을 가진 상위 $K$개의 토큰으로 구성한다. 이는 일정한 영역 크기를 강제한다.
2. TRE-P (뉴클리어스 기반): 누적 확률 질량이 임계값 $P$(예: 0.99)에 도달하는 최소한의 토큰 집합을 신뢰 영역으로 설정한다. 이는 모델의 확신도에 따라 적응하는 동적인 영역 크기를 생성한다.

손실 함수 공식화

$L_{TRE,t}$는 다음과 같이 계산된다:

1. 신뢰 영역 내 토큰들에 대한 로짓을 추출한다.
   2.과 이들 토근에 대해서만 소프트맥스(softmax)를 적용하여 국소 확률 분포 $\pi_{local}$을 계산한다.
2. 국소 엔트로피 $H(\pi_{local})$를 계산한다.
3. 전체 어휘 엔트로피 스케일과 정규화 강도를 맞추기 위해 스케일링 인자 $\gamma = \frac{\log |A|}{\log |\mathcal{A}_{TR}|}$를 적용한다.

최종 훈련 목적 함수는 표준 서로게이트 손실(예: PPO)과 스케일링된 TRE 손실을 결합한다:
$$L_{total,t}(\theta) = L_{surr,t}(\theta) + \alpha L_{TRE,t}(\theta)$$

주요 기여

1. 실패 모드에 대한 경험적 분석: 본 논문은 표준 엔트로피 정규화가 긴 호라이즌 생성에서 해롭다는 것을 보여주는 상세한 분석을 제공한다. 실험 결과, 완만한 정규화는 짧은 컨텍스트에서는 도움이 될 수 있으나, 꼬리 위험의 누적으로 인해 생성 예산이 늘어남에 따라 성능 붕괴를 초래함을 입증했다.
2. 신뢰 영역 엔트로피(TRE) 제안: 고확신 신뢰 영역 내에서만 독점적으로 탐색을 장려함으로써, 다양성을 유지하면서도 일관된 추론을 보존하는 새로운 정규화 프레임워크를 제안한다.
3. 포괄적 평가: 수학적 추론(MATH), 조합 탐색(Countdown), 선호도 정렬(HH) 작업에 대한 광범한 실험을 통해, TRE가 바닐라 PPO, 표준 엔트로피 정규화, 그리고 다른 최첨단 탐색 베이스라인(예: Forking-Tokens, KL-Cov)보다 일관되게 우수한 성능을 보임을 입증했다.

실험 결과

저자들은 Qwen2.5-1.5B-Instruct 및 Qwen2.5-7B-Instruct 모델을 사용하여 TRE를 평가하였다.

• 성능 향상: TRE(특히 TRE-P 변형)는 6개 평가 카테고리 중 5개에서 가장 좋은 결과를 얻었다.
  • MATH (1.5B): TRE-P는 58.28%의 Pass@1을 달성하여, Vanilla PPO(57.04%)와 강력한 KL-Cov 베이스라인(58.23%)을 모두 앞질렀다.
  • Countdown (1.5B): TRE-P는 66.96%의 Pass@1에 도달하여, Vanilla 대비 +2.84% 향상되었으며, KL-Cov(+2.38%)를 유의미하게 상회했다.
  • HH (1.5B): TRE-P는 +0.6364의 보상 이득을 보였는데, 이는 KL-Cov(+0.1531)가 달성한 이득보다 약 4.1배 더 큰 수치이다.
• 확장성: 이러한 개선은 7B 모델 규모에서도 일관되게 나타났으며, TRE-P는 정렬 및 탐색 작업에서 KL-Cov보다 2~3배 더 큰 이득을 보였다.
• 정책 역학(Policy Dynamics): 평균 피크 확률($\bar{P}_{max}$) 분석 결과, 표준 PPO는 빠르게 포화(과잉 확신)되는 반면, TRE는 더 낮고 다양한 확신 수준을 유지한다. 이는 결정론적인 경로로의 '붕괴'를 방지하며, 특히 추론 작업에서 효과적이다.
• 하이퍼파라미터 민감도:
  • TRE-K: $K=2$에서 최적의 성능이 관찰되었다. 더 큰 $K$ 값은 노이즈를 다시 도입하여 성능을 저하시켰다.
  • TRE-P: 임계값 $P=0.99$가 최적이었다. $P=0.9$는 너무 제한적이었고, $P=0.999$는 유효하지 않은 꼬리 토큰을 포함하여 점수를 떨어뜨렸다.

의의 및 주장

본 논문은 LLM에서 표준 엔트로피 정규화가 실패하는 이유가 탐색이라는 개념 자체의 결함이 아니라, 탐색의 *범위(scope)*가 맞지 않기 때문이라고 주장한다. 타당한 토큰들의 신뢰 영역 내로 탐색을 국한함으로써, TRE는 긴 호리즌 추론을 방해하는 누적된 꼬리 위험을 효과적으로 완화한다.

저자들은 TRE를 복잡한 추론과 긴 컨텍스트 생성이 필요한 작업에 있어 LLM을 정렬하기 위한 견고하고 확장 가능한 솔루션으로 제시한다. 또한, TRE가 7B 파라미터 규모까지 효과적으로 확장됨을 언급하면서도, 거대 모델 및 초장기 호라이즌(예: 128k+ 토큰)에서의 동작은 향후 연구 과제로 남겨두었다. 이 연구는 생성적 추론 체인의 무결성을 유지하기 위해서는 전역적인 균일성보다 구조화된 국소적 탐색이 더 우월하다는 점을 시사한다.