pith. sign in

arxiv: 2606.22958 · v1 · pith:LW3O66IInew · submitted 2026-06-22 · 💻 cs.LG · cs.CV

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

Ruolan Sun , Pawel Polak This is my paper

Pith reviewed 2026-06-26 08:58 UTC · model grok-4.3

classification 💻 cs.LG cs.CV
keywords inference-time alignmentdiffusion modelsflow-matchingMAP optimizationtext-to-image generationjoint optimizationpreference alignment
0
0 comments X

The pith

PG-MAP jointly optimizes text conditioning and image latents at inference time to align diffusion and flow-matching models.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper introduces PG-MAP as a training-free method that casts inference-time alignment as trajectory-level Gibbs-MAP and proximal energy optimization over both conditioning variables and latent states. It relies on a forward-consistency coupling to coordinate updates between these two modalities while staying compatible with diffusion and flow-matching transports. Experiments show consistent gains in PickScore and aesthetic metrics on diffusion backbones, with the method reducing to a latent-only form on flow-matching models that reaches 91.9 percent PickScore win rate. The approach can be paired with tuned classifier-free guidance and is supported by human preference evaluations over strong baselines.

Core claim

PG-MAP formulates inference-time alignment as a trajectory-level Gibbs-MAP / proximal energy optimization over the conditioning c and latent state z_t via a forward-consistency coupling, optionally guided by a frozen preference reward. This joint formulation enables coordinated updates across modalities while remaining compatible with both diffusion and flow-matching models through transport-specific adaptations.

What carries the argument

forward-consistency coupling that links conditioning and latent variables inside the joint proximal energy optimization

If this is right

  • Improves PickScore and Aesthetic metrics across diffusion backbones including SD 1.5 and SDXL.
  • Combines with tuned classifier-free guidance to reach the strongest overall performance.
  • Reduces to a latent-only variant on flow-matching models that records 91.9 percent PickScore and 75.7 percent HPS win rates against a static baseline.
  • Receives consistent human preference over tuned CFG and compute-matched universal guidance.
  • Oracle analysis indicates that the relative value of conditioning versus latent optimization varies by prompt type.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • A per-prompt selector that routes between conditioning and latent optimization could capture the headroom identified in the oracle analysis.
  • The transport-specific adaptations suggest the joint formulation may apply to additional generative model families beyond diffusion and flow-matching.

Load-bearing premise

The forward-consistency coupling enables coordinated joint updates across conditioning and latent variables without introducing new artifacts or inconsistencies.

What would settle it

An ablation that removes the forward-consistency coupling while preserving all other optimization components and shows the performance gains disappear would test the central mechanism.

Figures

Figures reproduced from arXiv: 2606.22958 by Pawel Polak, Ruolan Sun.

Figure 1
Figure 1. Figure 1: Joint PG-MAP exercises both axes at once on SDXL (same seed within each pair). Side annotations identify the per-prompt c- and z-side gains; zoom-in boxes mark them. Population-scale PartiPrompts win rates: Tab. 1; trajectory-level mechanism [PITH_FULL_IMAGE:figures/full_fig_p002_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: PG-MAP trajectory analysis on SDXL (50 DDIM, same seed within each row). Two specializations of Jt target different failure modes: (a)/(c) MAP-c moves only c to fix prompt align￾ment; (b)/(d) Reward-z moves only zt to lift perceptual quality. The opposite slopes of (c) (growing) and (d) (decaying) are a concrete signature of the non-stationary objective and the asymmetric, schedule-adaptive prior design (§… view at source ↗
Figure 3
Figure 3. Figure 3: 4× zoom on the maximum-diff 128×128 patch from the highest-HPS-gain prompt (“The Statue of Liberty in Minecraft”). Left: baseline (SD3.5 static). Center: UG-FM. Right: |UG￾baseline| × 8 amplified intensity heatmap. The perturbation localizes in textured / shaded regions, visible at 4–8× zoom. G Related work landscape and extended limitations G.1 Inference-time alignment landscape comparison A comparison-ma… view at source ↗
read the original abstract

Inference-time alignment of pretrained text-to-image models is typically performed along a single control axis, such as classifier-free guidance, attention editing, or reward-based latent perturbations. This limitation prevents modeling joint dependencies between conditioning and latent variables and hinders transfer across generative transports. We propose PG-MAP, a training-free framework that formulates inference-time alignment as a trajectory-level Gibbs-MAP / proximal energy optimization over the conditioning $c$ and latent state $z_t$ via a forward-consistency coupling, optionally guided by a frozen preference reward. This joint formulation enables coordinated updates across modalities while remaining compatible with both diffusion and flow-matching models through transport-specific adaptations. Across diffusion backbones (SD~1.5, SDXL), PG-MAP consistently improves alignment metrics such as PickScore and Aesthetic, and can be effectively combined with tuned classifier-free guidance to achieve the strongest overall performance. On flow-matching models (SD3.5-medium), the framework reduces to a latent-only variant, achieving $\mathbf{91.9\%}$ PickScore and $75.7\%$ HPS win rates against a static baseline, with controlled experiments ruling out noise-related artifacts. Human evaluations further confirm consistent preference over strong baselines, including tuned CFG and compute-matched universal guidance. Finally, an oracle-routing analysis shows that the relative importance of conditioning and latent optimization depends on prompt types, surfacing further headroom that a per-prompt selector could exploit.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

1 major / 0 minor

Summary. The manuscript presents PG-MAP, a training-free framework for inference-time alignment of diffusion and flow-matching text-to-image models. It formulates the alignment as joint Gibbs-MAP / proximal energy optimization over conditioning c and latent z_t using forward-consistency coupling, with transport-specific adaptations. The paper reports metric improvements on diffusion backbones (SD 1.5, SDXL) and high win rates on flow-matching (SD3.5-medium) with a latent-only variant, supported by human evaluations and an oracle-routing analysis.

Significance. Should the joint optimization prove effective across transports, the work would offer a notable contribution to inference-time alignment methods by enabling coordinated updates between conditioning and latent variables. Strengths include the human evaluations confirming preferences over baselines like tuned CFG and the oracle analysis highlighting prompt-dependent optimization importance. The reduction to latent-only on flow-matching, however, limits the demonstrated scope of the joint formulation.

major comments (1)
  1. [Abstract] Abstract: The central claim is that PG-MAP enables joint MAP optimization over both conditioning c and latent z_t via forward-consistency coupling. However, the abstract states that on flow-matching models (SD3.5-medium) the framework reduces to a latent-only variant while reporting 91.9% PickScore and 75.7% HPS win rates against a static baseline. This reduction means the headline results cannot be attributed to the joint conditioning-latent mechanism or the coupling term, and the controlled experiments (which only rule out noise artifacts) do not test whether the proximal energy or coupling functions for flow-matching.

Simulated Author's Rebuttal

1 responses · 0 unresolved

We thank the referee for identifying the need to clarify the scope of the joint optimization claim. The manuscript explicitly notes the reduction to latent-only on flow-matching models; we will revise the abstract and related sections to ensure the attribution of results is unambiguous while preserving the framework's transport-specific design.

read point-by-point responses

  1. Referee: [Abstract] Abstract: The central claim is that PG-MAP enables joint MAP optimization over both conditioning c and latent z_t via forward-consistency coupling. However, the abstract states that on flow-matching models (SD3.5-medium) the framework reduces to a latent-only variant while reporting 91.9% PickScore and 75.7% HPS win rates against a static baseline. This reduction means the headline results cannot be attributed to the joint conditioning-latent mechanism or the coupling term, and the controlled experiments (which only rule out noise artifacts) do not test whether the proximal energy or coupling functions for flow-matching.

    Authors: We agree that the abstract should more clearly separate the diffusion results (where joint conditioning-latent optimization is active) from the flow-matching results (where the framework reduces to latent-only). The reduction arises because the forward-consistency coupling and proximal energy terms for conditioning updates are transport-specific; on SD3.5-medium the conditioning component yields no additional gain, so the implementation drops it while retaining the latent MAP step. The headline win rates therefore reflect the latent-only variant, as already stated in the text. The controlled experiments on flow-matching were designed only to rule out sampling noise, not to validate the (unused) coupling functions. We will revise the abstract to foreground this distinction, move the flow-matching numbers into a dedicated sentence, and add a short methods paragraph explaining why the joint terms are inactive for this transport. The diffusion results (SD 1.5, SDXL) continue to demonstrate the full joint formulation. revision: yes

Circularity Check

0 steps flagged

No circularity; framework and results presented as independent empirical formulation

full rationale

The provided abstract introduces PG-MAP as a novel training-free joint MAP optimization framework using forward-consistency coupling, with transport-specific adaptations for diffusion and flow-matching models. No equations, derivations, or self-citations are shown that reduce claimed improvements (e.g., PickScore gains) to fitted inputs, self-definitions, or prior author results by construction. Results on SD 1.5, SDXL, and SD3.5-medium are reported as experimental outcomes from controlled tests, not as predictions forced by the method's own parameters. The reduction to latent-only on flow-matching is noted but does not create a definitional loop in the derivation chain itself. The paper's central claims remain self-contained against external benchmarks.

Axiom & Free-Parameter Ledger

0 free parameters · 0 axioms · 0 invented entities

Only the abstract is available; no free parameters, axioms, or invented entities are described in sufficient detail to populate the ledger.

pith-pipeline@v0.9.1-grok · 5789 in / 1251 out tokens · 23983 ms · 2026-06-26T08:58:30.791218+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

18 extracted references · 16 canonical work pages · 4 internal anchors

  1. [1]

    Jinho Chang, Jaemin Kim, and Jong Chul Ye

    arXiv:2402.14017. Jinho Chang, Jaemin Kim, and Jong Chul Ye. Training-free reward-guided image editing via trajectory optimal control. InInternational Conference on Learning Representations,

    work page arXiv

  2. [2]

    Training-Free Reward-Guided Image Editing via Trajectory Optimal Control

    arXiv:2509.25845. Hila Chefer, Yuval Alaluf, Yael Vinker, Lior Wolf, and Daniel Cohen-Or. Attend-and-excite: Attention-based semantic guidance for text-to-image diffusion models.ACM Transactions on Graphics, 42(4),

    work page internal anchor Pith review Pith/arXiv arXiv

  3. [3]

    arXiv:2505.18547. Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, Dustin Podell, Tim Dockhorn, Zion English, Kyle Lacey, Alex Goodwin, Yannik Marek, and Robin Rombach. Scaling rectified flow trans- formers for high-resolution image synthesis. InInternational ...

    work page arXiv

  4. [4]

    Ruiqi Feng, Chenglei Yu, Wenhao Deng, Peiyan Hu, and Tailin Wu

    arXiv:2406.04312. Ruiqi Feng, Chenglei Yu, Wenhao Deng, Peiyan Hu, and Tailin Wu. On the guidance of flow matching. InInternational Conference on Machine Learning,

    work page arXiv

  5. [5]

    Rinon Gal, Yuval Alaluf, Yuval Atzmon, Or Patashnik, Amit H

    arXiv:2502.02150. Rinon Gal, Yuval Alaluf, Yuval Atzmon, Or Patashnik, Amit H. Bermano, Gal Chechik, and Daniel Cohen-Or. An image is worth one word: Personalizing text-to-image generation using textual inversion. InInternational Conference on Learning Representations,

    work page arXiv

  6. [6]

    Prompt-to-Prompt Image Editing with Cross Attention Control

    Amir Hertz, Ron Mokady, Jay Tenenbaum, Kfir Aberman, Yael Pritch, and Daniel Cohen-Or. Prompt- to-prompt image editing with cross attention control.arXiv preprint arXiv:2208.01626,

    work page internal anchor Pith review Pith/arXiv arXiv

  7. [7]

    Classifier-Free Diffusion Guidance

    Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance.arXiv preprint arXiv:2207.12598,

    work page internal anchor Pith review Pith/arXiv arXiv

  8. [8]

    Peter Holderrieth, Uriel Singer, Tommi Jaakkola, Ricky T. Q. Chen, Yaron Lipman, and Brian Karrer. GLASS flows: Transition sampling for alignment of flow and diffusion models.arXiv preprint arXiv:2509.25170,

    work page arXiv

  9. [9]

    Yidong Ouyang, Liyan Xie, Hongyuan Zha, and Guang Cheng

    arXiv:2510.23974. Yidong Ouyang, Liyan Xie, Hongyuan Zha, and Guang Cheng. Alignment of diffusion model and flow matching for text-to-image generation.arXiv preprint arXiv:2602.00413,

    work page arXiv

  10. [10]

    10 Jiangweizhi Peng et al

    arXiv:2412.00100. 10 Jiangweizhi Peng et al. Safeguarding text-to-image generation via inference-time prompt-noise optimization.arXiv preprint arXiv:2412.03876,

    work page arXiv

  11. [11]

    DreamBooth: Fine tuning text-to-image diffusion models for subject-driven generation.arXiv preprint arXiv:2208.12242,

    Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, and Kfir Aberman. DreamBooth: Fine tuning text-to-image diffusion models for subject-driven generation.arXiv preprint arXiv:2208.12242,

    work page arXiv

  12. [12]

    arXiv preprint arXiv:2405.18881 (2024)

    arXiv:2405.18881. Bram Wallace, Meihua Dang, Rafael Rafailov, Linqi Zhou, Aaron Lou, Senthil Purushwalkam, Stefano Ermon, Caiming Xiong, Shafiq Joty, and Nikhil Naik. Diffusion model alignment using direct preference optimization. InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,

    work page arXiv

  13. [13]

    Human Preference Score v2: A Solid Benchmark for Evaluating Human Preferences of Text-to-Image Synthesis

    Xiaoshi Wu, Yiming Hao, Keqiang Sun, Yixiong Chen, Feng Zhu, Rui Zhao, and Hongsheng Li. Human preference score v2: A solid benchmark for evaluating human preferences of text-to-image synthesis.arXiv preprint arXiv:2306.09341,

    work page internal anchor Pith review Pith/arXiv arXiv

  14. [14]

    arXiv:2303.09833. A Mathematical foundations A.1 Reward models, gradient derivations, and reward chain rule Reward models.PickScore [Kirstain et al., 2023] is a CLIP-based scorer trained on the Pick-a-Pic dataset of human pairwise preferences ( 500k pairs). HPS v2 [Wu et al., 2023] similarly trains on hu- man preference data with an improved encoder; Imag...

    work page arXiv 2023

  15. [15]

    forward-consistency residual

    adds 11 text-faithfulness annotations on top of preference labels. The LAION aesthetic predictor [Schuhmann et al., 2022] is a small MLP head over CLIP features regressed against curated aesthetic ratings. All four are publicly available frozen models that accept an image x and prompt y and return a scalar Q(x, y)∈R, differentiable with respect to the ima...

    2022

  16. [16]

    Suppose at this point ∥Jc∥op ≤L c, ∥Jz∥op ≤L z, ∥rt∥ ≤R t, ∥∇cQ∥ ≤G Q c , ∥∇zt Q∥ ≤G Q z

    Proposition 2(Local stationary-point displacement bound).Let (c⋆ t , z⋆ t ) be an interior stationary point of Jt. Suppose at this point ∥Jc∥op ≤L c, ∥Jz∥op ≤L z, ∥rt∥ ≤R t, ∥∇cQ∥ ≤G Q c , ∥∇zt Q∥ ≤G Q z . Then ∥c⋆ t −µ t∥ ≤σ 2 c √at|s Lc Rt βt|s +λG Q c ,(8) ∥z⋆ t −z ddim t ∥ ≤σ z(t)2 (1+√at|sLz)Rt βt|s +λG Q z .(9) Proof. From the stationary fixed-point...

    work page arXiv 2088

  17. [17]

    Robustness on HPDv2

    Variation across λ is bounded by ≤1.0 pp on every metric, within bootstrap CI; the headline retainsλ=0.05as the default.Bold= highest in column. λPickScore HPS Aesthetic CLIPScore 0(MAP-cz)56.7% 47.5% 55.6% 48.8% 0.05(default) 56.4% 47.1% 56.2% 48.1% 0.157.7%47.9% 56.1%49.6% 0.2 56.0% 46.9%56.4%49.6% C.2 NoiseZoo: variance decomposition of DDIM-inverted S...

    work page arXiv 2024

  18. [18]

    The Statue of Liberty in Minecraft

    Figure 3: 4× zoom on the maximum-diff 128×128 patch from the highest-HPS-gain prompt (“The Statue of Liberty in Minecraft”). Left: baseline (SD3.5 static). Center: UG-FM. Right: |UG- baseline| ×8 amplified intensity heatmap. The perturbation localizes in textured / shaded regions, visible at4–8×zoom. G Related work landscape and extended limitations G.1 I...

    2023