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

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 parameters of the generator 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. On image generator alignment, FAV fine-tunes diverse few-step backbones, including GAN, drifting models, consistency models, and flow maps, scaling from ImageNet-256 to 1024 text-to-image synthesis.