5 Inference-time scaling via self-refinement
The chapter introduces inference-time scaling as a way to improve an LLM’s reasoning without retraining it, building on self-consistency and moving to self-refinement. Instead of sampling many answers and voting, self-refinement has the model critique a draft answer and revise it in one or more iterations. The chapter’s setup also emphasizes that, before refining answers, it is useful to have scoring methods that can compare outputs and decide whether a revision is actually better.
To support that process, the chapter develops two kinds of scorers. First, it shows a simple rule-based heuristic that rewards clearly extractable final answers, boxed formatting, and brevity, mainly as a practical baseline and tie-breaker. Then it explains token probabilities and log-probabilities, showing how a model can score an existing answer by measuring how strongly it prefers each next token. Because raw probabilities quickly become tiny when multiplied across a sequence, the chapter switches to log-space, where log-probabilities are summed for stability and easier comparison, and then uses the average log-probability of answer tokens as a length-normalized confidence score.
The final part puts everything together in an automated self-refinement loop that generates an initial answer, produces a critique, revises the answer, and optionally accepts the revision only if its score improves. The chapter demonstrates this workflow on a math problem and then evaluates it on MATH-500, finding that self-refinement can help but only modestly in this setting. In the reported experiments, heuristic scoring and average log-probability scoring sometimes helped, but they could also favor polished yet wrong answers, and plain self-refinement without scoring often performed best for the base model. The chapter closes by noting that stronger results may come from using different models or external judge-style scorers, and it sets up the transition from inference-time methods to training-based reasoning improvements.
A mental model of the topics covered in this book. This chapter continues stage 3 and focuses on inference-time techniques for improving reasoning without additional training. This chapter introduces self-refinement, where the model iteratively critiques and improves its own answers.
Three inference-time methods to improve reasoning covered in this book. The first two methods were covered in the previous chapter. This chapter covers the third method, self-refinement, where the model iteratively improves its own answers.
We build a simple rule-based score, compute token probabilities and log-probabilities, and then use these scores as part of a self-refinement method where the model iteratively improves its own answers.
This section implements a rule-based scorer to rank different answers generated by the pre-trained LLM.
Two generated responses reach the same correct answer but differ in their explanations. A scorer evaluates the responses and assigns a score to each response.
A simple rule-based length penalty used by our scorer. Longer explanations receive a smaller score contribution.
In this section, we move from a simple rule-based scorer to token-level probabilities. These probabilities form the basis for the logprob scoring method that we will use later in the self-refinement approach.
How an LLM selects the next token. The model converts the input text into token IDs, computes a score for every vocabulary token, and chooses the token with the highest score. The plot on the right shows the logit values for a subset of the vocabulary, with the token for Berlin having the highest score.
How we look up the logit scores for specific tokens. After passing the input text through the model, we obtain a logit value for every vocabulary token. We then convert the candidate tokens we want to score into token IDs and read off their corresponding logit values from the distribution.
Next-token scoring. The input text is converted into token IDs and fed to the LLM, which outputs logits for the next token. After applying a softmax, these logits become probabilities, where tokens like "Berlin" receive high probability and unlikely alternatives such as "Bridge" receive values near zero.
Computing token probability scores for a given sequence. For each position, we feed the preceding text into the model and read off the softmax probability of the next token. Multiplying these conditional probabilities yields the joint probability of the full sequence.
Extracting next-token probabilities. After converting the input text into token IDs, the model computes logits that are transformed into probabilities with a softmax function. Using index tensors for positions and true next tokens, we then get the model's computed probability for each next token.
Overview of how we move from token probabilities to token log-probabilities, which provide a numerically more stable basis for log-probability scoring used later in self-refinement.
Comparison of logits, softmax probabilities, and log-probabilities for a simple example. The log-probabilities preserve the ordering of the probabilities.
How logits are converted to probabilities and log-probabilities for next-token scoring. The correct next token ("Berlin") receives a high logit, which becomes a high probability and a less negative log-probability, while unlikely candidates like "Bridge" map to very small probabilities and large negative log-probabilities.
How token-level log-probabilities accumulate to form sequence log-probabilities. Each row shows the log-probability of the next token given the preceding text. Summing these values results in the joint log-probability of the full sequence
This section implements a logprob scorer, based on token log-probabilities, which we will use in the self-refinement method later in this chapter.
The modified logprob scoring procedure. The prompt tokens are excluded from the calculation, and only the log-probabilities of the answer tokens are collected. These values are then averaged to obtain a length-normalized score, which allows us to compare answers of different lengths fairly.
The final step in our workflow, where the logprob scorer developed earlier is used inside the self-refinement method.
The self-refinement process. The LLM first produces an initial answer to the prompt, then receives a critique prompt that asks it to analyze its own response and produce a short critique with a plan to refine the answer. In the final step, the model is given a refine prompt that contains the original question, its draft answer, and the critique, and it generates a revised answer that incorporates the suggested improvements.
The self-refinement loop with optional scoring. The model first produces an initial answer, then critiques it, and generates a revised answer based on the critique. Both answers can be evaluated with the scoring functions from this chapter (for example, logprob scoring), and the revised answer is only accepted if its score improves on the previous one.
Overview of the book's progression from basic LLM usage to inference-time reasoning methods and finally to training-based techniques. This chapter concludes the inference-scaling methods without additional training, and the next chapters introduce approaches that update model weights to further improve reasoning performance.
Summary
- Self-refinement extends the inference-time scaling ideas from the previous chapter by iteratively critiquing and improving a single answer instead of relying on multiple independent samples as in self-consistency.
- A simple rule-based scoring function ranks model outputs by rewarding extractable final answers and shorter, more economical completions.
- Next-token scoring quantifies model confidence by converting logits into normalized probabilities and combining these into a sequence-level likelihood.
- Log-probabilities replace raw probabilities to avoid numerical underflow and to turn products over many tokens into stable sums and averages.
- These scoring functions are more generally useful beyond self-refinement, for example, for breaking ties in self-consistency or implementing Best-of-N selection strategies.
- The self-refinement procedure consists of three stages: generating an initial draft, producing a short critique and fix plan, and generating a revised answer.
- A reusable refinement function automates the self-refinement workflow with multiple iterations (refinement rounds), and it can use a score-based acceptance to keep only revisions that do not degrade a computed answer score using one of the scoring functions we developed earlier.
FAQ
What is inference-time scaling, and how does self-refinement differ from self-consistency?
Inference-time scaling improves model accuracy without further training by using extra compute at generation time. Self-consistency generates multiple answers and chooses the final one by majority vote, while self-refinement generates one answer and iteratively critiques and improves it.
Why does this chapter introduce scoring functions before building self-refinement?
Scoring functions are needed to compare and rank candidate answers during refinement. The chapter first builds simple rule-based and probability-based scorers so the self-refinement loop can decide whether a revised answer is better than the previous one.
What does the simple heuristic scorer reward?
The heuristic scorer rewards answers that are easy to extract and concise. It gives a bonus for a boxed final answer, a smaller bonus for an extractable number, and a brevity bonus that mildly favors shorter responses.
Why is brevity included in the heuristic score?
Brevity is included because, when two answers are similarly good, shorter ones are cheaper to generate and often preferred. The length term is designed to break ties or slightly favor compact answers without forcing overly short outputs.
What are token probabilities used for in this chapter?
Token probabilities are used to estimate how confident the model is in a proposed answer. By looking at the probability assigned to each next token, the chapter builds a sequence-level score that reflects the model’s own preferences.
Why do raw joint probabilities become difficult to use?
Raw joint probabilities are products of many numbers between 0 and 1, so they quickly become extremely small. This can lead to numerical underflow and makes the scores hard to compare, especially for longer sequences.
How do log-probabilities solve the numerical stability problem?
Log-probabilities replace multiplication with addition, since the log of a product is the sum of the logs. This keeps values in a more manageable numeric range and preserves the same ranking of answers.
What is the difference between token probabilities and average log-probability scoring?
Token probabilities score the whole sequence by multiplying per-token probabilities, while average log-probability scoring averages the log-probabilities of only the answer tokens. The average form makes scores more comparable across answers of different lengths.
How does the self-refinement loop work step by step?
The model first generates an initial answer, then produces a critique of that answer, and finally generates a revised answer using the original question, draft answer, and critique. This process can repeat for multiple iterations, optionally using a scorer to keep only improvements.
Does self-refinement always improve results?
No. The chapter shows that self-refinement can help, but it can also make answers worse in later iterations or accept confident but incorrect revisions. Its usefulness depends on the model, the task, and the scoring method, so experimentation is important.
Build a Reasoning Model (From Scratch) ebook for free
