Self-taught reasoning via post-rationalization fine-tuning approximates policy gradient optimization

In what ended up being the precursor to RLVR, I explore how the STaR method bootstraps system 2 reasoning in LLMs by fine-tuning on self-generated rationales and post-hoc rationalizations.

The idea

Standard pretrained and fine-tuned language models operate like fast system 1 thinkers, jumping straight to answers via direct word associations without stepping through deliberate reasoning. We can slow them down into system 2 thinkers by requiring them to generate a rationale before answering, but prompting alone has limits. By iteratively generating rationales, filtering those that lead to correct answers, and post-rationalizing failed questions with answer hints to create extra training samples, we construct a fine-tuning dataset that teaches the model to reason. Fine-tuning the base model on this growing set of successful rationale-answer pairs iteratively improves both reasoning quality and overall accuracy.

The mechanism

We define the expected reward $J(\theta)$ of the model policy as the expected correctness of predicted rationale-answer pairs $(\hat{r}, \hat{y})$ given input $x$ :

J(\theta) = \mathbb{E}_{(\hat{r}, \hat{y}) \sim P_\theta(\cdot \mid x)} \left[ \mathbb{I}(\hat{y} = y) \right]

This objective measures the expected task success under the model's generated reasoning chains and final answers.

Expanding the expectation into an explicit summation over candidate rationales and outputs gives:

J(\theta) = \sum_{\hat{r}, \hat{y}} P_\theta(\hat{r}, \hat{y} \mid x) \mathbb{I}(\hat{y} = y)

This form explicitly weights each rationale-answer trajectory by its joint probability under model parameters $\theta$ .

Applying the log-derivative trick $\nabla_\theta P_\theta = P_\theta \nabla_\theta \log P_\theta$ to differentiate through the summation yields:

\nabla_\theta J(\theta) = \sum_{\hat{r}, \hat{y}} P_\theta(\hat{r}, \hat{y} \mid x) \nabla_\theta \log P_\theta(\hat{r}, \hat{y} \mid x) \mathbb{I}(\hat{y} = y)

This converts the gradient of the trajectory probabilities into a likelihood ratio gradient multiplied by the scalar reward signal.

Rewriting the expression back into an expectation yields the policy gradient loss:

\nabla_\theta J(\theta) = \mathbb{E}_{(\hat{r}, \hat{y}) \sim P_\theta(\cdot \mid x)} \left[ \nabla_\theta \log P_\theta(\hat{r}, \hat{y} \mid x) \mathbb{I}(\hat{y} = y) \right]

This proves that standard supervised fine-tuning on filtered, correct reasoning trajectories directly approximates policy gradient reinforcement learning.

Worth knowing

Fine-tuning in each STaR iteration must start from the original base model rather than the previously fine-tuned checkpoint, preventing sample autocorrelation that violates the IID training assumption.
In practice, the optimal number of STaR iterations seems to be incredibly variable. Sometimes 3 iterations is enough, but often it takes >10. In the paper, the authors used ~36 iterations for GSM8K. Even the number of training steps per iteration is a hyperparameter and needs to be carefully tuned.
Standard rationale generation cannot improve performance on problems the model initially misses; post-rationalization provides the answer as a hint to generate backward rationale training targets for those failed cases.
STaR plans entirely in token space rather than an abstract representation space, leaving long reasoning sequences susceptible to compounding/reinforcing errors.

Use it when / don't use it when

Use it when

You have dataset inputs and ground-truth answers, but lack ground-truth intermediate reasoning steps.
You want to boost system 2 reasoning accuracy in language models without human-annotated chain-of-thought data.
You are fine-tuning a base language model on complex reasoning benchmarks like mathematics or commonsense QA.

Don't use it when

Ground-truth final answers are unavailable or cannot be automatically verified with an indicator function.
You can use a more modern method like RLVR instead.