Katkat
Back

Explain Importance Sampling in LLM Training

GoogleMachine Learning Engineer · Technical Screen
HardMachine Learning EngineerTechnical ScreenMachine Learning

Last updated

Tests understanding of importance sampling in modern LLM fine-tuning, focusing on sequence-level likelihood ratios, exponential variance growth, support conditions, and practical mitigations like clipping and diagnostics.

Scenario

In modern reinforcement learning from human feedback (RLHF) and off-policy policy optimization (such as PPO or off-policy rejection sampling) for Large Language Models (LLMs), training pipelines frequently evaluate or optimize expectations under a target policy pi_theta(y | x) using sequences y = (y_1, ..., y_T) sampled from a behavior policy pi_beta(y | x) (e.g., an earlier checkpoint or an exploratory model).

You are tasked with formally analyzing how importance sampling (IS) is applied to token generation, diagnosing failure modes, and proposing variance-reduction mechanisms.

Tasks

1. Likelihood Ratio and Support Conditions

  • Formulate the importance sampling estimator for the expected sequence reward E_{y ~ pi_theta}[R(x, y)] given N independent sequence completions generated from pi_beta(y | x).
  • State the exact support condition required for the estimator to remain unbiased. What happens in an autoregressive vocabulary setting when a token has zero probability under pi_beta but positive probability under pi_theta?

2. Variance Growth across Sequence Length

  • Express the sequence-level importance weight w(y | x) as a product of token-level conditional probability ratios r_t = pi_theta(y_t | x, y_ 0, derive how the variance of the sequence-level weight Var(w(y | x)) scales with sequence length T.
  • Explain why this makes vanilla sequence-level importance sampling intractable for long-context generation.

3. Practical Mitigations and Clipping

  • Compare sequence-level importance weights against per-token importance weights as used in algorithms like PPO or actor-critic policy gradients.
  • Explain the mathematical mechanism of weight clipping (e.g., clip(r_t, 1 - epsilon, 1 + epsilon)) or truncated importance sampling. Discuss the explicit trade-off introduced between estimator bias and variance.

4. Monitoring and Diagnostics

  • Define the Effective Sample Size (ESS) metric based on normalized importance weights.
  • Explain how you would use ESS and policy divergence metrics (such as token-level KL divergence) in production to trigger policy rollout refreshes or dynamically shrink learning rate steps.

Hints

Hint 1

Recall that the sequence probability decomposes autoregressively as the product of per-token probabilities, meaning the sequence-level importance weight is a product of T ratios.

Hint 2

For independent random variables X_i with mean 1, the variance of their product can be calculated using E[X_i^2] = 1 + sigma^2.

Hint 3

Consider how clipping bounds the maximum contribution of any single token update at the expense of introducing bias away from the true policy gradient direction.

Follow-up questions

  • How does Self-Normalized Importance Sampling (SNIS) compare to standard IS in terms of bias and mean squared error when the normalization constant is unknown or weights have heavy tails?
  • Why does standard PPO in LLMs often drop the product across timesteps altogether and apply token-level clipping independently at each position t?
  • In multi-step RL algorithms like V-trace or Retrace, how are importance sampling ratios truncated at both sequence and token levels to preserve contraction mappings?