Inference-Time Alignment of Diffusion Models via Trust-Region Iterative Twisted Sequential Monte Carlo
- URL: http://arxiv.org/abs/2605.25123v1
- Date: Sun, 24 May 2026 15:08:52 GMT
- Title: Inference-Time Alignment of Diffusion Models via Trust-Region Iterative Twisted Sequential Monte Carlo
- Abstract summary: We study inference-time alignment for diffusion-based generative models.<n>We propose a trust-region framework for learning twisting functions in SMC-based inference-time alignment.
- Score: 15.105260031413698
- License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
- Abstract: We study inference-time alignment for diffusion-based generative models, aiming to steer a base model toward high-reward outputs without updating its weights. Recent Sequential Monte Carlo (SMC)-based steering methods approximate reward-tilted target distributions in a principled way, but their proposals remain largely tied to the base sampler. Since reward information is mainly used after propagation through particle reweighting and resampling, these methods can require large particle budgets and suffer from weight degeneracy and high-variance estimates. One way to reduce variance and improve particle efficiency is to iteratively learn twisting functions that provide look-ahead guidance, as in twisted SMC. However, existing learnable twisting methods are developed mainly for classical sequential inference and can be unstable when applied to diffusion-based alignment with high-dimensional state spaces and terminal, noisy, or black-box rewards. We propose Trust-Region Iterative Twisted Sequential Monte Carlo (TRI-TSMC), a trust-region framework for learning twisting functions in SMC-based inference-time alignment. Each iteration computes an exact KL-constrained update in path space, which admits a closed-form solution by tempered importance reweighting, and projects this target back to the parameterized twisted family by weighted maximum likelihood. Theoretically, we formalize the value-function interpretation of the optimal twisting function and show that it yields a zero-variance sampler. We prove that the trust-region update follows an escort path toward the target distribution, that the weighted maximum-likelihood update is a forward-KL projection, and that the path reduces residual importance-weight variance. Empirically, TRI-TSMC improves primary alignment objectives on discrete diffusion text generation and text-to-image generation under matched inference-time budgets.
Related papers
- Designing Reinforcement Learning for Diffusion Models: A Unified Path-Space View [17.57225063689823]
Reinforcement learning post-training provides a way to align diffusion models with human preferences and task-specific rewards.<n>Current RL algorithms for diffusion models remain fragmented.<n>This paper shows that these seemingly different losses arise from a single path-space principle.
arXiv Detail & Related papers (2026-08-14T16:11:55Z) - Expanding SPHERE-JEPA: A Family of Statistical Regularizers for the Hypersphere [11.087393998848254]
In Self-Supervised Learning (SSL), preventing representation collapse by explicitly enforcing a uniform distribution on the unit hypersphere has proven to be effective.<n>Current frameworks typically rely on sliced statistical regularizers such as SIGReg (used in LeJEPA) and SUSReg (used in SPHEREJEPA)<n>This injects projection variance into the training gradients, destabilizing optimization, and hindering convergence.<n>In this work, we first show that analytically integrating out these random projections yields a deterministic Maximum Mean Discrepancy (MMD), bypassing the variance of sliced methods.
arXiv Detail & Related papers (2026-06-16T07:10:05Z) - Contrastive Distribution Matching for Amortized Sequential Monte Carlo in Discrete Diffusion [33.94204857658877]
We introduce Contrastive Distribution Matching (CDM), a novel framework that amortizes the cost of SMC inference by learning a parameterized twist function via positive and negative samples.<n>In practice, evaluating our learned twist function incurs less than 5% additional computational overhead compared to a single forward pass of the base model.<n>We validate the effectiveness and versatility of our approach across a diverse range of applications, including toxic text generation, regulatory DNA sequence design, protein designability, and diffusion large language model alignment.
arXiv Detail & Related papers (2026-05-22T08:06:52Z) - Self-Distilled Trajectory-Aware Boltzmann Modeling: Bridging the Training-Inference Discrepancy in Diffusion Language Models [65.89572755202245]
Diffusion Language Models (DLMs) offer stronger global awareness and highly parallel generation.<n>Post-training DLMs with standard Negative Evidence Lower Bound (NELBO)-based supervised fine-tuning remains inefficient.<n>We propose a self-distilled trajectory-based post-training framework that aligns training with the easy-to-hard structure of inference.
arXiv Detail & Related papers (2026-05-12T09:39:06Z) - Diffusion Controller: Framework, Algorithms and Parameterization [54.82539154511621]
We present a unified control-theoretic view that casts reverse diffusion sampling as state-only control within (generalized) linearly-solvable Markov Decision Processes.<n>Under this framework, control acts by reweighting the pretrained reverse-time transition kernels, balancing terminal objectives against an $f$-divergence cost.<n>Experiments on Stable Diffusion v1.4 show consistent gains in preference-alignment win rates and improved quality-efficiency trade-offs.
arXiv Detail & Related papers (2026-03-07T01:49:59Z) - Learnable Chernoff Baselines for Inference-Time Alignment [64.81256817158851]
We introduce Learnable Chernoff Baselines as a method for efficiently and approximately sampling from exponentially tilted kernels.<n>We establish total-variation guarantees to the ideal aligned model, and demonstrate in both continuous and discrete diffusion settings that LCB sampling closely matches ideal rejection sampling.
arXiv Detail & Related papers (2026-02-08T00:09:40Z) - Self-Rewarding Sequential Monte Carlo for Masked Diffusion Language Models [58.946955321428845]
This work presents self-rewarding sequential Monte Carlo (SMC)<n>Our algorithm stems from the observation that most existing MDLMs rely on a confidence-based sampling strategy.<n>We introduce the trajectory-level confidence as a self-rewarding signal for assigning particle importance weights.
arXiv Detail & Related papers (2026-02-02T09:21:45Z) - Inference-Time Scaling of Diffusion Language Models with Particle Gibbs Sampling [70.8832906871441]
We study how to steer generation toward desired rewards without retraining the models.<n>Prior methods typically resample or filter within a single denoising trajectory, optimizing rewards step-by-step without trajectory-level refinement.<n>We introduce particle Gibbs sampling for diffusion language models (PG-DLM), a novel inference-time algorithm enabling trajectory-level refinement while preserving generation perplexity.
arXiv Detail & Related papers (2025-07-11T08:00:47Z) - Training-Free Stein Diffusion Guidance: Posterior Correction for Sampling Beyond High-Density Regions [46.59494117137471]
Training free diffusion guidance provides a flexible way to leverage off-the-shelf classifiers without additional training.<n>We introduce Stein Diffusion Guidance (SDG), a novel training-free framework grounded in a surrogate SOC objective.<n>Experiments on molecular low-density sampling tasks suggest that SDG consistently surpasses standard training-free guidance methods.
arXiv Detail & Related papers (2025-07-07T21:14:27Z) - Sequential Kalman Tuning of the $t$-preconditioned Crank-Nicolson algorithm: efficient, adaptive and gradient-free inference for Bayesian inverse problems [1.3654846342364308]
We propose an adaptive implementation of EKI and Flow Annealed Kalman Inversion.
EKI is only exact in the regime of Gaussian target measures and linear forward models.
We show significant improvements in the rate of convergence compared to adaptation within standard SMC.
arXiv Detail & Related papers (2024-07-10T15:56:30Z) - Sequential Monte Carlo for Inclusive KL Minimization in Amortized Variational Inference [3.126959812401426]
We propose SMC-Wake, a procedure for fitting an amortized variational approximation that uses sequential Monte Carlo samplers to estimate the gradient of the inclusive KL divergence.
In experiments with both simulated and real datasets, SMC-Wake fits variational distributions that approximate the posterior more accurately than existing methods.
arXiv Detail & Related papers (2024-03-15T18:13:48Z) - SIXO: Smoothing Inference with Twisted Objectives [8.049531918823758]
We introduce SIXO, a method that learns targets that approximate the smoothing distributions.
We then use SMC with these learned targets to define a variational objective for model and proposal learning.
arXiv Detail & Related papers (2022-06-13T07:46:35Z) - Distributionally Robust Models with Parametric Likelihood Ratios [123.05074253513935]
Three simple ideas allow us to train models with DRO using a broader class of parametric likelihood ratios.
We find that models trained with the resulting parametric adversaries are consistently more robust to subpopulation shifts when compared to other DRO approaches.
arXiv Detail & Related papers (2022-04-13T12:43:12Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.