One-sentence Summary

Vector Institute introduces Twisted Sequential Monte Carlo for probabilistic inference in LLMs, employing a novel contrastive method to learn twist functions that direct inference-time computation toward promising partial sequences, and proposes bidirectional SMC bounds for evaluating inference accuracy, demonstrating effectiveness on automated red-teaming, sentiment-controlled generation, and infilling tasks.

Key Contributions

• The paper formulates capability and safety techniques such as RLHF, automated red-teaming, prompt engineering, and infilling as sampling from unnormalized target distributions, and introduces twisted Sequential Monte Carlo with learned twist functions to focus inference-time computation on promising partial sequences.
• The paper proposes a contrastive method for learning twist functions, establishing connections to soft reinforcement learning and positioning the approach as a principled probabilistic extension of best-of-K decoding and related value-based sampling schemes.
• The paper introduces bidirectional Sequential Monte Carlo bounds on the log partition function for estimating KL divergence between inference and target distributions, demonstrating twisted SMC for sampling undesirable pretrained outputs, generating reviews with varied sentiment, and infilling tasks.

Introduction

Many language model tasks, including RLHF, red-teaming, and reasoning, can be seen as steering generated text toward a specified property by scoring full sequences. The authors frame this as probabilistic inference, sampling from a base language model modulated by a potential function and estimating an intractable partition function over all possible responses. Prior sampling approaches are limited because the desired full-sequence distribution is non-causal, so generating tokens sequentially requires unknown future-aware marginals, and earlier SMC work mostly uses fixed or short-horizon guidance rather than learned twist functions. The authors leverage twisted sequential Monte Carlo and learn intermediate twist functions, including a contrastive twist learning method, to guide generation toward promising partial sequences and provide partition-function bounds for evaluating inference quality.

Method

The authors leverage a twisted sequential Monte Carlo framework to sample from language distributions defined by a pretrained base language model and a terminal scoring function. Let the base autoregressive model assign sequence probability $p_0(\mathbf{s}_{1:T})$p0​(s1:T​), and let the desired target distribution be defined by a terminal potential $\phi(\mathbf{s}_{1:T})$ϕ(s1:T​):

Here $\sigma$σ is the inference target, $\phi$ϕ encodes the final objective, such as reward, likelihood, classification probability, or another terminal scoring signal, and $\mathcal{Z}_\sigma$Zσ​ is the partition function. The central difficulty is that direct sampling from $\sigma$σ is intractable because $\phi$ϕ is terminal, while generation proceeds token by token. The authors therefore introduce intermediate target distributions that guide early token decisions with information from the eventual terminal objective.

The method represents intermediate targets using twist functions $\psi_t(\mathbf{s}_{1:t})$ψt​(s1:t​), which modulate the base model to approximate the true marginal of the target distribution at time $t$t. For $t < T$t<T, the twisted intermediate target is defined as

while the final target remains

The role of the twist is to summarize how future potential can be incorporated into present sampling decisions. If the twists were ideal, the intermediate targets would match the true target marginals:

The optimal twist satisfies a backward marginalization condition:

with the corresponding recursion

This condition shows that the optimal twist evaluates the expected terminal potential under future continuations generated by the base model. In other words, $\psi_t(\mathbf{s}_{1:t})$ψt​(s1:t​) estimates how promising a partial sequence is for ultimately achieving high terminal potential.

At each generation step, the authors run a sequential Monte Carlo procedure with $K$K partially generated sequences. Given a proposal $q(s_t \mid \mathbf{s}_{1:t-1})$q(st​∣s1:t−1​) and the twisted intermediate targets, the incremental importance weight at step $t$t is

These weights correct the mismatch between the base model, the proposal, and the twisted intermediate targets. The authors highlight two natural proposal choices. The first is to use the base model directly as the proposal:

This is simple and tractable, but it may require a large number of particles $K$K to approximate the terminal target well, especially when the terminal potential is sharp or highly selective.

The second choice is a twist-induced proposal. For a given set of twisted intermediate targets, the one-step variance-minimizing proposal has the form

Equivalently,

where

This proposal is tractable in transformer architectures for $t < T$t<T because the normalization sums over the discrete vocabulary $\mathcal{V}$V of next tokens. The full twisted proposal over a sequence is

At the final step, the true terminal potential $\phi(\mathbf{s}_{1:T})$ϕ(s1:T​) may require a separate network, such as a reward model, classifier, or likelihood evaluator. The authors therefore use an approximate final twist $\psi_T(\mathbf{s}_{1:T}) \approx \phi(\mathbf{s}_{1:T})$ψT​(s1:T​)≈ϕ(s1:T​) inside the proposal:

With this proposal, the incremental weights for $t < T$t<T become

and at the final step,

The key practical benefit is that the proposed next tokens are already biased toward continuations that the twist believes will achieve high terminal potential. The final weight then corrects any remaining gap between the approximate twist and the true potential.

The authors also connect this formulation to conditional inference. If the target is a posterior obtained by conditioning on an observation variable $o_T$oT​, the target is

In this case, the optimal twist has a direct probabilistic meaning:

Thus, for conditional generation tasks, the twist learns the probability that a partial sequence can still lead to the observed target outcome. The unconditional terminal potential setting can be recovered by treating a binary observation as the potential. This perspective also supports exact posterior sampling in simulated settings through ancestral sampling from the joint distribution and conditioning on the simulated observation.

The authors further interpret the twist functions through reinforcement learning. When the terminal potential is written as an exponentiated reward,

the twist corresponds to an exponentiated soft value function:

From this viewpoint, the twist behaves like a critic that estimates the value of partial sequences, while the proposal behaves like an actor that generates token continuations. The optimal twist recursion mirrors a soft Bellman-like equation, even when rewards are present only at the terminal step.

The main training method for the twists is contrastive twist learning. The authors train a parameterized family of twists $\psi_t^\theta$ψtθ​ by minimizing the average KL divergence between the desired target marginal and the twisted intermediate marginal:

The gradient of this objective has a contrastive structure:

The negative part of the gradient encourages the twist to increase scores under the current twisted distribution $\pi_t^\theta$πtθ​. The positive part encourages the twist to increase scores on partial sequences that better represent the true target marginal $\sigma(\mathbf{s}_{1:t})$σ(s1:t​). This resembles contrastive learning or energy-based model training, where positive samples push the energy down, while negative samples push the model’s distribution toward realistic but less target-aligned examples.

For negative sampling, the authors approximate the expectation under $\pi_t^\theta$πtθ​ using sequential importance sampling or sequential Monte Carlo. One can sample from a proposal such as the base model $p_0$p0​ or the twist-induced proposal $q^\pi$qπ, then reweight the generated partial sequences using the unnormalized intermediate target weights. In a full SMC workflow, resampling across $K$K particles provides a lower-variance approximation to the negative samples.

For positive sampling, the method relies on partial sequences that are drawn from, or weighted to approximate, the true target marginals. When exact target samples are available, for example through ancestral sampling from a tractable joint observation model, the exact posterior sample can be used directly. The authors describe truncating a complete exact sequence to a prefix of length $t$t, since that prefix corresponds to a sample from the marginal $\sigma(\mathbf{s}_{1:t})$σ(s1:t​). This gives a clean positive sample for the twist gradient at step $t$t.

When exact samples are not available, the authors use approximate positive sampling. They run SMC or SIS with $K$K particles, but crucially they weight these particles using the true terminal potential $\phi(\mathbf{s}_{1:T})$ϕ(s1:T​), not merely the learned twist. The resulting weights approximate the target marginal distribution over sequences. The same set of particles and weights can then be used to estimate positive gradient contributions across different timesteps. Rejection sampling can also provide exact positive samples when an upper bound on the proposal-to-target likelihood ratio is known.

The training workflow can be summarized as a loop over generation steps. The authors draw or approximate positive partial sequences that reflect the true target marginal. They also draw negative partial sequences from the current proposal and reweight them under the current twisted intermediate target. The twist network is updated so that it assigns higher scores to positive partial sequences and lower scores to negative partial sequences. Over time, the intermediate twisted targets move closer to the true marginals of the final target distribution.

The authors also provide tools for evaluating the quality of this inference procedure. A central quantity is the log partition function $\log \mathcal{Z}_\sigma$logZσ​, because it appears in KL divergence estimates and in likelihood-based diagnostics. They derive bidirectional stochastic bounds for $\log \mathcal{Z}_\sigma$logZσ​ using the extended-state view of SMC. Let $S$S denote the full extended state of particles, tokens, weights, and resampling indices. The lower bound is

and the upper bound is

These satisfy

The lower bound corresponds to sampling the extended SMC system from the proposal. The upper bound requires sampling from the corresponding target in the extended state space. This can be done when at least one exact target sample is available. During the extended SMC process, the exact target sample is protected so that its partial sequence remains in the particle set. The remaining particles follow the usual proposal and resampling dynamics. As $K$K increases, the bounds become exact, allowing the authors to sandwich $\log \mathcal{Z}_\sigma$logZσ​ and assess how closely the SMC procedure approximates the true target distribution.

Experiment

The experiments evaluate the framework by comparing upper and lower log partition bounds and KL divergences between proposal samples and target distributions across toxic story generation, sentiment review generation, and infilling. They show that twisted SMC can substantially improve sampling efficiency and estimation quality for rare events compared with standard importance sampling, while also serving as a tool to assess other inference methods. The results indicate that proposal quality depends on the training objective and sample availability, with CTL performing best under approximate positive sampling and DPG performing best when exact target samples are available.

In the conditional infilling setting, DPG achieves the best expected divergence scores in both KL directions, benefiting from access to exact target samples. Contrastive learning performs competitively but less consistently than DPG and is near SIXO, while FUDGE and RL are less favorable across the evaluated divergences. PPO is weaker for the target-to-proposal direction because it trains only from proposal samples. DPG is strongest for both proposal-to-target and target-to-proposal KL divergences in infilling. PPO performs relatively poorly when measuring how well the proposal covers the posterior.

In the conditional infilling setting, divergence scores are evaluated in both KL directions to assess how well each method aligns proposals with target and posterior samples. DPG achieves the strongest performance in both directions, aided by access to exact target samples, while contrastive learning remains competitive but less consistent than DPG. FUDGE and RL are less favorable overall, and PPO is weaker, particularly when the evaluation emphasizes coverage of the posterior.