기술 요약: 시작하는 방식이 추론하는 방식이다: 프리픽스 튜닝된 사전 분포를 통한 RLVR 의 탐색 유도

1. 문제 제기

검증 가능한 보상을 활용한 강화 학습 (RLVR) 은 수학 문제 해결 및 정리 증명과 같은 작업에서 대형 언어 모델 (LLM) 의 추론 능력을 향상시키는 강력한 패러다임으로 부상했습니다. 그러나 저자들은 현재 RLVR 접근법의 중요한 한계를 지적합니다: 엔트로피 붕괴입니다.

RLVR 은 종종 단일 롤아웃 정확도 (Pass@1) 를 향상시키지만, 성공적인 추론 궤적의 범위를 확장하는 데는 실패하는 경우가 많습니다. 이는 모델이 좁은 범위의 추론 경로로 수렴하는 출력 다양성의 감소로 나타납니다. 엔트로피 정규화나 내재적 보상과 같은 기존 탐색 기술들은 종종 기본 모델의 분포를 고정된 것으로 간주하고 RL 목적 함수나 샘플링 절차를 수정하려 시도합니다. 본 논문은 이러한 방법들이 기본 모델의 사전 분포에 의해 RL 이 강력하게 형성되기 때문에 불충분하다고 주장합니다. 또한, 수동적인 탐색 기술은 종종 생성 품질을 저하시켜 효과적인 추론에 기여하지 않는 노이즈가 포함된 롤아웃을 초래합니다.

핵심적인 과제는 정책이 중복되거나 저품질의 궤적으로 붕괴되지 않고, 전체 모델 파라미터를 업데이트하는 계산 비용을 들이지 않으면서도 효과적인 탐색—다양하고 작업 관련성이 높으며 고품질인 추론 행동을 발견하는 것—을 장려하는 것입니다.

2. 방법론: IMAX 프레임워크

저자들은 **정보 최대화 증강 탐색 (Information-Maximizing Augmented eXploration, IMAX)**을 제안합니다. 이는 LLM 파라미터가 아닌 소프트 프롬프트 튜닝으로 최적화 대상을 이동시키는 경량 프레임워크입니다. 고정된 사전 분포 위에 RL 프로세스를 수정하는 대신, IMAX 는 추론 궤적에 대한 모델의 사전 분포를 능동적으로 재구성하는 소프트 프리픽스 풀을 학습합니다.

2.1 프리픽스 조건부 추론

핵심 메커니즘은 $C$개의 소프트 프리픽스 (연속적인 가상 토큰) $\mathcal{E} = \{e_1, \dots, e_C\}$로 구성된 풀을 포함합니다. 주어진 입력 프롬프트 $x$에 대해, 프리픽스 $e_z$가 입력 임베딩 시퀀스에 선행됩니다. 백본 LLM 은 **동결 (frozen)**된 채 유지되며, 정책 $\pi_E(y|x, z)$는 프롬프트와 선택된 프리픽스 식별자 $z$ 모두에 조건부로 설정됩니다.

• 초기화: 프리픽스는 외부 LLM 에 의해 생성된 의미 있는 텍스트 프롬프트 (예: "변수를 명시적으로 도입하세요", "경우로 나누세요") 풀에서 초기화됩니다. 이를 통해 프리픽스들이 서로 다른 추론 전략으로 시작하도록 보장합니다.
• 제어: 각 프리픽스는 "제어 노브" 역할을 하여 고유한 롤아웃 분포를 유도하고 모델을 추론 공간의 다른 영역 (예: 대수적 vs. 도식적 vs. 코드 유사 솔루션) 으로 유도합니다.

2.2 정보 최대화 (InfoMax) 보상

학습된 프리픽스 풀이 중복된 제어 (서로 다른 프리픽스가 동일한 궤적을 유도함) 로 붕괴되는 것을 방지하기 위해, 저자들은 InfoMax 보상을 도입합니다. 목표는 프롬프트 $X$가 조건일 때 프리픽스 식별자 $Z$와 생성된 궤적 $Y$ 사이의 상호 정보 (Mutual Information, MI) 를 최대화하는 것입니다:
$$I(Z; Y | X) = H(Y | X) - H(Y | X, Z)$$
직관적으로 이는 동일한 질문에 대해 서로 다른 프리픽스가 구별 가능하고 유효한 추론 행동을 생성하도록 장려합니다.

2.3 변분 하한 및 학습

상호 정보를 직접 계산하는 것은 불가능합니다. 저자들은 주어진 프롬프트 $x$에 대해 특정 응답 $y$를 생성한 프리픽스를 예측하도록 훈련된 경량 분류기인 변분 사후분포 (variational posterior) $q_\phi(z | x, y)$를 사용하여 변분 하한을 유도합니다.

• 역방향 형태: 그들은 MI 의 역방향 분해를 활용합니다: $I(Z; Y | X) = H(Z | X) - H(Z | X, Y)$.
• 보상 공식화: InfoMax 보상은 $\log q_\phi(z | x, y)$로 정의됩니다. 프리픽스 풀의 전체 목적 함수는 다음과 같습니다:
  $$\max_E \mathbb{E} [R(x, y) + \beta \log q_\phi(z | x, y)]$$
  여기서 $R(x, y)$는 표준 검증 가능 보상 (정확도) 이고 $\beta$는 하이퍼파라미터입니다.

2.4 교차 최적화

학습 절차는 교차 최적화 (좌표 상승 변분 추론) 방식을 따릅니다:

1. 사후분포 ($q_\phi$) 업데이트: 프리픽스 풀을 고정하고, 프롬프트 - 응답 쌍 (소프트 프리픽스 없이 생성되어 정보 누출을 방지함) 에서 프리픽스 식별자를 예측하도록 분류기 $q_\phi$를 업데이트합니다.
2. 프리픽스 풀 ($E$) 업데이트: $q_\phi$를 고정하고, InfoMax 보상이 증강된 그룹 상대 정책 최적화 (GRPO) 알고리즘을 사용하여 소프트 프리픽스를 업데이트합니다.
3. 층화 샘플링: 학습 중 각 프롬프트는 프리픽스 공간 내의 균일한 커버리지를 보장하기 위해 GRPO 그룹 내에서 서로 다른 프리픽스 식별자와 짝지어집니다.

3. 주요 기여

1. 경량 RLVR 패러다임: 본 논문은 전체 모델 파라미터가 아닌 소프트 프롬프트 풀을 최적화함으로써 추론을 개선하는 방법을 제시합니다. 이는 전체 파인튜닝의 계산 오버헤드 없이 동결된 모델의 조건부 생성 분포를 재구성합니다.
2. 다양성을 위한 InfoMax 보상: 저자들은 서로 다른 프리픽스가 구별 가능하고 작업 관련성이 높은 추론 궤적을 유도하도록 명시적으로 장려하는 응답 수준 InfoMax 보상을 제안하여, 프리픽스 풀이 중복된 제어 수단으로 붕괴되는 것을 방지합니다.
3. 실증적 검증: 세 가지 모델 규모 (Qwen2.5-1.5B, Qwen3-4B, Qwen3-8B) 에 걸쳐 수학 추론 벤치마크 (GSM8K, MATH-500, Minerva, AMC) 에서 수행된 광범위한 실험을 통해 IMAX 가 추론 품질과 다양성 모두를 향상시킨다는 것을 입증했습니다.

4. 실험 결과

저자들은 풀 정책 GRPO, 텍스트 프롬프트 초기화, 그리고 기타 다양성 지향 RLVR 방법 (DIVER, MERCI) 을 포함한 여러 베이스라인에 대해 IMAX 를 평가했습니다.

• 성능 향상: IMAX 는 표준 프리픽스 튜닝 RLVR 베이스라인을 일관되게 능가했습니다.
  • Qwen3-4B에서 IMAX 는 특정 벤치마크에서 표준 베이스라인 대비 Pass@4 를 최대 **11.60%**까지, Avg@4 를 **10.57%**까지 향상시켰습니다.
  • Qwen3-8B에서 DAPO+IMAX 는 GSM8K, MATH-500, Minerva 에서 DAPO-prefix 보다 유의미한 개선을 보였습니다.
• 다양성과 커버리지:
  • t-SNE 분석: 응답 임베딩의 시각화는 IMAX 에 의해 유도된 프리픽스가 다른 방법들에 비해 더 분리된 클러스터를 형성하여 더 강력한 그룹 내 다양성을 나타냈음을 보여주었습니다.
  • 추론 구성 요소: 히트맵은 서로 다른 프리픽스가 성공적으로 서로 다른 추론 패턴 (예: 프리픽스 1 은 대수적 조작을 선호, 프리픽스 2 는 경우 분할 및 검증을 선호) 을 유도하여 의미론적 초기화와 일치함을 보여주었습니다.
  • 오라클 성능: 여러 프리픽스의 결합된 성능 (오라클) 이 가장 좋은 단일 프리픽스를 초과하여, 프리픽스들이 문제 공간의 상호 보완적인 성공 영역을 커버하고 있음을 확인시켜 주었습니다.
• 효율성: IMAX 는 경량 프리픽스 임베딩과 분류기 헤드만 업데이트하므로 전체 파라미터 GRPO 에 비해 학습 비용 (GPU 시간) 이 3.0 배~6.0 배 적게 소요됩니다.

5. 중요성과 주장

본 논문은 모델이 추론 과정을 시작하는 방식 (사전 분포) 이 근본적으로 어떻게 추론하는지를 결정한다고 주장합니다. 소프트 프롬프트를 재사용 가능한 제어 변수로 취급함으로써 IMAX 는 RLVR 의 탐색 문제에 대한 확장 가능한 솔루션을 제공합니다.

• 알고리즘 독립성: 이 프레임워크는 기존 RLVR 파이프라인 (예: GRPO, DAPO) 에 원활하게 통합되도록 설계되었습니다.
• 구조화된 다양성: 단순히 무작위성을 증가시키는 방법과 달리, IMAX 는 구조화된 추론 다양성을 장려하여 탐색이 의미론적으로 구별 가능하고 고품질의 경로로 이어지도록 보장합니다.
• 약간의 한계: 저자들은 이 방법이 주로 추론 작업을 위해 설계되었으며 일반 지시 따르기 작업에서는 제한적인 향상만 보인다고 인정합니다. 또한, 효과성은 초기 텍스트 프롬프트의 품질에 의존하며, 이는 도메인 간에 수동 설계나 재튜닝이 필요할 수 있습니다. 학습된 프리픽스는 추론 중 지속적이지만 작은 오버헤드를 도입합니다.

요약하자면, IMAX 는 소프트 프리픽스를 통해 생성 과정의 "시작"을 최적화하는 것이 LLM 에서 탐색을 유도하고 추론 능력을 향상시키는 매우 효과적인 전략임을 보여주며, 풀 정책 강화 학습에 대한 계산적으로 효율적인 대안을 제시합니다.