Probabilistic Inference in Language Models via Twisted Sequential Monte Carlo

Stephen Zhao; Rob Brekelmans; Alireza Makhzani; Roger Grosse

Source-linked AI summary

Probabilistic Inference in Language Models via Twisted Sequential Monte Carlo

Stephen Zhao, Rob Brekelmans, Alireza Makhzani, Roger Grosse

arXiv:2404.17546v1Apr 26, 2024cs.LGcs.AIcs.CLstat.ML

TL;DR

Many language-model learning and inference tasks can be formulated as inference over sequence-level target distributions, but direct sampling is difficult because the target depends on full sequences. The paper uses learned twist functions with SMC to prioritize promising prefixes and introduces bidirectional log-partition bounds to evaluate inference quality. It demonstrates the framework across sampling and evaluation tasks, while noting that optimal twists require intractable future marginalization.

Problem
Many language-model learning and inference tasks seek samples satisfying a sequence-level property or reward, motivating probabilistic inference over unnormalized target distributions.
Source: §1 · p. 1
Method
Twisted SMC learns intermediate twist functions that estimate future target value, focuses computation on promising partial sequences, and uses contrastive twist learning for terminal-potential targets.
Source: §1 · p. 2, §3 · p. 4, §4.1 · p. 6
Results
Twisted SMC supplies bidirectional bounds on log Zσ for estimating KL divergence in both directions and evaluating language-model inference quality.
Source: §1 · p. 2, §5.1 · p. 8
Takeaways & Limitations
The framework supports sampling and evaluation across undesirable-output generation, varied-sentiment reviews, infilling, and other controlled-generation settings.
Source: §8 · p. 10
Optimal twist functions are unavailable because computing them requires marginalizing over future timesteps, so the paper uses learned approximations.
Source: §3.1 · p. 4

Abstract

from arXiv · show

Numerous capability and safety techniques of Large Language Models (LLMs), including RLHF, automated red-teaming, prompt engineering, and infilling, can be cast as sampling from an unnormalized target distribution defined by a given reward or potential function over the full sequence. In this work, we leverage the rich toolkit of Sequential Monte Carlo (SMC) for these probabilistic inference problems. In particular, we use learned twist functions to estimate the expected future value of the potential at each timestep, which enables us to focus inference-time computation on promising partial sequences. We propose a novel contrastive method for learning the twist functions, and establish connections with the rich literature of soft reinforcement learning. As a complementary application of our twisted SMC framework, we present methods for evaluating the accuracy of language model inference techniques using novel bidirectional SMC bounds on the log partition function. These bounds can be used to estimate the KL divergence between the inference and target distributions in both directions. We apply our inference evaluation techniques to show that twisted SMC is effective for sampling undesirable outputs from a pretrained model (a useful component of harmlessness training and automated red-teaming), generating reviews with varied sentiment, and performing infilling tasks.

1. Introduction

The paper frames controlled language-model generation as probabilistic inference over sequence-level target distributions and develops twisted SMC to sample and evaluate these targets. It learns intermediate twist functions to focus computation on promising prefixes and introduces bidirectional partition-function bounds for diagnosing inference quality.

Twisted Sequential Monte Carlo for Language Modeling

Learned twist functions approximate target marginals over partial sequences, focusing each generation step on prefixes with promising future potential.The optimal twists require intractable future marginalization, motivating learned approximations.Source: §1 · pp. 1–2
Twisted SMC provides a general framework for sampling and evaluating language-model inference when target distributions are defined by terminal potentials.The framework is intended for capability and safety tasks, including RLHF, red-teaming, reasoning, and controlled generation.Source: §1 · pp. 1–2

Contrastive Twist Learning

Contrastive twist learning estimates intermediate twists using methods inspired by energy-based modeling and density-ratio estimation, while connecting twisted SMC to soft reinforcement learning.The paper adapts existing twisted SMC methods to the language-modeling setting.Source: §1 · p. 2

Evaluating Inference in Language Models

Bidirectional SMC bounds on log Zσ evaluate inference accuracy by bounding KL divergence between inference and target distributions in both directions.These tools can diagnose mode-dropping in methods such as PPO, which optimizes the mode-seeking DKL(q ∥σ).Source: §1 · p. 2

Experiments

Experiments show twisted SMC improves controlled generation and provides insight into inference quality across the paper’s reported settings.The experiments cover sampling and evaluation, including undesirable-output sampling, varied-sentiment reviews, and infilling tasks as described in the paper context.Source: §1 · p. 2

2. Background

The background introduces importance sampling and sequential Monte Carlo as estimators and samplers for an intractable partition function. SMC improves on whole-sequence resampling by maintaining and selectively cloning promising partial sequences through intermediate targets.

2.1. Simple Importance Sampling

Simple importance sampling estimates Zσ unbiasedly from proposal-distribution importance weights, while normalized weights support approximate target expectations and samples.Approximation quality depends on how closely the proposal matches σ; exact inference gives zero importance-weight variance.Source: §2.1 · p. 2

SMC-PROPOSAL

Twisted SMC sampling uses proposal draws, incremental weights, and optional resampling to maintain K partial sequences and extend selected prefixes.High-weight sequences can be cloned and used as prefixes for the next proposal step.Source: §2.1 · p. 3, SMC-PROPOSAL · p. 3, §2.2 · p. 3

2.2. Sequential Monte Carlo

Unlike SIS, SMC can resample partial sequences before completion, cloning promising prefixes so later sampling better covers high-probability target regions.Without resampling, SMC reduces to SIS; resampling times may be selected adaptively using criteria such as ESS.Source: Figure 2 · p. 3, §2.2 · pp. 3–4
SMC decomposes inference into intermediate target distributions and estimates the final partition function by composing incremental partition-function-ratio estimates.The estimator remains unbiased when the final intermediate target equals σ, regardless of the intermediate targets and proposal.Source: §2.2 · p. 3

3. Twisted Sequential Monte Carlo for Language Modeling

The framework applies twisted SMC to language-model targets defined by terminal potentials, learning twists that approximate target marginals and support sequential inference. It also develops tractable proposals, conditional targets, and probabilistic connections to soft reinforcement learning.

3. Twisted Sequential Monte Carlo for Language Modeling

Twisted SMC addresses terminal-potential language-model targets by learning intermediate twists that approximate true target marginals and focus computation on promising prefixes.The final target remains the intended distribution, while learned twists summarize future information relevant to sequential generation.Source: §3 · p. 4, §3.1 · p. 4

3.1. Twist Functions

Optimal twists correspond to future target information, but they require marginalizing over future timesteps and therefore must be approximated by learned functions.The optimal intermediate targets are the target marginals σ(s1:t), and their twists satisfy a recursive characterization.Source: §3.1 · p. 4

3.2. Proposal Distribution

The twist-induced proposal minimizes one-step incremental-weight variance and can be made tractable for transformer architectures, while the base model proposal may require large K.At the final timestep, the proposal uses an approximation to the terminal potential because that potential may require a separate reward model or classifier.Source: §3.2 · pp. 4–5

3.3. Conditional Target Distributions

Conditional targets treat sequences as posteriors given an observation, allowing twists and proposals to condition on that observation and enabling exact target samples when joint sampling is tractable.The BDMC trick uses exact posterior samples for training or evaluation in this conditional formulation.Source: §3.3 · p. 5

3.4. Connections with Reinforcement Learning

The framework includes soft RL as a special case while extending beyond RL objectives to general target distributions, probabilistic SMC resampling, and density-ratio-based twist learning.In the soft-RL interpretation, twists are analogous to critics and proposals to actors; the framework also supports intermediate potentials or rewards.Source: §3.4 · pp. 5–6

4. Learning the Twist Functions

The paper introduces contrastive twist learning to estimate intermediate twist functions and discusses exact or approximate sampling strategies, alongside related twist-learning methods.

4.1. Contrastive Twist Learning

Contrastive twist learning minimizes separate KL divergences to target marginals, encouraging mass coverage so early errors do not aggressively prune promising sequences.The method is motivated by energy-based modeling and density-ratio estimation.Source: §4.1 · p. 6
CTL’s gradient resembles energy-based-model maximum likelihood, combining positive and negative sampling to update twist functions.Exact target samples can supply the positive term when available, while approximate techniques are used for both sampling terms.Source: §4.1 · p. 6

4.1.1. APPROXIMATE NEGATIVE SAMPLING

Negative sampling estimates expectations with SIS using proposals such as the base model or twist-induced proposal, and may also use SMC resampling.The challenge is the intractable normalization of the target marginal.Source: §4.1.1 · p. 6

4.1.2. (APPROXIMATE) POSITIVE SAMPLING

Positive samples may come from exact posterior samples, rejection sampling, or approximate SMC/SIS sampling when exact sampling is unavailable or inefficient.Exact rejection sampling requires a known likelihood-ratio bound M; otherwise approximate or quasi-rejection methods may be needed.Source: §4.1.2 · pp. 6–7

4.2. Twist Learning Methods from Related Work

Related twist-learning approaches include soft Q-learning, path-consistency losses, SIXO, and FUDGE, connecting twist estimation to soft reinforcement learning and classification.Table 1 summarizes losses for twist and proposal learning.Source: §4.2 · p. 7, Table 1 · p. 7

5. Evaluating Inference in Language Models

The framework evaluates language-model inference by estimating the target partition function and deriving bidirectional KL bounds. With an exact target sample, bidirectional SMC bounds can become tight as the number of particles increases.

5. Evaluating Inference in Language Models

A single exact target sample enables upper as well as lower bounds on log Zσ, and the bounds can tightly sandwich it as K increases.These bounds support reliable conclusions about inference quality.Source: §5 · p. 7, §5.2 · p. 8

5.1. Applications of log Zσ Estimation

Estimating log Zσ allows KL divergence to the target to be evaluated in both directions for tractable inference distributions q.The reverse direction requires target samples, which may be available through the BDMC trick for Bayesian posteriors.Source: §5.1 · pp. 7–8
For twisted SMC, tight lower or upper log-partition bounds provide upper bounds on forward or reverse KL divergence for the marginal approximate sample distribution.This enables evaluation even though the marginal distribution of the returned SMC sample is not tractable.Source: §5.1 · p. 8

5.2. Bidirectional SMC Bounds on log Zσ

The proposed bidirectional SMC bound is novel and can evaluate inference beyond language modeling by bounding log Zσ in an extended SMC state space.The extended state includes tokens and resampling indices, with lower and upper bounds derived for the corresponding proposal and target distributions.Source: §5.2 · p. 8
Constructing the upper bound requires one exact target sample, whose partial sequence is guaranteed to be cloned during resampling while remaining tokens are proposed.The lower-bound expectation instead samples from the ordinary SMC proposal.Source: §5.2 · p. 8
A small gap between the bounds implies that the single approximate SMC sample is close to the target in symmetrized KL divergence.The bounds become exact as K approaches infinity for any proposal.Source: §5.2 · p. 8

6. Related Work

The related work situates twisted SMC among controlled decoding, variational proposals, and reinforcement-learning-based generation methods.

6. Related Work

Prior language-model work studies diverse or controlled decoding, while Lew et al. (2023) apply SMC without learning twist functions or proposals.The paper distinguishes its learned-twist framework from this prior SMC approach.Source: §6 · p. 8
The bidirectional bound gap corresponds to the sum of forward and reverse KL divergences in the extended SMC state space.This provides the formal connection between the bounds and symmetrized divergence.Source: §6 · p. 8
Soft reinforcement-learning value functions play a role analogous to twist functions for SMC resampling, connecting PPO-based methods and value-guided search.Related work uses Monte Carlo Tree Search with PPO value estimates or discriminator-driven guidance.Source: §6 · p. 9

7. Experiments

The experiments evaluate twisted SMC for rare-event sampling and for bidirectional KL-based comparison of inference methods across toxicity, sentiment, and infilling tasks.

7.1. Comparing SIS and SMC for log Zσ Estimation

Orders of magnitude fewer samples suffice for tight log Zσ estimation and near-exact toxic-output sampling with the twist-induced proposal, whereas base-proposal SIS generally fails without resampling.With the base proposal, SMC resampling eventually tightens both bounds and produces small KL divergence; with the twist-induced proposal, both SIS and SMC achieve tight bounds.Source: §7.1 · p. 9, Figure 3 · p. 9

7.2. Evaluating Twist-Induced or Variational Proposals

The evaluation compares PPO, DPG, CTL, SIXO, and FUDGE using forward and reverse KL divergences estimated from SMC log-partition bounds.SMC serves both as an inference method with tractable qπ and as an evaluation method for other inference techniques.Source: §7.2 · pp. 9–10

7.2.1. GENERATING TOXIC STORIES

PPO performs best for DKL(q ∥σ), while CTL performs best for DKL(σ ∥q) in toxic-story generation.The target uses a classifier probability of toxicity, providing a smoother learning signal while retaining exact sampling through rejection sampling.Source: §7.2.1 · p. 10

7.2.2. GENERATION WITH VARIED SENTIMENT

CTL performs best in both KL directions for generating reviews with varied sentiment, although all methods achieve low KL divergences.The target is a classifier-based distribution over 1-star reviews, and exact samples are available by rejection sampling.Source: §7.2.2 · p. 10

7.2.3. INFILLING

DPG performs best in both KL directions for infilling, while CTL is comparable to SIXO and PPO performs relatively poorly on DKL(σoT ∥qoT).The experiment averages KL divergences over 2000 continuation-conditioned tasks with T = 15 and c = 10.Source: §7.2.3 · p. 10
The experiments suggest preferring DPG when exact target samples are available and CTL when only approximate positive sampling is available.This conclusion follows the toxicity, sentiment, and infilling comparisons.Source: §7.2.3 · p. 10, §7.2.1 · p. 10, §7.2.2 · p. 10

8. Conclusion

The paper presents twisted SMC as a probabilistic inference framework for language-model capability and safety tasks, combining sampling, contrastive twist learning, and bidirectional evaluation bounds.

8. Conclusion

Twisted SMC is presented as a principled framework for sampling and evaluating language-model target distributions across varied experimental settings.The paper reports quantitative and qualitative evidence for both sampling and evaluation.Source: §8 · p. 10

Impact Statement

The impact statement emphasizes both the potential benefits and misuse risks of controlled generation, while the appendix contains examples and proof material.

Impact Statement

Controlled generation may improve safety, automated red-teaming, and estimation of rare behaviors, but the same techniques can also magnify harms when misused.The statement frames these capabilities as inherently double-edged.Source: Impact Statement · p. 11

Appendix

The appendix includes examples of target posteriors and proof material supporting the paper’s theoretical results.The supplied appendix passages identify Table 5 and proofs of propositions and the CTL gradient.Source: Table 5 · p. 15, A. Proofs · p. 15

A.1. Proof for Optimal Intermediate Target Distributions

This appendix develops the theoretical foundations and connections of twisted SMC, covering optimal intermediate targets, proposal and twist-learning objectives, conditional and soft-RL formulations, and implementation choices. It also relates the framework to controlled decoding and related future-value methods, including SIXO and FUDGE.

A.1. Proof for Optimal Intermediate Target Distributions

Optimal twisted SMC produces a partition-function estimator with zero bias and zero variance.The result follows because optimal multi-step importance weights telescope to a constant independent of sampled particles.Source: §A.1 · p. 16
Optimality is equivalent to matching every intermediate target to the corresponding target marginal and satisfying the associated twist and proposal recursions.These conditions are necessary and sufficient, and the optimal intermediate distributions are σ(s1:t) for all t.Source: §A.1 · pp. 16–17

A.2. Proof of Twist-Induced Proposal

The twist-induced proposal minimizes one-step importance-weight variance, while its multi-step generalization minimizes variance when resampling c steps ahead.The multi-step proposal’s denominator is generally intractable for c > 1 in language-modeling applications.Source: §A.2 · p. 17

A.3. Derivation of CTL Gradient

The appendix derives the negative CTL gradient with respect to θ using the log-derivative identity and the fact that log Zψt is sequence-independent.The derivation applies the identity ∇θf(θ) = f(θ)∇θ log f(θ) to the twisted target expression.Source: §A.3 · p. 18

B. SMC with Intermediate Potentials and Connection with Soft Reinforcement Learning

With intermediate potentials, twisted SMC generalizes beyond terminal-only rewards and connects its probabilistic formulation to soft reinforcement learning.The framework supports intermediate potentials, while the soft-RL special case uses a terminal potential such as ϕ(s1:T ) = eβrT (s1:T ).Source: B. SMC with Intermediate Potentials… · p. 18

B.1. Twisted SMC with Intermediate Potentials

For intermediate potentials, optimal twists recover target marginals, factor as ψt(s1:t) = ϕt(s1:t)Φt(s1:t), and distinguish future-value Φt from Q-like twist functions ψt.Φt summarizes expected future potentials, whereas ψt incorporates both the immediate potential and future value; rescaling either by a sequence-independent constant leaves marginals unchanged.Source: §B.1 · p. 19

B.2. Conditional Twisted SMC

Conditional twists represent future observations through conditional likelihoods, recovering unconditional targets for fixed observations and supporting infilling.In the infilling example, future tokens are observations generated by the base language model, and the conditional formulation also yields classification losses for twist learning.Source: §B.2 · pp. 20–21
Conditional twisted SMC includes unconditional targets as a special case by fixing the conditioning observations and recovering σ(s1:T) ∝ p0(s1:T)ϕT(s1:T).The formulation also matches soft-RL notation when binary optimality observations encode exponentiated rewards.Source: §B.2 · p. 22

B.3. Connection with Soft Reinforcement Learning

Soft reinforcement learning is represented as an MDP over partial sequences, with next-token actions, deterministic append transitions, and reward rt(s1:t).The target potential is ϕt(s1:t) = eβrt(s1:t), and the optimal soft value corresponds to a scaled log partition function.Source: §B.3 · p. 22
Optimal intermediate soft values decompose the regularized optimization over next-token conditionals, inducing intermediate target marginals and twist-based one-step proposals.Approximate Vt or Qt functions provide corresponding twisting targets, while their induced normalization constants act as soft-value estimates.Source: §B.3 · p. 23

B.4. Remarks on Parameterization

The choice between ψ and Φ twist parameterizations depends on potential-evaluation cost: Φ supports exact intermediate potentials over |V| next tokens, whereas ψ is preferable when potentials are evaluated over only K sequences.At the final timestep, a learned ψT enables tractable proposal sampling while the terminal potential is evaluated on K sampled sequences for unbiased importance weighting.Source: §B.4 · p. 24

C.1. Soft Q-Learning (RL) and Path Consistency Losses from Log Importance Weights

Twist-learning losses can be derived by minimizing variance or squared log importance weights, with optimality corresponding to constant importance weights.The resulting framework includes soft Q-learning and path consistency learning losses from the soft-RL literature.Source: §C.1 · pp. 24–25

C.1.1. SOFT Q-LEARNING AND RL BASELINE

For ψ-parameterized targets, summed one-step squared log-weight consistency recovers the familiar soft Q-learning loss, with errors of the form E[β2(rt + Vt − Qt)2].Without intermediate rewards, the corresponding RL baseline enforces the twist recursion in log space at intermediate timesteps.Source: §C.1.1 · pp. 25–26

C.1.2. PATH CONSISTENCY LEARNING (FOR TWIST LEARNING)

Using value-parameterized targets, one-step log-consistency losses recover path consistency learning, and multi-step log-weight losses recover multi-step PCL.The derivation treats PCL as an importance-sampling consequence rather than its original KKT-based derivation; jointly learning proposals is left to future work.Source: §C.1.2 · p. 26

C.2. Controlled Decoding Losses via Optimal Twist Identities (Mudgal et al., 2023)

Proposition B.1 (or Proposition 3.2 and Eq. (13) in the main text) notes that the optimal twists satisfy the following relationships.Source: §C.2 · p. 27
Controlled decoding losses enforce optimal twist identities: CD-Q yields zero-variance one-step importance weights, while CD-FUDGE matches the gradient of an intractable squared-error objective (Mudgal et al., 2023).CD-FUDGE requires base-model rollouts because multi-step marginalization is generally intractable; CD-Q instead uses a twist-induced proposal.Source: §C.2 · p. 27

C.3. SIXO: Smoothing Inference with Twisted Objectives (Lawson et al., 2022)

SIXO learns likelihood-ratio twists with noise-contrastive classification between conditional target prefixes and base-model prefixes.Exact conditional sampling can use the BDMC trick, while ancestral sampling and shuffled observations provide tractable estimates for the required positive and negative terms.Source: §C.3 · p. 28

C.4. FUDGE: Future Discriminators (Yang & Klein, 2021)

FUDGE directly trains a future discriminator to predict conditional observation likelihoods, using cross-entropy classification and base-model rollouts for correct targets (Yang & Klein, 2021).The framework supports unconditional twists and conditional twists for infilling, with the latter amortizing over observation variables and requiring independent base-model rollouts in the objective.Source: §C.4 · pp. 29–30

D.1. Proposal Sampling in Mudgal et al. (2023)

The paper contrasts proposal-learning and twisted-SMC objectives, showing where their targets, regularization, and decoding procedures diverge. It also surveys proposal-learning approaches and their relationships to SMC.

D.1. Proposal Sampling in Mudgal et al. (2023)

The CD and soft-RL objectives share one-step KL regularization, but soft values additionally account for future regularization through expected terminal rewards.As β →0, the soft value recovers expected base-model reward, whereas the CD objective mixes β=0 future evaluation with β>0 current-step regularization.Source: §D.1 · p. 31
Mudgal et al.’s one-step proposal uses twist-like values, but its composition is not shown to sample the marginals of a natural target distribution.The paper distinguishes policy evaluation under the base model from proposals that induce the desired SMC targets.Source: §D.1 · p. 31

D.2. Blockwise Greedy Decoding in Mudgal et al. (2023)

Blockwise arg max decoding of partial completions does not produce samples from the target distribution, whereas the SMC framework uses probabilistic resampling based on importance weights.The blockwise method generalizes best-of-K decoding, but the paper recommends c- or M-step resampling for target matching.Source: §D.2 · p. 32

E. Proposal Learning Methods

SMC proposal-learning methods can be used either for particle sampling with intermediate targets or directly in full-sequence importance-sampling bounds.The paper discusses DPG, path consistency learning, and policy-gradient formulations as related ways to learn controlled-generation proposals.Source: E. Proposal Learning Methods · p. 32

E.1. Path Consistency Learning for Controlled Generation

Path consistency learning treats learned soft values as twisting targets, while its proposal-learning formulation does not itself consider resampling according to those targets.This distinguishes an SMC-centered objective, where twisting targets are primary, from proposal learning for controlled generation.Source: §E.1 · p. 32

E.2. Policy Gradient Methods

Traditional PPO-style RLHF uses a mode-seeking KL objective that does not exploit exact target samples and can reduce fine-tuned-model diversity.The paper also derives mass-covering policy-gradient objectives and relates DPG to sequential maximum-likelihood learning for energy-based proposals.Source: §E.2 · p. 33, §E.3 · pp. 33–34

E.3. Policy Gradient with Mass-Covering / Maximum Likelihood KL Divergence

DPG learns a single sequence-level KL objective, whereas CTL optimizes separate KL divergences for each intermediate marginal and uses different sampling signals.DPG uses positive prefixes from target samples or approximate weighted sampling; CTL uses approximate negative sampling under each intermediate target.Source: §E.3 · p. 34

E.3.1. NAIVE USE OF PROPOSAL LEARNING TO DEFINE TWISTED SMC TARGETS

The paper shows that naively converting learned proposal energies into twisted-SMC targets can produce incorrect intermediate marginals or uninformative resampling. It contrasts this failure with CTL and introduces the extended-state-space construction used for bidirectional bounds.

E.3.1. NAIVE USE OF PROPOSAL LEARNING TO DEFINE TWISTED SMC TARGETS

Proposal-learning energies should not be reused as target twists: even optimal proposals can induce intermediate targets that differ from the desired target marginals.A prefix-dependent normalization factor can preserve optimal one-step proposals while breaking the SMC condition πt(s1:t)=σ(s1:t).Source: §E.3.1 · pp. 34–35
Empirically, resampling with naively induced targets produced worse SMC log Zσ bounds than SIS or IWAE using the learned proposal directly.This demonstrates that proposal quality alone does not guarantee valid or effective twisted-SMC targets.Source: §E.3.1 · p. 35
The CTL objective avoids this mismatch: at its global optimum, its intermediate targets equal σ(s1:t) and satisfy the necessary and sufficient conditions for optimal twisted SMC.This contrasts CTL’s target-focused optimization with naive proposal-energy reuse.Source: §E.3.1 · p. 35

E.3.2. SMC WITH NORMALIZED TARGETS INDUCED BY LEARNED PROPOSAL LEADS TO UNIFORM WEIGHTS

Normalizing targets induced from learned proposals yields uniform resampling weights before the final step, so SMC samples have the same marginals as proposal-driven SIS.Thus, the construction is well behaved algebraically but provides no meaningful intermediate resampling advantage.Source: §E.3.2 · p. 35

F. Bidirectional SMC

Bidirectional SMC represents proposal and target distributions on an extended state space, producing unbiased partition-function estimates and lower and upper bounds on log Zσ.The construction uses K sequences, resampling histories, intermediate targets, and a final target πT=σ.Source: F. Bidirectional SMC · p. 36

SMC-TARGET

This section presents the SMC target and twist parameterizations, covering exact targeting, partition-function and KL-bound estimation, and standardized proposal training. It also examines computational tradeoffs and performance across parameterization choices, including infilling.

SMC-TARGET

The extended target fixes one index path to an exact target sample, while the remaining K −1 particles use the proposal and undergo weighted resampling.This ensures at least one exact target sequence appears at every timestep and may be cloned when its importance weight is high.Source: SMC-TARGET · p. 38
Bidirectional SMC bounds lower- and upper-bound log Zσ, with gaps equal to DKL(qSMC(S) ∥σSMC(S)) and DKL(σSMC(S) ∥qSMC(S)), respectively.The extended-state construction also yields unbiased estimators of Zσ.Source: SMC-TARGET · pp. 39–40

G.1. Common Details Across Experiments

All methods use matched initialization and custom implementations; twist modifiers begin near zero so initial proposal samples approximately follow the base model p0.When available, exact posterior samples are used, and RL twist learning combines q and exact σ samples.Source: §G.1 · pp. 40–41

G.2. Choices of Twist Parameterization

Twist functions use a shared parameterization across prefixes rather than separate models for each timestep, reducing memory demands and potentially improving learning speed.The shared design mirrors the base language model’s single architecture across timesteps.Source: §G.2 · p. 41

G.2.1. LINEAR HEAD

The linear-head parameterization adds little computation but is capacity constrained, achieving worse KL divergences than other parameterizations.It keeps the base model fixed and trains only a new linear head.Source: §G.2.1 · p. 41

G.2.2. MLP HEAD

The 3-layer ReLU MLP head costs more than a linear head but remains inexpensive relative to the transformer and generally performs well.It is used in the toxicity-threshold and sentiment experiments; a separate transformer twist is also considered.Source: §G.2.2 · p. 41, §G.2.3 · p. 41

G.2.3. SEPARATE TRANSFORMER FOR THE TWIST

The separate-transformer parameterization is significantly more costly, does not always outperform an MLP head, and was retained for toxicity classification because it performed well there.Both separate-transformer variants roughly double forward-pass computation and substantially increase backward-pass memory and training time, while sampling remains the main bottleneck.Source: §G.2.3 · p. 41, §G.2.4 · p. 42

G.2.4. SEPARATE TRANSFORMER FOR THE TWIST, WITH MLP HEAD

The separate transformer with an MLP head performed materially better than all other parameterizations for 15-token infilling, particularly with DPG, so it was used for all infilling experiments.This variant is the most computationally costly approach described.Source: §G.2.4 · pp. 41–42

G.3. Comments on Our Choices of Experiment Settings

The experiments are designed to expose the advantage of twisted SMC over base-model SIS in difficult threshold sampling while using probabilistic potentials for baseline comparisons. Thresholded objectives create a sparse-learning-signal challenge.

G.3. Comments on Our Choices of Experiment Settings

The toxicity-threshold experiment is deliberately difficult for the base model and is used to demonstrate improvement from CTL-learned twists over base-model SIS.The setting is chosen so that drawing threshold-satisfying samples with p0 is extremely unlikely.Source: §G.3 · p. 42
With approximate positive sampling and a thresholded target, all importance weights are zero until at least one sampled sequence meets the threshold.Consequently, early twist-learning iterations receive no learning signal when sampling from p0 or near-initialized twists.Source: §G.3 · p. 42
For baseline comparisons, the paper uses probability-valued potentials and notes that the threshold and probability settings share the same underlying framework.This permits evaluating both directional KL divergences or plotting log Zσ bounds as a function of K across the settings.Source: §G.3 · p. 42

G.4. Experiment-Specific Details

These experiments specify task-specific twist architectures, sampling procedures, evaluation estimators, and training controls for toxicity, sentiment, and infilling. Exact posterior samples enable conditional KL evaluation, while computational cost and sample availability constrain practical use.

G.4. Experiment-Specific Details

Approximate positive sampling is used in the 10-token SIS/SMC comparison, while collecting 2,000 exact posterior samples offline takes several hours and practical settings may allow fewer samples.The toxicity setup generates 20 output tokens with 100 SMC particles; sentiment generates 10 tokens and uses an MLP head with 1024 hidden units per layer.Source: §G.4 · pp. 42–43
KL estimates use sandwich bounds on log Zσ, with K = 500 samples and shared median estimates to control comparison noise across methods and seeds.The bounds are generally tight; inaccuracies affect absolute KL values but not relative differences among learning methods.Source: §G.4 · p. 42
Infilling conditions the twist on continuation tokens by broadcasting their final embedding into the MLP head alongside prefix embeddings.This lets the head produce outputs that depend on the conditioning tokens.Source: §G.4 · p. 43
Conditional infilling KL divergences average one-sample estimates over 2000 continuation contexts, with confidence intervals generally between 0.20 and 0.30.A single exact conditional target sample comes from the BDMC trick, and one conditional proposal sample is drawn per context.Source: §G.4 · p. 43
Infilling training uses at least two negative samples per conditioning context for CTL, unlike other methods that can use one, and adopts approximate positive samples because they perform slightly better.The fixed batch size is split across contexts and negatives for CTL, with 25 contexts and 4 negatives in the main infilling experiments.Source: §G.4 · pp. 43–44

H.1. Qualitative Results

The qualitative and ablation results show that twisted SMC can control toxicity and sentiment, while infilling quality tracks KL divergence across methods. Exact posterior samples generally improve learning, but their rejection-sampling cost limits practicality and generated reviews remain somewhat incoherent when the base model is incoherent.

H.1. Qualitative Results

Twisted SMC produces a coherent TinyStories-style story with approximately 88% toxicity-classifier probability when exact posterior samples are unavailable.The example uses β = 10 and approximate positive sampling, supporting control toward undesirable outputs.Source: §H.1 · p. 44
Twisted SMC learns reviews associated with each rating class under β = 100, although the generated reviews are slightly incoherent because the base GPT2-Medium model is often incoherent.The generations are qualitatively more coherent than uncontrolled p0 outputs.Source: §H.1 · p. 44
Infilling quality follows KL divergence: DPG respects grammar and topic best, SIXO performs worse, and CTL has the highest KL divergence with incorrect grammar and weaker topicality.The comparison is reported for DPG, SIXO, and CTL in that order.Source: §H.1 · p. 44

H.2. Infilling with Fewer Tokens

With two output tokens and one conditioning token, DPG and RL learn best, while FUDGE learns notably more slowly and PPO performs poorly on D_KL(σ ∥ q).Most methods achieve low KL divergence in both directions in this easier setting.Source: §H.2 · p. 45, Table 8 · p. 45

H.3. Approximate vs. Exact Posterior Sampling

For a fixed number of training samples, exact σ samples improve KL performance for most methods, but rejection sampling makes exact-sample training substantially slower in wall-clock time.Figure 5 shows exact-sample learning reaching lower KL divergences more reliably, while exact results terminate earlier because fewer updates fit into the same time.Source: §H.3 · p. 45, Figure 5 · p. 46