Hierarchical Variational Policies for Reward-Guided Diffusion

2026-05-23 · Source: cs.LG updates on arXiv.org · Field: Technology & Digital — Artificial Intelligence & Machine Learning, Computer Vision · Depth: Expert, extended

Summary

The paper introduces Hierarchical Variational Policies (HVP), a principled framework designed to adapt pretrained diffusion models for downstream objectives, such as inverse problems, while substantially reducing inference costs. This approach formulates test-time adaptation as a hierarchical variational model, amortizing control into a lightweight yet expressive stochastic policy. HVP naturally supports few-step diffusion sampling, maintaining high sample quality even with large step sizes. The framework includes two main variants: Amortized HVP (AHVP), a fully amortized sampler that achieves a strong quality–speed tradeoff, and Semi-Amortized HVP (SHVP), which combines amortized proposals with limited test-time optimization to reach state-of-the-art perceptual quality. For instance, on 4x super-resolution, AHVP delivers better perceptual quality with over 5x faster inference compared to leading baselines.

Key takeaway

For machine learning engineers developing diffusion-based solutions for inverse problems, you should evaluate Hierarchical Variational Policies (HVP) to significantly reduce inference costs while maintaining or improving sample quality. AHVP offers over 5x faster inference for tasks like 4x super-resolution, making it ideal for real-time or high-throughput applications. If your project demands the absolute best perceptual quality, SHVP provides state-of-the-art results with only a modest increase in computational expense.

Key insights

Hierarchical Variational Policies amortize diffusion model guidance, enabling faster, high-quality sampling for inverse problems via learned stochastic control.

Principles

• Amortizing guidance shifts computational cost from inference to a one-time training phase.
• Hierarchical variational models with latent variables enhance policy flexibility for multimodal posteriors.
• Stochastic policies offer more flexible control than deterministic ones, especially for complex datasets.

Method

The framework involves a two-stage procedure: first, learn an initial noise distribution maximizing reward, then train per-step stochastic controllers to capture residual structure. This replaces expensive inner-loop optimization with a single forward pass.

In practice

• Apply AHVP for inverse problems requiring fast, high-quality image generation.
• Use SHVP when state-of-the-art perceptual quality is paramount, accepting moderate additional cost.
• Consider two-stage training to recover both coarse structure and fine-grained details.

Topics

• Diffusion Models
• Variational Inference
• Stochastic Control
• Inverse Problems
• Image Super-Resolution
• Amortized Inference

Code references

• czi-ai/oc-guidance

Best for: Research Scientist, AI Scientist, Machine Learning Engineer, Computer Vision Engineer

Related on AIssential

• Stage-wise Distortion-Perception Traversal in Zero-shot Inverse Problems with Diffusion Models — A stage-wise diffusion model framework enables flexible distortion-perception tradeoff traversal in zero-shot inverse problems.
• Variational Test-time Optimization for Diffusion Synchronization — A new optimal control framework enables principled, test-time diffusion synchronization for collaborative generation without retraining.
• Fast and Robust Likelihood-Guided Diffusion Posterior Sampling with Amortized Variational Inference — Amortized variational inference improves diffusion posterior sampling efficiency while retaining robustness to unseen degradations.
• Estimation of instrument and noise parameters for inverse problem based on prior diffusion model — A Gibbs-Diffusion Posterior Sampling (G-DPS) based method efficiently estimates observation parameters and images in inverse problems with diffusion priors.
• From Noise to Control: Parameterized Diffusion Policies — PDP transforms diffusion from stochastic diversity into a precise, optimizable tool for behavior steering via a learned manifold.
• Large-scale empirical tuning and comparison of default optimizers for variational inference — A selection of 5 adaptive stochastic optimizers reliably approximates optimal performance for black-box variational inference without extensive tuning.

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by cs.LG updates on arXiv.org.