Research map / Probabilistic and scientific modeling / Variational and generative sampling

Aligning Few-Step Generative Model via Amortizing Sample-Based Variational Inference

NeurIPS 2026 · Probabilistic and scientific modeling › Variational and generative sampling · approach: Variational inference theory and gradients

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes FAV, a sample-based alignment framework that approximates reward-tilted sampling with SVGD (KDE-estimated reference score) and amortizes the Stein transport into few-step generators via fixed-point regression, evaluated on offline/offline-to-online RL and image generator alignment.

Problem

Aligning few-step generative models is hard because existing alignment methods require tractable likelihoods, specific ODE/SDE solvers, or model-specific formulations.

Approach

Cast alignment as sampling from a reward-tilted distribution anchored to a reference; approximate it with SVGD using only generator samples, estimate the reference score by Gaussian KDE, and distill the SVGD transport into generator weights via a stop-gradient fixed-point regression loss (optionally in a pretrained encoder's representation space).

Key results

• On 56 offline RL tasks, FAV averages 54 vs. 50 for FAV-Adaptive and 48 for the best baseline FQL.
• On 30 offline-to-online RL tasks, FAV reaches 69 average after 1M online steps vs. 54 for FQL and 44 for RLPD.
• FAV generates aligned images roughly 180x to 280x faster than inference-time methods (0.1 s vs. 17.9-28.6 s) with higher aesthetic score (6.61 vs. 6.52 for ReNO-50).

Surprising findings

• FAV-Adaptive, with a bandwidth set by Scott's rule and no tuning, already outperforms all baselines on average offline RL.
• Baselines such as Adjoint Matching and policy gradient degrade in the few-step regime on the toy MeanFlow setting.

Limitations

• KDE-based score estimation is only consistent asymptotically, so finite reference samples introduce approximation error in high dimensions.
• Image experiments rely on transport in a pretrained encoder's representation space, which does not fully resolve the high-dimensional approximation error.
• Baselines such as Regularized REINFORCE and Adjoint Matching are not directly applicable to one-step generators.

Experimental setup

Datasets: OGBench D4RL AntMaze ImageNet-256 SneakyPrompt DrawBench 8-Gaussians

Benchmarks: OGBench D4RL AntMaze 8-Gaussians toy KL ImageNet-256 aesthetic score SANA-Sprint 1024^2 HPSv2/NSFW

Metrics: average return KL divergence aesthetic score HPSv2 ImageReward DreamSim diversity CLIP diversity inference time

Baselines: IQL ReBRAC DSRL FAWAC IFQL QAM SRPO CAC FQL RLPD DRaFT Flow-GRPO

Builds on: Stein variational gradient descent FQL QAM MeanFlow DRaFT

Closest papers by meaning

• Learnable Chernoff Baselines for Provable Inference-Time Alignment (NeurIPS 2026 · similarity 0.86)
• Stitched Value Model for Diffusion Alignment (NeurIPS 2026 · similarity 0.86)
• Controlled Generation with Equivariant Variational Flow Matching (ICML 2025 · similarity 0.86)
• Aligning Model Properties via Conformal Risk Control (NeurIPS 2024 · similarity 0.86)
• Composition and Alignment of Diffusion Models using Constrained Learning (NeurIPS 2025 · similarity 0.86)
• Statistically Optimal Generative Modeling with Maximum Deviation from the Empirical Distribution (ICML 2024 · similarity 0.86)