기술 요약: 언어 모델을 위한 표현 기반 탐색 (Representation-Based Exploration)

문제 정의

강화 학습(RL)은 수학적 추론 및 코드 생성과 같이 검증 가능한 보상이 존재하는 작업에서 언어 모델(LM)의 성능을 향상시키는 데 유망한 가능성을 보여주었습니다. 그러나 현재의 RL 사후 학습(post-training) 기술이 가진 결정적인 한계는, 이 기술들이 실제로 새로운 행동을 발견하는 것인지, 아니면 단순히 베이스 모델에 이미 존재하는 능력을 (낮은 확률로 존재하던 것을) "날카롭게 만드는(sharpening/amplifying)" 것에 불과한지에 대한 불확실성입니다. 기존의 증거들은 표준적인 RL 레시피가 사전 학습된 모델에 결여된 능력을 깨우는 데 실패하고, 대신 이미 알려진 솔루션으로 수렴하는 경향이 있음을 시사합니다. 또한, 이를 완화하기 위한 데이터 큐레이션 및 알고리즘적 개입은 특히 복잡하고 개방적인 도메인에서 병목 현상이 되고 있습니다.

본 논문은 의도적인 탐색(deliberate exploration)—모델이 진정으로 새롭고 다양한 행동을 발견하도록 명시적으로 인센티브를 부여하는 것—이 언어 모델 추론을 위한 RL의 잠재력을 실현하는 데 필수적이라고 주장합니다. 핵심 과제는 조합론적으로 거대한 언어 공간 내에서 다양성과 새로움을 확장 가능한 방식으로 정량화하는 동시에, 효율적인 탐색을 위해 사전 학습된 모델에 내재된 사전 지식을 활용하는 것입니다.

방법론

저자들은 사전 학습된 언어 모델의 은닉 상태(hidden states)로부터 유도된 타원형 보너스(elliptical bonuses)를 사용하는 표현 기반 탐색(representation-based exploration) 전략을 제안합니다. 이 방법론은 두 가지 측면의 접근 방식을 통해 평가됩니다.

1. 추론 시 선택 (Inference-Time Selection)

다양성의 역할을 최적화 및 일반화와 분리하기 위해, 저자들은 먼저 추론 시 선택 문제를 정의합니다. 고정된 모델 $\pi$에 의해 생성된 $N$개의 후보 응답 풀이 주어졌을 때, 목표는 최대한 다양하면서도 정답을 포함할 확률이 높은 $k$개의 응답 서브셋을 선택하는 것입니다.

• 메커니즘 (RepExp): 이 알고리즘은 타원형 보너스를 최대화하도록 응답을 반복적으로 선택합니다. 응답 $y$에 대해, 모델의 마지막 레이어 은닉 상태를 응답의 토큰들에 대해 평균 내어 특징 벡터 $\bar{h}_\theta(x, y)$를 구성합니다.
• 다양성 지표: 이전에 선택된 응답들의 집합에 대한 후보 $y$의 새로움은 $div(y) = \bar{h}_\theta(x, y)^\top \Sigma^{-1} \bar{h}_\theta(x, y)$로 정의되며, 여기서 $\Sigma$는 선택된 특징들의 공분산 행렬입니다. 역공분산은 Woodbury 항등식을 사용하여 효율적으로 업데이트됩니다.
• 목표: 이 접근 방식은 검증기(verifier)에 쿼리를 던질 다양한 응답의 "코어셋(coreset)"을 선택하여, 검증기 효율성(verifier efficiency)(정답을 찾는 데 필요한 샘플 수를 줄이는 것)을 개선하는 것을 목표로 합니다.

2. RL 사후 학습 (RL Post-Training)

저자들은 이 탐색 전략을 GRPO(Group Relative Policy Optimization) 사후 학습 파이프라인에 통합합니다.

• 보상 증강: 특정 서브셋을 선택하는 대신, 타원형 보너스를 외적 보상(extrinsic reward)에 직접 더합니다. 롤아웃 $y_i$에 대해, 보상은 $r(x, y_i) + \beta \cdot \bar{h}_\theta(x, y_i)^\top \Sigma^{-1} \bar{h}_\theta(x, y_i)$가 됩니다.
• 구현: 공분산 행렬 $\Sigma$는 현재 배치 내의 롤아웃들과 비교하여 새로움을 측정하기 위해 각 훈련 배치마다 재초기화됩니다. 특징 벡터는 확장성을 보장하기 위해 희소 무작위 투영(sparse random projection)을 통해 낮은 차원(예: 512)으로 투영됩니다.

3. 토큰 수준 확장 (예비 연구)

개념 증명으로서, 저자들은 자기회귀 생성 과정 자체를 수정하는 방안도 탐구합니다. 각 생성 단계에서 타원형 보너스를 통해 토큰 수준의 로짓(logits)을 섭동(perturbing)함으로써, 사후 선택이 아닌 생성 과정 중에 다양성을 유도하는 것을 목표로 합니다.

주요 기여 및 결과

1. 개선된 검증기 효율성 (추론 시)

본 논문은 표현 기반 탐색이 광범 \text{범위의 모델}(Phi, Llama, Qwen)과 태스크(MATH, GSM8K, MBPP+, Game of 24)에서 검증기 효율성을 유의미하게 향상시킨다는 것을 입증합니다.

• 결과: Qwen-2.5-14b-Instruct와 같은 모델의 경우, 무작위 샘플링과 비교했을 때 정답을 찾는 데 필요한 샘플 수에서 50% 이상의 개선을 달anim했습니다.
• 모델 강도와의 상관관계: 탐색의 이점은 모델의 강도에 따라 달라집니다. 더 강력한 모델(예: Qwen-32B)은 크게 혜택을 받는 반면, 더 약한 모델은 개선이 없거나 오히려 성능이 저하될 수 있습니다.
• 난이도 민감도: 이 방법은 가장 어려운 문제들(예: MATH 데이터셋의 상위 20% 난이도 문제)에서 가장 큰 효율성 이득을 제공합니다.

2. 향상된 사후 학습 성능 및 다양성

표현 기반 탐색을 GRPO 사후 학습에 통합하면 표준 GRPO 및 "Unlikeliness" 베이스라인보다 우수한 성능을 나타냅니다.

• Pass@k 개선: 이 방법은 큰 $k$ 값에 대해 pass@k 비율을 유지하거나 향상시키는 반면, 표준 GRPO는 큰 $k$에 대해 베이스 모델 대비 성능이 저하되는 현상("다양성 붕괴", diversity collapse)을 보이는 경우가 많습니다.
• 샘플 효율성: AIME 2024 벤치마크에서, RepExp를 사용한 사후 학습된 Qwen-2.5-7b-Instruct는 표준 GRPO의 pass@256에 필적하는 pass@80을 달성하여, 테스트 타임 샘플 효율성 측면에서 3배의 향상을 보여주었습니다.
• 안티-샤프닝 (Anti-Sharpening): 베이스 모델의 로그 가능도(log-likelihood) 분석 결과, RepExp는 표준 GRPO가 베이스 모델의 분포를 날카롭게 만드는 경향이 있는 것과 달리, 훨씬 더 새로운(낮은 가능도를 가진) 응답을 생성함을 보여줍니다. 이는 이 방법이 "샤프닝 레짐(sharpening regime)"에서 성공적으로 탈출했음을 시사합니다.

3. 고비용 검증에 대한 적용

단백질 서열 생성(습식 실험을 통한 검증 비용이 매우 높은 분야)에 대한 예비 실험에서, 이 방법은 무작위 샘플링 대비 검증 효율성을 최대 56.9배 향상시켜, 검증 비용이 높은 도메인에서의 잠재적 유용성을 입증했습니다.

의의 및 주장

본 논문은 원칙적인 표현 기반 다양성 개념에 의해 가이드되는 의도적인 탐색이, 단순한 샤프닝을 넘어 언어 모델의 새로운 행동을 발견할 수 있는 실질적인 경로라고 주장합니다.

• 접근 방식의 참신함: 보조 네트워크(예: Random Network Distillation)나 복잡한 데이터 큐레이션에 의존하는 기존 연구와 달리, 이 방법은 사전 학습된 모델 자체의 은닉 상태를 활용하므로 추가적인 학습 메커니즘이 필요하지 않습니다.
• 검증 전략: 저자들은 다양성의 이점과 다른 RL 메커니즘을 분리하기 위한 두 단계 평가(추론 시 선택 후 사후 학습)의 가치를 강조합니다.
• 제한적 범위: 저자들은 현재 구현이 "단순하고 원칙적인" 체계임을 인정합니다. 이들이 모든 탐색 문제를 해결했다고 주장하는 것이 아니라, 이 접근 방식이 더 정교한 전략들의 트레이드오프를 이해하기 위한 강력한 베이스라인과 프레임워크를 제공한다고 제안합니다. 또한, 토큰 수준의 탐색이 유망하지만 현재는 최적화가 필요한 높은 계산 비용이 발생한다는 점을 언급했습니다.

요약하자면, 본 논문은 표현 기반 보너스를 통해 다양한 행동의 발견을 명시적으로 장려함으로써, 언어 모델이 표준적인 RL 사후 학습이 놓치는 새로운 솔루션을 찾도록 유도하여 추론 능력을 더욱 효과적으로 확장할 수 있다고 주장합니다.