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

PG-MAP (Preference-Guided Adaptive MAP) addresses the limitations of existing inference-time alignment methods, which typically optimize either the conditioning (text embeddings) or the latent state in isolation. The authors formulate a unified, training-free objective that treats both the conditioning $c$ and the latent state $z_t$ as variables to be optimized at each denoising step. This is achieved through a forward-consistency coupling that links the two modalities via the model's own denoising kernel, allowing for coordinated updates that adapt to the specific transport dynamics of either diffusion or flow-matching models.

The framework introduces a schedule-adaptive trust region and a transport-dependent active set. By re-optimizing $c$ and $z_t$ at every step, the model can refine global layout and local texture simultaneously. For diffusion models, the framework performs full joint refinement. For flow-matching models (like SD3.5-medium), the authors demonstrate that the optimal strategy reduces to a latent-only variant (UG-FM) due to the high dimensionality of the conditioning representation and the specific Jacobian properties of the flow-matching ODE.

Across multiple backbones (SD 1.5, SDXL, and SD3.5-medium), PG-MAP consistently outperforms static baselines and existing single-axis methods on metrics such as PickScore, Aesthetic, and HPS. Human evaluations confirm that these automated gains translate into a clear perceptual preference, with PG-MAP winning significantly against compute-matched baselines. The authors also provide an oracle-routing analysis, suggesting that the relative importance of conditioning versus latent optimization is prompt-dependent, indicating further potential for per-prompt selection strategies.

This work provides a unified, flexible framework for inference-time control that works across different generative architectures without requiring model retraining. By demonstrating that joint optimization is more effective than single-axis control, it offers a robust tool for practitioners to steer text-to-image models toward higher aesthetic and alignment standards using only a frozen reward signal.