🤖 machine learning

Where Should RL Post-Training Compute Go? Model Size, Search, Learning, and Feedback

This paper introduces a FLOP-accounting framework and the RACE diagnostic protocol to demonstrate that optimal RL post-training strategies under a fixed budget are not universal but depend on a complex interplay between model size, search depth, learning updates, and feedback mechanisms, urging researchers to report detailed compute allocations rather than just total FLOPs.

Original authors: Patrick Wilhelm, Odej Kao

Published 2026-07-16

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Patrick Wilhelm, Odej Kao

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a chef trying to create the world's most delicious soup, but you have a strict limit on how much electricity your kitchen stove can use. You could spend that electricity on a giant, expensive pot that holds a lot of soup at once, or you could use a smaller pot but stir it for hours. You could also spend energy on tasting the soup constantly to see if it needs salt, or on chopping vegetables faster. In the world of artificial intelligence, specifically when teaching computer programs to think and solve problems, scientists face a similar puzzle. They have a fixed "budget" of computer power (measured in a unit called FLOPs, which is just a count of how many math calculations the computer can do). The big question is: if you have a set amount of power, where should you spend it? Should you use a bigger, smarter brain for the AI? Should you let the AI try more guesses before learning? Or should you spend more power on a "teacher" that grades the AI's work? Getting this balance wrong means the AI might learn slowly or not at all, even if you have plenty of power.

This paper, titled "Where Should RL Post-Training Compute Go?", dives into exactly that question. The authors, Patrick Wilhelm and Odej Kao, treat the computer's energy budget like a limited amount of money that needs to be split between three main things: search (letting the AI try out many different answers), learning (actually updating the AI's brain based on those answers), and feedback (using a separate program to grade how good those answers are). They found that there is no single "best" way to spend the money. Instead, the perfect recipe changes depending on the size of the AI's brain, the type of teacher grading it, and what kind of test you want the AI to pass later.

The researchers discovered that simply saying "we used 100 units of power" isn't enough to understand how an AI was trained. It's like saying a chef used 100 watts of electricity without telling you if they spent it on the stove, the blender, or the lights. If you have a small AI brain, you might get the best results by spending most of your power on letting it try many different answers (search). But if you have a giant, powerful brain, that same amount of power might only buy you a few tries, so you might need to spend more on the actual learning steps. They also found that the "teacher" matters a lot. If the teacher is a simple rule-based checker (like a math key), almost all the power goes to the AI trying things. But if the teacher is another complex AI program, a huge chunk of the power budget gets eaten up just by the teacher doing its job, leaving less for the student AI to learn.

The paper suggests that the "best" way to train an AI depends entirely on what you are trying to achieve. If you want the AI to get a high score on a specific practice test, one mix of search and learning might work best. But if you want it to solve a completely new problem later, a different mix might be better. To help scientists figure this out without wasting millions of dollars on trial runs, the authors introduced a tool called RACE. Think of RACE as a quick, cheap taste-test. It runs a tiny, low-power experiment to guess where the "sweet spot" is for spending the budget, so researchers know which full-scale experiments are worth doing. The authors are careful to say this isn't a magic guarantee that the AI will get smarter; it's just a diagnostic tool to help them pick the right direction before they commit their full resources.

In the end, the main takeaway is that "more power" isn't the only answer. The smartest move is to be a careful accountant. Whether you are training a robot to walk or a computer to solve math problems, you need to report exactly how you split your energy between the brain size, the practice attempts, and the grading system. Because, as the authors show, the same amount of energy can buy you a completely different kind of intelligence depending on how you spend it.

Technical Summary: Where Should RL Post-Training Compute Go?

Problem Statement
Reinforcement Learning (RL) post-training is increasingly utilized to adapt foundation models for reasoning, planning, and feedback-driven robot-learning pipelines. However, current scaling analyses typically summarize the cost of these processes using a single aggregate metric: total FLOPs. This paper argues that such reporting is under-specified. Under a fixed post-training budget, the allocation of compute is a constrained resource-allocation problem involving four coupled dimensions:

Model Size: The capacity of the policy being trained.
Search/Rollout: The generation of candidate trajectories.
Learning/Updates: The gradient-based optimization of the policy.
Feedback/Reward Evaluation: The inference cost of reward models or verifiers.

The central question addressed is: Under a fixed FLOP budget, how should compute be distributed among these components to maximize performance? The authors posit that larger policies consume more per-token compute, thereby purchasing fewer rollouts and update steps under the same budget. Furthermore, the introduction of learned feedback mechanisms (e.g., Process Reward Models) introduces a new compute cost that competes with rollout and update budgets, a trade-off obscured by total-FLOP reporting.

Methodology
The study introduces a FLOP-accounting framework specifically for GRPO (Group Relative Policy Optimization) post-training. The total compute ( $C_{total}$ ) is decomposed into three distinct terms:
$C_{total} = C_{search} + C_{learning} + C_{reward}$
Where:

$C_{search}$ : Autoregressive rollout generation.
$C_{learning}$ : Policy optimization (accounting for LoRA adapter parameters and frozen backbone costs).
$C_{reward}$ : Reward or verifier inference.

The primary allocation variable is the update fraction ( $\rho$ ), defined as the ratio of learning compute to total compute. The authors conduct experiments using LoRA-adapted Qwen2.5 policies on the Polaris-53K mathematical reasoning dataset. The experimental design includes:

Stage A: Varying model sizes (1.5B, 3B, 7B) to study search-learning allocation.
Stage B & D: Comparing different reward systems (sparse, structured, dense, proxy-PRM, and real PRM) to assess how feedback cost shifts allocation frontiers.
Evaluation: Metrics include native training reward, downstream pass@1 accuracy on GSM8K and MATH-500, and a "common judge" process-quality score.

To address the high cost of exhaustive search, the authors introduce RACE (Reward-Aware Compute Allocation), a diagnostic pilot-grid protocol. RACE runs small IsoFLOP pilot grids to estimate local reward-allocation responses and identify whether a setting belongs to a "search-heavy," "update-heavy," or "reward-heavy" regime. The authors explicitly state that RACE is a diagnostic tool for regime identification, not a guarantee of held-out improvement.

Key Contributions

FLOP Accounting Framework: A formal decomposition of RL post-training compute into rollout, update, and reward-evaluation components, enabling the comparison of runs that are mechanistically different despite having identical total FLOPs.
Coupled Model-Size and Allocation: Empirical evidence showing that model choice and training allocation are inseparable. Under a fixed budget, larger policies inherently buy fewer updates and rollouts, altering the optimal allocation frontier.
Reward- and Target-Dependence: Demonstration that the "best" allocation depends on the reward system (rule-based vs. PRM) and the evaluation target (native reward vs. downstream accuracy vs. process quality). For instance, PRM-based feedback shifts the optimal update fraction ( $\rho$ ) lower because a portion of the budget is consumed by reward-model inference.
RACE Protocol: A diagnostic method for identifying allocation regimes via small pilot grids, intended to guide validation runs rather than replace them.

Results

Conditional Allocation Frontiers: The optimal allocation is not a fixed point but a conditional frontier. Search-heavy settings underutilize learning capacity, while update-heavy settings become exploration-limited.
Reward System Impact: Rule-based rewards (sparse, structured, dense) tend to favor high-update regimes ( $\rho \approx 0.72$ ). In contrast, PRM-style feedback (proxy and real) shifts the optimal native-reward frontier toward lower update fractions ( $\rho \approx 0.44–0.45$ ) because the fixed budget must now cover reward-model inference.
Target Disagreement: Different evaluation targets select different optimal allocations.
- Native Reward: Favors the allocation that maximizes the specific feedback signal used during training.
- Downstream Pass@1: Often favors high-update regimes (e.g., the $a180$ setting) across reward systems.
- Common Judge: Often peaks at mid-update regimes (e.g., $a100$ ) for dense and PRM rewards.
- This disagreement indicates that a recipe maximizing one metric may not maximize another.
RACE Performance: RACE successfully identifies the correct regime (high vs. low update) in validation settings with high accuracy (1.0 same-regime rate) but does not guarantee a global improvement in held-out reward over baselines.

Significance and Claims
The paper modestly claims that RL post-training is a constrained resource-allocation problem rather than a simple scaling law. The authors argue that reporting only total FLOPs hides the mechanistic trade-offs between model capacity, experience generation, and feedback evaluation.

The significance lies in the recommendation for the community to report allocation-aware compute: total FLOPs alongside the specific breakdown of rollout, update, and reward-evaluation FLOPs, and the resulting update fraction $\rho$ . This transparency is crucial for robot-learning systems, where improving a reasoning or feedback module via RL post-training involves deciding whether to invest limited resources in larger models, more extensive search, or stronger feedback mechanisms. The paper concludes that there is no universal "best" recipe; the optimal allocation depends jointly on model size, reward design, and the specific evaluation target being optimized.

Technical Summary: Where Should RL Post-Training Compute Go?

More like this