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

Researchers introduce FAV, a novel framework for aligning few-step generative models that requires only sample access to generators and reference distributions. The method uses Stein Variational Gradient Descent to cast alignment as sampling from reward-tilted distributions, demonstrating superior performance across robotic manipulation tasks and scaling to high-resolution image synthesis.

FAV represents a significant advancement in generative model alignment by removing restrictive assumptions that have constrained prior approaches. Traditional alignment frameworks require tractable likelihoods, specific solver architectures, or predetermined model families—limitations that substantially narrow their applicability. This research bypasses these constraints through a sample-based variational inference methodology, enabling alignment across diverse architectures including GANs, consistency models, and flow maps.

The framework's core innovation lies in amortizing Stein Variational Gradient Descent particle updates directly into generator parameters via fixed-point regression. This technical approach elegantly transforms the alignment problem into one of sampling from reward-tilted distributions, making the methodology broadly applicable. The distinction between sample-based access and model-specific requirements addresses a fundamental bottleneck in generative AI development.

Empirical validation spans two distinct domains with noteworthy results. In robotic manipulation, FAV surpasses existing policy extraction baselines across 86 combined offline and offline-to-online reinforcement learning tasks, demonstrating practical utility for robotics applications where alignment quality directly impacts task performance. For image generation, the framework successfully fine-tunes multiple few-step architectures and scales to high-resolution synthesis at 1024² resolution, indicating its versatility across generative domains.

The availability of open-source code accelerates potential adoption within the research community. This work addresses a genuine technical gap in making advanced alignment techniques accessible to practitioners using diverse model architectures. The general-purpose nature of FAV positions it as a foundational tool for future generative model development, particularly as few-step models become increasingly prevalent due to their computational efficiency advantages.

• →FAV enables alignment of few-step generative models without requiring tractable likelihoods or specific solver architectures.
• →The framework outperforms baseline methods on 86 robotic manipulation tasks, demonstrating practical reinforcement learning improvements.
• →Successfully scales to high-resolution image synthesis at 1024² across multiple backbone architectures including GANs and flow models.
• →Sample-based variational inference approach removes structural constraints that previously limited alignment method applicability.
• →Open-source implementation facilitates broader adoption and integration into existing generative model development pipelines.