Importance sampling

Learn about one policy from data another policy generated, by reweighting each sample. Almost every modern RL trick for LLMs is built on this.

Key idea
On-policy vs off-policy data, and the importance weight new ÷ old
Real world
A phone poll reaches too many retirees. Pollsters don't throw the answers away; they count each retiree's answer for less and everyone else's for more, so the result matches the real population.
In an LLM
Answers generated by last step's weights, or by a separate fast inference engine, are used to train the current weights. Each token is reweighted by how much more or less likely the current model makes it.

On-policy and off-policy

Policy gradient assumes the samples came from the policy you are updating. That data is on-policy. As soon as the policy changes, even by one gradient step, older samples become off-policy: they show what the old policy did, not the new one.

Throwing them away is wasteful, because generating answers is the expensive part of LLM training. Importance sampling lets you keep them. Weight each sample by

w=probability under the policy you’re trainingprobability under the policy that generated itw = \frac{\text{probability under the policy you're training}}{\text{probability under the policy that generated it}}

Samples the new policy likes more than the old one count for more. Samples it likes less count for less.

example · reweighting a poll

A poll asks whether people support a new park. 60% of the people reached are over 60, but they are only 30% of the town.

GroupShare of pollShare of townSupportWeight
Over 6060%30%70%0.3 ÷ 0.6 = 0.5
Under 6040%70%40%0.7 ÷ 0.4 = 1.75

Plain average of the poll. Wrong: it describes the people reached.

0.6×0.7+0.4×0.4=58%0.6 \times 0.7 + 0.4 \times 0.4 = 58\%

Weighted. That is exactly the town's support.

0.6×0.5×0.7+0.4×1.75×0.4=49%0.6 \times 0.5 \times 0.7 + 0.4 \times 1.75 \times 0.4 = 49\%
example · reweighting tokens

The model wrote an answer that got advantage +1. Since then, the weights have moved:

TokenProbability when sampledProbability nowWeight
So0.600.601.0
Wait0.100.151.5
Alternatively0.020.126.0

Each token's gradient is multiplied by its weight. Alternatively is now six times likelier than when it was sampled, so this old sample says a lot about the current policy, perhaps too much. That is where capping comes in.

Try it

Probability of each answer, and its weight new ÷ old
short guess r = 0.1×0.32
one-line answer r = 0.4×0.87
worked solution r = 0.8×2.38
worked + checked r = 1.0×6.46
Red weights are above the cap, so the capped estimate counts them as ×3.0.
true value0.64
no correction0.28
importance sampling0.58
capped at 30.43
true valueplain averageimportance-weightedweights capped at C
Repeat the whole thing 150 times: where each method’s final estimate lands
True value 0.64. The plain average centres on 0.30: it measures the old policy. Weighting centres on the truth but spreads widely, because “worked + checked” counts ×6.5. Capping at 3 narrows the spread but centres on 0.46, 0.17 too low: the capped answers are the ones the new policy likes most.
Answers are sampled from the old policy (grey). The estimate we want is the average reward under the new policy (purple).

Why weights explode for long answers

The probability of a whole answer is the product of its token probabilities, so the weight of a whole answer is the product of the token weights. Tiny differences on each token multiply up:

  • 1,000 tokens each 1% more likely: 1.011000≈21,0001.01^{1000} \approx 21{,}000.
  • 1,000 tokens each 1% less likely: 0.991000≈0.000040.99^{1000} \approx 0.00004.

That's why LLM methods almost never use the raw sequence weight. They use per-token weights (PPO), clip or cap them (CISPO, TIS), or average the log-weights over the tokens (GSPO).

Where it shows up

  • PPO: the ratio ρ is an importance weight, so a batch can be reused for several gradient steps.
  • Asynchronous training: answers written by an older copy of the model, a few updates behind.
  • Sampler–trainer mismatch: the fast inference engine and the trainer disagree slightly about probabilities, even with identical weights.
going deeper

The identity

For a target policy π\pi and a behaviour policy μ\mu that generated the data:

Ea∼π[f(a)]=∑aπ(a)f(a)=∑aμ(a)π(a)μ(a)f(a)=Ea∼μ[π(a)μ(a)f(a)]\mathbb{E}_{a \sim \pi}\big[f(a)\big] = \sum_a \pi(a) f(a) = \sum_a \mu(a) \frac{\pi(a)}{\mu(a)} f(a) = \mathbb{E}_{a \sim \mu}\Big[\frac{\pi(a)}{\mu(a)} f(a)\Big]

The estimate is unbiased but its variance grows with Eμ[(π/μ)2]\mathbb{E}_\mu[(\pi/\mu)^2], which explodes when π\pi puts weight where μ\mu rarely samples.

Truncated importance sampling uses min⁡(π/μ, C)\min(\pi/\mu,\ C): biased, but with bounded variance. For sequences, the policy gradient is usually approximated per token,

∑tπθ(ot∣st)μ(ot∣st)A^t∇log⁡πθ(ot∣st)\sum_t \frac{\pi_\theta(o_t \mid s_t)}{\mu(o_t \mid s_t)} \hat A_t \nabla \log \pi_\theta(o_t \mid s_t)

which ignores how earlier tokens changed the states later tokens see. It's a good approximation while π\pi and μ\mu stay close, which is what clipping enforces.

Related

  • ppo - Reuse each batch of experience for several updates without letting one lucky sample swing the policy too far.
  • mismatch - The fast engine that generates answers and the code that trains on them compute slightly different probabilities for the same tokens. Left alone, this quietly destabilises RL.