BandPO: 확률 인지 경계를 통한 LLM 강화학습의 신뢰 영역과 비율 클리핑 연결

기존 PPO나 GRPO의 고정된 클리핑 방식은 확률이 낮은 행동의 탐색을 억제하여 모델이 새로운 전략을 배우지 못하게 만드는 한계가 있었다. BandPO는 수학적 신뢰 영역을 동적 경계로 변환하여 낮은 확률의 행동도 충분히 업데이트될 수 있도록 허용함으로써 학습의 안정성과 탐색 효율을 동시에 잡았다.

PPO와 같은 강화학습 알고리즘은 정책 업데이트 시 급격한 변화를 막기 위해 확률 비율(new/old)을 일정 범위(예: 0.8~1.2)로 제한하는 클리핑(Clipping) 기법을 사용한다. 하지만 이 방식은 Softmax 출력 확률이 매우 낮은 행동(Tail actions)의 경우, 허용되는 확률 변화량($\Delta \pi$)이 기존 확률에 비례해 매우 작아지는 문제를 낳는다. 결과적으로 유망한 새로운 전략이 발견되어도 모델이 이를 충분히 학습하지 못하고 특정 답변에만 매몰되는 엔트로피 붕괴(Entropy Collapse)가 발생한다.

BandPO는 이 문제를 해결하기 위해 '신뢰 영역(Trust Region)'이라는 기하학적 개념을 다시 도입한다. 단순히 비율을 고정하는 대신, 확률 분포 간의 거리(f-divergence)를 일정 수준 이하로 유지한다는 제약 조건을 각 토큰의 확률값에 맞게 동적으로 변환한다. 이를 통해 확률이 낮은 행동에 대해서는 더 넓은 업데이트 마진을 부여하고, 확률이 높은 행동은 엄격하게 관리하여 안정성을 유지한다.

이 접근법은 수학적으로 볼록 최적화(Convex Optimization) 문제로 정의되며, 각 상황에 맞는 최적의 클리핑 경계를 실시간으로 계산한다. 실험 결과, 수학 문제 풀이와 같은 복잡한 추론 작업에서 모델이 조기에 단순한 패턴에 갇히지 않고 더 나은 해결책을 지속적으로 탐색할 수 있게 함으로써 성능을 크게 향상시켰다.

BandPO는 GRPO 프레임워크를 기반으로 하며, 기존의 고정 클리핑 연산자를 Band 연산자로 대체한다. Band 연산자는 f-divergence로 정의된 고차원 신뢰 영역을 개별 행동의 확률 비율에 대한 스칼라 구간으로 투영하는 역할을 수행한다.

수학적으로는 $D_f(Q\|P) \leq \delta$ 제약 조건을 만족하면서 확률 비율 $r = Q(a)/P(a)$의 최대값과 최소값을 찾는 최적화 문제로 정형화된다. [입력값: 이전 정책 확률 $P(a)$, 신뢰 영역 반경 $\delta$, f-divergence 종류] → [연산: 라그랑주 승수법을 통한 단변량 루트 찾기] → [결과: 동적 상한 $\bar{r}$ 및 하한 $\underline{r}$] → [의미: 해당 토큰 확률에 최적화된 클리핑 범위 결정].

Total Variation(TV) 발산의 경우 $r = 1 \pm \delta/p$라는 선형적 폐쇄형 해를 가지며, Pearson χ² 발산은 $r = 1 \pm \sqrt{\delta(1-p)/p}$라는 제곱근 형태의 해를 가진다. KL 발산은 폐쇄형 해가 없으므로 이분법(Bisection) 기반의 수치적 솔버를 사용하여 전역 최적해를 구한다.

Qwen2.5-3B-Instruct 모델을 사용한 AMC2023 벤치마크에서 BandPO(KL, 0.05)는 55.17%의 mean@32 점수를 기록하여 기본 GRPO(45.94%) 대비 약 10%p에 가까운 성능 향상을 보였다. 특히 AIME 2024/2025 등 고난도 수학 벤치마크에서도 기존의 휴리스틱 방식인 Clip-Higher보다 일관되게 우수한 성과를 거두었다.

엔트로피 분석 결과, 기본 GRPO는 학습 초기(약 340단계)에 엔트로피가 급격히 감소하며 특정 답변으로 수렴하는 경향을 보였으나, BandPO는 학습 전반에 걸쳐 높은 엔트로피를 유지하며 탐색 능력을 보존했다. 이는 낮은 확률의 행동에 대한 클리핑 발생 빈도를 거의 0에 가깝게 줄인 결과로 분석된다.

BandPO의 핵심은 f-divergence 제약 조건을 단변량 최적화 문제로 축소한 Lemma 1(Uniform Complement Rescaling)에 있다. 타겟 행동 이외의 모든 행동에 대해 확률 비율을 일정하게 유지하는 것이 최적해임을 증명함으로써, 고차원 심플렉스 상의 문제를 단일 스칼라 $r$에 대한 루트 찾기 문제로 단순화했다.

제안된 Band 연산자는 $Band_{f,\delta}(r; a, P) \triangleq clip(r, \underline{r}_{f,\delta}(a; P), \bar{r}_{f,\delta}(a; P))$로 정의된다. 여기서 경계값들은 $g_f(p, r) = \delta$ 방정식의 해이며, $g_f$는 볼록 함수(Convex function)이므로 이분법을 통해 안정적인 수치해를 보장한다.

이론적 분석을 통해 확률 $p \to 0$일 때 상한 $\bar{r} \to \infty$가 됨을 보여, 낮은 확률의 행동이 조기에 클리핑되는 병목 현상을 해결함을 입증했다. 또한 심플렉스 포화(Simplex Saturation) 조건을 정의하여 수치적 불안정성을 방지하는 클램핑 로직을 포함했다.

수치적 솔버(Numerical Solver) 사용에 따른 계산 오버헤드가 존재한다. 특히 KL 발산의 경우 반복적인 루트 찾기 알고리즘이 필요하여 단순 클리핑 대비 연산량이 증가하며, 이를 해결하기 위해 CUDA 가속 이분법이나 룩업 테이블(Lookup Table) 활용이 권장된다. 또한 모든 토큰에 동일한 신뢰 영역 반경 $\delta$를 적용하는 정적 가정을 사용하므로, 토큰별 중요도에 따른 적응형 반경 설정 연구가 필요하다.