PaperlayerView exact revision on arXiv

Source-linked AI summary

Posterior sampling by source-space MCMC via prior-based few-step transport maps

Hoang Phuc Hau Luu, Marcelo Hartmann, Zhongjian Wang

arXiv:2610.01034v1stat.MLcs.AIcs.LG

TL;DR

Sample-only priors make posterior MCMC difficult because their densities and scores are unavailable, and guarantees for learned priors do not by themselves cover fixed few-step maps. The paper uses iMF transports for source-space sampling and establishes posterior stability bounds; its experiments report efficient synthetic sampling and CLIP-guided ImageNet steering, while hybrid SPT improves synthetic efficiency and one-step iMF reduces ImageNet sampling costs relative to multi-step SiT.

  • Problem

    Sample-only priors can represent complex structures but lack directly available densities and scores for posterior MCMC, while posterior stability for learned transports remains an open guarantee.

    Source: §1 · pp. 1–2

  • Method

    The paper represents sample-only priors with one- or few-step iMF maps and samples in Gaussian source space using SPT with pCN and optional split HMC.

    Source: §1 · pp. 2–3

  • Results

    Hybrid SPT improves efficiency on synthetic benchmarks, while one-step iMF reduces CLIP-guided ImageNet sampling costs relative to multi-step SiT.

    Source: §5 · p. 10

  • Takeaways & Limitations

    The framework supports posterior approximation for synthetic settings and steering a pretrained image prior toward text-specified preferences.

    Source: §1 · p. 3

    For original MeanFlow, the parameter-dependent objective gap prevents the paper’s training-suboptimality and approximation-error decomposition in general.

    Source: §3.2 · p. 6

Abstract

from arXiv · show

Bayesian inference increasingly uses informative but implicit priors represented only by samples, such as historical ensembles, simulator outputs, and pretrained generative models. The same computational problem appears in the test-time guidance task (generalized Bayes), where an explicit positive weight, e.g., an exponentiated reward, tilts an implicit prior. We develop a framework for source-space generalized Bayesian inference that combines inexpensive few-step prior transports with posterior stability guarantees. Specifically, we represent the prior using a one- or few-step improved MeanFlow (iMF) map and perform posterior sampling in its Gaussian source space. We establish Wasserstein error bounds between the exact and learned posteriors in terms of the joint population iMF and auxiliary-velocity loss, decomposed into training suboptimality and model-class approximation error. In the iMF source space, we adopt parallel tempering with preconditioned Crank-Nicolson updates and introduce a hybrid variant that incorporates split Hamiltonian Monte Carlo to improve sampling efficiency. Synthetic experiments show that the proposed framework can approximate posterior distributions accurately and efficiently, while CLIP-guided ImageNet experiments demonstrate its ability to steer a pretrained iMF image prior toward text-specified preferences.

1 Introduction

The paper addresses inference with sample-only priors, whose densities and scores are unavailable, and the unresolved question of how transport-learning error affects posterior accuracy. It proposes few-step iMF source-space inference with stability guarantees and MCMC methods for generalized Bayesian tasks.

1 Introduction

  • Sample-only priors can encode complex structures, but their unavailable densities and scores obstruct posterior MCMC; generalized Bayes also covers reward-weighted tasks.For example, a text prompt can define a reward through CLIP similarity.Source: §1 · pp. 1–2
  • Source-space inference represents the learned data prior through a Gaussian reference and samples its posterior, though informative weights can restore difficult geometry.The likelihood or weight may reintroduce curvature, stiffness, anisotropy, or multimodality.Source: §1 · p. 2
  • Posterior stability asks how prior-transport error controls posterior discrepancy, but existing guarantees do not establish stability for fixed few-step implementations.Flow-matching losses control average error, while numerical integration adds discretization error not controlled by the training objective alone.Source: §1 · pp. 2–3
  • The paper uses iMF for native one- or few-step transport and derives stability guarantees from its population training objective; original MeanFlow does not directly provide the same guarantee.The approach targets both computational efficiency and posterior stability for implemented maps.Source: §1 · p. 3
  • The framework combines SPT with pCN and optionally split HMC, and reports accurate, efficient synthetic posterior approximation plus CLIP-guided steering of a pretrained ImageNet prior.Split HMC is used at the coldest temperature when likelihood or reward gradients are available.Source: §1 · p. 3

Other related work

2 Background

Flow matching learns a velocity field through regression and generates samples by numerically integrating its ODE. MeanFlow and iMF instead learn average velocities, with iMF adding an auxiliary velocity objective and supporting one- or few-step sampling.

2.1 Flow matching

  • Flow matching regresses a time-dependent marginal velocity from source-target interpolations, with the samplewise interpolation velocity providing a tractable training target.Generation integrates the learned velocity-field ODE with a numerical solver.Source: §2.1 · p. 4

2.2 MeanFlow and improved MeanFlow

  • MeanFlow parameterizes average velocity and uses the MeanFlow identity to construct its training objective.The original objective substitutes learned derivatives for unknown derivatives, and its tractable form uses the samplewise interpolation velocity.Source: §2.2 · p. 4
  • Original MeanFlow’s tractable and marginal objectives are generally inequivalent because their gap depends on the model parameters.iMF replaces the interpolation velocity in the derivative term with a learnable marginal velocity.Source: §2.2 · p. 5
  • iMF jointly trains its average-velocity map and auxiliary velocity head, but uses only the average-velocity map during sampling.A fixed time schedule defines the few-step map; the auxiliary head is explicitly supervised through flow-matching loss.Source: §2.2 · p. 5

3 Source-space sampling via improved MeanFlow

The paper performs generalized posterior sampling by pulling the target back through a learned prior transport to Gaussian source space, then studies how transport-learning error affects posterior accuracy. It uses improved MeanFlow for few-step maps and source-space tempering samplers, including a hybrid variant.

3.1 Bayesian inference through a learned prior transport

  • An exact prior transport maps Gaussian source samples to the prior, so reweighting the source by the transported likelihood and mapping posterior draws back gives exact data-space posterior samples.The paper seeks to control the Wasserstein discrepancy between the learned and ideal source posteriors through the transport’s training error.Source: §3.1 · p. 5

3.2 Improved MeanFlow source and its posterior stability

  • The joint iMF and auxiliary-velocity objective differs from marginal risk only by a parameter-independent constant, enabling risk decomposition into training suboptimality and model-class approximation error.The original MeanFlow objective generally lacks this decomposition because its objective-risk gap depends on the parameters.Source: §3.2 · p. 6
  • The few-step iMF transport error bound vanishes when δ = 0 without requiring step sizes to approach zero, and includes a one-step guarantee.The bound links transport error to population training suboptimality and model-class approximation error through δ.Source: §3.2 · p. 7
  • Under the stated regularity and weight conditions, posterior Wasserstein error is bounded when the transport error is sufficiently small, connecting the learned map’s stability to posterior stability.The results require conditions including a bound on the weight’s Lipschitz constant relative to the transport error and normalization constant.Source: §3.2 · pp. 7–8

3.3 Source parallel tempering with improved MeanFlow

  • SPT tempers the pulled-back target and swaps adjacent replicas; the hybrid uses split HMC at the cold replica and pCN at hotter replicas, with geometric convergence under stated conditions.The cold marginal’s W2 distance to the learned posterior is bounded by Cρ^n for some C < ∞ and ρ ∈ (0, 1).Source: §3.3 · p. 8

4 Experiments

The experiments evaluate synthetic posterior sampling and CLIP-guided ImageNet generation. The hybrid sampler improves synthetic sampling efficiency, while the iMF approaches produce prompt-aligned ImageNet samples within the reported comparisons.

4.1 Synthetic experiments

  • Synthetic comparisons use shared prior-training data of 20,000 samples and return 6,000 samples per method; iMF uses a six-step transport, versus 100-step RK4 for FM.The methods are iMF+SPT+pCN, iMF+SPT+hybrid, FM+SPT+pCN, and DPS.Source: §4.1 · p. 9
  • FM+SPT+pCN sampling takes 33.7× and 34.5× the iMF hybrid’s time on banana and sine, respectively, while DPS shows substantially larger distributional errors and systematic histogram deviations.The FM comparison uses its configured 100-step RK4 transport.Source: §4.1 · p. 9

4.2 CLIP-guided ImageNet generation

  • For class-conditional 256×256 ImageNet prompts, iMF samplers retain 8,000 samples per method; SiT is evaluated within a matched per-prompt sampling-time budget exceeding 10 hours.The comparison includes best-of-K with K = 12, retaining the highest-CLIP-reward image among independent prior samples.Source: §4.2 · pp. 9–10

4.1 Synthetic experiments

  • 3.24× and 4.44× mean coordinate ESS gains on banana and sine over iMF+SPT+pCN; the hybrid also has the lowest mean sliced Wasserstein and binned TV errors, with small sliced-Wasserstein gaps versus seed variability.Its coordinate ESS per second is 2.25× and 3.13× higher on banana and sine, respectively, than iMF+SPT+pCN.Source: §4.1 · p. 9, Table 1 · p. 10

4.2 CLIP-guided ImageNet generation

  • The two iMF methods achieve comparable rewards and both outperform SiT+SPT+pCN; their samples more often match the specified autumn setting and snow-capped-mountain backgrounds.The hybrid has no clear reward or sample-quality improvement over iMF+SPT+pCN in these examples, and best-of-K achieves the lowest rewards.Source: §4.2 · p. 10

5 Conclusion

The conclusion reports synthetic efficiency gains from hybrid SPT and lower CLIP-guided ImageNet sampling costs for one-step iMF than multi-step SiT. It also summarizes posterior-stability guarantees and identifies capacity-dependent approximation error as an open question.

5 Conclusion

  • The framework links training suboptimality and model-class approximation error to posterior stability and end-to-end error bounds under stated assumptions.Quantifying how approximation error decreases with capacity while controlling regularity constants remains open.Source: §5 · p. 10
  • Hybrid SPT improves efficiency on synthetic benchmarks, while one-step iMF reduces CLIP-guided ImageNet sampling costs relative to multi-step SiT.Table 2 reports CLIP reward and measured sampling cost.Source: §5 · p. 10, Table 2 · p. 11
  • Figures 1 and 2 respectively pair the class labels golden retriever and bald eagle with photograph prompts depicting a retriever running through fallen autumn leaves and an eagle soaring above snow-capped mountains.Source: Figure 1 · p. 11, Figure 2 · p. 11

AI use statement

  • The authors report using generative AI for language editing, literature search and summarization, implementation and testing, mathematical work, and manuscript review; they state they take responsibility for the final manuscript and artifacts.They state that generative AI was not used to generate synthetic datasets.Source: AI use statement · p. 12

A.1 Equivalence between LiMF and RiMF

  • The proof uses E[V | W_T, R, T] = v(W_T, T) to identify the conditional mean velocity in the Aθ expression.Source: §A.1 · p. 15

A.2 Proof of relation (14)

A.3 Details on Remark 3.4

  • Finite affine-SiLU networks with direct spatial and time inputs have bounded first and second input derivatives, with resulting constants that may depend on network parameters.Source: §A.3 · pp. 15–16

A.4 Details on Remark 3.6

  • The weighted condition controls the Jacobian acting on velocity error averaged over space and time, not a uniform operator-norm bound; a uniform spatial Lipschitz bound is sufficient.Source: §A.4 · p. 16

A.5 Proof of Theorem 3.7

The proof derives the fixed-schedule transport bound by connecting joint marginal risk to local errors and composing the step maps.

A.5 Proof of Theorem 3.7

  • Joint marginal risk is related to fixed-time slices by averaging over intervals of width h, with the δ = 0 case obtained by letting h decrease to zero.Source: §A.5 · pp. 16–17
  • The local transport error qθ combines uθ − v with a time-interval correction involving time and spatial derivatives of uθ.Source: §A.5 · p. 17
  • For a multi-step schedule, the proof replaces exact maps one at a time and bounds the composition; for K = 1, no Lipschitz condition on the learned step map is needed.The composition argument uses the triangle inequality and Cauchy–Schwarz.Source: §A.5 · pp. 17–19

A.6 Proof of Theorem 3.9

The proof establishes source-space and data-space posterior stability by bounding normalized density differences and propagating transport error.

A.6 Proof of Theorem 3.9

  • If L_y e ≤ Z_y/2, the weight assumptions ensure the learned normalizing constant remains positive, with Zθ,y ≥ Z_y/2.Source: §A.6 · p. 19
  • The proof decomposes normalized densities and couples their common and remaining parts to bound the source-space Wasserstein error.The resulting source second moments are finite because normalized weights are bounded relative to γ.Source: §A.6 · p. 19
  • Data-space stability follows by inserting (Sθ)#νy between the exact and learned posteriors, applying Lipschitz propagation, and substituting the earlier transport bound.Source: §A.6 · p. 19

A.7 Geometric convergence of the source samplers

Under bounded-weight and fixed-sampler assumptions, the appendix establishes geometric convergence bounds for SPT with pCN and hybrid updates, then relates sampler convergence to cold-marginal and end-to-end bounds.

A.7 Geometric convergence of the source samplers

  • Under bounded source weights and fixed production settings, Theorem A.2 gives geometric χ² convergence bounds for iMF+SPT+pCN and iMF+SPT+hybrid.Source: §A.7 · p. 20
  • For hybrid updates, every two production sweeps contract the centered norm by at most q_hyb, and any fixed cold-target-invariant kernel is allowed without a separate split-HMC spectral-gap assumption.Source: §A.7 · pp. 20–21
  • The cold-marginal Wasserstein bounds follow by combining joint χ² convergence with conditional Jensen and transport inequalities; the hybrid rate has form Cρ^n with ρ = √q_hyb.Source: §A.7 · p. 22
  • Corollary A.3 extends sampler convergence to end-to-end convergence bounds under its stated assumptions, including a Lipschitz learned transport and a bound on δ.Source: §A.7 · p. 22

B.1 Empirical loss-to-map-error relationship

On the two-dimensional banana prior, marginal joint risk generally falls alongside squared transport-map error for one- and six-step iMF generation. The observed relationship is not strictly monotonic and supports, but does not establish, the loss-to-map-error connection.

B.1 Empirical loss-to-map-error relationship

  • Numerical-reference sensitivity checks indicate reference error is negligible at the observed map-error scale.Doubling quadrature order and tightening ODE tolerances changed reference maps by 8.3 × 10^-15 and 8.7 × 10^-9 RMS, respectively.Source: §B.1 · p. 22
  • Figure 3 plots marginal joint risk against squared transport-map error across training checkpoints for one- and six-step generation, with final checkpoints marked by stars.Error bars are one Monte Carlo standard error conditional on the fitted model and numerical reference, not variability across training runs.Source: §B.1 · p. 22, Figure 3 · p. 23
  • The experiment uses held-out interpolation samples and a numerical reference to estimate marginal risk, and evaluates squared map error for one and six uniformly spaced generation steps.The evaluation draws are independent of the training dataset and fixed across checkpoints.Source: §B.1 · p. 23
  • As marginal risk bδ fell from 2.743 to 0.0653, squared map error fell from 0.554 to 0.0198 for one-step generation and 0.00732 for six-step generation.The relationship was not strictly monotonic across checkpoints.Source: §B.1 · p. 23

B.2 Synthetic experiments

In synthetic banana- and sine-prior experiments, iMF- and FM-based source-space parallel tempering recover the main posterior well, while DPS shows systematic deviations. The section compares learned priors and posterior histograms and specifies the sampler and evaluation setup.

B.2 Synthetic experiments

  • All three prior-learning models capture the main geometric structure of the banana and sine priors.The comparison includes iMF, flow matching, and diffusion models trained on the same fixed dataset of 20,000 prior samples.Source: §B.2 · p. 24, Figure 4 · p. 24
  • The two iMF samplers use a six-step transport, while FM+SPT+pCN uses 100 RK4 steps; the hybrid replaces the coldest pCN update with six-step split HMC.SPT uses ten parallel ladders, 24 inverse temperatures, and retains 6,000 samples per method and seed after adaptation and burn-in.Source: §B.2 · p. 24
  • Both iMF- and FM-based SPT recover the main posterior well, while DPS exhibits systematic deviations from the numerical reference distributions.Posterior sample histograms are shown for the banana and sine priors.Source: §B.2 · p. 24, Figure 5 · p. 25, Figure 6 · p. 25
  • The evaluation defines mean coordinate ESS and ESS per second, and also uses sliced Wasserstein and binned total variation to compare posterior samples with references.The section presents posterior histograms against numerical references for both priors.Source: §B.2 · p. 25, Figure 5 · p. 25, Figure 6 · p. 25

B.3.1 Experiment details

The ImageNet comparison uses one-step pretrained iMF and a 125-step SiT flow-matching prior, with SPT and pCN or hybrid updates. Sampling settings are matched through a prompt-specific budget for SiT.

B.3.1 Experiment details

  • The comparison uses pretrained iMF-XL/24 with one transport step and frozen SiT-XL/25 with 125 Heun steps.The iMF setup uses ω = 8 and guidance interval [0.4, 0.65]; SiT uses classifier-free guidance scale 4.Source: §B.3.1 · p. 26
  • Each class-prompt pair uses eight independent ladders, each with eight replicas, and retains 8,000 cold states across the iMF methods.The iMF samplers use 250 adaptation sweeps, 250 burn-in sweeps, and 1,000 retained sweeps per ladder.Source: §B.3.1 · p. 26
  • SiT receives a prompt-specific sampling budget equal to the measured iMF+SPT+hybrid sampling time, plus one hour for overhead.Its available sampling time is split into adaptation, burn-in, and retained sampling in proportions 1/6, 1/6, and 2/3.Source: §B.3.1 · p. 26

B.3.2 Full results

The full ImageNet results report CLIP rewards and sampling costs for two eagle prompts, a mink prompt, a castle prompt, and two aircraft-carrier prompts. Figures 7–18 show generated samples from the methods.

B.3.2 Full results

  • For the bald-eagle fish prompt, Table 3 lists Best-of-K 0.3247/2.59/–, SiT+SPT+pCN 0.3407/10.78/376.60, iMF+SPT+pCN 0.3625/2.46/5.90, and iMF+SPT+hybrid 0.3596/11.17/26.81.The table reports CLIP reward over retained samples and sampling cost.Source: §B.3.2 · pp. 26–27, Table 3 · p. 26
  • For the bald-eagle snow prompt, Table 3 lists Best-of-K 0.2984/2.58/–, SiT+SPT+pCN 0.3139/10.74/371.58, iMF+SPT+pCN 0.3727/2.41/5.77, and iMF+SPT+hybrid 0.3713/11.16/26.78.The section also displays samples for the prompt in Figures 7–18.Source: §B.3.2 · pp. 26–27, Table 3 · p. 26
  • For the mink stream prompt, Table 3 lists Best-of-K 0.3629/2.75/–, SiT+SPT+pCN 0.3946/10.72/371.04, iMF+SPT+pCN 0.4190/2.45/5.89, and iMF+SPT+hybrid 0.4112/11.12/26.70.The table reports CLIP reward over retained samples and sampling cost.Source: §B.3.2 · pp. 26–27, Table 3 · p. 26
  • For the castle sunrise prompt, Table 3 lists Best-of-K 0.3023/2.46/–, SiT+SPT+pCN 0.3476/10.96/372.13, iMF+SPT+pCN 0.3820/2.52/6.05, and iMF+SPT+hybrid 0.3849/11.40/27.35.The section includes sample figures for the prompt among Figures 7–18.Source: §B.3.2 · pp. 26–27, Table 3 · p. 26
  • For the aircraft-carrier sunset prompt, Table 3 lists Best-of-K 0.2879/2.49/–, SiT+SPT+pCN 0.3286/10.38/373.72, iMF+SPT+pCN 0.3667/2.42/5.80, and iMF+SPT+hybrid 0.3659/10.90/26.16.The table reports CLIP reward over retained samples and sampling cost.Source: §B.3.2 · pp. 26–27, Table 3 · p. 26
  • For the aircraft-carrier watercolor prompt, Table 3 lists Best-of-K 0.2956/2.52/–, SiT+SPT+pCN 0.3055/10.63/367.95, iMF+SPT+pCN 0.3736/2.49/5.97, and iMF+SPT+hybrid 0.3774/11.07/26.55.Generated samples for the prompt appear in Figures 7–18.Source: §B.3.2 · pp. 26, 28, Table 3 · p. 26

Related papers

Loading 2610.01034v1…