Reflect, Retry, Reward: Self-Improving LLMs via Reinforcement Learning

Paper Details

Title: Reflect, Retry, Reward: Self-Improving LLMs via Reinforcement Learning

Authors: Shelly Bensal, Umar Jamil, Christopher Bryant, Melisa Russak, Kiran Kamble, Dmytro Mozolevskyi, Muayad Ali, and Waseem AlShikh, all affiliated with Writer, Inc.

Version: arXiv 2505.24726v1, submitted on May 30, 2025 in cs.CL.

Link: https://arxiv.org/abs/2505.24726

Code status: The paper does not provide a public code repository. Searches by title, distinctive method sentences, and author names also did not surface an official implementation. The implementation notes point to an extension of TRL's GRPOTrainer, vLLM rejection sampling, and task-specific validators.

One-Sentence Summary

The paper trains the failure response, not the final answer: when a first attempt fails, the model reflects; when a reflected retry succeeds, GRPO assigns reward only to the reflection tokens, improving verifiable function calling and equation construction under binary supervision.

Motivation

Large language models now perform well across mathematics, code, reasoning, and tool use, yet their competence is still uneven. Success on one prompt does not guarantee success on a similar prompt, and strength on a public benchmark does not imply stability on a newly deployed verifiable task. The direct remedy is to collect those failures and fine-tune on them, but that assumes the failure distribution has been gathered and that correct demonstrations are available.

Synthetic data offers another route: sample demonstrations from a stronger teacher and train a smaller student. That route weakens when the strongest available model also fails often, when data cannot be shared, or when the task depends on a live tool environment. The paper therefore asks a narrower and more practical question: if the only reliable supervision is success or failure, can a model still improve?

Prompt-time self-reflection has long been popular. The model answers, critiques its own attempt, and revises. This requires no training set, but its value depends heavily on the reflection prompt. Models do not always recognize their own errors, and repeated reflection can introduce noise for easy prompts or already-strong models.

Training-based self-correction attempts to internalize that behavior. Prior systems fine-tune on correction trajectories or formulate the process as multi-turn reinforcement learning. Many still rely on teacher models for data or supervision. The authors instead impose a stricter condition: no external teacher, no reward for the correct final answer, and only a binary validator that determines whether the first attempt failed and whether the second attempt succeeded.

The central move is to shift the learning target from solving a particular task to producing an effective reflection after failure. Task-specific RL optimizes the final response; the resulting skill may remain tied to the training distribution. Reflection after verified failure is a structural behavior that can, in principle, be reused wherever success is checkable.

Preliminaries: Why GRPO Fits the Setting

Group Relative Policy Optimization, introduced in DeepSeekMath, is an outcome-level RL method. It avoids a separate value network. For one query, it samples a group of completions and estimates each completion's advantage relative to the group's rewards. That suits sparse terminal rewards: whether an equation evaluates to a target, whether code executes, or whether a tool call satisfies the schema.

For a query $q$, GRPO samples $G$ completions $y_i$ from the previous policy $\theta_{\text{old}}$. If the verifier assigns $v(q,y_i) \in \{0,1\}$, the group-relative advantage is:

$$\hat A_i=\frac{v(q,y_i)-\operatorname{mean}_{j\in[1,G]} v(q,y_j)}{\operatorname{std}_{j\in[1,G]} v(q,y_j)+\epsilon}$$

Here, $\epsilon$ stabilizes division. The formula converts binary outcomes into a relative signal: successful members of the group receive positive advantage and unsuccessful members receive negative advantage. The policy update then raises the probability of successful completions relative to the group.

The complication in this paper is that generation has two stages. The model first answers; after failure, it generates a reflection; then it retries. If the success of the retry were credited to the whole sequence, the update would also strengthen retry-answer tokens. The authors want only the reflection to learn, so they need a reward mask that confines credit to that intermediate text.

Method

1. Verifiability defines the state machine

The method starts with a binary validator:

$$v(q,y)=\mathbb{1}\{y \text{ satisfies } C(q)\}$$

Here, $q$ is the query, $y$ is the model output, and $C(q)$ denotes the task constraints. In function calling, $C$ checks the tool, parameter names, parameter values, and format. In Countdown, it checks that all provided numbers are used exactly once and that the equation evaluates to the target. In code, it may simply check execution. The validator need not produce a score.

Let the first generation be $y_0 \sim \pi_\theta(\cdot \mid q)$. If $v(q,y_0)=1$, training does nothing for that example. If $v(q,y_0)=0$, the model receives a short failure prompt and produces a reflection $r$:

$$r \sim \pi_\theta(\cdot \mid q,y_0,c_{\text{fail}})$$

The original task and reflection are then placed in context for a second attempt $y_1$:

$$y_1 \sim \pi_\theta(\cdot \mid q,y_0,r)$$

These are not three independent tasks. Failure switches the state to reflection; reflection changes the retry context; the retry outcome determines whether the reflection is credited.

Figure 1: A failed first answer triggers a reflection. If the reflected retry succeeds, only the reflection tokens are rewarded.

2. Reward lands on reflection tokens only

Let $s=v(q,y_1)$ denote the retry outcome. A positive learning signal exists only when $s=1$. The token-level reward structure can be expressed as:

$$\hat A_i^{(t)} = \hat A_i \cdot \mathbb{1}\{t \in R_i\}, \quad R_i=\{t:\text{token }t\text{ belongs to reflection }r_i\}$$

Here, $\hat A_i^{(t)}$ is the advantage for token $t$ in trajectory $i$, and $R_i$ is the set of token indices belonging to the reflection. If the retry fails, the reflection receives no positive reinforcement. If the retry succeeds, the first answer, retry answer, and all non-reflection tokens have advantage zero.

This is stronger than saying the model learns from reflection. It does not update because it wrote the right tool call; it updates because it wrote a passage that turned a verified failure into a verified success. The authors exclude the final answer from the learning signal to reduce memorization of a particular Countdown equation or APIGen tool parameterization and to keep the learned behavior as task-agnostic as possible.

The token-level policy gradient can be written as:

$$\nabla_\theta J(\theta)=\mathbb{E}\left[\sum_{t\in R}\hat A \, \nabla_\theta \log \pi_\theta(r_t\mid q,y_0,r_{{<t}})\right]$$

Here, $r_{{<t}}$ is the reflection prefix generated before token $r_t$. Only the conditional probabilities of reflection tokens receive gradient. The GRPO ratio and clipping can operate as usual in TRL, but the mask compresses the update to the reflection segment.

Adding the usual reference-policy KL term gives:

$$J_{\text{KL}}(\theta)=\mathbb{E}\left[\sum_{t\in R}\hat A \, \nabla_\theta \log \pi_\theta(r_t\mid \cdot)\right]-\beta\,D_{\mathrm{KL}}\!\left[\pi_\theta\|\pi_{\text{ref}}\right]$$

The authors set the KL coefficient $\beta=0.001$. This is small enough to preserve exploration while limiting drift away from the reference model. Binary success supplies the direction of learning; the KL regularizer supplies stability.

3. A dataset of failures

Training did not sweep the full task sets. The authors first sampled up to 64 responses per query and retained only queries for which every sampled response failed the validator. Larger models fail less often, so they generally need more samples to expose enough failure queries. vLLM with prefix caching accelerated this rejection-sampling pass.

This has three consequences. First, queries that the model already solves cannot contribute to reflection training, so skipping them saves compute. Second, nondeterministic sampling lets occasionally failing queries enter training even when their success rate is high; a query solved 80% of the time still carries a learnable 20% failure mode. Third, a failure-only set makes convergence easier to interpret, because every training example lies near a verified boundary.

The authors emphasize that this was an efficiency and analysis choice, functionally equivalent to learning from an online stream of successes and failures. The important point is not that successes were artificially excluded, but that clearly solvable queries were filtered out so RL could concentrate near the decision boundary.

4. Multi-step GRPO

TRL's GRPOTrainer normally maps one generation to one reward. The authors extend _prepare_inputs with a second_step. After the first completion, the trainer does not immediately update. It preserves the masks already computed by GRPO and performs another completion generation. Because training uses the failure dataset, the first completion corresponds to the reflection, so the existing mask identifies exactly the reflection tokens.

This implementation also adds a structure to reward-function inputs so downstream rewards can understand the first completion and subsequent steps. The authors note that the multi-step design can accept more complex downstream reward mechanisms, but the experiments deliberately use only binary success. The architecture remains general while the actual reward remains minimal.

Training ran for up to 1,750 steps with an effective batch of 256 failures, though most models converged earlier and were evaluated at their convergence point. The final hyperparameters used a 5e-7 learning rate, 0.03 warmup ratio, cosine annealing, and KL coefficient 0.001. Models were trained on 4 to 8 H100 GPUs. The study stops at 1.5B to 8B because GRPO's compute and scalability remain known bottlenecks for larger models.

Experiments

APIGen function calling

APIGen contains 60,000 function-calling examples; the test split has 12,000. Each item provides a user query, candidate tools with JSON parameter schemas, and the correct call. Across the dataset there are 4,211 unique tools, an average of 2.3 parameters per tool, and an average of 2.8 candidate tools per query, with a minimum of 1 and maximum of 8. Success requires the correct tool, parameters, values, and format.

To avoid data contamination, the authors evaluate models released before APIGen: Qwen2 1.5B/7B Instruct, Llama3.1 8B Instruct, and Phi3.5-mini Instruct. They also report untrained 70B-scale baselines: Qwen2 72B, Llama3.1 70B, and Palmyra-X4. Because model families use different tool-calling templates, they tested formats and chose the strongest baseline rather than a weak baseline that would inflate gains.

Averaged across the small models, one vanilla reflection adds 4.5 percentage points. After training, most one-attempt scores exceed their vanilla two-attempt scores, and a reflected retry adds another 4.7 percentage points on average. The headline comparison is Qwen2 7B trained to 77.3% on the second attempt versus Qwen2 72B vanilla at 76.6%.

How to read the two columns

The APIGen and Countdown tables report four numbers for each trained model: first and second attempt before training, and first and second attempt after training. The vanilla first column measures base competence. The vanilla second column measures how much a prompted reflection helps without any parameter update. The trained first column is especially important because it does not involve an explicit reflection path; if it rises, the reflection training has altered behavior that shows up before any retry. The trained second column then measures the combined value of the learned reflection policy and the retry itself.

That distinction prevents a common misreading. The method does not merely test whether asking the model to reflect is useful. It tests whether reinforcement learning can make the reflection itself more useful. The gap between vanilla second and trained second isolates the effect of training under the same two-attempt procedure. The gap between vanilla first and trained first is less direct but reveals that the intervention changes more than an optional correction subroutine.

Another useful comparison is with the 70B-class baselines. They are not trained in the paper, so they are not evidence that small models beat equally trained large models. They are evidence that a well-targeted training objective can close or reverse a large parameter gap on two verifiable tasks. The phrase small beats large should therefore be read with its precise conditions: small trained against large vanilla, on tasks where a binary verifier exists.

APIGen	Vanilla 1st	Vanilla 2nd	Trained 1st	Trained 2nd
Qwen2 1.5B Instruct	32.6%	34.8%	48.6%	52.9%
Qwen2 7B Instruct	66.4%	69.4%	72.2%	77.3%
Llama3.1 8B Instruct	64.9%	70.9%	68.7%	74.9%
Phi3.5-mini Instruct	47.5%	50.2%	52.9%	56.0%
Qwen2 72B Instruct	73.7%	76.6%	—	—
Llama3.1 70B Instruct	66.8%	76.9%	—	—
Palmyra-X4 73B	79.9%	83.5%	—	—

Table 1: Accuracy on the 12,000-example APIGen test set. Training rewards only reflection tokens, yet first-attempt accuracy also rises.

Figure 3: The original layout of Table 1, useful for checking the four-way comparison of vanilla and trained models with and without reflection.

Countdown equation construction

Countdown comes from the TinyZero project. The training side has 450,000 problem lists; the test set has 15,000. Given 3 to 4 numbers and a target, the model must construct an equation with basic arithmetic, using every provided number exactly once. The validator needs no gold equation: it checks validity, number usage, and whether the expression evaluates to the target.

If $\mathbf{x}$ is the vector of provided numbers, $\mathbf{o}$ is the generated operator sequence, and $T$ is the target, the Countdown verifier can be stated as:

$$v(\mathbf{x},\mathbf{o},T)=\mathbb{1}\left\{\operatorname{eval}(e(\mathbf{x},\mathbf{o}))=T \wedge \operatorname{operands}(e(\mathbf{x},\mathbf{o}))=\operatorname{multiset}(\mathbf{x})\right\}$$

The second conjunct means the multiset of operands actually used by the generated expression must match the multiset of provided numbers. This prevents the model from inventing extra numbers, dropping inconvenient ones, or reusing a convenient number.

Countdown gains are larger. Vanilla reflection improves accuracy by 5.3 percentage points on average, while reflection after training adds 8.6 percentage points on the second retry; the paper reports an average overall gain of 16.0%. Llama3.1 8B rises from 2.2% to 17.8% after a trained retry. Qwen2.5 1.5B rises from 6.0% to 45.0%. Qwen2.5 7B reaches 50.3% after training, exceeding Qwen2.5 72B vanilla's 49.9%.

Why the baseline matters

The Countdown results cannot be interpreted by absolute percentages alone. A model starting at 2.2% has much more room to improve than a model starting at 31.7%. This is why the largest relative jumps belong to Llama models, while the Qwen models start from much stronger equation-writing behavior. The common pattern is that the objective finds a useful gradient near the failure boundary, but the distance to that boundary determines how visible the improvement becomes.

The Countdown verifier also makes the task less ambiguous than many open-ended reasoning benchmarks. It does not ask whether an explanation sounds plausible. It asks whether the equation is valid, whether the number multiset matches the input, and whether evaluation equals the target. Because those checks are cheap and deterministic, the reward signal is not contaminated by human preference. That makes Countdown a good test bed for the central claim, although it also means the result may not automatically transfer to tasks whose validators are noisy or adversarial.

The APIGen task, by contrast, uses gold answers. This is stricter than checking whether a call returns HTTP 200: a wrong API can return 200, and a correct API can still fail because of authentication or environmental state. The paper's validator avoids that operational nondeterminism, but it shifts the task toward exact agreement with a canonical call. Readers should treat it as a schema-and-value reasoning benchmark rather than a full live-agent benchmark.

What the first-attempt gain suggests

Across both experiments, the trained one-attempt column improves even though the update mask excludes final-answer tokens. The authors hypothesize that learning to reflect improves reasoning more generally. A cautious interpretation is that the training changes the model's internal representation of failure-prone contexts: when it sees a difficult tool schema or a Countdown target, it may allocate attention differently even without producing explicit reflective text.

This distinction matters for deployment. If all gains required an extra generation, the method would trade accuracy for latency and cost. Because first-attempt accuracy also improves, the trained policy can be used without paying the retry overhead on easy cases. On hard cases, the reflection and retry remain available. The result is an adaptive procedure: most inputs terminate immediately, and only verified failures invoke the more expensive path.

That adaptivity is also what separates the design from always-reflect prompting. An unconditional reflection prompt forces every query through an extra generation, including easy ones, and may introduce unnecessary revisions. Here, the validator gates the reflection. The paper's mechanism preserves the success path and spends computation only where the binary signal says it is needed.

Failure-only sampling in practice

The rejection-sampling stage deserves as much attention as the RL update. A query that succeeds on every sample provides no information about how to cross from failure to success. Removing it saves both GPU time and optimization noise. Conversely, a query with a low but nonzero success rate may be less useful than a query that always fails, because easy successes may dominate the group and dilute the failure signal.

The 64-sample cap is a practical compromise. It ensures that even strong 7B-8B models yield enough hard queries while avoiding unbounded sampling on a query that a strong model occasionally solves. In a production system, the cap would likely become a budget: sample until a fixed number of failures is collected, until a maximum cost is reached, or until the validator has seen enough diverse failures for the intended task family.

Prefix caching makes this tractable because the query and tool schema are common across samples. The variable part is the sampled completion. This is especially relevant for APIGen, where candidate tool descriptions can be long. Reusing the prompt prefix avoids recomputing the same context for dozens of candidates.

Countdown	Vanilla 1st	Vanilla 2nd	Trained 1st	Trained 2nd
Qwen2.5 1.5B Instruct	6.0%	10.2%	34.9%	45.0%
Qwen2.5 3B Instruct	18.8%	29.0%	33.9%	47.3%
Qwen2.5 7B Instruct	31.7%	38.0%	41.6%	50.3%
Llama3.1 8B Instruct	2.2%	4.6%	8.8%	17.8%
Llama3.2 3B Instruct	2.1%	3.0%	8.8%	13.8%
Palmyra 1.7B	26.8%	31.8%	33.3%	38.6%
Qwen2.5 32B Instruct	38.6%	45.1%	—	—
Qwen2.5 72B Instruct	45.2%	49.9%	—	—
Llama3.1 70B Instruct	17.3%	25.5%	—	—
Palmyra-X4 73B	46.8%	51.6%	—	—

Table 2: Accuracy on the 15,000-example Countdown test set. Models with lower baselines have more room to learn from failure.

Figure 2: Reflections from the same model before and after training. The trained reflection is shorter, clearer, and less repetitive.

Reflections become shorter

The qualitative example in Figure 2 is one of the paper's most interesting observations. Vanilla reflections tend to be verbose, confusing, and redundant. Trained reflections are shorter and more direct. This runs against the common intuition that longer chain-of-thought text improves reasoning.

The RL objective explains the trend without implying that brevity is universally better. Reward is not given for thinking longer; it is given for a reflection that helps the next attempt pass the validator. Verbose text that increases downstream noise is not reinforced. Compact error localization that reliably changes the retry context is selected repeatedly. The model learns the reward-relevant prefix, not a human rhetorical preference.

Error analysis

The appendix groups first-attempt errors. In APIGen, errors are tool-choice, parameter, and format errors. Qwen2 1.5B reduces tool-choice errors from 33.9% to 21.3% and format errors from 8.2% to 4.8%, while parameter errors remain at 25.3%. Qwen2 7B improves parameter errors from 27.2% to 22.7%. The small models mainly improve tool choice; 7B-8B models mainly improve parameter selection.

APIGen error analysis	Tool-choice error	Parameter error	Format error
Qwen2 1.5B: vanilla → trained	33.9% → 21.3%	25.3% → 25.3%	8.2% → 4.8%
Qwen2 7B: vanilla → trained	5.2% → 3.5%	27.2% → 22.7%	1.2% → 1.6%
Llama3.1 8B: vanilla → trained	4.6% → 3.7%	28.5% → 25.6%	2.0% → 2.0%
Phi3.5-mini: vanilla → trained	18.8% → 17.6%	24.8% → 23.9%	8.9% → 5.7%

Table 3: First-attempt error categories on APIGen. Training does not eliminate parameter errors, but it moves the dominant bottleneck.

Low catastrophic forgetting

Because the target is reflection rather than a task answer, the authors evaluate MMLU-Pro, GSM8K, HellaSwag, and MATH. Table 4 shows that most changes are below one percentage point, and some metrics improve. Qwen2.5 1.5B moves from 30.9% to 31.5% on MMLU-Pro and from 25.5% to 26.3% on MATH. The small KL coefficient and reward mask do not collapse the model into a single-task distribution.

General benchmark	MMLU-Pro	GSM8K	HellaSwag	MATH
Qwen2 1.5B: vanilla → trained	22.1% → 21.9%	59.6% → 59.5%	66.0% → 66.1%	21.2% → 20.9%
Qwen2 7B: vanilla → trained	43.4% → 42.8%	77.6% → 77.0%	80.7% → 80.5%	34.0% → 33.7%
Qwen2.5 1.5B: vanilla → trained	30.9% → 31.5%	59.4% → 58.0%	68.3% → 68.3%	25.5% → 26.3%
Qwen2.5 7B: vanilla → trained	44.8% → 44.2%	84.2% → 82.1%	80.5% → 80.5%	42.3% → 40.3%
Llama3.1 8B: vanilla → trained	41.0% → 40.9%	78.0% → 77.7%	79.3% → 79.3%	31.1% → 31.5%

Table 4: General-capability checks. Almost all degradation is below one percentage point.

Why the Diagram Is Method-Specific

The flowchart captures the two branches that make the paper different from ordinary retry RL. Reflection is entered only after a failed first attempt, and reflection receives positive advantage only when the reflected retry succeeds. First-attempt success and retry failure leave the reflection distribution unchanged; the retry answer itself is not reinforced.

flowchart TB
    Q[Query q] --> Y0[First answer y0]
    Y0 --> V0{Validator v q y0}
    V0 -->|Success 1| Stop[No policy update]
    V0 -->|Failure 0| Fail[Add failure context]
    Fail --> R[Generate reflection r]
    R --> Y1[Retry answer y1]
    Y1 --> V1{Validator v q y1}
    V1 -->|Success 1| Mask[Set advantage mask to R]
    Mask --> GRPO[GRPO updates reflection tokens only]
    V1 -->|Failure 0| NoReward[No positive reward for r]
    GRPO --> Next[Improved reflection policy]

Limitations

First, the authors note that not every task admits a clean binary validator. Open-ended writing, planning, safety review, or subjective judgment cannot always be reduced to zero or one. Gold labels can help, and a larger model can act as judge, but that reintroduces reliability questions: judge errors will be exploited by RL and can produce reward hacking.

Second, the authors also report that the method does not work for every model or task. The model must already have some basic ability to perform the task, reflect, and learn. Llama3.2 3B Instruct failed to learn self-correction on function calling. Smaller 0.5B and 1B models struggled with both accuracy and reflection, so their results were not reported.

Third, from an evaluation-design perspective, the APIGen validator relies on exact gold calls rather than a live API returning a useful response. That makes the benchmark controllable, but it may undervalue semantically equivalent calls with different formatting. Real tool environments add state, authentication, and side effects.

Fourth, the procedure uses only two attempts. The paper does not study whether another round of reflection continues to help, whether reflection chains remain stable, or whether reflection-only reward scales to larger models and harder tasks. GRPO's compute cost also left models above 8B untrained in this study.

Engineering Interpretation

The method is most attractive when an agent already has an execution result but not a gold demonstration. A database agent knows whether a query errored. A code agent knows whether tests failed. A calculator or tool-calling agent knows whether the returned object matched the schema. In those settings, the validator is often not an extra model but an existing runtime check. The paper turns that check into a learning boundary without requiring a human to write the correct trajectory.

The implementation posture is also pragmatic. Instead of inventing a new trainer, the authors patch TRL's input preparation to perform a second completion while retaining the original mask. That keeps the GRPO machinery mostly unchanged and moves the novel part to where it belongs: the boundary between first generation, reflection, and retry. For teams maintaining agent stacks, this is easier to reproduce than a bespoke multi-turn RL environment.

At the same time, the reward mask creates an audit requirement. One must be certain which tokens belong to the reflection. If an implementation accidentally includes retry-answer tokens, it becomes ordinary task RL. If it accidentally excludes part of the reflection, the policy learns only part of the diagnostic behavior. The paper's conceptual clarity does not remove the need for precise token accounting.

The method also suggests a different way to manage failure logs. Rather than storing only the failed response, store the verifier outcome, the first context, the reflection prompt, the reflection text, and the retry context. This makes it possible to replay the exact transition and compare different reflection policies later. Without that provenance, a promising reflection may be indistinguishable from a lucky retry.

The observed shorter reflections also have a useful operational implication. A reflection that helps a retry does not need to be a polished essay. It can be a compact list of omitted constraints, wrong tool assumptions, or numerical mistakes. This aligns with agentic systems, where the next turn's value often comes from precise constraints rather than fluent prose.

Conclusion

The contribution is not a new RL algorithm but a different partition of credit. Conventional task RL attributes success to the whole output. Here, success is attributed to the interval between a verified failure and a successful retry. The binary validator only says that a boundary was crossed; the reflection text says how. With that interface, average function-calling accuracy improves by 9.0% and Countdown by 16.0%. Trained 7B models surpass untrained 72B models from the same family on both tasks.

For practitioners, the practical recipe is concrete: when success can be checked automatically, sample from a real failure stream, give the model one reflection opportunity, and use group-relative advantage to update the reflection tokens. No large teacher annotation is needed, and each task's final answer does not have to become a training target. Future work must answer three questions: does reflection training transfer across tasks, can complex task graphs be decomposed into multiple verifiable boundaries, and is a reflection-only mask still sufficient when the causal contribution of reflection and final answer is harder to separate?

The broader takeaway for LLM systems is that reward placement can matter more than adding another prompting heuristic. Here, the authors do not make the model reason louder. They define the exact moment where it should reason differently and make that moment the only place where reinforcement accumulates. The binary signal is weak, but it is placed with unusual precision.

When the only supervision is a zero or a one, a model can still improve: if failure is followed by success, credit the reflection that carried it across the boundary.