Knapsack RL: 예산 할당 최적화를 통한 대형 언어 모델 탐색 개방

초록

대형 언어 모델(LLMs)은 강화 학습을 통해 스스로 개선할 수 있으며, 이 과정에서 더 나은 해결책을 탐색하고 발견하기 위해 궤적을 생성합니다. 그러나 이러한 탐색 과정은 계산 비용이 많이 들어, 현재의 방법론에서는 각 작업에 제한된 탐색 예산을 할당할 수밖에 없습니다. 이러한 균일한 할당 방식은 문제적인 극단적인 사례를 초래합니다: 쉬운 작업은 지속적으로 성공하는 반면, 어려운 작업은 지속적으로 실패하며, 이 둘 모두 널리 사용되는 그룹 상대 정책 최적화(GRPO)의 학습 업데이트 동안 기울기가 제로가 됩니다. 우리는 이 문제를 탐색 예산 할당의 관점에서 접근합니다. 각 작업의 탐색을 고유한 "가치"와 "비용"을 가진 "항목"으로 간주함으로써, 고전적인 배낭 문제와의 연결을 설정합니다. 이 공식을 통해 모델의 현재 학습 상태에 따라 자원을 적응적으로 분배하는 최적의 할당 규칙을 도출할 수 있습니다. 이 방법을 GRPO에 적용했을 때, 학습 중에 비제로 정책 기울기의 효과적인 비율이 20-40% 증가했습니다. 이는 계산적 "무료 점심"으로 작용하여, 학습이 포화된 작업에서 가장 영향력 있는 작업으로 탐색 예산을 재할당할 수 있게 합니다. 이를 통해 특히 어려운 문제에 대해 균일한 할당 방식에서는 계산적으로 불가능했던 상당히 큰 예산(예: 93회 롤아웃)을 할당할 수 있습니다. 이러한 개선은 수학적 추론 벤치마크에서 의미 있는 성과로 이어졌으며, 평균적으로 2-4점의 향상과 특정 작업에서 최대 9점의 성과를 보였습니다. 특히, 전통적인 균일 할당 방식으로 비슷한 성능을 달성하려면 약 2배의 계산 자원이 필요할 것입니다.

English

Large Language Models (LLMs) can self-improve through reinforcement learning, where they generate trajectories to explore and discover better solutions. However, this exploration process is computationally expensive, often forcing current methods to assign limited exploration budgets to each task. This uniform allocation creates problematic edge cases: easy tasks consistently succeed while difficult tasks consistently fail, both producing zero gradients during training updates for the widely used Group Relative Policy Optimization (GRPO). We address this problem from the lens of exploration budget allocation. Viewing each task's exploration as an "item" with a distinct "value" and "cost", we establish a connection to the classical knapsack problem. This formulation allows us to derive an optimal assignment rule that adaptively distributes resources based on the model's current learning status. When applied to GRPO, our method increases the effective ratio of non-zero policy gradients by 20-40% during training. Acting as a computational "free lunch", our approach could reallocate exploration budgets from tasks where learning is saturated to those where it is most impactful. This enables significantly larger budgets (e.g., 93 rollouts) for especially challenging problems, which would be computationally prohibitive under a uniform allocation. These improvements translate to meaningful gains on mathematical reasoning benchmarks, with average improvements of 2-4 points and peak gains of 9 points on specific tasks. Notably, achieving comparable performance with traditional homogeneous allocation would require about 2x the computational resources.

Knapsack RL: 예산 할당 최적화를 통한 대형 언어 모델 탐색 개방

Knapsack RL: Unlocking Exploration of LLMs via Optimizing Budget Allocation

초록

Support