🤖 machine learning

PG-MAP: Joint MAP Optimization for Inference-Time Alignment of Diffusion and Flow-Matching Models

The paper introduces PG-MAP, a training-free framework that enhances inference-time alignment for both diffusion and flow-matching models by formulating the process as a joint Gibbs-MAP optimization over conditioning and latent variables, thereby achieving superior performance across various metrics and human evaluations compared to existing single-axis methods.

Original authors: Ruolan Sun, Pawel Polak

Published 2026-06-23

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ruolan Sun, Pawel Polak

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a master chef trying to cook a perfect dish based on a customer's order. In the world of AI image generation, the "chef" is a complex model (like Stable Diffusion), the "order" is the text prompt (e.g., "a red panda astronaut"), and the "dish" is the final image.

Usually, when the chef makes a mistake, you can only fix it in one way at a time:

Change the order: You rewrite the text prompt to be more specific (like telling the chef, "Make sure the panda is red").
Tweak the ingredients: You adjust the raw mixture while it's cooking to fix the texture or lighting (like adding more salt or adjusting the heat).

Existing methods force you to choose one of these paths. If you fix the text, you might ruin the texture. If you fix the texture, you might lose the meaning of the prompt.

PG-MAP is a new "super-chef assistant" that does both at the same time, dynamically, while the dish is being cooked.

The Core Idea: A Dynamic Dance

The paper introduces PG-MAP (Preference-Guided Adaptive MAP). Instead of making a single decision at the start or tweaking things once at the end, PG-MAP treats the image creation process as a continuous journey.

Think of the image generation process like a long, winding road trip from a foggy starting point (random noise) to a clear destination (the final image).

The Problem: On this road, the "fog" (noise) is thick at the beginning and clears up at the end. Existing methods try to steer the car using a static map or a single adjustment.
The PG-MAP Solution: PG-MAP acts like a GPS that constantly recalculates the route step-by-step. It looks at the current state of the image and asks two questions simultaneously:
1. "Is the text description still matching what we are seeing?" (Conditioning, or the c-side).
2. "Does the visual quality look good?" (Latent state, or the z-side).

It adjusts the "text description" and the "visual ingredients" together, ensuring they work in harmony rather than fighting each other.

How It Works: The "Forward-Consistency" Rule

The paper uses a fancy term called "forward-consistency coupling." Here is a simple analogy:

Imagine you are walking through a dark forest (the generation process). You have a flashlight (the AI model) that shows you where you are right now.

Old methods would say, "I know where I started, so I'll just keep walking in a straight line," or "I'll just adjust my path once at the end."
PG-MAP says, "Let's look at where the flashlight is pointing right now. If the light suggests we are drifting off the path, we will gently nudge our direction (the text) and our footing (the image details) at the same time to get back on track."

It does this by calculating a "score" (a reward) for every tiny step of the journey. If the image looks like a "red panda" but the fur looks blurry, it sharpens the fur. If the fur looks great but the animal looks like a human, it tweaks the text embedding to force the panda shape back.

Two Different Types of Roads

The paper makes a fascinating discovery: the best way to drive depends on the type of road you are on.

Diffusion Models (The bumpy, foggy road): On older models (like SD 1.5 or SDXL), the road is full of noise. PG-MAP works best here by adjusting both the text and the image details early in the process when the fog is thick.
Flow-Matching Models (The smooth, straight highway): On newer, faster models (like SD3.5), the road is much smoother. Here, PG-MAP realizes that tweaking the text is unnecessary (the road is too stable). Instead, it focuses entirely on fine-tuning the image details at the very end of the trip.

This adaptability is unique; most other tools try to use the same strategy for both types of roads.

The Results: Better Pictures, Less Guesswork

The authors tested this on thousands of prompts.

Visual Quality: The images were sharper, had better lighting, and more accurate details (like a sword grip or a feather texture).
Prompt Accuracy: The images matched the text descriptions better (e.g., the "red panda" actually looked like a panda, not a human).
Human Preference: When real people were asked to choose between the old way and PG-MAP, they picked PG-MAP about 60% to 67% of the time.

The "Oracle" Secret

The paper also ran a "what-if" experiment. They asked: "If we could magically know the perfect way to fix every single specific prompt, how much better could we get?"
They found that different types of prompts need different fixes.

Short, specific prompts (like "a red cube") need more help with the text.
Atmospheric prompts (like "a sunset over the ocean") need more help with the visual texture.

PG-MAP is a single tool that handles both, but the paper suggests that in the future, a smart "traffic controller" could automatically choose the best setting for each specific request, potentially making the images even better.

Summary

PG-MAP is a training-free tool that makes AI image generators smarter by letting them adjust both the meaning of the prompt and the look of the image simultaneously, step-by-step, as the image is being created. It adapts its strategy based on the type of AI model it's using, resulting in images that are both more accurate to the text and more beautiful to look at.

Technical Summary: PG-MAP: Joint MAP Optimization for Inference-Time Alignment of Diffusion and Flow-Matching Models

Problem Statement
Current inference-time alignment methods for pretrained text-to-image models (Diffusion and Flow-Matching) typically operate along a single control axis. Conditioning-side methods (e.g., attention editing, prompt embedding learning) modify the text embedding $c$ , while latent-side methods (e.g., reward-guided latent perturbations) modify the latent state $z_t$ . Existing frameworks fail to model the joint dependencies between conditioning and latent variables, treating updates as additive rather than coordinated. Furthermore, these methods are often static, fixing control strategies at the start of the trajectory ( $z_T$ ) or offline, ignoring the dynamic information needs of the denoiser as the trajectory evolves from high-noise (global layout) to low-noise (local detail) regimes. This limitation hinders the transfer of alignment techniques across different generative transports (Diffusion vs. Flow-Matching).

Methodology: PG-MAP
The authors propose PG-MAP (Preference-Guided Adaptive MAP), a training-free framework that formulates inference-time alignment as a trajectory-level Gibbs-MAP / proximal energy optimization. The core innovation is a forward-consistency coupling that jointly optimizes the conditioning $c$ and the latent state $z_t$ at every denoising step.

Objective Function: The framework defines a step-dependent energy function $J_t(c, z_t)$ comprising three components:
1. Forward-Consistency Residual: A term $\ell_t(c, z_t)$ that penalizes the deviation between the current latent $z_t$ and the model's own prediction of the previous step's state, ensuring the trajectory remains consistent with the denoiser's forward kernel.
2. Gaussian Anchoring Priors: Terms $R_c(c)$ and $R_z(z_t)$ that anchor the optimized variables to their unperturbed values ( $c_0$ and the DDIM trajectory point $z^{ddim}_t$ ). Crucially, the latent prior scale $\sigma_z(t)$ is schedule-adaptive, tracking the marginal noise scale of the diffusion process ( $\sigma_z(t) = \gamma\sqrt{1-\bar{\alpha}_t}$ ) to provide a scale-invariant trust region that tightens as the trajectory approaches the data manifold.
3. Preference Reward Tilt: A term $\lambda Q(\hat{x}_0, y)$ that guides the optimization toward human preference (e.g., PickScore, HPS) based on the model's per-step preview of the clean image.
Adaptive Active Set: The framework employs a step-dependent active set $A_t \subseteq \{c, z_t\}$ and a reward gate $\lambda_t$ .
- Diffusion Models (SD 1.5, SDXL): The framework jointly refines both $c$ and $z_t$ , primarily in the high-noise window, allowing the conditioning and latent variables to inform each other as the trajectory unfolds.
- Flow-Matching Models (SD3.5): Due to the deterministic nature of the ODE and the high capacity of the concatenated text encoders, the framework reduces to a latent-only variant (UG-FM). Here, the active set collapses to $\{z_t\}$ restricted to the data-side window (low noise), as conditioning updates are found to be ineffective and noise-side perturbations are overly amplified by the ODE dynamics.
Optimization: The refined pair $(c^\star_t, z^\star_t)$ is obtained via a small number of joint ascent steps (typically $K=2$ ) using gradients derived from the energy function. This process is composable with existing methods like Classifier-Free Guidance (CFG), as CFG modifies the vector field while PG-MAP moves the query point under a fixed vector field.

Key Contributions

Joint (c, $z_t$ ) MAP Framework: The first inference-time framework to couple conditioning and latent variables through the denoiser's forward kernel, simultaneously addressing composition failures (c-side) and texture failures (z-side).
Unified Objective: The objective $J_t$ analytically recovers prior single-axis methods (MAP-c, Reward-z) and Universal Guidance as special cases, while CFG remains a composable external mechanism rather than a special case.
Schedule-Adaptive Optimization: Introduces a time-dependent trust region and active set that adapts to the noise schedule, loosening constraints at high noise and tightening them near the data end.
Transport-Dependent Analysis: Demonstrates that the optimal active set depends on the generative transport. Diffusion benefits from joint refinement at high noise, whereas Flow-Matching benefits from latent-only refinement at the data end, a hypothesis supported by local perturbation analysis and empirical validation.

Results

Diffusion Backbones (SD 1.5, SDXL): PG-MAP consistently improves alignment metrics. On SDXL, it achieves a 56.4% PickScore win rate and 56.2% Aesthetic win rate against a static baseline. Combining PG-MAP with tuned CFG yields the strongest overall performance, reaching 64.6% HPS and 56.5% Aesthetic.
Flow-Matching Backbone (SD3.5-medium): The latent-only variant (UG-FM) achieves 91.9% PickScore and 75.7% HPS win rates against the static baseline. Controlled experiments rule out noise-related artifacts, confirming the gains are due to structured, low-frequency latent perturbations.
Human Evaluation: On SDXL, PG-MAP is preferred over static baselines, tuned CFG, and compute-matched Universal Guidance in pairwise human evaluations (60–67% win rates).
Oracle Routing: An oracle analysis reveals that different prompt types (e.g., attribute-binding vs. atmospheric scenes) benefit from different optimization variants. A per-prompt selector could theoretically improve metrics by an additional 5–14 percentage points, indicating significant headroom for future routing mechanisms.

Significance and Claims
The paper claims that PG-MAP represents a shift from static, single-axis guidance heuristics toward dynamic, trajectory-aware optimization as a default design principle for inference-time alignment. By formulating alignment as a proximal energy optimization with forward-consistency coupling, the framework enables coordinated updates across modalities that are compatible with both diffusion and flow-matching transports. The authors emphasize that the "load-bearing" ingredients are the joint coupling and non-stationary scheduling, rather than simply larger step sizes or stronger reward signals. The work provides a unified theoretical and empirical foundation for understanding how conditioning and latent variables interact during the generation process, offering a flexible, training-free approach to improve generative model alignment.