Projection Sampling

Why People Wanted It

An expert trace may contain the right reasoning, but the base model may be unlikely to produce that exact trajectory on its own. Rejection sampling would be an easy way to get model-generated correct traces if the base model could reliably generate them; in the regime this method targets, that is not feasible. So the real question is: can you start from an expert trace, preserve its correctness-relevant information, and move the samples toward what the base model itself is more likely to produce?

The Core Move

Start an information-preserving Metropolis–Hastings chain from an expert trace, using proposals that remain within the set of information-equivalent trajectories; the target is the base model restricted to that set. In the paper’s notation, the target is $p_C(\mathbf{x}) \propto p(\mathbf{x})\mathbf{1}[\mathbf{x}\in C]$.

Here $C$ is the set of trajectories that retain the required information, and the indicator excludes trajectories outside it. The method calls the resulting procedure projection sampling. The information-projection target is the closest-to-base distribution under that constraint; the focus here is how to sample toward it.

Quiz

What makes a Metropolis–Hastings proposal information-preserving in projection sampling?

It replaces the expert trace with an unconstrained base-model sample before every step.

It accepts every candidate generated by the base model.

It proposes only trajectories that remain in the information-preserving set.

It changes the target to the expert policy itself.

$ check

Why Information-Preserving Proposals Can Target the Base Model

Now the version for readers who want more of the machinery: restricting proposals to correct trajectories does not mean the chain must target a different distribution. Suppose the current trajectory is in $C$, and the proposal kernel $\kappa_C$ only proposes candidates in $C$. The Metropolis–Hastings acceptance ratio uses the target density and the forward and reverse proposal probabilities:

$$\alpha(\mathbf{x},\mathbf{x}') = \min\left(1,\frac{p(\mathbf{x}')\kappa_C(\mathbf{x}\mid\mathbf{x}')}{p(\mathbf{x})\kappa_C(\mathbf{x}'\mid\mathbf{x})}\right).$$

Because both candidates are in $C$, the indicator in the restricted target is one for each, so the restriction contributes no additional factor to this ratio. A proposal outside $C$ would have zero target mass and therefore zero acceptance probability. This is why the paper’s proposition relates MH with irreducible proposals targeting $p_C$ to MH with information-preserving proposals targeting $p$.

The chain still needs suitable reachability and return behavior: the stated conditions are irreducibility and aperiodicity. A proposal that preserves information but cannot reach relevant trajectories does not, by that property alone, guarantee convergence to the desired target.

Why Repeated Steps Reduce the Policy Gap

Here is the part that gets technical: the claimed progressive reduction follows from the Markov transition’s stationary distribution and the data processing inequality, rather than from each individual proposal being more likely under the base model. Let $\Psi$ denote one MH step, and let $\pi_k$ be the candidate distribution after $k$ steps. Since $p_C$ is stationary, applying the same Markov kernel to $\pi_k$ and $p_C$ cannot increase their KL divergence:

$$\mathrm{KL}(\pi_{k+1}\|p_C) \leq \mathrm{KL}(\pi_k\|p_C).$$

In words, each chain step does not increase the divergence to the stationary target. On trajectories in $C$, the restricted target differs from $p$ by a normalization constant, so the corresponding KL divergences to $p_C$ and to $p$ differ by a fixed additive constant. Thus the policy gap to the base model is non-increasing across steps, while the candidates remain information-preserving. This is a monotonicity claim, not a guarantee that every individual trace improves or that a finite run has reached the target.

Why Blockwise Sampling Fits Autoregressive Models

Now for the implementation-level deep dive: direct full-sequence MH would repeatedly regenerate long token sequences, which is described as infeasible and can also mix slowly in high-dimensional spaces. The method instead defines restrictions of the target over prefixes of increasing length, in units of block size $B$. It samples within one prefix-length distribution and uses that result to initialize sampling at the next length. The blockwise schedule is a way to make autoregressive refinement tractable; it is not a claim that the full sequence is resampled in one operation.

Why the Proposal Uses an Expert-Guided Partial Trace

Here is the other technical ingredient: the expert trajectory enters through the proposal kernel. Given a candidate, the procedure selects an index uniformly from the sequence positions, truncates the candidate there, and prompts the base model with the query, expert trajectory, and partial response. The prompt asks the model to continue from that partial response in a way consistent with the expert solution and to reach a complete answer. Conditioning the base model on this prompt defines the proposal distribution.

The expert trace guides the proposed continuation, while the partial response gives the model a point from which to continue. In practice, the text says this prompt can sufficiently generate correct trajectories from expert demonstrations, satisfying the information-preservation constraint. The MH acceptance step remains important: the proposal is not itself the target distribution.

Minimal Causal Code

This schematic loop makes the accept-or-stay update visible. log_target evaluates the base-model log probability for candidates in the information-preserving set, and log_proposal(a, b) means the log probability of proposing a from b. The prompt-based proposal is represented by propose.

import math
import random

current = expert_trace
for step in range(num_steps):
    candidate = propose(query, expert_trace, current)
    if not information_preserving(candidate, expert_trace):
        continue

    log_ratio = (
        log_target(candidate) + log_proposal(current, candidate)
        - log_target(current) - log_proposal(candidate, current)
    )
    if math.log(random.random()) < min(0.0, log_ratio):
        current = candidate

A rejected candidate leaves current unchanged, which is the self-transition in MH. The code is illustrative: the paper’s sequence-length restrictions and blockwise schedule are additional parts of its autoregressive algorithm.

What You Gain, and What You Lose

You begin with expert trajectories and use a base-model proposal to produce information-preserving candidates that move toward the base model’s restricted distribution. The paper states that the policy gap is non-increasing across MCMC steps, not that every sample becomes better or that finite computation guarantees convergence.

The tradeoff is computational and algorithmic. Full-sequence MH entails repeated LLM inference and may mix slowly; blockwise refinement addresses the sequence-length burden but introduces a block-size choice. Larger blocks lower the stated sampling cost, while smaller blocks provide higher resolution and less approximation error. The procedure is a one-time cost before finetuning, rather than an inference-time sampling cost. The proposal also depends on the model following the expert-guided continuation prompt sufficiently well to preserve the required information.

Quiz

Why does projection sampling refine autoregressive sequences in blocks?

It guarantees that every proposed full sequence is correct.

It samples at increasing prefix lengths and uses each result to initialize the next stage.

It makes the proposal independent of the expert trajectory.

It eliminates the Metropolis–Hastings accept-or-reject step.

$ check