🤖 AI

Where Rollouts Begin: Low-Load, High-Leverage First-Token Diversification for RLVR

This paper introduces REFT, a lightweight RLVR method that enhances rollout diversity and improves reasoning performance by uniformly sampling from the policy's top- $N$ candidates for the first token after the reasoning marker, thereby broadening exploration without altering the correctness signal.

Original authors: Soeun Kim, Albert No

Published 2026-05-28

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Soeun Kim, Albert No

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

The Core Problem: The "First Step" Trap

Imagine you are teaching a student how to solve complex math problems. Instead of giving them the answer key, you ask them to try solving the problem eight different times (this is called a "rollout group"). You then grade their answers and tell them which attempts were right and which were wrong. The student uses this feedback to get smarter.

The problem? The student is incredibly stubborn about how they start. Every single time, they begin their solution with the exact same phrase, like "Let’s see..." or "First,".

Because they always start the same way, their eight attempts look very similar. If they get the answer wrong, they probably got it wrong in all eight attempts because they were stuck in the same "thinking lane." The teacher (the AI trainer) can’t give useful feedback because there’s no variety to compare.

Most AI researchers thought, "Well, the first word doesn't matter much. It’s just a polite opener. Let’s focus on changing the middle or end of the solution to create variety."

This paper argues that this is a mistake. It suggests that the very first word is actually a high-leverage point—a small change that opens up huge new possibilities.

The Discovery: The "Doorway" Effect

The researchers looked closely at the first word the AI generates after it decides to start reasoning (marked by a special tag like <think>). They found two surprising things:

The AI is Overconfident: The AI is extremely sure about its first word. It picks the same opener 50–60% of the time, even though there are many other valid ways to start.
Correctness Doesn’t Care: Whether the AI starts with "Let," "First," or even a rare word like "Consider," the chance of getting the final math answer right is almost the same.

The Analogy: Think of the first word as a doorway into a house.

The AI always uses the front door because it’s the most obvious one.
But there are 19 other doors (side doors, back doors, garage doors) that lead into the same house.
The AI ignores these other doors because it’s biased toward the front door.
However, once you step through any of those doors, you can explore different rooms (reasoning paths) inside the house. By only using the front door, the AI is missing out on exploring the rest of the house.

The Solution: REFT (Rollout Exploration with First-Token Diversification)

The authors created a simple trick called REFT. It doesn’t change the AI’s brain or its grading system. It only changes how the AI is asked to start its attempts.

How it works:

The AI looks at its top 20 possible first words.
Instead of picking the most popular one every time, REFT says: "Pick 4 different words from this list at random."
The AI then writes 8 solutions, but it forces itself to start with those 4 different words (2 solutions per word).

The Analogy:
Instead of sending 8 students through the front door, you send 2 students through the front door, 2 through the side door, 2 through the back door, and 2 through the garage.

The students still use the same logic once they are inside.
But because they entered from different places, they explore different parts of the house.
This gives the teacher a much wider variety of attempts to grade, making the learning process more effective.

Why This Works Better Than Other Methods

Usually, to get more variety, researchers use a method called "Temperature." This is like shaking the dice before rolling them. It makes the AI more random everywhere.

The Problem with Temperature: It makes the AI random in the middle of the solution too. If the AI gets random in the middle, it might make silly mistakes or nonsensical leaps. It’s like shaking the dice so hard that the students start walking through walls instead of just exploring different rooms.
The Benefit of REFT: REFT only shakes the dice for the first step. The rest of the solution is still logical and careful. It’s like ensuring the students enter through different doors, but once inside, they walk carefully and logically.

The Results

The researchers tested this on several AI models (from small 0.5 billion parameter models to larger 7 billion parameter models) and various math datasets.

Better Accuracy: The AI got more correct answers overall.
Better Coverage: When the AI was allowed to try 8 or 64 times, it was much more likely to find at least one correct solution.
Less "All-Wrong" Groups: In standard training, sometimes all 8 attempts are wrong, giving the teacher no useful feedback. With REFT, because the attempts are more diverse, it’s less likely that all of them fail. At least one attempt usually finds a correct path, giving the AI something good to learn from.

Summary

The paper flips a common assumption on its head. We usually think the "boring" first step of a reasoning process doesn't matter. But this paper shows that the first step is a powerful lever. By simply forcing the AI to try different starting phrases, we can unlock a wider variety of reasoning paths without making the AI less logical. It’s a small, cheap change that leads to smarter, more robust AI reasoning.

Technical Summary: Where Rollouts Begin: Low-Load, High-Leverage First-Token Diversification for RLVR

Problem Statement
Reinforcement Learning with Verifiable Rewards (RLVR) improves the reasoning capabilities of language models by sampling candidate reasoning traces and scoring them with an automatic verifier, rather than relying on labeled trajectories. In grouped RLVR methods like GRPO and DAPO, the policy is updated based on a finite set of rollouts for the same prompt. The diversity of these rollouts is a central bottleneck: if rollouts are semantically redundant or receive identical rewards, the group provides little contrastive signal for learning. Existing methods typically broaden exploration by adjusting temperature, branching at high-entropy reasoning steps, or selecting diverse final answers. These approaches share an assumption that early, low-entropy prefix tokens are not valuable sites for exploration. This paper challenges that assumption, identifying the first token after the reasoning marker (<think>) as an overlooked position for broadening rollout diversity.

Methodology: REFT
The authors introduce REFT (Rollout Exploration with First-Token Diversification), a lightweight modification to the rollout sampling process in RLVR. The method is grounded in the observation that the first token is "low-load" (it is typically a discourse opener like "First" or "Let" rather than a substantive reasoning step) but "high-leverage" (it conditions the distribution of all subsequent tokens).

The authors diagnose that the policy’s distribution over the first token is sharply peaked, yet the correctness of the final answer remains nearly flat across the top-20 alternative first tokens. This decoupling suggests that low-probability first tokens are not necessarily poor reasoning choices but rather under-sampled routes to viable solutions.

REFT operates as follows:

Candidate Selection: For a given prompt, the policy identifies the top- $N$ valid first-token candidates.
Uniform Sampling: Instead of sampling all $G$ rollouts from the original skewed distribution, REFT uniformly samples $K$ distinct first tokens from the top- $N$ candidates.
Budget Allocation: The rollout budget is evenly allocated across these $K$ selected tokens. For each selected first token, the remaining continuation is generated using the standard decoder.
Integration: All other components of the RLVR pipeline—including the verifier, reward calculation, advantage estimation, and RL objective (e.g., GRPO or DAPO)—remain unchanged. REFT is thus a drop-in replacement for the first decision of the rollout sampler, not a new RL algorithm.

Key Contributions

Diagnostic Insight: The paper identifies a structural asymmetry in RLVR: the first-token distribution is sharply concentrated, but this concentration is weakly correlated with correctness. Standard RLVR training (e.g., GRPO) exacerbates this by over-crediting the most frequent first token, further narrowing exploration.
REFT Algorithm: A simple, efficient method to diversify rollouts by explicitly covering multiple plausible first-token openings, thereby exposing the policy to broader continuation regions without injecting noise into later, more critical reasoning steps.
Analysis of Diversity: The paper demonstrates that REFT increases semantic diversity and answer diversity during training, reduces the frequency of "all-wrong" rollout groups (a specific type of zero-variance group), and prevents the over-sharpening of the first-token prior during training.

Results
Experiments were conducted on four base models (Qwen2.5-0.5B/3B/7B-Instruct and Llama3.2-3B-Instruct) across three difficulty regimes (GSM8K, BigMath-Easy, BigMath) using GRPO and DAPO baselines.

Performance Gains: REFT consistently improved aggregate Pass@1, Pass@8, and Pass@64 metrics over baselines. For example, on Qwen2.5-3B-Instruct with GRPO, Math Avg. Pass@1 increased from 28.91 to 32.13, and Pass@8 from 48.39 to 52.38.
Scalability: Gains were observed across model scales from 0.5B to 7B parameters, indicating the method is not contingent on large model capacity.
Budget Efficiency: REFT improved performance under matched rollout budgets. It also improved Pass@64, showing that the gains in single-sample correctness (Pass@1) did not come at the cost of larger-budget discoverability.
Ablation: The method is robust to hyperparameter choices. A candidate width of $N=20$ and selected token count of $K=4$ provided optimal balance, avoiding the dilution of budget seen with wider pools ( $N=50, 100$ ) or insufficient replication ( $K=1$ ).

Significance
The paper claims that REFT offers a simple inversion of usual exploration intuition in RLVR. Useful diversity can be recovered not only from high-entropy reasoning forks but also from low-load prefix choices whose probability is sharply biased and weakly tied to correctness. By targeting the first token, REFT avoids the pitfalls of high-temperature sampling (which injects harmful noise into later reasoning steps) and provides controlled coverage of alternative reasoning paths. The results suggest that reallocating rollout budget away from repeated visits to a weakly supported discourse preference toward plausible first-token routes yields more useful reward contrast and improved reasoning performance.