Command Palette
Search for a command to run...
Poly-EPO: 탐색적 추론 모델 학습
Poly-EPO: 탐색적 추론 모델 학습
Ifdita Hasan Orney Jubayer Ibn Hamid Shreya S Ramanujam Shirley Wu Hengyuan Hu Noah Goodman Dorsa Sadigh Chelsea Finn
초록
경험 기반 학습에서 탐험(Exploration)은 핵심적인 요소로, 에이전트가 복잡한 문제에 대한 해결책을 찾게 하고, 새로운 상황에 일반화하며, 테스트 시간 컴퓨테이션(Test-time compute)을 통해 성능을 확장할 수 있도록 합니다. 본 논문에서는 사후 훈련(Post-training)된 언어 모델(LMs)에 대해 적극적 탐험(Optimistic exploration)을 명시적으로 장려하고 탐험과 활용(Exploitation) 간의 시너지를 촉진하는 프레임워크를 제안합니다. 이 프레임워크의 핵심 아이디어는 보상 함수(Reward function) 하에서 집단적으로 정확하고 추론 전략(Reasoning strategies) 측면에서 탐색적인 응답 집합(Set of responses)을 생성하도록 언어 모델을 훈련시키는 것입니다. 먼저, 임의의 목적 함수(Objective functions) 하에서 집합 강화 학습(Set RL)을 통해 언어 모델을 최적화하기 위한 일반적인 방법을 개발했으며, 이는 이점 계산(Advantage computation)을 수정함으로써 표준 강화 학습 알고리즘(Standard RL algorithms)을 해당 설정에 어떻게 적응시킬 수 있는지를 보여줍니다. 이어 탐험과 활용을 명시적으로 시너지를 이루는 목적 함수를 통해 이 프레임워크를 구체화한 다채색 탐색 정책 최적화(Polychromatic Exploratory Policy Optimization, POLY-EPO)를 제안합니다. 다양한 추론 벤치마크(Reasoning benchmarks)에 걸친 실험 결과, POLY-EPO가 높은 pass@k 커버리지(Pass@k coverage)를 통해 일반화 능력을 향상시키고, 모델 생성물(Model generations) 간의 다양성을 더 잘 보존하며, 테스트 시간 컴퓨테이션과 효과적으로 확장됨을 확인했습니다.
One-sentence Summary
The authors propose POLY-EPO, a post-training framework for language models utilizing set reinforcement learning with modified advantage computation to generate response sets that are collectively accurate under the reward function and employ exploratory reasoning, thereby synergizing exploration and exploitation to improve generalization via higher pass@k coverage, preserve generation diversity, and scale effectively with test-time compute across reasoning benchmarks.
Key Contributions
- A framework for post-training language models is presented to explicitly encourage optimistic exploration while promoting synergy between exploration and exploitation. This approach trains models to generate sets of responses that are collectively accurate under the reward function and exploratory in their reasoning strategies.
- The work develops a general recipe for optimizing language models with set reinforcement learning under arbitrary objective functions by adapting standard algorithms through a modification to the advantage computation. This shows how standard reinforcement learning algorithms can be adapted to this setting.
- Polychromatic Exploratory Policy Optimization (POLY-EPO) is introduced as an instantiation of the framework that encodes synergy directly in the advantage function through covariance between average reward and diversity. Experiments across reasoning benchmarks demonstrate improved generalization, higher pass@k coverage, and effective scaling with test-time compute.
Introduction
Exploration is critical for language models to solve complex reasoning tasks and scale effectively with test-time compute. However, standard reinforcement learning fine-tuning often collapses generation diversity onto narrow high-reward behaviors. Prior methods attempt to address this by adding exploration bonuses, yet they treat exploration and exploitation as separate objectives requiring careful hyperparameter tuning. The authors introduce Polychromatic Exploratory Policy Optimization (POLY-EPO), a framework that leverages set reinforcement learning to optimize sets of responses collectively. This method explicitly synergizes exploration and exploitation by encoding the covariance between average reward and strategy diversity directly into the advantage function. As a result, the approach encourages optimistic exploration of novel reasoning strategies and improves generalization metrics like pass@k.
Method
The authors introduce Polychromatic Exploratory Policy Optimization (POLY-EPO), a method designed to optimize language models using a set reinforcement learning framework. This approach generalizes standard reinforcement learning by assigning rewards to sets of sampled generations rather than to each generation independently. The core objective is to balance exploration and exploitation, encouraging the model to discover diverse reasoning strategies while maintaining accuracy.
The framework begins with a set-level objective function. For a given prompt, the policy samples a set of n generations. A set-level reward function f is applied to the entire set, coupling all generations under a shared learning signal. To optimize this objective efficiently, the authors derive a tractable gradient estimator. For each prompt, N independent generations are sampled, where N is greater than the set size n. From these samples, K sets of size n are constructed combinatorially. A set-level score is computed for each constructed set, and a baseline is established as the average of these scores across all sets. The advantage for a specific set is the difference between its score and this baseline.
To integrate this with standard reinforcement learning algorithms, a marginal set advantage is defined for each individual generation. This value is calculated as the average advantage of all sets that contain the specific generation. This construction allows the method to use existing policy gradient algorithms, such as PPO or GRPO, by replacing the standard advantage function with the marginal set advantage. The marginal set advantage for a generation y is defined as:
Amarg♯(x,y;f):=∣G(y)∣1G∈G(y)∑A♯(x,G;f)where G(y) represents the collection of all sets containing the generation y. The resulting estimator is unbiased and scales efficiently with the number of samples.
The specific objective optimized in POLY-EPO is the polychromatic objective. This function is defined as the product of the mean reward of the generations in the set and a diversity measure. The mathematical formulation ensures that a set must optimize both exploration and exploitation simultaneously, as the score depends on the product of these two factors. The polychromatic objective is expressed as:
fpoly(x,y1,⋯,yn)=n1i=1∑nr(x,yi)⋅d(x,y1,⋯,yn)If either the reward or the diversity is low, the overall score is diminished.
Diversity is measured using a language model judge to cluster responses according to their underlying reasoning strategies. The judge groups generations based on macro and micro strategies while ignoring superficial textual differences. The diversity of a set is calculated as the number of distinct clusters represented in the set divided by the set size. Degenerate generations that exhibit reward hacking or nonsensical behavior are isolated into a dedicated cluster and excluded from the diversity calculation to ensure robustness.
The final training algorithm follows a standard on-policy update loop. The model samples generations for a batch of prompts, constructs sets, computes the polychromatic scores, and derives the marginal set advantages. These advantages are then used to update the policy parameters. This process prioritizes simplicity and scalability, leveraging infrastructure already present in modern post-training pipelines while introducing the benefits of set-level credit assignment.
Experiment
The experiments evaluate POLY-EPO on mathematical reasoning benchmarks and synthetic domains with verifiable rewards to assess whether its exploration mechanism improves reasoning performance. In mathematical tasks, POLY-EPO demonstrates superior diversity in reasoning strategies and better utilization of test-time compute compared to baselines that suffer from strategy collapse. Synthetic domain results further confirm that POLY-EPO discovers significantly more distinct successful strategies than standard methods which quickly converge to single solutions, validating that the proposed objective effectively balances exploration and exploitation to enhance generalization.