LoopRPT: 루프형 언어 모델을 위한 강화 학습 기반 사전 학습 프레임워크

기존 LLM은 겉으로 보이는 텍스트를 통해서만 생각하지만, 루프형 모델은 내부에서 여러 번 되풀이하며 속으로 생각한다. 이 논문은 루프형 모델이 내부 연산 과정에서 직접 보상을 받아 더 효율적으로 추론하도록 만드는 새로운 학습법을 제시하여, 적은 연산으로도 복잡한 문제를 풀 수 있게 한다.

일반적인 Transformer는 입력 시퀀스를 한 번 통과하며 다음 토큰을 예측한다. 하지만 복잡한 추론이 필요한 경우, 모델은 명시적인 텍스트(Chain-of-Thought)를 생성해야만 생각할 시간을 벌 수 있다. 루프형 언어 모델(LoopLM)은 이와 달리 동일한 레이어 블록을 여러 번 반복 통과하며 내부 표현을 정교화하는 방식으로 속으로 생각한다.

LoopRPT는 이 내부 반복 과정에 직접 보상을 주는 방식을 택한다. 기존의 강화 학습이 최종 출력 토큰의 정답 여부만 따졌다면, LoopRPT는 각 루프 단계마다 이 단계에서 얼마나 정답에 가까워졌는가를 측정한다. 이를 위해 안정적인 기준점인 EMA 교사 모델을 두고, 학생 모델이 교사보다 더 빠르고 정확하게 정답 확률을 높이면 보상을 준다.

결과적으로 모델은 쉬운 토큰에서는 빠르게 루프를 빠져나오고, 어려운 토큰에서는 필요한 만큼 충분히 내부 연산을 수행하도록 학습된다. 이는 고정된 연산량을 사용하는 기존 모델보다 훨씬 유연하며, 동일한 정확도를 유지하면서도 평균 연산 단계를 획기적으로 줄이는 효과를 가져온다.

루프형 아키텍처를 기반으로 하며, 각 단계 k에서 잠재 상태 h를 업데이트한다. 종료 게이트를 통해 각 단계에서 나갈 확률을 정의하며, 이는 누적 분포 함수를 통해 특정 임계값을 넘을 때 종료되는 구조다.

엔트로피 기반 선택 기법을 사용한다. 교사 모델의 출력 분포에서 엔트로피 H = -sum(p * log(p))를 계산하여 상위 rho%의 어려운 토큰만 학습에 활용한다. [교사 모델의 예측 불확실성 수치 입력 → 상위 백분위 필터링 → 학습 대상 토큰 결정 → 학습 효율성 증대 의미]

단계별 보상 R(k) = Delta_acc(k) - C(k)를 정의한다. 여기서 Delta_acc는 교사 모델의 기준 로그 확률 대비 학생 모델의 개선량이며, C(k)는 연산 단계가 길어질수록 부과되는 시간 페널티다. [단계별 로그 확률 차이 계산 → 시간 페널티 차감 → 최종 보상 산출 → 정확도와 효율성의 균형 조절 의미]

노이즈 잠재 롤아웃을 통해 정책 경사를 계산한다. 잠재 상태에 N(0, sigma^2 * I) 노이즈를 더해 여러 번의 실행을 수행하고, 그룹 정규화된 이득을 사용하여 종료 정책을 최적화한다. [노이즈 섞인 잠재 상태 입력 → 여러 경로의 보상 평균 및 표준편차 계산 → 이득 산출 → 안정적인 정책 학습 의미]

Ouro-2.6B 모델에서 LoopRPT 적용 시, 어려운 토큰에 대한 정확도가 기존 대비 +3.58% 향상되었으며, 평균 추론 단계는 3.51단계에서 2.28단계로 약 35% 감소했다. 이는 정확도와 연산량 사이의 파레토 지배를 달성했음을 의미한다.

GSM8K 수학 벤치마크에서 2.6B 규모 모델의 정확도가 81.76에서 85.36으로 상승했으며, MBPP+ 코딩 테스트에서도 +2.91점의 성능 향상을 기록했다. 이는 사전 학습 단계에서의 강화 학습이 복잡한 하위 작업으로 잘 전이됨을 보여준다.

Ablation Study 결과, 노이즈 롤아웃과 시간 페널티, 토큰 선택 전략 중 하나라도 제거할 경우 정확도가 떨어지거나 추론 단계가 늘어나는 등 성능 저하가 관찰되어 각 구성 요소의 유효성이 입증됐다.

LoopRPT는 루프형 모델의 매개변수 공유 백본을 통해 잠재 상태를 반복 갱신하는 구조를 최적화한다. 각 단계는 독립적인 예측 로짓과 종료 게이트 값을 생성하여 유연한 추론 깊이를 지원한다.

핵심은 EMA 교사 모델을 활용한 밀집 보상 설계다. 교사는 학생보다 느리게 업데이트되며 안정적인 기준점 역할을 수행하여, 강화 학습 특유의 타겟 체이싱 문제를 완화하고 학습 안정성을 높인다.

손실 함수는 정책 경사 손실, 단계 가중치 적용 NTP 손실, 엔트로피 보너스, 그리고 KL 발산 제약의 가중 합으로 구성된다. 특히 단계 가중치 손실은 보상이 높은 단계에 더 큰 가중치를 두어 백본이 조기에 정답을 찾도록 유도한다.

노이즈 주입 기법은 결정론적인 루프 경로에 확률적 탐색 가능성을 부여한다. 이는 정책 경사 학습 시 분산을 줄이고, 다양한 잠재 경로에서도 모델이 강건하게 작동하도록 돕는 정규화 효과를 제공한다.

루프형 모델의 잠재 추론이 깊어질수록 성능이 항상 단조 증가하지 않고 오히려 떨어질 수 있는 현상이 관찰됐다. 이는 후반부 단계가 토큰 수준의 교차 엔트로피를 단조적으로 개선하도록 명시적으로 최적화되지 않았기 때문일 수 있다.