Aligning Few-Step Generative Models by
Amortizing Sample-based Variational Inference

Jaewoo Lee*¹² Hyeongyu Kang*¹ Dohyun Kim¹ Kyuil Sim¹ Woocheol Shin¹ Minsu Kim¹³ Taeyoung Yun¹ Jeongjae Lee¹ Sanghyeok Choi⁴ Tabitha Edith Lee³⁵ Jongchul Ye†¹ Jinkyoo Park†¹⁶
¹KAIST ²MongooseAI ³Mila – Quebec AI Institute ⁴University of Edinburgh ⁵Université de Montréal ⁶Omelet
Paper Code
RL alignment results
Image generator alignment results

TL;DR

  • Problem: Existing alignment methods rely on tractable likelihoods, specific solvers, or model-specific formulation — which few-step generators (GAN, drifting, consistency, flow map) often do not provide.
  • Solution: Approximate the reward-tilted distribution $q^*(x) \propto p_\text{ref}(x) \exp(\beta \cdot r(x))$ via Stein Variational Gradient Descent (SVGD) with a KDE, and amortize the result of the SVGD transport into the generator.
  • Advantages:
    • General: Only sample access to the generator and $p_\text{ref}$ — no likelihood, solver, or architecture assumptions.
    • Scalable: Single recipe across 2D toy settings, 56 offline + 30 offline-to-online RL tasks, ImageNet-256, and 1024² text-to-image synthesis.
    • Fast: Amortization preserves few-step inference — 180–280× faster than inference-time alignment baselines.
    • Robust: Reference anchoring and the explicit diversity term in the Stein transport mitigate reward over-optimization.
    • Effective: New SOTA on RL tasks and Pareto frontier for reward–diversity and reward–preference trade-offs in image alignment.

Challenges

  • Generative model alignment is essential for visual content generation and effective robotic manipulation.
  • Existing methods rely on tractable denoising likelihoods, specific ODE/SDE solvers, or backbone-specific formulations.
  • These assumptions break down for modern few-step generative models, such as consistency models, flow maps, and drifting models.

Methods

FAV — (Few-step Generative Models Alignment via Sample-based Variational Inference)

FAV loss algorithm
FAV mechanism diagram
Illustration of Stein transport.
  • To sidestep restrictive assumptions, we design FAV to require only sample access to the generator and the reference distribution.
  • We set the target as a reward-tilted distribution anchored to a reference distribution:
\[q^*(x)\propto p_\text{ref}(x) \exp(r(x))\]
  • To approximate this target, we employ SVGD as a particle-based inference method: it transports samples from the implicit generator toward $q^*(x)$ without evaluating the likelihood of the generator.
  • The Stein transport vector field toward the reward-tilted distribution cleanly decomposes into three interpretable terms:
\[\phi^*_{q_\theta, q^*}(x) = \mathbb{E}_{x' \sim q_\theta} \Big[ \underbrace{k(x', x)\, \nabla_{x'} \log p_{\text{ref}}(x')}_{\text{Prior Alignment}} + \underbrace{\beta \cdot k(x', x)\, \nabla_{x'} r(x')}_{\text{Reward Guidance}} + \underbrace{\nabla_{x'} k(x', x)}_{\text{Diversity Enforcement}} \Big].\]
  • For the prior alignment term, which includes the intractable score $\nabla_{x’} \log p_{\text{ref}}(x’)$, we employ sample-based KDE with an RBF kernel of bandwidth $\sigma$.
  • This yields the following transport field, which only requires samples from the generator and the reference distribution:
\[\hat{\phi}_{q_\theta, q^*_\sigma}^*(x) = \mathbb{E}_{\substack{x' \sim q_\theta \\ x^{\text{ref}} \sim p_{\text{ref}}}} \left[ k_\sigma(x', x) \left( \frac{\tilde{k}_\sigma(x', x^{\text{ref}})(x^{\text{ref}} - x')}{\sigma^2} + \beta \cdot \nabla_{x'} r(x') + \frac{x - x'}{\sigma^2} \right) \right],\]
  • To preserve fast inference, we amortize the transport into the generator with a fixed-point regression loss:
\[\mathcal L(\theta)=\mathbb E_{\epsilon\sim \mathcal N}[||\text{stopgrad}(f_\theta(\epsilon)+\hat\phi^*_{q_\theta,q^*_\sigma }(f_\theta(\epsilon)))-f_\theta(\epsilon)||^2_2].\]
  • For high-dimensional inputs, FAV operates in the latent space of a pre-trained encoder, mitigating the curse of dimensionality.

📝 Takeaway. By amortizing sample-based variational inference with KDE-based nonparametric score estimation, FAV can align any few-step generator without structural assumptions.

Experiments - Toy Setting

2D toy experiments
FAV enables reward-tilted sampling across diverse pre-trained few-step backbones, from 1-step to 16-step.

Experiments - Reinforcement Learning

Offline RL results
Offline RL results.
Offline-to-online RL results
Offline-to-online RL results.
  • We directly pre-train FAV as a single-step generative policy, without a separate BC pre-training stage.
  • FAV achieves the best average performance across 56 offline RL tasks and 30 offline-to-online RL tasks.
  • FAV-Adaptive, which automatically sets the kernel bandwidth $\sigma$ using Scott’s rule, also outperforms the baselines.

Experiments - Image Generator Alignment

ImageNet main results
FAV achieves Pareto-frontier trade-offs on ImageNet 256 under aesthetic-score optimization.
ImageNet qualitative results
Qualitative results on ImageNet 256.
ImageNet non-differentiable reward
FAV-B optimizes gradient-free black-box rewards using a zeroth-order gradient estimator.
FAV vs test-time scaling
FAV preserves the fast generation capability of few-step generative models.
Sana qualitative results
FAV aligns a high-resolution text-to-image model toward human preference and safety objectives.

Citation

@article{lee2026aligning,
  title={Aligning Few-Step Generative Models by Amortizing Sample-based Variational Inference},
  author={Lee, Jaewoo and Kang, Hyeongyu and Kim, Dohyun and Sim, Kyuil and Shin, Woocheol and Kim, Minsu and Yun, Taeyoung and Lee, Jeongjae and Choi, Sanghyeok and Lee, Tabitha Edith and Ye, Jongchul and Park, Jinkyoo},
  journal={arXiv preprint arXiv:2605.26552},
  year={2026}
}