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

Jaewoo Lee*^1,2, Hyeongyu Kang*¹, Dohyun Kim¹, Kyuil Sim¹, Woocheol Shin¹, Minsu Kim^1,3, Taeyoung Yun¹, Jeongjae Lee¹, Sanghyeok Choi⁴, Tabitha Edith Lee^3,5, Jong Chul Ye†¹, Jinkyoo Park†^1,6

¹KAIST   ²MongooseAI   ³Mila – Quebec AI Institute   ⁴University of Edinburgh   ⁵Université de Montréal   ⁶Omelet

_{* Equal contribution    † Corresponding author}

Abstract

Aligning a few-step generative model is challenging, since existing alignment frameworks typically rely on restrictive assumptions: a tractable likelihood, a specific ODE/SDE solver, or a particular model family. We introduce FAV (Few-step Generative Models Alignment via Sample-based Variational Inference), a general alignment framework that requires only sample access to the generator and the reference distribution. We cast alignment as sampling from a reward-tilted distribution anchored to a reference distribution. We leverage Stein Variational Gradient Descent as a sample-based variational inference scheme and amortize its particle updates into the generator parameters via fixed-point regression. We evaluate FAV on two domains: robotics manipulation and image generator alignment. On generative policy alignment for robotic manipulation, FAV outperforms prevailing policy extraction baselines across 56 offline and 30 offline-to-online RL tasks. For image generator alignment, FAV fine-tunes diverse few-step backbones, including GAN, drifting model, consistency models, and flow maps, scaling from ImageNet-256 to 1024² text-to-image synthesis.

Tasks

Task	Directory
Reinforcement Learning (offline & offline-to-online)	fav-offrl/
2D Toy Setting	fav-toy/
Conditional Image Generation	fav-conditional-image/
Text-to-Image	fav-text-to-image/

Citation

If you find this repository helpful, please cite our work:

@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, Jong Chul and Park, Jinkyoo},
  journal={arXiv preprint arXiv:2605.26552},
  year={2026}
}