通过单步隐式采样实现摊销推断
摘要
本文提出 OSIS（单步隐式采样器），利用短时间 MCMC 动力学为隐式建模分布定义局部目标，实现单步摊销采样。作者分析了收敛性质，并表明该方法在采样效率上优于多步摊销采样器和非摊销 MCMC 基线。
arXiv:2610.12679v1 Announce Type: new
Abstract: Efficient sampling from Bayesian posterior distributions in high dimensions remains a hard problem, despite advances in theory and algorithms. Multiple steps are usually required to draw unbiased samples, for example, iterations of an MCMC algorithm or integration steps in a diffusion-based amortised sampler. However, scalable one-step amortised inference remains a challenging open problem. In this work, we build upon recent progress in data-driven generative modelling to offer a possible solution for one-step amortised sampling. Our one-step implicit sampler (OSIS) uses short-time MCMC dynamics to define a local target for the evolution of an implicitly modelled distribution. In the asymptotic limit, the modelled distribution is a fixed point of the MCMC kernel and thus samples from the target. We analyse the convergence properties of OSIS and show connections with other generative modelling and sampling algorithms. We study various design choices and show that OSIS performs well compared to multi-step amortised samplers and non-amortised MCMC baselines while being far more efficient at sampling time.
查看缓存全文
缓存时间: 2026/10/12 06:42
# Amortised inference throughone-step implicit sampling Source: [https://arxiv.org/html/2610.12679](https://arxiv.org/html/2610.12679) Kirill TamogashevArran CarterAffiliation:University of EdinburghSanghyeok ChoiAffiliation:University of EdinburghStefan BauerEsmeralda S\. WhitammerAffiliation:Technical University of MunichAffiliation:Munich Center for Machine LearningAffiliation:CIFAR Fellow\*Equal contribution ###### Abstract Efficient sampling from Bayesian posterior distributions in high dimensions remains a hard problem, despite advances in theory and algorithms\. Multiple steps are usually required to draw unbiased samples, for example, iterations of an MCMC algorithm or integration steps in a diffusion\-based amortised sampler\. However, scalable one\-step amortised inference remains a challenging open problem\. In this work, we build upon recent progress in data\-driven generative modelling to offer a possible solution for one\-step amortised sampling\. Our*one\-step implicit sampler*\(OSIS\) uses short\-time MCMC dynamics to define a local target for the evolution of an implicitly modelled distribution\. In the asymptotic limit, the modelled distribution is a fixed point of the MCMC kernel and thus samples from the target\. We analyse the convergence properties ofOSISand show connections with other generative modelling and sampling algorithms\. We study various design choices and show thatOSISperforms well compared to multi\-step amortised samplers and non\-amortised MCMC baselines while being far more efficient at sampling time\. \*\*footnotetext:Correspondence to[vincent\.paulinef@gmail\.com](mailto:[email protected])and[kirill\.tamogashev@ed\.ac\.uk](mailto:[email protected])\.## 1Introduction ##### The sampling problem\. Many problems in physics, biology, statistics, and machine learning\([Welling and Teh, 2011](https://arxiv.org/html/2610.12679#bib.bib1);[Wirnsberger et al\., 2020](https://arxiv.org/html/2610.12679#bib.bib30);[Noé et al\., 2019](https://arxiv.org/html/2610.12679#bib.bib12);[Holdijk et al\., 2023](https://arxiv.org/html/2610.12679#bib.bib31)\)require drawing samples from a distribution whose density is known up to a normalising constant\. Formally, this means sampling from a distribution defined via an*energy function*ℰ:ℝd→ℝ\\mathcal\{E\}:\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}, with density π\(x\)=e−ℰ\(x\)/Z,Z=∫ℝde−ℰ\(x\)𝑑x\.\\pi\(x\)=\{e^\{\-\\mathcal\{E\}\(x\)\}\}/Z,\\quad Z=\{\\int\_\{\\mathbb\{R\}^\{d\}\}e^\{\-\\mathcal\{E\}\(x\)\}\\,\\mathrm\{d\}x\}\.\(1\)A typical instance of this problem is Bayesian inference: a posteriorπ\(x∣y\)∝p\(x\)p\(y∣x\)\\pi\(x\\mid y\)\\propto p\(x\)p\(y\\mid x\)over a variablexxgiven an observationyyhas the form in \([1](https://arxiv.org/html/2610.12679#S1.E1)\) withℰ\(x\)=−logp\(x\)−logp\(y∣x\)\\mathcal\{E\}\(x\)=\-\\log p\(x\)\-\\log p\(y\\mid x\)\. The classic way to sample from a distribution as in \([1](https://arxiv.org/html/2610.12679#S1.E1)\) is by using \(training\-free\) Monte Carlo algorithms\. Many such algorithms draw samples by simulating an ergodic chain usingπ\\pi\-invariant kernels for a sufficiently long time\. Popular Monte Carlo algorithms include Metropolis\-Hastings\([Metropolis et al\., 1953](https://arxiv.org/html/2610.12679#bib.bib36);[Hastings, 1970](https://arxiv.org/html/2610.12679#bib.bib37)\), unadjusted and Metropolis\-adjusted Langevin\([Rossky et al\., 1978](https://arxiv.org/html/2610.12679#bib.bib88);[Besag, 1994](https://arxiv.org/html/2610.12679#bib.bib87);[Roberts and Tweedie, 1996](https://arxiv.org/html/2610.12679#bib.bib14)\), Hamiltonian Monte Carlo\([Duane et al\., 1987](https://arxiv.org/html/2610.12679#bib.bib13);[Neal, 2011](https://arxiv.org/html/2610.12679#bib.bib38)\), and sequential Monte Carlo\([Del Moral et al\., 2006](https://arxiv.org/html/2610.12679#bib.bib6)\)\. The use of these methods comes at the cost of running chains for many steps, or with many simulated particles, to obtain high\-quality samples from the target\. Generative modelling algorithms that learn to imitate unbiased samples from a distribution –*i\.e\.*, approximate an empirical measure – scale quite successfully, as evidenced by many large\-scale applications of deep learning\([Kong et al\., 2021](https://arxiv.org/html/2610.12679#bib.bib68);[Hoogeboom et al\., 2022](https://arxiv.org/html/2610.12679#bib.bib67);[Sahoo et al\., 2024](https://arxiv.org/html/2610.12679#bib.bib66)\)\. Meanwhile, generative modelling algorithms that rely only on a density function, without unbiased samples, continue to lag behind, and scaling them remains a challenge despite the deep significance of this problem\. ##### Amortised sampling\. Deep learning has brought about new possibilities for introducing learnable components into sampling algorithms\. Such components allow to spend more time and compute for training, yet gain speed when sampling and exploit generalisable structure in the target\. Amortisation techniques include learning MCMC transition kernels\([Song et al\., 2017](https://arxiv.org/html/2610.12679#bib.bib53);[Samsonov et al\., 2022](https://arxiv.org/html/2610.12679#bib.bib69);[Wang et al\., 2025](https://arxiv.org/html/2610.12679#bib.bib70)\), amortising MCMC dynamics\([Li et al\., 2017](https://arxiv.org/html/2610.12679#bib.bib50)\), and training density estimators to sample from unnormalised distributions\([Rezende and Mohamed, 2015](https://arxiv.org/html/2610.12679#bib.bib45);[Midgley et al\., 2023](https://arxiv.org/html/2610.12679#bib.bib7)\)\. The most successful family of amortised samplers adopts diffusion models\([Sohl\-Dickstein et al\., 2015](https://arxiv.org/html/2610.12679#bib.bib2);[Ho et al\., 2020](https://arxiv.org/html/2610.12679#bib.bib3);[Song et al\., 2021](https://arxiv.org/html/2610.12679#bib.bib4)\)as approximators of the target distribution\([Zhang and Chen, 2022](https://arxiv.org/html/2610.12679#bib.bib39);[Vargas et al\., 2023](https://arxiv.org/html/2610.12679#bib.bib8);[Richter et al\., 2024](https://arxiv.org/html/2610.12679#bib.bib42);[Sendera et al\., 2024](https://arxiv.org/html/2610.12679#bib.bib10);[Chen et al\., 2025a](https://arxiv.org/html/2610.12679#bib.bib43);[Albergo and Vanden\-Eijnden, 2025](https://arxiv.org/html/2610.12679#bib.bib44);[Berner et al\., 2026](https://arxiv.org/html/2610.12679#bib.bib11);[Havens et al\., 2025](https://arxiv.org/html/2610.12679#bib.bib64),*e\.g\.*,\)\. Diffusion samplers introduce a learnable stochastic process to model the transition dynamics from a known distribution to a target, usually by approximately reversing a fixed noising process\. In contrast to diffusion models trained using a variational bound on the likelihood of empirical samples, diffusion samplers are trained using stochastic optimal control objectives\([Zhang and Chen, 2022](https://arxiv.org/html/2610.12679#bib.bib39);[Berner et al\., 2024](https://arxiv.org/html/2610.12679#bib.bib9)\)or path\-space divergences\([Richter et al\., 2020](https://arxiv.org/html/2610.12679#bib.bib40);[Lahlou et al\., 2023](https://arxiv.org/html/2610.12679#bib.bib41)\), among other choices\. Recent work has found and exploited connections between diffusion samplers and amortisation of SMC\([Choi et al\., 2026](https://arxiv.org/html/2610.12679#bib.bib52)\)\. 𝒟\(ρθ,ρh\)\\mathcal\{D\}\(\\rho\_\{\\theta\},\\,\\rho\_\{h\}\)KK𝒟\(ρθ,ρh\)\\mathcal\{D\}\(\\rho\_\{\\theta\},\\,\\rho\_\{h\}\)KK𝒟\(ρθ,ρh\)\\mathcal\{D\}\(\\rho\_\{\\theta\},\\,\\rho\_\{h\}\)KKx∼ρθ≔\(Gθ\)\#qx\\sim\\rho\_\{\\theta\}\\coloneqq\(G\_\{\\theta\}\)\_\{\\\#\}qx∼ρθ≔\(Gθ\)\#qx\\sim\\rho\_\{\\theta\}\\coloneqq\(G\_\{\\theta\}\)\_\{\\\#\}qx∼ρθ≔\(Gθ\)\#qx\\sim\\rho\_\{\\theta\}\\coloneqq\(G\_\{\\theta\}\)\_\{\\\#\}qz∼qz\\sim qTrainingSourceqqx∼ρh≔Kρθx\\sim\\rho\_\{h\}\\coloneqq K\\rho\_\{\\theta\}x∼ρh≔Kρθx\\sim\\rho\_\{h\}\\coloneqq K\\rho\_\{\\theta\}x∼ρh≔Kρθx\\sim\\rho\_\{h\}\\coloneqq K\\rho\_\{\\theta\}Targetπ\\pi Figure 1:OSISevolves an implicit modelled distributionρθ\\rho\_\{\\theta\}towards a targetπ\\piavailable through its unnormalised density \(green contours\)\.GθG\_\{\\theta\}maps sampleszzthe source distributionqqto samplesx=Gθ\(z\)x=G\_\{\\theta\}\(z\)from the modelled distributionρθ≔\(Gθ\)\#q\\rho\_\{\\theta\}\\coloneqq\(G\_\{\\theta\}\)\_\{\\\#\}q\(blue KDE contours\)\. This pushforward distribution is evolved by aπ\\pi\-invariant MCMC kernelKK, givingρh≔Kρθ\\rho\_\{h\}\\coloneqq K\\rho\_\{\{\\theta\}\}\. The parametersθ\\thetaare updated to minimise an empirical divergence betweenρθ\\rho\_\{\\theta\}andρh\\rho\_\{h\}\. Upon convergence,ρθ≈π\\rho\_\{\\theta\}\\approx\\pi\. ##### One\-step amortised sampling\. Although diffusion samplers allow for a considerable inference speed\-up by avoiding long\-chain MCMC, they require simulation of a neural stochastic dynamical system during both training and sampling\. Amortised sampling with a*single*neural function evaluation, however, remains a challenge\. Although it is possible to use normalising flows for one\-step sampling\([Rezende and Mohamed, 2015](https://arxiv.org/html/2610.12679#bib.bib45)\), they require the invertibility of a neural network and efficient computation of its Jacobian, both of which remain significant constraints\. More recent work shows that such methods scale poorly and are prone to mode collapse\([He et al\., 2025](https://arxiv.org/html/2610.12679#bib.bib54)\)\. These methods can further be placed within a Monte Carlo outer loop and augmented with MCMC sampling\([Midgley et al\., 2023](https://arxiv.org/html/2610.12679#bib.bib7)\)\. In an attempt to design an efficient one\-step sampler, we take a different approach, drawing inspiration from the novel paradigm in generative modelling termed drifting models\([Deng et al\., 2026](https://arxiv.org/html/2610.12679#bib.bib33);[Gretton et al\., 2026](https://arxiv.org/html/2610.12679#bib.bib55)\)\. Drifting models evolve a implicit distribution towards an empirical target by defining a local direction for its evolution via a*drifting field*estimated from batches of target and generated samples\. Such an approach, cannot be directly applied to sampling, as it requires access to samples from the target distribution to compute the drifting field\. ##### Our contributions\. We bypass this limitation by using an MCMC kernel invariant to the target distribution to implicitly define a drifting field\. This leads us to the first main contribution: 1. \(1\)We propose an algorithm for training a one\-step implicit sampler \(OSIS\) using an objective that enforces invariance of the sampler’s output distribution to a target\-invariant MCMC kernel \([Section2\.1](https://arxiv.org/html/2610.12679#S2.SS1)\)\. There have been prior attempts to design a one\-step amortised sampling algorithm, using kernel density approximations to circumvent data unavailability\([Cao et al\., 2026](https://arxiv.org/html/2610.12679#bib.bib32)\)and using adversarial training\([Li et al\., 2017](https://arxiv.org/html/2610.12679#bib.bib50)\)\. We discuss the connections betweenOSIS, drifting models, computational Wasserstein gradient flows, and this prior work in[Section3](https://arxiv.org/html/2610.12679#S3)\. We also study the theoretical properties ofOSIS, showing global convergence under assumptions similar to those in the MCMC literature: 1. 1\.We prove convergence ofOSISunder asymptotic conditions and constraints on the target \([Section4](https://arxiv.org/html/2610.12679#S4)\)\. To achieve scalability and efficiency in practice, we study the empirical performance ofOSISunder different design choices, such as choice of MCMC kernel, empirical divergence, and other techniques: 1. 1\.We study design choices forOSIS\([§5](https://arxiv.org/html/2610.12679#S5)\) and their influence on performance \([Table3](https://arxiv.org/html/2610.12679#S6.T3)\)\. Finally, we benchmark our sampler on amortised inference problems ranging from synthetic targets to posterior sampling under large\-scale image model priors, leading to our final contribution: 1. 1\.We show thatOSISperforms on par with multi\-step baselines, while being more efficient \([§6\.1](https://arxiv.org/html/2610.12679#S6.SS1)\) and remaining scalable \([§6\.2](https://arxiv.org/html/2610.12679#S6.SS2)\)\. ## 2Proposed method: Training by distilling MCMC refinements ### 2\.1Preliminaries Algorithm 1OSIS using a divergence𝒟\\mathcal\{D\}1:Prior distribution qq, π\\pi\-invariant kernel KK, generator GθG\_\{\\theta\}, divergence 𝒟\\mathcal\{D\}\. 2:whilenot convergeddo 3: ρθ←Gθ\#q\\rho\_\{\\theta\}\\leftarrow G\_\{\\theta\\,\\\#\}q 4: ⊳\\trianglerightApply the MCMC kernel 5: ρh←Kρsg\[θ\]\\rho\_\{h\}\\leftarrow K\\rho\_\{\\text\{sg\}\\left\[\\theta\\right\]\} 6: ⊳\\trianglerightCompute loss using𝒟\\mathcal\{D\} 7: ℒ\(θ\)←𝒟\(ρθ,ρh\)\\mathcal\{L\}\(\\theta\)\\leftarrow\\mathcal\{D\}\\left\(\\rho\_\{\\theta\},\\rho\_\{h\}\\right\) 8:Gradient update on θ\\thetawith ∇θℒ\\nabla\_\{\\theta\}\\mathcal\{L\} 9:endwhile 10:return GθG\_\{\\theta\} ##### Problem setting: one\-step amortised sampling\. Suppose thatG:ℝdz→ℝdG:\\mathbb\{R\}^\{d\_\{z\}\}\\to\\mathbb\{R\}^\{d\}is a measurable function andqqis a simple distribution \(*e\.g\.*,𝒩\(0,I\)\\mathcal\{N\}\(0,I\)\) overℝdz\\mathbb\{R\}^\{d\_\{z\}\}, wheredzd\_\{z\}is not necessarily equal todd\. Define a random variablez∼qz\\sim qandx=G\(z\)x=G\(z\)\. The distribution ofxxis the pushforward ofqqbyGG, denotedρ=G\#q\\rho=G\_\{\\\#\}q\. IfGGis not \(efficiently\) invertible or differentiable, the density ofρ\\rhois intractable to compute\. Given a target distributionπ\\pioverℝd\\mathbb\{R\}^\{d\}defined by an unnormalised density as in \([1](https://arxiv.org/html/2610.12679#S1.E1)\), the goal of one\-step amortised sampling is to findGGsuch thatπ≈G\#q\\pi\\approx G\_\{\\\#\}q\. If the functionGGdepends on parametersθ\\theta, we denote itGθG\_\{\\theta\}and abbreviateρθ=\(Gθ\)\#q\\rho\_\{\\theta\}=\(G\_\{\\theta\}\)\_\{\\\#\}q\. The goal is then to findθ\\thetasuch thatρθ≈π\\rho\_\{\\theta\}\\approx\\pi\. ##### Stationary kernels\. Suppose thatKKis a stochastic kernel fromℝd\\mathbb\{R\}^\{d\}toℝd\\mathbb\{R\}^\{d\}\(*i\.e\.*, a conditional distributionK\(y∣x\)K\(y\\mid x\)satisfying measurability conditions\)\. One can construct a Markov chain\(xk\)k=0∞\(x\_\{k\}\)\_\{k=0\}^\{\\infty\}viaxk\+1∼K\(xk\+1∣xk\)x\_\{k\+1\}\\sim K\(x\_\{k\+1\}\\mid x\_\{k\}\)\. Markov chain Monte Carlo algorithms selectKKsuch that, regardless of the distribution ofx0x\_\{0\}, the chain converges to a unique stationary distributionπ\\piwith respect to whichKKis invariant,*i\.e\.*,Kπ\(𝑑y\)≔∫K\(𝑑y∣x\)π\(x\)𝑑x=π\(𝑑y\)K\\pi\(\\mathrm\{d\}y\)\\coloneqq\\int K\(\\mathrm\{d\}y\\mid x\)\\pi\(x\)\\,\\mathrm\{d\}x=\\pi\(\\mathrm\{d\}y\)\([Douc et al\., 2018](https://arxiv.org/html/2610.12679#bib.bib34)\)\. Many choices forKKuse a discrete\-time approximation to Langevin dynamics, such as ULA, MALA\([Rossky et al\., 1978](https://arxiv.org/html/2610.12679#bib.bib88);[Besag, 1994](https://arxiv.org/html/2610.12679#bib.bib87);[Roberts and Tweedie, 1996](https://arxiv.org/html/2610.12679#bib.bib14)\), HMC\([Duane et al\., 1987](https://arxiv.org/html/2610.12679#bib.bib13);[Neal, 2011](https://arxiv.org/html/2610.12679#bib.bib38)\), and improved Langevin integrators\([Leimkuhler and Matthews, 2015](https://arxiv.org/html/2610.12679#bib.bib72)\)\. ### 2\.2Training objectives for one\-step sampling Training of a one\-step sampler, by some algorithm, produces a sequence of parameter iterates\(θk\)k≥0\(\\theta\_\{k\}\)\_\{k\\geq 0\}, and thus an evolution of the modelled distributionρθk≔\(Gθk\)\#q\\rho\_\{\\theta\_\{k\}\}\\coloneqq\(G\_\{\\theta\_\{k\}\}\)\_\{\\\#\}q\. The training algorithm should be designed to ensure thatρθk\\rho\_\{\\theta\_\{k\}\}converges toπ\\piask→∞k\\to\\infty\. One way to express such convergence is through invariance of the limiting distribution with respect to aπ\\pi\-invariant stochastic kernelKK\. ##### Divergence\-based loss\. Given aπ\\pi\-invariant kernelKK, a one\-step samplerρθ=\(Gθ\)\#q\\rho\_\{\\theta\}=\(G\_\{\\theta\}\)\_\{\\\#\}qshould satisfyρθ≈Kρθ\\rho\_\{\\theta\}\\approx K\\rho\_\{\\theta\}upon convergence\. We propose to turn this into a training objective: ℒ\(θ\)=𝒟\(ρθ,Kρsg\[θ\]\),\\mathcal\{L\}\(\\theta\)=\\mathcal\{D\}\\left\(\\rho\_\{\\theta\},\\,K\\rho\_\{\\text\{sg\}\\left\[\\theta\\right\]\}\\right\),\(2\)where𝒟\\mathcal\{D\}is some measure of divergence between distributions,Kρ\(𝑑y\)≔∫K\(𝑑y∣x\)ρ\(𝑑x\)K\\rho\(\\mathrm\{d\}y\)\\coloneqq\\int K\(\\mathrm\{d\}y\\mid x\)\\rho\(\\mathrm\{d\}x\), andsg\[⋅\]\\text\{sg\}\\left\[\\cdot\\right\]is the stop\-gradient operator\. The complete algorithm, given in[Algorithm1](https://arxiv.org/html/2610.12679#alg1), leaves room for two design choices: the kernelKKand the divergence𝒟\\mathcal\{D\}\. Both should be efficient to compute, and the distance𝒟\\mathcal\{D\}should be differentiable with respect to \(empirical approximations of\)ρθ\\rho\_\{\\theta\}\. We discuss these choices in[§C\.1](https://arxiv.org/html/2610.12679#A3.SS1)and[§C\.2](https://arxiv.org/html/2610.12679#A3.SS2), respectively\. Additional aspects of training are discussed in[§5](https://arxiv.org/html/2610.12679#S5)and variants of[Algorithm1](https://arxiv.org/html/2610.12679#alg1)in[AppendixA](https://arxiv.org/html/2610.12679#A1)\. ## 3Connections and related work ##### Connection to other implicit samplers\. The connection between generative modelling and implicit models, those without a parametrisation of density, goes back to[Mohamed and Lakshminarayanan \(2016\)](https://arxiv.org/html/2610.12679#bib.bib73)\. Subsequent works explored algorithms that use MCMC to train implicit samplers:[Feng et al\. \(2017\)](https://arxiv.org/html/2610.12679#bib.bib74)amortises Stein variational gradient descent\([Liu and Wang, 2016](https://arxiv.org/html/2610.12679#bib.bib26), SVGD;\), while[Titsias \(2017\)](https://arxiv.org/html/2610.12679#bib.bib75)fits an invertible map to sample from the distribution induced by a short MCMC chain\. Related approaches include[Hoffman \(2017\)](https://arxiv.org/html/2610.12679#bib.bib76),[Ruiz and Titsias \(2019\)](https://arxiv.org/html/2610.12679#bib.bib77),[Gabrié et al\. \(2022\)](https://arxiv.org/html/2610.12679#bib.bib78), and[Yi and Liu \(2023\)](https://arxiv.org/html/2610.12679#bib.bib90)\. The closest precursor toOSISis[Li et al\. \(2017\)](https://arxiv.org/html/2610.12679#bib.bib50)\. That algorithm initialises an MCMC chain from an implicit approximation network, applies a transition kernel, and updates the network with an adversarial loss matching the initial samples to their few\-step refinements\.OSISdiffers from this line of work in three key ways\. First, unlike[Li et al\. \(2017\)](https://arxiv.org/html/2610.12679#bib.bib50),OSISdoes not rely on an adversarial loss, but remains stable with a fixed choice of empirical divergence\. Second,OSISis compatible with any MCMC kernel and any divergence between distributions, and we compare several choices of each \(see[Table3](https://arxiv.org/html/2610.12679#S6.T3)for kernels and[Tables1](https://arxiv.org/html/2610.12679#S6.T1)and[1](https://arxiv.org/html/2610.12679#S6.T1)for divergences\)\. Third, we link the convergence of these methods to the convergence of the underlying MCMC kernel and establish a connection to Wasserstein gradient flows that was not known before in the sampling context \(see[Section4](https://arxiv.org/html/2610.12679#S4)\)\. ##### Connection to drifting models\. OSIScan be linked to drifting models\([Deng et al\., 2026](https://arxiv.org/html/2610.12679#bib.bib33)\), which define a local target for the evolution of a modelled distribution to guide it towards a target given empirically through data samples\. AlthoughOSISand drifting models solve different problems \(approximating an empirical data distribution and an energy\-based distribution, respectively\), they share some similarities\. In[Deng et al\. \(2026\)](https://arxiv.org/html/2610.12679#bib.bib33), stationarity of an implicitly defined generative model is expressed through a drifting field, a vector field defining an evolution of particles under which the target distribution is stationary under asymptotic conditions\. The field is defined similarly to the gradient of a kernel\-based divergence and thus evolves model\-generated samples towards the data distribution over the course of training\. Unlike in diffusion and flow\-based models, which approximate a time\-dependent evolution, this field is a non\-stationary training target, rather than part of the generation process\. More recent works propose other ways to define this drifting field, some of which can be identified with gradients of divergences between distributions, making the training process simulate a gradient flow\([He et al\., 2026](https://arxiv.org/html/2610.12679#bib.bib56);[Gretton et al\., 2026](https://arxiv.org/html/2610.12679#bib.bib55);[Han et al\., 2026](https://arxiv.org/html/2610.12679#bib.bib51)\)\. As we discuss below,OSISshares parts of this connection; however, drifting models and their successors*approximate the vector field giving the gradient flow*, whereasOSIS*uses the vector field giving the gradient flow to define the local target for optimisation*\. [Cao et al\. \(2026\)](https://arxiv.org/html/2610.12679#bib.bib32)propose several ways to adapt drifting models for sampling\. Their method relies on either Monte Carlo simulation or second\-order derivatives to approximate the gradient of the energy, both of which scale poorly with dimensionality\. In contrast,OSISdefines the drift implicitly by computing the divergence between two probability distributions\. ##### Connection to computational Wasserstein gradient flows\. A Wasserstein gradient flow \(WGF\) describes a continuous evolution\(ρt\)t≥0\(\\rho\_\{t\}\)\_\{t\\geq 0\}on the space of absolutely continuous probability distributions with finite second moments𝒫2\(ℝd\)\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\)\. The flow moves a distributionρ\\rhoalong the gradients of a functional𝒥:𝒫2\(ℝd\)→ℝ\\mathcal\{J\}:\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\)\\to\\mathbb\{R\}\. The flow admits two equivalent descriptions: as an evolution of particlesxt∼ρtx\_\{t\}\\sim\\rho\_\{t\}\([3](https://arxiv.org/html/2610.12679#S3.E3)\) and the corresponding continuity equation, the evolution of the densityρt\\rho\_\{t\}\([4](https://arxiv.org/html/2610.12679#S3.E4)\): ddtxt=−∇δ𝒥δρt\(xt\),\\frac\{\\mathrm\{d\}\}\{\\mathrm\{d\}t\}\{x\_\{t\}\}=\-\\nabla\\frac\{\\delta\\mathcal\{J\}\}\{\\delta\\rho\_\{t\}\}\(x\_\{t\}\),\(3\) ∂∂tρt=∇⋅\(ρt∇δ𝒥δρt\),\\frac\{\\partial\}\{\\partial t\}\\rho\_\{t\}=\\nabla\\cdot\\Big\(\\rho\_\{t\}\\,\\nabla\\frac\{\\delta\\mathcal\{J\}\}\{\\delta\\rho\_\{t\}\}\\Big\),\(4\) where initial conditions are given byx0∼ρ0x\_\{0\}\\sim\\rho\_\{0\}andδ𝒥δρ:ℝd→ℝ\\frac\{\\delta\\mathcal\{J\}\}\{\\delta\\rho\}:\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}denotes the first variation of𝒥\\cal J\. The flow given in \([4](https://arxiv.org/html/2610.12679#S3.E4)\) admits a discrete\-time approximation scheme by proximal steps in the Wasserstein geometry, known as the Jordan\-Kinderlehrer\-Otto \(JKO\) scheme\([Jordan et al\., 1998](https://arxiv.org/html/2610.12679#bib.bib22)\): ρk\+1=argmin𝜌\{2η𝒥\(ρ\)\+𝒲22\(ρ,ρk\)\}\\rho\_\{k\+1\}=\\underset\{\\rho\}\{\\operatorname\{argmin\}\}\\Big\\\{2\\eta\\mathcal\{J\}\(\\rho\)\+\\mathcal\{W\}^\{2\}\_\{2\}\(\\rho,\\rho\_\{k\}\)\\Big\\\}\(5\)that converges to the continuous\-time flow asη→0\\eta\\to 0\. Refer to[Ambrosio et al\. \(2005\)](https://arxiv.org/html/2610.12679#bib.bib21)for a rigorous treatment or to[Santambrogio \(2017\)](https://arxiv.org/html/2610.12679#bib.bib61)for a succinct introduction\. The scheme in \([5](https://arxiv.org/html/2610.12679#S3.E5)\) offers a tempting way to solve optimisation problems on the space of probability distributions, and recent years have seen multiple attempts to convert it into a scalable machine learning algorithm\.[Mokrov et al\. \(2021\)](https://arxiv.org/html/2610.12679#bib.bib57)proposes to learn WGF by parametrising each step of the JKO scheme \([5](https://arxiv.org/html/2610.12679#S3.E5)\) via an input\-convex neural network\. An alternative approach was introduced in[Galashov et al\. \(2025\)](https://arxiv.org/html/2610.12679#bib.bib59)\. Some more recent works\([Chen et al\., 2025b](https://arxiv.org/html/2610.12679#bib.bib60);[Tian et al\., 2026](https://arxiv.org/html/2610.12679#bib.bib58)\)offer ways to compute WGF using different variations of MMD as the functional𝒥\\mathcal\{J\}\. Finally,[Gretton et al\. \(2026\)](https://arxiv.org/html/2610.12679#bib.bib55)provides an interpretation of drifting \([Section3](https://arxiv.org/html/2610.12679#S3)\) as WGF\.[He et al\. \(2026\)](https://arxiv.org/html/2610.12679#bib.bib56);[Han et al\. \(2026\)](https://arxiv.org/html/2610.12679#bib.bib51)show that Sinkhorn divergence to the target distribution can be used as the functional and can achieve strong results on image generation\. Notably,*all methods that attempt to neurally approximate WGF to this day fundamentally rely on the availability of unbiased samples*\. ## 4Theoretical analysis ofOSIS In this section we analyse the convergence ofOSISunder idealised assumptions: 1. \(1\)We takeKKto be continuous Langevin diffusion and𝒟\\mathcal\{D\}to be12𝒲22\\frac\{1\}\{2\}\\mathcal\{W\}\_\{2\}^\{2\}\. In practice, this assumption is satisfied in the infinite\-particle limit of empirical estimation of the divergence and with perfect numerical integration in the Monte Carlo updates\. 2. \(2\)We assume that optimisation of the modelled distribution is performed in Wasserstein space, rather than the parameter space \(θ\\theta\)\. This can be achieved in practice by replacing the gradient step on𝒲22\\mathcal\{W\}\_\{2\}^\{2\}with respect toθ\\thetaby a proximal gradient step in the𝒲22\\mathcal\{W\}\_\{2\}^\{2\}geometry and assuming the parametric family and optimisation algorithm are sufficiently expressive to perform this step exactly\. Our analysis uses that the evolution of density given by Langevin dynamics is the WGF of the KL divergence functional; we show thatOSISinherits the convergence properties of Langevin dynamics\([Durmus and Moulines, 2017](https://arxiv.org/html/2610.12679#bib.bib71);[Chewi et al\., 2025](https://arxiv.org/html/2610.12679#bib.bib16),*e\.g\.*,\)\. ### 4\.1Setting π\\piρk\\rho\_\{k\}ρk\+1\\rho\_\{k\+1\}ρh\\rho\_\{h\}OSISstepWasserstein geodesic \(dynamic OT\)Langevin \(KL descent\)12𝒲22\(⋅,π\)\\frac\{1\}\{2\}\\mathcal\{W\}\_\{2\}^\{2\}\(\\cdot,\\pi\)KL\(⋅∥π\)\\mathrm\{KL\}\(\\cdot\\\|\\pi\) Figure 2:One idealised OSIS iteration\. KL and Wasserstein distance share the minimiserπ\\pi\. However, KL decreases along the Langevin path \(red\), whereas Wasserstein decreases along the geodesic \(grey\)\. OSIS takes a step along the Wasserstein gradient of12𝒲22\(⋅,ρh\)\\frac\{1\}\{2\}\\mathcal\{W\}\_\{2\}^\{2\}\(\\cdot,\\rho\_\{h\}\)towardsρk\+1\\rho\_\{k\+1\}\(blue\)\.We first formalise the algorithm described in[Section2\.1](https://arxiv.org/html/2610.12679#S2.SS1)\. ##### Notation and assumptions\. We assume that all distributions are elements of𝒫2\(ℝd\)\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\), the space of absolutely continuous distributions onℝd\\mathbb\{R\}^\{d\}with finite second moment, and we identify them with their densities\. We assume that the targetπ∝e−ℰ\\pi\\propto e^\{\-\\cal E\}is defined by a twice continuously differentiable energyℰ:ℝd→ℝ\\mathcal\{E\}:\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}\. Let\{Xt\}t≥0\\\{X\_\{t\}\\\}\_\{t\\geq 0\}be a Langevin process, defined by an Itô SDE: dXt=−∇ℰ\(Xt\)dt\+2dBt,\\mathrm\{d\}X\_\{t\}=\-\\nabla\\mathcal\{E\}\(X\_\{t\}\)\\,\\mathrm\{d\}t\+\\sqrt\{2\}\\,\\mathrm\{d\}B\_\{t\},\(6\)where\{Bt\}t≥0\\\{B\_\{t\}\\\}\_\{t\\geq 0\}isdd\-dimensional Brownian motion\. Forh≥0h\\geq 0, we defineKhK^\{h\}to be the stochastic integration kernel of the process \([6](https://arxiv.org/html/2610.12679#S4.E6)\),*i\.e\.*,Kh\(A∣x\)≔ℙ\(Xt\+h∈A∣Xt=x\)K^\{h\}\(A\\mid x\)\\coloneqq\\mathbb\{P\}\(X\_\{t\+h\}\\in A\\mid X\_\{t\}=x\), which isπ\\pi\-invariant\. For a distributionρ\\rhoand an absolutely continuous kernelKKwe denote byKρK\\rhothe pushforward distribution, which has densityKρ\(y\)≔∫K\(y∣x\)ρ\(x\)𝑑xK\\rho\(y\)\\coloneqq\\int K\(y\\mid x\)\\rho\(x\)\\mathrm\{d\}x\. As before, we denote byT\#ρT\_\{\\\#\}\\rhothe pushforward of a distributionρ\\rhoby a measurable functionTT\. Integration is assumed to be taken over the entireℝd\\mathbb\{R\}^\{d\}unless stated otherwise\. We use\|⋅\|\|\\cdot\|to denote Euclidean distance and define 2\-Wasserstein distance and KL divergence between distributionsρ\\rhoandμ\\muas follows: 𝒲22\(ρ,μ\)=infT\#ρ=μ\{∫\|T\(x\)−x\|2ρ\(x\)dx\},KL\(ρ∥μ\)=∫logρ\(x\)μ\(x\)ρ\(x\)dx\.\\mathcal\{W\}^\{2\}\_\{2\}\(\\rho,\\mu\)=\\inf\_\{T\_\{\\\#\}\\rho=\\mu\}\\left\\\{\\int\|T\(x\)\-x\|^\{2\}\\rho\(x\)\\mathrm\{d\}\{x\}\\right\\\},\\quad\\operatorname\{KL\}\\left\(\\rho\\\!\\parallel\\\!\\mu\\right\)=\\int\\log\\frac\{\\rho\(x\)\}\{\\mu\(x\)\}\\rho\(x\)\\mathrm\{d\}\{x\}\.We use the Monge formulation of 2\-Wasserstein distance, since for absolutely continuous distributions, the Monge optimal transport maps exist and are unique\. Other formulations are discussed in[SectionB\.1](https://arxiv.org/html/2610.12679#A2.SS1)\. ##### One step ofOSISin Wasserstein space\. We now formalise one step ofOSIS, visualised in[Figure2](https://arxiv.org/html/2610.12679#S4.F2)\. At each training iterationk≥0k\\geq 0,OSIShas the initial distributionρk≔ρθk=\(Gθk\)\#q\\rho\_\{k\}\\coloneqq\\rho\_\{\\theta\_\{k\}\}=\(G\_\{\\theta\_\{k\}\}\)\_\{\\\#\}q\. First, we useρhk≔Khkρk\\rho\_\{h\_\{k\}\}\\coloneqq K^\{h\_\{k\}\}\\rho\_\{k\}to denote the distribution after a continuous evolution ofρk\\rho\_\{k\}by a stochastic kernelKhkK^\{h\_\{k\}\}for the timehk\>0h\_\{k\}\>0\.Second, we compute the 2\-Wasserstein OT mapThkT^\{h\_\{k\}\}betweenρk\\rho\_\{k\}andρhk\\rho\_\{h\_\{k\}\}\. Withρhk\\rho\_\{h\_\{k\}\}fixed, the functional12𝒲22\(⋅,ρhk\)\\frac\{1\}\{2\}\\mathcal\{W\}\_\{2\}^\{2\}\(\\cdot,\\rho\_\{h\_\{k\}\}\)has Wasserstein gradientId−Thk\\operatorname\{Id\}\-T^\{h\_\{k\}\}atρk\\rho\_\{k\}\([Alfonsi and Jourdain, 2020](https://arxiv.org/html/2610.12679#bib.bib49), Theorem 2\.2\)\.Third, we setTηkhk≔ηkThk\+\(1−ηk\)IdT^\{h\_\{k\}\}\_\{\\eta\_\{k\}\}\\coloneqq\\eta\_\{k\}T^\{h\_\{k\}\}\+\(1\-\\eta\_\{k\}\)\\operatorname\{Id\}withηk∈\(0,1\]\\eta\_\{k\}\\in\(0,1\], which is the result of taking a gradient step on12𝒲22\(⋅,ρhk\)\\frac\{1\}\{2\}\\mathcal\{W\}\_\{2\}^\{2\}\(\\cdot,\\rho\_\{h\_\{k\}\}\)with step sizeηk\\eta\_\{k\}\. Finally, one idealised iteration ofOSISis expressed as ρk\+1≔\(Tηkhk\)\#ρk\.\\rho\_\{k\+1\}\\coloneqq\(T^\{h\_\{k\}\}\_\{\\eta\_\{k\}\}\)\_\{\\\#\}\\rho\_\{k\}\.\(7\) ### 4\.2Main results We prove convergence of the procedure \([7](https://arxiv.org/html/2610.12679#S4.E7)\) under various assumptions onπ\\pi\. Proofs appear in[SectionB\.2](https://arxiv.org/html/2610.12679#A2.SS2)\. ##### Convergence for strongly log\-concave target\. We consider the case whereπ\\piisλ\\lambda\-strongly log\-concave \(λ\\lambda\-SLC\), that is,∇2ℰ⪰λI\\nabla^\{2\}\\mathcal\{E\}\\succeq\\lambda Ifor someλ\>0\\lambda\>0\. main\-pratenddefaultcategory\.tex main\-pratenddefaultcategory\.tex main\-pratenddefaultcategory\.tex main\-pratenddefaultcategory\.tex main\-pratenddefaultcategory\.tex ###### Theorem 4\.1\(Contraction of one step ofOSISunderλ\\lambda\-SLCπ\\pi\)\. Letπ\\pibeλ\\lambda\-SLC\. Letρk,ρhk,Thk,Tηkhk\\rho\_\{k\},\\rho\_\{h\_\{k\}\},T^\{h\_\{k\}\},T\_\{\\eta\_\{k\}\}^\{h\_\{k\}\},ρk\+1\\rho\_\{k\+1\}be as in[Section4\.1](https://arxiv.org/html/2610.12679#S4.SS1)\. For anyhk\>0h\_\{k\}\>0,ηk∈\(0,1\]\\eta\_\{k\}\\in\(0,1\], the following bounds hold: 𝒲2\(ρk\+1,π\)\\displaystyle\\mathcal\{W\}\_\{2\}\(\\rho\_\{k\+1\},\\pi\)≤\(\(1−ηk\)\(1\+e−λhk\)\+e−λhk\)𝒲2\(ρk,π\);\\displaystyle\\leq\\left\(\(1\-\\eta\_\{k\}\)\(1\+e^\{\-\\lambda h\_\{k\}\}\)\+e^\{\-\\lambda h\_\{k\}\}\\right\)\\mathcal\{W\}\_\{2\}\(\\rho\_\{k\},\\pi\);\(8\)KL\(ρk\+1∥π\)\\displaystyle\\operatorname\{KL\}\\left\(\\rho\_\{k\+1\}\\\!\\parallel\\\!\\pi\\right\)≤\(1−ηk\(1−e−2λhk\)\)KL\(ρk∥π\)\.\\displaystyle\\leq\(1\-\\eta\_\{k\}\(1\-e^\{\-2\\lambda h\_\{k\}\}\)\)\\operatorname\{KL\}\\left\(\\rho\_\{k\}\\\!\\parallel\\\!\\pi\\right\)\.\(9\) main\-pratenddefaultcategory\.tex main\-pratenddefaultcategory\.tex ###### Theorem 4\.2\(Convergence ofOSISunderλ\\lambda\-SLCπ\\pi\)\. Under the assumptions of[Theorem4\.1](https://arxiv.org/html/2610.12679#S4.Thmtheorem1), let\(ρk\)k≥0\(\\rho\_\{k\}\)\_\{k\\geq 0\}be the sequence of probability distributions produced by theOSISiterations\. Assuminghk\>0h\_\{k\}\>0andηk∈\(0,1\]\\eta\_\{k\}\\in\(0,1\]for allk≥0k\\geq 0, the following bounds hold for anyN≥1N\\geq 1: KL\(ρN∥π\)≤e−∑k=0N−1ηk\(1−e−2λhk\)KL\(ρ0∥π\)\.\\operatorname\{KL\}\\left\(\\rho\_\{N\}\\\!\\parallel\\\!\\pi\\right\)\\leq e^\{\-\\sum\_\{k=0\}^\{N\-1\}\\eta\_\{k\}\(1\-e^\{\-2\\lambda h\_\{k\}\}\)\}\\operatorname\{KL\}\\left\(\\rho\_\{0\}\\\!\\parallel\\\!\\pi\\right\)\.\(10\)Moreover, ifh¯=infk≥0hk\>0\\bar\{h\}=\\inf\_\{k\\geq 0\}h\_\{k\}\>0,η¯=infk≥0ηk\>0\\bar\{\\eta\}=\\inf\_\{k\\geq 0\}\\eta\_\{k\}\>0, andη¯\>2/\(1\+exp\(λh¯\)\)\\bar\{\\eta\}\>2/\(1\+\\exp\(\\lambda\\bar\{h\}\)\), then: 𝒲2\(ρN,π\)≤\(\(1−η¯\)\(1\+e−λh¯\)\+e−λh¯\)N𝒲2\(ρ0,π\);\\mathcal\{W\}\_\{2\}\(\\rho\_\{N\},\\pi\)\\leq\\big\(\(1\-\\bar\{\\eta\}\)\(1\+e^\{\-\\lambda\\bar\{h\}\}\)\+e^\{\-\\lambda\\bar\{h\}\}\\big\)^\{N\}\\mathcal\{W\}\_\{2\}\(\\rho\_\{0\},\\pi\);\(11\) main\-pratenddefaultcategory\.tex main\-pratenddefaultcategory\.tex As a result, we have shown a geometric convergence rate forOSISunder aλ\\lambda\-SLC assumption onπ\\pi\. ##### Beyond strong log\-concavity\. The same update also converges for non\-log\-concaveπ\\pisatisfying anα\\alpha\-log\-Sobolev inequality \(α\\alpha\-LSI\) and a lower curvature bound∇2ℰ⪰−κI\\nabla^\{2\}\\mathcal\{E\}\\succeq\-\\kappa I\. Hereα\>0\\alpha\>0andα\\alpha\-LSI means2αKL\(ρ∥π\)≤∫\|∇log\(ρ/π\)\|2ρdx2\\alpha\\operatorname\{KL\}\\left\(\\rho\\\!\\parallel\\\!\\pi\\right\)\\leq\\int\|\\nabla\\log\(\\rho/\\pi\)\|^\{2\}\\rho\\,\\mathrm\{d\}xfor allρ\\rho\. In the appendix \(\), we prove convergence rates, which in particular imply: main\-pratenddefaultcategory\.tex main\-pratenddefaultcategory\.tex main\-pratenddefaultcategory\.tex main\-pratenddefaultcategory\.tex ###### Corollary 4\.3\(Geometric convergence under LSI and bounded negative curvature\)\. Under the notations of[§4\.1](https://arxiv.org/html/2610.12679#S4.SS1), and theα\\alpha\-LSI and curvature assumptions above, with fixedhk=h\>0h\_\{k\}=h\>0andηk=η∈\(0,1\]\\eta\_\{k\}=\\eta\\in\(0,1\]withκh\(1−η\)<2\\kappa h\(1\-\\eta\)<2,𝒲22\(ρk,π\)\\mathcal\{W\}^\{2\}\_\{2\}\(\\rho\_\{k\},\\pi\)andKL\(ρk∥π\)\\operatorname\{KL\}\\left\(\\rho\_\{k\}\\\!\\parallel\\\!\\pi\\right\)converge geometrically to 0\. main\-pratenddefaultcategory\.tex main\-pratenddefaultcategory\.tex ## 5Practical considerations The idealised setting, which is detailed in[Algorithm1](https://arxiv.org/html/2610.12679#alg1)and theoretically analysed in[Section4](https://arxiv.org/html/2610.12679#S4)cannot be implemented in practice\. Instead, we approximate the distributions and the divergence between them 8by their empirical counterparts and replace the continuous\-time Langevin by a discrete\-time MCMC simulation\. The next paragraphs detail various design choices forOSISrelevant in practice\. An extended introduction to the practical considerations of the algorithm is provided in[AppendixC](https://arxiv.org/html/2610.12679#A3)\. ##### Refinement kernel\. OSIScan be instantiated with different MCMC kernels\. The theory concerns exact Langevin, but we use the MALA kernel in practice and consider other choices in[Table3](https://arxiv.org/html/2610.12679#S6.T3)\. ##### Empirical divergence\. It is not possible to use infinite populations as considered in the theory, insteadOSIScomputes finite\-batch empirical divergences between generated samples and their kernel\-refined versions\. Our presented theory considers a squared 2\-Wasserstein loss, butOSIScan be extended to other discrepancies between distributions\. For the experiments on the synthetic targets \([Section6\.1](https://arxiv.org/html/2610.12679#S6.SS1)\) we use sliced Wasserstein distances\([Bonneel et al\., 2015](https://arxiv.org/html/2610.12679#bib.bib17), SWD;\), and for the experiments on Bayesian posteriors \([Section6\.2](https://arxiv.org/html/2610.12679#S6.SS2)\) we use minibatch OT with squared\-Euclidean cost, which converges to exact 2\-Wasserstein in the limit\([Fatras et al\., 2021](https://arxiv.org/html/2610.12679#bib.bib47)\)\. ##### Annealing\. Kernel refinements can initially target the tempered distributionπβ∝exp\(−βℰ\)\\pi\_\{\\beta\}\\propto\\exp\(\-\\beta\\mathcal\{E\}\)withβt\\beta\_\{t\}and then annealβ\\betato 1 over the course of training\. The starting temperature and annealing schedule balance early exploration with learning the final target\. ##### Complete algorithm\. To facilitate training stability, we make several gradient steps after each MCMC refinement\. The complete training procedure detailing the use of empirical divergence, finite\-step MCMC kernel, annealing, and multiple gradient steps per iteration is presented in[Algorithm3](https://arxiv.org/html/2610.12679#alg3)\. ## 6Experiments ### 6\.1Synthetic benchmarks ##### Target distributions\. We evaluateOSISon standard sampling benchmarks: Neal’s funnel \(d=10d=10\)\([Neal, 2003](https://arxiv.org/html/2610.12679#bib.bib46)\), ManyWell \(d=32,64d=32,64\)\([Noé et al\., 2019](https://arxiv.org/html/2610.12679#bib.bib12)\), and 40GMM \(d=2,10d=2,10\)\([Blessing et al\., 2024](https://arxiv.org/html/2610.12679#bib.bib5)\)\. We also consider benchmarks in higher dimension: NICE\-MNIST \(d=196d=196\) and NICE\-FMNIST \(d=728d=728\)\([Blessing et al\., 2024](https://arxiv.org/html/2610.12679#bib.bib5)\)\. Table 1:Sampler comparison on GMM, Funnel, and ManyWell targets\. Entries are mean±\\pmstandard deviation over five seeds\. Second column shows number of NN and energy evaluations per sampleTable 2:Sampler comparison on NICE targets\. Conventions follow[Table1](https://arxiv.org/html/2610.12679#S6.T1)\. ##### Baselines\. We evaluateOSISagainst several amortised samplers: the normalising\-flow\-based FAB\([Midgley et al\., 2023](https://arxiv.org/html/2610.12679#bib.bib7)\)and diffusion samplers trained with different algorithms – PIS\([Zhang and Chen, 2022](https://arxiv.org/html/2610.12679#bib.bib39)\), DIS\([Berner et al\., 2024](https://arxiv.org/html/2610.12679#bib.bib9)\), and off\-policy methods TB\+\+MALA\([Sendera et al\., 2024](https://arxiv.org/html/2610.12679#bib.bib10)\)and SubTB\+\+SMC\([Choi et al\., 2026](https://arxiv.org/html/2610.12679#bib.bib52)\)\. Since FAB can either be used as a one\-step sampler or placed in an annealed importance sampling loop, we compareOSISto the one\-step version of FAB as well\. Additionally, we compareOSISto SMC\([Del Moral et al\., 2006](https://arxiv.org/html/2610.12679#bib.bib6)\), a training\-free sampling algorithm\. Figure 3:Sampling time comparison on ManyWell \(median of250250runs\)\. DIS is trained for different numbers of denoising steps\.OSISachieves𝒲22\\mathcal\{W\}\_\{2\}^\{2\}comparable to the baselines while sampling up to1000×1000\\timesfaster\.Figure 4:Visualisation of samples generated byOSISand the baselines for 40GMM \(d=2d=2\), Funnel \(d=10d=10\), ManyWell \(d=32d=32\), MNIST \(d=196d=196\)\. Funnel and ManyWell are projected ontoℝ2\\mathbb\{R\}^\{2\}\. ##### Results\. [Table1](https://arxiv.org/html/2610.12679#S6.T1)shows sample\-based metrics on synthetic targets, with 2\-Wasserstein and MMD evaluated between 10000 draws from a trainedOSISmodel and an equal number of oracle samples from the target\.OSISis not consistently outperformed by any of the MCMC, normalising flow, or diffusion\-based baselines despite requiring only one neural function evaluation and no energy evaluations to draw a sample\. While the diffusion\-based samplers trained by off\-policy methods\([Sendera et al\., 2024](https://arxiv.org/html/2610.12679#bib.bib10);[Choi et al\., 2026](https://arxiv.org/html/2610.12679#bib.bib52)\)are the strongest overall, they require many sampling steps to perform well,*cf\.*the analysis in[Berner et al\. \(2026\)](https://arxiv.org/html/2610.12679#bib.bib11)\. On the higher\-dimensional targets – given by normalising flow densities trained on MNIST and FMNIST treated as energy\-based models –OSISagain performs on par with FAB and stronger than diffusion\-based samplers \([Table2](https://arxiv.org/html/2610.12679#S6.T2)\)\. [Figure4](https://arxiv.org/html/2610.12679#S6.F4)shows samples from various models on 40GMM \(d=2d=2\), Funnel \(d=10d=10\), ManyWell \(d=32d=32\), and MNIST \(d=2d=2\)\. It is clear from the figure thatOSISperforms comparably in terms of sampling quality while being significantly faster at inference\. ##### Sampling efficiency\. The main advantage of one\-step generative models is low sampling cost\.OSIS, a one\-step sampler, is much more efficient at sampling time than most baselines, whose sampling procedures require hundreds of function evaluations\. To quantify this exactly, in[Figure3](https://arxiv.org/html/2610.12679#S6.F3)we show sampling time per batch of20002000samples on ManyWell \(d=32d=32\), measured on an NVIDIA GH200 GPU, excluding compilation and warm\-up\.OSISyields a significant cost reduction with comparable or improved metrics \([Table1](https://arxiv.org/html/2610.12679#S6.T1)\)\. To assess the trade\-off between number of sampling steps and sampling quality, we also train and evaluate the diffusion\-based sampler DIS with44,1616,6464, and128128denoising steps\. Although DIS performance improves with the number of denoising steps,OSISconsistently outperforms at strictly lower sampling time\. Figure 5:Uncurated samples obtained using outsourced sampling \([Section6\.2](https://arxiv.org/html/2610.12679#S6.SS2)\)\. The samplers correspond to the ImageReward\-based target conditioned on the given prompts\. More examples in[SectionE\.6](https://arxiv.org/html/2610.12679#A5.SS6)\.Table 3:OSIS kernel ablation withM=1M=1inner iterations \([Algorithm3](https://arxiv.org/html/2610.12679#alg3)\); details in[AppendixD](https://arxiv.org/html/2610.12679#A4)\. ##### Ablation of kernel choice\. [Table3](https://arxiv.org/html/2610.12679#S6.T3)compares different choices of the kernel inOSIS\(see[§C\.1](https://arxiv.org/html/2610.12679#A3.SS1)for a description of the kernels\)\. MALA, ULA, and OBABO tend to preform the strongest, and the higher\-order splitting schemes for underdamped Langevin \(OBABO and BAOAB\) often improve over MALA, approaching the diffusion sampler baselines\. However, based on the experiments,OSISdoes not consistently benefit from sophisticated MCMC kernels\. Additionally, we compare with the RWMH kernel, which does not use energy gradients\. Although it does not provide the best performance, it remains competitive and can serve as an alternative when target\-energy gradients are unavailable\. Note that these methods are compared in terms ofOSIStraining steps, not at matched computational cost, and different kernels require different numbers of energy and gradient evaluations per step\. Further ablations are provided in[§E\.4](https://arxiv.org/html/2610.12679#A5.SS4)\. ### 6\.2Bayesian posterior sampling under deep generative model priors We show the scalability ofOSISby applying it to the problem of sampling from a Bayesian posterior under a pretrained generative model prior in latent space, known otherwise as*outsourced sampling*\([Venkatraman et al\., 2025](https://arxiv.org/html/2610.12679#bib.bib79)\)\. Suppose that a pretrained generative model is expressed as a deterministic transformation of a latent variable:x=g\(z\)x=g\(z\), wherez∼q\(z\)z\\sim q\(z\)follows a known base distribution\. The model’s distribution overxx,p=g\#qp=g\_\{\\\#\}q, is treated as a Bayesian prior\. Ifrris a likelihood or reward function, we wish to sample from the posterior, or reward\-tilted distribution \(which has density proportional top\(x\)r\(x\)p\(x\)r\(x\)ifpphas a density\)\. It can be shown that this posterior is the pushforward byggof the latent\-space distribution with density proportional toq\(z\)r\(g\(z\)\)q\(z\)r\(g\(z\)\), reducing sampling from the posterior inxx\-space to sampling from this distribution over the latentzz\. Past work adopting such a formulation\([Venkatraman et al\., 2025](https://arxiv.org/html/2610.12679#bib.bib79);[Tang et al\., 2025](https://arxiv.org/html/2610.12679#bib.bib83);[Kalaivanan et al\., 2025](https://arxiv.org/html/2610.12679#bib.bib85);[Om et al\., 2025](https://arxiv.org/html/2610.12679#bib.bib86);[Wang et al\., 2026](https://arxiv.org/html/2610.12679#bib.bib84);[Tamogashev and Whitammer, 2026](https://arxiv.org/html/2610.12679#bib.bib82)\)used various amortised and Monte Carlo methods to sample from the latent\-space posterior; here, we trainOSISsamplers\. We conduct two sets of experiments: 1. \(1\)withggthe decoder of a VAE\([Kingma and Welling, 2014](https://arxiv.org/html/2610.12679#bib.bib80)\)trained on MNIST with 16\-dimensional latent, with a pretrained classifier’s likelihoodpclassifier\(c∣x\)p\_\{\\text\{classifier\}\}\(c\\mid x\)used to define the reward; 2. \(2\)withggthe generator of a R3GAN\([Huang et al\., 2024](https://arxiv.org/html/2610.12679#bib.bib48)\)trained on FFHQ\-64 with 64\-dimensional latent, with exponentiated ImageReward\([Xu et al\., 2023](https://arxiv.org/html/2610.12679#bib.bib81)\)as a reward,r\(x\)=exp\(βRImageReward\(x,c\)\)r\(x\)=\\exp\(\\beta R\_\{\\text\{ImageReward\}\}\(x,c\)\)for a given captioncc\. Details are given in[AppendixA](https://arxiv.org/html/2610.12679#A1)\. ##### MNIST with classifier reward\. For the experiments with a VAE pretrained on MNIST, we use a mixture of different classes as the reward:r\(x\)=∑c=09wcp\(c∣x\)r\(x\)=\\sum\_\{c=0\}^\{9\}w\_\{c\}p\(c\\mid x\)\. Samples are shown in[SectionE\.5](https://arxiv.org/html/2610.12679#A5.SS5)\. Table 4:Comparison on the outsourced sampling problem\. ##### FFHQ\-64 with ImageReward\. We sample from the reward\-reweighted distribution, which is defined by ImageReward conditioned on a prompt\. For these experiments, we consider the four prompts shown in[Figure5](https://arxiv.org/html/2610.12679#S6.F5), which shows the decoded samples from a trainedOSISmodel\. For the prompta young blond girl, we compareOSISto a diffusion sampler \([Sendera et al\. \(2024\)](https://arxiv.org/html/2610.12679#bib.bib10), as used for this problem in[Venkatraman et al\. \(2025\)](https://arxiv.org/html/2610.12679#bib.bib79)\), to an MCMC algorithm \(MALA with adaptive step size\), and to unbiased samples obtained by an expensive rejection sampling; see[Table4](https://arxiv.org/html/2610.12679#S6.T4)\.OSISperforms similarly – higher mean reward, lower FID – to the diffusion\-based amortised baseline \(TB\+\+MALA\) while requiring only one function evaluation per sample\.OSISfurther improves with MCMC refinement\. Visualisations in[SectionE\.6](https://arxiv.org/html/2610.12679#A5.SS6)show comparable quality of samples\. ## 7Conclusion This paper has proposed an algorithm for one\-step amortised sampling based on matching the distribution induced by a generator with its MCMC refinement\. With appropriate techniques to make the refinement and divergence optimisation practical,OSISperforms surprisingly well compared to strong samplers that require many neural net and energy evaluations to draw each sample\. The most immediate challenge is to improve the sampling quality on highly multimodal targets, where diffusion samplers have been successful due to the geometric mixing induced by the noise, which is absent with tempering alone\. The theoretical analysis should also be extended to include temperature annealing and inexact integration of the dynamics within various refinement kernels\. #### Acknowledgements The work of AC is supported by the UKRI EPSRC through the CDT in Machine Learning Systems hosted in the School of Informatics, University of Edinburgh \(EP/Y03516X/1\), as well as joint sponsorship from Level E Research\. The work of SC and ESW is supported by the Advanced Research and Invention Agency \(ARIA\)\. SB and ESW acknowledge support from the CIFAR Learning in Machines and Brains programme\. This work has been partially supported by the Munich Center for Machine Learning \(MCML\) and G\-Research\. This work was enabled by the computational resources of the Edinburgh International Data Facility \(EIDF\), the Edinburgh Generative AI Laboratory \(GAIL\), and the John von Neumann Institute for Computing \(NIC\) on the GCS Supercomputer JUPITER—JUWELS at Jülich Supercomputing Centre \(JSC\)\. ## References - Albergo and Vanden\-Eijnden \(2025\)M\. S\. Albergo and E\. Vanden\-EijndenNETS: a non\-equilibrium transport sampler\.International Conference on Machine Learning \(ICML\)\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px2.p2.1)\. - Alfonsi and Jourdain \(2020\)A\. Alfonsi and B\. JourdainSquared quadratic Wasserstein distance: optimal couplings and Lions differentiability\.ESAIM: Probability and Statistics\.Cited by:[§4\.1](https://arxiv.org/html/2610.12679#S4.SS1.SSS0.Px2.p2.1)\. - Ambrosioet al\.\(2005\)L\. Ambrosio, N\. Gigli, and G\. SavaréGradient flows: in metric spaces and in the space of probability measures\.Springer\.Cited by:[§A\.1](https://arxiv.org/html/2610.12679#A1.SS1.p1.1),[§B\.1](https://arxiv.org/html/2610.12679#A2.SS1.SSS0.Px2.p2.2),[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px3.p3.2)\. - Andrieu and Thoms \(2008\)C\. Andrieu and J\. ThomsA tutorial on adaptive MCMC\.Statistics and computing\.Cited by:[§C\.1](https://arxiv.org/html/2610.12679#A3.SS1.SSS0.Px8.p1.1)\. - Berneret al\.\(2026\)J\. Berner, L\. Richter, M\. Sendera, J\. Rector\-Brooks, and N\. MalkinFrom discrete\-time policies to continuous\-time diffusion samplers: asymptotic equivalences and faster training\.Transactions on Machine Learning Research \(TMLR\)\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px2.p2.1),[§6\.1](https://arxiv.org/html/2610.12679#S6.SS1.SSS0.Px3.p1.1)\. - Berneret al\.\(2024\)J\. Berner, L\. Richter, and K\. UllrichAn optimal control perspective on diffusion\-based generative modeling\.Transactions on Machine Learning Research \(TMLR\)\.Cited by:[§D\.1](https://arxiv.org/html/2610.12679#A4.SS1.SSS0.Px2.p1.1),[Table 10](https://arxiv.org/html/2610.12679#A5.T10.2.1.11.1),[Table 9](https://arxiv.org/html/2610.12679#A5.T9.2.1.11.1),[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px2.p2.1),[§6\.1](https://arxiv.org/html/2610.12679#S6.SS1.SSS0.Px2.p1.1)\. - Besag \(1994\)J\. BesagComments on “Representations of knowledge in complex systems” by U\. Grenander and M\.I\. Miller\.Journal of the Royal Statistical Society: Series B \(Statistical Methodology\)56,pp\. 591–592\.Cited by:[§C\.1](https://arxiv.org/html/2610.12679#A3.SS1.SSS0.Px3.p1.1),[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px1.p2.1),[§2\.1](https://arxiv.org/html/2610.12679#S2.SS1.SSS0.Px2.p1.1)\. - Blessinget al\.\(2024\)D\. Blessing, X\. Jia, J\. Esslinger, F\. Vargas, and G\. NeumannBeyond ELBOs: a large\-scale evaluation of variational methods for sampling\.International Conference on Machine Learning \(ICML\)\.Cited by:[§D\.1](https://arxiv.org/html/2610.12679#A4.SS1.SSS0.Px1.p1.1),[§D\.1](https://arxiv.org/html/2610.12679#A4.SS1.SSS0.Px2.p1.1),[§D\.1](https://arxiv.org/html/2610.12679#A4.SS1.SSS0.Px2.p2.1),[§6\.1](https://arxiv.org/html/2610.12679#S6.SS1.SSS0.Px1.p1.1)\. - Bonneelet al\.\(2015\)N\. Bonneel, J\. Rabin, G\. Peyré, and H\. PfisterSliced and Radon Wasserstein barycenters of measures\.Journal of Mathematical Imaging and Vision\.Cited by:[§C\.2](https://arxiv.org/html/2610.12679#A3.SS2.SSS0.Px3.p1.1),[§D\.1](https://arxiv.org/html/2610.12679#A4.SS1.SSS0.Px3.p1.1),[§5](https://arxiv.org/html/2610.12679#S5.SS0.SSS0.Px2.p1.1)\. - Caoet al\.\(2026\)W\. Cao, K\. Yan, and L\. ZhaoOne\-step sampler for Boltzmann distributions via drifting\.arXiv preprint arXiv:2603\.17579\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px4.p1.2),[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px2.p3.1)\. - Chenet al\.\(2025a\)J\. Chen, L\. Richter, J\. Berner, D\. Blessing, G\. Neumann, and A\. AnandkumarSequential controlled Langevin diffusions\.International Conference on Learning Representations \(ICLR\)\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px2.p2.1)\. - Chenet al\.\(2025b\)Z\. Chen, A\. Mustafi, P\. Glaser, A\. Korba, A\. Gretton, and B\. K\. Sriperumbudur\(De\)\-regularized maximum mean discrepancy gradient flow\.Journal of Machine Learning Research\.Cited by:[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px3.p4.1)\. - Chewiet al\.\(2025\)S\. Chewi, M\. A\. Erdogdu, M\. Li, R\. Shen, and M\. S\. ZhangAnalysis of Langevin Monte Carlo from Poincaré to log\-Sobolev\.Foundations of Computational Mathematics\.Cited by:[§4](https://arxiv.org/html/2610.12679#S4.p1.2)\. - Choiet al\.\(2026\)S\. Choi, S\. Mittal, V\. Elvira, J\. Park, and E\. S\. WhitammerReinforced sequential Monte Carlo for amortised sampling\.International Conference on Machine Learning \(ICML\)\.Cited by:[§D\.1](https://arxiv.org/html/2610.12679#A4.SS1.SSS0.Px2.p3.1),[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px2.p2.1),[§6\.1](https://arxiv.org/html/2610.12679#S6.SS1.SSS0.Px2.p1.1),[§6\.1](https://arxiv.org/html/2610.12679#S6.SS1.SSS0.Px3.p1.1)\. - Cuturi \(2013\)M\. CuturiSinkhorn distances: lightspeed computation of optimal transport\.Neural Information Processing Systems \(NeurIPS\)\.Cited by:[§C\.2](https://arxiv.org/html/2610.12679#A3.SS2.SSS0.Px4.p1.1)\. - Del Moralet al\.\(2006\)P\. Del Moral, A\. Doucet, and A\. JasraSequential Monte Carlo samplers\.Journal of the Royal Statistical Society: Series B \(Statistical Methodology\)\.Cited by:[§D\.1](https://arxiv.org/html/2610.12679#A4.SS1.SSS0.Px2.p1.1),[Table 10](https://arxiv.org/html/2610.12679#A5.T10.2.1.10.1),[Table 9](https://arxiv.org/html/2610.12679#A5.T9.2.1.10.1),[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px1.p2.1),[§6\.1](https://arxiv.org/html/2610.12679#S6.SS1.SSS0.Px2.p1.1)\. - Denget al\.\(2026\)M\. Deng, H\. Li, T\. Li, Y\. Du, and K\. HeGenerative modeling via drifting\.arXiv preprint arXiv:2602\.04770\.Cited by:[§A\.1](https://arxiv.org/html/2610.12679#A1.SS1.p1.1),[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px3.p2.1),[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px2.p1.1)\. - Doucet al\.\(2018\)R\. Douc, E\. Moulines, P\. Priouret, and P\. SoulierMarkov chains\.Springer\.Cited by:[§2\.1](https://arxiv.org/html/2610.12679#S2.SS1.SSS0.Px2.p1.1)\. - Duaneet al\.\(1987\)S\. Duane, A\. D\. Kennedy, B\. J\. Pendleton, and D\. RowethHybrid Monte Carlo\.Physics Letters B\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px1.p2.1),[§2\.1](https://arxiv.org/html/2610.12679#S2.SS1.SSS0.Px2.p1.1)\. - Durmus and Moulines \(2017\)A\. Durmus and E\. MoulinesNonasymptotic convergence analysis for the unadjusted Langevin algorithm\.The Annals of Applied Probability\.Cited by:[§4](https://arxiv.org/html/2610.12679#S4.p1.2)\. - Fatraset al\.\(2021\)K\. Fatras, Y\. Zine, S\. Majewski, R\. Flamary, R\. Gribonval, and N\. CourtyMinibatch optimal transport distances; analysis and applications\.arXiv preprint arXiv:2101\.01792\.Cited by:[§C\.2](https://arxiv.org/html/2610.12679#A3.SS2.SSS0.Px2.p1.2),[§5](https://arxiv.org/html/2610.12679#S5.SS0.SSS0.Px2.p1.1)\. - Fenget al\.\(2017\)Y\. Feng, D\. Wang, and Q\. LiuLearning to draw samples with amortized Stein variational gradient descent\.Conference on Uncertainty in Artificial Intelligence \(UAI\)\.Cited by:[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px1.p1.1)\. - Feydyet al\.\(2019\)J\. Feydy, T\. Séjourné, F\. Vialard, S\. Amari, A\. Trouvé, and G\. PeyréInterpolating between optimal transport and MMD using Sinkhorn divergences\.International Conference on Artificial Intelligence and Statistics \(AISTATS\)\.Cited by:[§C\.2](https://arxiv.org/html/2610.12679#A3.SS2.SSS0.Px4.p1.3)\. - Flamaryet al\.\(2021\)R\. Flamary, N\. Courty, A\. Gramfort, M\. Z\. Alaya, A\. Boisbunon, S\. Chambon, L\. Chapel, A\. Corenflos, K\. Fatras, N\. Fournier,et al\.Pot: python optimal transport\.Journal of Machine Learning Research\.Cited by:[§D\.1](https://arxiv.org/html/2610.12679#A4.SS1.SSS0.Px1.p2.1)\. - Foglianiet al\.\(2026\)L\. Fogliani, B\. Loureiro, and M\. GabriéAnnealing in variational inference mitigates mode collapse: A theoretical study on Gaussian mixtures\.arXiv preprint arXiv:2602\.12923\.Cited by:[§D\.1](https://arxiv.org/html/2610.12679#A4.SS1.SSS0.Px3.p2.1)\. - Gabriéet al\.\(2022\)M\. Gabrié, G\. M\. Rotskoff, and E\. Vanden\-EijndenAdaptive Monte Carlo augmented with normalizing flows\.Proceedings of the National Academy of Sciences119\(10\),pp\. e2109420119\.Cited by:[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px1.p1.1)\. - Galashovet al\.\(2025\)A\. Galashov, V\. De Bortoli, and A\. GrettonDeep MMD gradient flow without adversarial training\.International Conference on Learning Representations \(ICLR\)\.Cited by:[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px3.p4.1)\. - Grettonet al\.\(2012\)A\. Gretton, K\. M\. Borgwardt, M\. J\. Rasch, B\. Schölkopf, and A\. SmolaA kernel two\-sample test\.Journal of Machine Learning Research\.Cited by:[§C\.2](https://arxiv.org/html/2610.12679#A3.SS2.SSS0.Px1.p1.1)\. - Grettonet al\.\(2026\)A\. Gretton, L\. K\. Wenliang, A\. Galashov, J\. Thornton, V\. De Bortoli, and A\. DoucetOn the Wasserstein gradient flow interpretation of drifting models\.arXiv preprint arXiv:2605\.05118\.Cited by:[§A\.1](https://arxiv.org/html/2610.12679#A1.SS1.p1.1),[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px3.p2.1),[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px2.p2.1),[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px3.p4.1)\. - Hanet al\.\(2026\)J\. Han, P\. Li, Q\. Guo, R\. Xu, S\. Ermon, and E\. J\. CandèsOne\-step generative modeling via Wasserstein gradient flows\.arXiv preprint arXiv:2605\.11755\.Cited by:[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px2.p2.1),[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px3.p4.1)\. - Hastings \(1970\)W\. K\. HastingsMonte Carlo sampling methods using Markov chains and their applications\.\.Cited by:[§C\.1](https://arxiv.org/html/2610.12679#A3.SS1.SSS0.Px1.p1.1),[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px1.p2.1)\. - Havenset al\.\(2025\)A\. Havens, B\. K\. Miller, B\. Yan, C\. Domingo\-Enrich, A\. Sriram, B\. Wood, D\. Levine, B\. Hu, B\. Amos, B\. Karrer, X\. Fu, G\. Liu, and R\. T\. Q\. ChenAdjoint sampling: highly scalable diffusion samplers via adjoint matching\.International Conference on Machine Learining \(ICML\)\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px2.p2.1)\. - Heet al\.\(2025\)J\. He, Y\. Du, F\. Vargas, D\. Zhang, S\. Padhy, R\. OuYang, C\. Gomes, and J\. M\. Hernández\-LobatoNo trick, no treat: pursuits and challenges towards simulation\-free training of neural samplers\.Frontiers in Probabilistic Inference workshop at International Conference on Machine Learining \(ICML\)\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px3.p1.1)\. - Heet al\.\(2016\)K\. He, X\. Zhang, S\. Ren, and J\. SunDeep residual learning for image recognition\.Computer Vision and Pattern Recognition \(CVPR\)\.Cited by:[§D\.1](https://arxiv.org/html/2610.12679#A4.SS1.SSS0.Px3.p2.1)\. - Heet al\.\(2026\)P\. He, O\. Khangaonkar, H\. Pirsiavash, Y\. Bai, and S\. KolouriSinkhorn\-drifting generative models\.arXiv preprint arXiv:2603\.12366\.Cited by:[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px2.p2.1),[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px3.p4.1)\. - Hoet al\.\(2020\)J\. Ho, A\. Jain, and P\. AbbeelDenoising diffusion probabilistic models\.Neural Information Processing Systems \(NeurIPS\)\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px2.p2.1)\. - Hoffman \(2017\)M\. D\. HoffmanLearning deep latent Gaussian models with Markov chain Monte Carlo\.International Conference on Machine Learning \(ICML\)\.Cited by:[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px1.p1.1)\. - Holdijket al\.\(2023\)L\. Holdijk, Y\. Du, F\. Hooft, P\. Jaini, B\. Ensing, and M\. WellingStochastic optimal control for collective variable free sampling of molecular transition paths\.Neural Information Processing Systems \(NeurIPS\)\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px1.p1.1)\. - Hoogeboomet al\.\(2022\)E\. Hoogeboom, V\. G\. Satorras, C\. Vignac, and M\. WellingEquivariant diffusion for molecule generation in 3d\.International Conference on Machine Learning \(ICML\)\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px1.p3.1)\. - Huanget al\.\(2024\)Y\. Huang, A\. Gokaslan, V\. Kuleshov, and J\. TompkinThe GAN is dead; long live the GAN\! A modern GAN baseline\.Neural Information Processing Systems \(NeurIPS\)\.Cited by:[§D\.2](https://arxiv.org/html/2610.12679#A4.SS2.SSS0.Px1.p2.1),[item \(2\)](https://arxiv.org/html/2610.12679#S6.I1.i2.p1.1)\. - Jordanet al\.\(1998\)R\. Jordan, D\. Kinderlehrer, and F\. OttoThe variational formulation of the Fokker–Planck equation\.SIAM journal on mathematical analysis\.Cited by:[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px3.p3.1)\. - Kalaivananet al\.\(2025\)A\. Kalaivanan, Z\. Zhao, J\. Sjölund, and F\. LindstenESS\-Flow: training\-free guidance of flow\-based models as inference in source space\.arXiv preprint arXiv:2510\.05849\.Cited by:[§6\.2](https://arxiv.org/html/2610.12679#S6.SS2.p2.1)\. - Kingma and Welling \(2014\)D\. P\. Kingma and M\. WellingAuto\-encoding variational Bayes\.International Conference on Learning Representations \(ICLR\)\.Cited by:[§D\.2](https://arxiv.org/html/2610.12679#A4.SS2.SSS0.Px1.p1.1),[item \(1\)](https://arxiv.org/html/2610.12679#S6.I1.i1.p1.1)\. - Kolouriet al\.\(2019\)S\. Kolouri, K\. Nadjahi, U\. Simsekli, R\. Badeau, and G\. RohdeGeneralized sliced Wasserstein distances\.Neural Information Processing Systems \(NeurIPS\)\.Cited by:[§C\.2](https://arxiv.org/html/2610.12679#A3.SS2.SSS0.Px3.p1.1)\. - Konget al\.\(2021\)Z\. Kong, W\. Ping, J\. Huang, K\. Zhao, and B\. CatanzaroDiffwave: a versatile diffusion model for audio synthesis\.International Conference on Learning Representations \(ICLR\)\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px1.p3.1)\. - Lahlouet al\.\(2023\)S\. Lahlou, T\. Deleu, P\. Lemos, D\. Zhang, A\. Volokhova, A\. Hernández\-Garcıa, L\. N\. Ezzine, Y\. Bengio, and N\. MalkinA theory of continuous generative flow networks\.International Conference on Machine Learning \(ICML\)\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px2.p2.1)\. - Leimkuhler and Matthews \(2013a\)B\. Leimkuhler and C\. MatthewsRational construction of stochastic numerical methods for molecular sampling\.Applied Mathematics Research eXpress\.Cited by:[§C\.1](https://arxiv.org/html/2610.12679#A3.SS1.SSS0.Px4.p1.1)\. - Leimkuhler and Matthews \(2013b\)B\. Leimkuhler and C\. MatthewsRobust and efficient configurational molecular sampling via Langevin dynamics\.The Journal of Chemical Physics138\(17\)\.Cited by:[§C\.1](https://arxiv.org/html/2610.12679#A3.SS1.SSS0.Px4.p2.2)\. - Leimkuhler and Matthews \(2015\)B\. Leimkuhler and C\. MatthewsMolecular dynamics: with deterministic and stochastic numerical methods\.Springer\.Cited by:[§2\.1](https://arxiv.org/html/2610.12679#S2.SS1.SSS0.Px2.p1.1)\. - Liet al\.\(2017\)Y\. Li, R\. E\. Turner, and Q\. LiuApproximate inference with amortised MCMC\.arXiv preprint arXiv:1702\.08343\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px2.p1.1),[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px4.p1.2),[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px1.p2.1)\. - Liu and Wang \(2016\)Q\. Liu and D\. WangStein variational gradient descent: a general purpose Bayesian inference algorithm\.Neural Information Processing Systems \(NeurIPS\)\.Cited by:[§C\.1](https://arxiv.org/html/2610.12679#A3.SS1.SSS0.Px6.p1.1),[§C\.1](https://arxiv.org/html/2610.12679#A3.SS1.p1.1),[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px1.p1.1)\. - Loshchilov and Hutter \(2019\)I\. Loshchilov and F\. HutterDecoupled weight decay regularization\.International Conference on Learning Representations \(ICLR\)\.Cited by:[§D\.2](https://arxiv.org/html/2610.12679#A4.SS2.SSS0.Px3.p1.1)\. - Metropoliset al\.\(1953\)N\. Metropolis, A\. W\. Rosenbluth, M\. N\. Rosenbluth, A\. H\. Teller, and E\. TellerEquation of state calculations by fast computing machines\.The Journal of Chemical Physics\.Cited by:[§C\.1](https://arxiv.org/html/2610.12679#A3.SS1.SSS0.Px1.p1.1),[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px1.p2.1)\. - Midgleyet al\.\(2023\)L\. I\. Midgley, V\. Stimper, G\. N\. Simm, B\. Schölkopf, and J\. M\. Hernández\-LobatoFlow annealed importance sampling bootstrap\.International Conference on Learning Representations \(ICLR\)\.Cited by:[§D\.1](https://arxiv.org/html/2610.12679#A4.SS1.SSS0.Px2.p1.1),[Table 10](https://arxiv.org/html/2610.12679#A5.T10.2.1.8.1),[Table 10](https://arxiv.org/html/2610.12679#A5.T10.2.1.9.1),[Table 9](https://arxiv.org/html/2610.12679#A5.T9.2.1.8.1),[Table 9](https://arxiv.org/html/2610.12679#A5.T9.2.1.9.1),[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px2.p1.1),[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px3.p1.1),[§6\.1](https://arxiv.org/html/2610.12679#S6.SS1.SSS0.Px2.p1.1)\. - Mohamed and Lakshminarayanan \(2016\)S\. Mohamed and B\. LakshminarayananLearning in implicit generative models\.arXiv preprint arXiv:1610\.03483\.Cited by:[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px1.p1.1)\. - Mokrovet al\.\(2021\)P\. Mokrov, A\. Korotin, L\. Li, A\. Genevay, J\. M\. Solomon, and E\. BurnaevLarge\-scale Wasserstein gradient flows\.Neural Information Processing Systems \(NeurIPS\)\.Cited by:[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px3.p4.1)\. - Nadjahiet al\.\(2020\)K\. Nadjahi, A\. Durmus, L\. Chizat, S\. Kolouri, S\. Shahrampour, and U\. SimsekliStatistical and topological properties of sliced probability divergences\.Neural Information Processing Systems \(NeurIPS\)\.Cited by:[§C\.2](https://arxiv.org/html/2610.12679#A3.SS2.SSS0.Px3.p1.1)\. - Neal \(2003\)R\. M\. NealSlice sampling\.The Annals of Statistics\.Cited by:[§6\.1](https://arxiv.org/html/2610.12679#S6.SS1.SSS0.Px1.p1.1)\. - Neal \(2011\)R\. M\. NealMCMC using Hamiltonian dynamics\.Handbook of Markov chain Monte Carlo\.Cited by:[§C\.1](https://arxiv.org/html/2610.12679#A3.SS1.SSS0.Px5.p1.1),[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px1.p2.1),[§2\.1](https://arxiv.org/html/2610.12679#S2.SS1.SSS0.Px2.p1.1)\. - Noéet al\.\(2019\)F\. Noé, S\. Olsson, J\. Köhler, and H\. WuBoltzmann generators: sampling equilibrium states of many\-body systems with deep learning\.Science\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px1.p1.1),[§6\.1](https://arxiv.org/html/2610.12679#S6.SS1.SSS0.Px1.p1.1)\. - Omet al\.\(2025\)K\. Om, K\. Sim, T\. Yun, H\. Kang, and J\. ParkPosterior inference in latent space for scalable constrained black\-box optimization\.arXiv preprint arXiv:2507\.00480\.Cited by:[§6\.2](https://arxiv.org/html/2610.12679#S6.SS2.p2.1)\. - Rezende and Mohamed \(2015\)D\. Rezende and S\. MohamedVariational inference with normalizing flows\.International Conference on Machine Learning \(ICML\)\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px2.p1.1),[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px3.p1.1)\. - Richteret al\.\(2024\)L\. Richter, J\. Berner, and G\. LiuImproved sampling via learned diffusions\.International Conference on Learning Representations \(ICLR\)\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px2.p2.1)\. - Richteret al\.\(2020\)L\. Richter, A\. Boustati, N\. Nüsken, F\. J\. R\. Ruiz, and Ö\. D\. AkyildizVarGrad: a low\-variance gradient estimator for variational inference\.Neural Information Processing Systems \(NeurIPS\)\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px2.p2.1)\. - Roberts and Tweedie \(1996\)G\. O\. Roberts and R\. L\. TweedieExponential convergence of Langevin distributions and their discrete approximations\.\.Cited by:[§C\.1](https://arxiv.org/html/2610.12679#A3.SS1.SSS0.Px3.p1.1),[§C\.1](https://arxiv.org/html/2610.12679#A3.SS1.SSS0.Px3.p1.2),[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px1.p2.1),[§2\.1](https://arxiv.org/html/2610.12679#S2.SS1.SSS0.Px2.p1.1)\. - Ronnebergeret al\.\(2015\)O\. Ronneberger, P\. Fischer, and T\. BroxU\-net: convolutional networks for biomedical image segmentation\.Medical Image Computing and Computer\-Assisted Intervention \(MICCAI\)\.Cited by:[§D\.1](https://arxiv.org/html/2610.12679#A4.SS1.SSS0.Px3.p2.1)\. - Rosskyet al\.\(1978\)P\. J\. Rossky, J\. D\. Doll, and H\. L\. FriedmanBrownian dynamics as smart Monte Carlo simulation\.The Journal of Chemical Physics69,pp\. 4628–4633\.Cited by:[§C\.1](https://arxiv.org/html/2610.12679#A3.SS1.SSS0.Px3.p1.1),[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px1.p2.1),[§2\.1](https://arxiv.org/html/2610.12679#S2.SS1.SSS0.Px2.p1.1)\. - Ruiz and Titsias \(2019\)F\. Ruiz and M\. TitsiasA contrastive divergence for combining variational inference and MCMC\.International Conference on Machine Learning \(ICML\)\.Cited by:[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px1.p1.1)\. - Sahooet al\.\(2024\)S\. S\. Sahoo, M\. Arriola, Y\. Schiff, A\. Gokaslan, E\. Marroquin, J\. T\. Chiu, A\. Rush, and V\. KuleshovSimple and effective masked diffusion language models\.Neural Information Processing Systems \(NeurIPS\)\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px1.p3.1)\. - Samsonovet al\.\(2022\)S\. Samsonov, E\. Lagutin, M\. Gabrié, A\. Durmus, A\. Naumov, and E\. MoulinesLocal\-global MCMC kernels: the best of both worlds\.Neural Information Processing Systems \(NeurIPS\)\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px2.p1.1)\. - Santambrogio \(2015\)F\. SantambrogioOptimal transport for applied mathematicians: calculus of variations, PDEs, and modeling, volume 87 of progress in nonlinear differential equations and their applications\.\.Cited by:[§B\.2](https://arxiv.org/html/2610.12679#A2.SS2.p1.1)\. - Santambrogio \(2017\)F\. Santambrogio\{\\\{euclidean, Metric, and wasserstein\}\\\}gradient flows: an overview\.Bulletin of Mathematical Sciences7\(1\),pp\. 87–154\.Cited by:[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px3.p3.2)\. - Senderaet al\.\(2024\)M\. Sendera, M\. Kim, S\. Mittal, P\. Lemos, L\. Scimeca, J\. Rector\-Brooks, A\. Adam, Y\. Bengio, and E\. S\. WhitammerImproved off\-policy training of diffusion samplers\.Neural Information Processing Systems \(NeurIPS\)\.Cited by:[§D\.1](https://arxiv.org/html/2610.12679#A4.SS1.SSS0.Px2.p3.1),[Figure 9](https://arxiv.org/html/2610.12679#A5.F9),[Figure 9](https://arxiv.org/html/2610.12679#A5.F9.2.3.1),[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px2.p2.1),[§6\.1](https://arxiv.org/html/2610.12679#S6.SS1.SSS0.Px2.p1.1),[§6\.1](https://arxiv.org/html/2610.12679#S6.SS1.SSS0.Px3.p1.1),[§6\.2](https://arxiv.org/html/2610.12679#S6.SS2.SSS0.Px2.p1.1)\. - Sohl\-Dicksteinet al\.\(2015\)J\. Sohl\-Dickstein, E\. Weiss, N\. Maheswaranathan, and S\. GanguliDeep unsupervised learning using nonequilibrium thermodynamics\.International Conference on Machine Learning \(ICML\)\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px2.p2.1)\. - Songet al\.\(2017\)J\. Song, S\. Zhao, and S\. ErmonA\-NICE\-MC: adversarial training for MCMC\.Neural Information Processing Systems \(NeurIPS\)\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px2.p1.1)\. - Songet al\.\(2021\)Y\. Song, J\. Sohl\-Dickstein, D\. P\. Kingma, A\. Kumar, S\. Ermon, and B\. PooleScore\-based generative modeling through stochastic differential equations\.International Conference on Learning Representations \(ICLR\)\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px2.p2.1)\. - Strang \(1968\)G\. StrangOn the construction and comparison of difference schemes\.SIAM journal on numerical analysis5\(3\),pp\. 506–517\.Cited by:[§C\.1](https://arxiv.org/html/2610.12679#A3.SS1.SSS0.Px4.p2.3)\. - Tamogashev and Whitammer \(2026\)K\. Tamogashev and E\. S\. WhitammerData\-to\-energy stochastic dynamics\.International Conference on Learning Representations \(ICLR\)\.Cited by:[§6\.2](https://arxiv.org/html/2610.12679#S6.SS2.p2.1)\. - Tanget al\.\(2025\)Z\. Tang, J\. Peng, J\. Tang, M\. Hong, F\. Wang, and T\. ChangInference\-time alignment of diffusion models with direct noise optimization\.International Conference on Machine Learning \(ICML\)\.Cited by:[§6\.2](https://arxiv.org/html/2610.12679#S6.SS2.p2.1)\. - Tianet al\.\(2026\)C\. Tian, B\. K\. Sriperumbudur, A\. Gretton, and Z\. ChenSobolev regularized MMD gradient flow\.arXiv preprint arXiv:2605\.11884\.Cited by:[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px3.p4.1)\. - Titsias \(2017\)M\. K\. TitsiasLearning model reparametrizations: implicit variational inference by fitting MCMC distributions\.arXiv preprint arXiv:1708\.01529\.Cited by:[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px1.p1.1)\. - Vargaset al\.\(2023\)F\. Vargas, W\. Grathwohl, and A\. DoucetDenoising diffusion samplers\.International Conference on Learning Representations \(ICLR\)\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px2.p2.1)\. - Venkatramanet al\.\(2025\)S\. Venkatraman, M\. Hasan, M\. Kim, L\. Scimeca, M\. Sendera, Y\. Bengio, G\. Berseth, and N\. MalkinOutsourced diffusion sampling: efficient posterior inference in latent spaces of generative models\.International Conference on Machine Learning \(ICML\)\.Cited by:[§6\.2](https://arxiv.org/html/2610.12679#S6.SS2.SSS0.Px2.p1.1),[§6\.2](https://arxiv.org/html/2610.12679#S6.SS2.p1.1),[§6\.2](https://arxiv.org/html/2610.12679#S6.SS2.p2.1)\. - Wanget al\.\(2025\)C\. Wang, W\. Chen, H\. Kanagawa, and C\. OatesReinforcement learning for adaptive MCMC\.International Conference on Artificial Intelligence and Statistics \(AISTATS\)\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px2.p1.1)\. - Wanget al\.\(2026\)Z\. Wang, A\. Harting, M\. Barreau, M\. Zavlanos, and K\. H\. JohanssonSource\-guided flow matching\.International Conference on Learning Representations \(ICLR\)\.Cited by:[§6\.2](https://arxiv.org/html/2610.12679#S6.SS2.p2.1)\. - Welling and Teh \(2011\)M\. Welling and Y\. W\. TehBayesian learning via stochastic gradient Langevin dynamics\.International Conference on Machine Learning \(ICML\)\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px1.p1.1)\. - Wirnsbergeret al\.\(2020\)P\. Wirnsberger, A\. J\. Ballard, G\. Papamakarios, S\. Abercrombie, S\. Racanière, A\. Pritzel, D\. Jimenez Rezende, and C\. BlundellTargeted free energy estimation via learned mappings\.The Journal of Chemical Physics\.Cited by:[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px1.p1.1)\. - Xuet al\.\(2023\)J\. Xu, X\. Liu, Y\. Wu, Y\. Tong, Q\. Li, M\. Ding, J\. Tang, and Y\. DongImageReward: learning and evaluating human preferences for text\-to\-image generation\.Neural Information Processing Systems \(NeurIPS\)\.Cited by:[§D\.2](https://arxiv.org/html/2610.12679#A4.SS2.SSS0.Px1.p2.1),[item \(2\)](https://arxiv.org/html/2610.12679#S6.I1.i2.p1.1)\. - Yi and Liu \(2023\)M\. Yi and S\. LiuSliced Wasserstein variational inference\.Asian Conference on Machine Learning \(ACML\)\.Cited by:[§3](https://arxiv.org/html/2610.12679#S3.SS0.SSS0.Px1.p1.1)\. - Zhang and Chen \(2022\)Q\. Zhang and Y\. ChenPath integral sampler: a stochastic control approach for sampling\.International Conference on Learning Representations \(ICLR\)\.Cited by:[§D\.1](https://arxiv.org/html/2610.12679#A4.SS1.SSS0.Px2.p1.1),[Table 10](https://arxiv.org/html/2610.12679#A5.T10.2.1.12.1),[Table 9](https://arxiv.org/html/2610.12679#A5.T9.2.1.12.1),[§1](https://arxiv.org/html/2610.12679#S1.SS0.SSS0.Px2.p2.1),[§6\.1](https://arxiv.org/html/2610.12679#S6.SS1.SSS0.Px2.p1.1)\. ## Appendix AAdditional discussion of the algorithm ### A\.1Vector\-field\-based loss It is possible to an consider alternative training procedure forOSIS, which is similar to the one used for drifting models\([Deng et al\., 2026](https://arxiv.org/html/2610.12679#bib.bib33);[Gretton et al\., 2026](https://arxiv.org/html/2610.12679#bib.bib55)\)\. Gradient descent on the parametersθ\\thetaalong the objective in \([2](https://arxiv.org/html/2610.12679#S2.E2)\) approximates the Wasserstein gradient descent of the distributionρ\\rhoalong the*Wasserstein gradients of the divergence𝒟\\mathcal\{D\}*\(see[Section3](https://arxiv.org/html/2610.12679#S3)\)\. The theory of Wasserstein geometry\([Ambrosio et al\., 2005](https://arxiv.org/html/2610.12679#bib.bib21)\)tells us that the evolution of a distribution can be identified with the evolution of particles that follow this distribution along a gradient field\. In this case, the evolution of particles follows a vector fieldVVsuch thatV=−∇δ𝒟δρθV=\-\\nabla\\frac\{\\delta\\mathcal\{D\}\}\{\\delta\\rho\_\{\\theta\}\}, with the term on the right being the gradient of the first variation of𝒟\\mathcal\{D\}\. This observation allows us to use an alternative training objective: ℒ\(θ\)=𝔼z∼q\[\|\|Gθ\(z\)−sg\[Gθ\(z\)\+ηV\(Gθ\(z\)\)\]‖2\],\\displaystyle\\mathcal\{L\}\(\\theta\)=\\mathbb\{E\}\_\{z\\sim q\}\\left\[\\left\|\|G\_\{\\theta\}\(z\)\-\\text\{sg\}\\left\[G\_\{\\theta\}\(z\)\+\\eta V\(G\_\{\\theta\}\(z\)\)\\right\]\\right\\\|^\{2\}\\right\],\(12\)where nowVV, rather than the divergence𝒟\\mathcal\{D\}, is estimated empirically using samples fromρθ\\rho\_\{\\theta\}\. The algorithm corresponding to this objective is presented in[Algorithm2](https://arxiv.org/html/2610.12679#alg2)and experiments are in[SectionE\.3](https://arxiv.org/html/2610.12679#A5.SS3)\. ### A\.2Full algorithm [Algorithms1](https://arxiv.org/html/2610.12679#alg1)and[2](https://arxiv.org/html/2610.12679#alg2)describe the idealised version ofOSIS\. However, as discussed in[Section5](https://arxiv.org/html/2610.12679#S5), there are many design choices that can help improve performance of the algorithm in practice\.[Algorithms3](https://arxiv.org/html/2610.12679#alg3)and[4](https://arxiv.org/html/2610.12679#alg4)present the complete version ofOSISand include tempering, approximation of a Markov chain and the divergence with their empirical analogues, as well as the use of a replay buffer and inner iterations\. We note that this work focuses on the variant ofOSISdescribed in[Algorithms1](https://arxiv.org/html/2610.12679#alg1)and[3](https://arxiv.org/html/2610.12679#alg3)\. Pseudocode in[Algorithms2](https://arxiv.org/html/2610.12679#alg2)and[4](https://arxiv.org/html/2610.12679#alg4)is provided for completeness\. Algorithm 2OSIS using a vector fieldVV1:Prior distributon qq, π\\pi\-invariant kernel KπK\_\{\\pi\}, number of chain steps MM, generator GθG\_\{\\theta\}, vector field V\(⋅,⋅\)V\_\{\(\\cdot,\\cdot\)\}, step size η\\eta\. 2:whilenot convergeddo 3: ρθ←Gθ\#q\\rho\_\{\\theta\}\\leftarrow G\_\{\\theta\\,\\\#\}q 4: ⊳\\trianglerightRun a chain for timehh 5: ρh←Khρsg\[θ\]\\rho\_\{h\}\\leftarrow K^\{h\}\\rho\_\{\\text\{sg\}\\left\[\\theta\\right\]\} 6: ⊳\\trianglerightCompute loss usingVV 7: ℒ\(θ\)←𝔼z∼q\(z\),x=Gθ\(z\)‖Gθ\(z\)−sg\[x\+Vρsg\[θ\],ρh\(x\)\]‖2\\mathcal\{L\}\(\\theta\)\\leftarrow\\mathbb\{E\}\_\{z\\sim q\(z\),x=G\_\{\\theta\}\(z\)\}\\\|G\_\{\\theta\}\(z\)\-\\text\{sg\}\\left\[x\+V\_\{\\rho\_\{\\text\{sg\}\\left\[\\theta\\right\]\},\\,\\rho\_\{h\}\}\(x\)\\right\]\\\|^\{2\} 8:Gradient update on θ\\thetawith ∇θℒ\\nabla\_\{\\theta\}\\mathcal\{L\} 9:endwhile 10:return GθG\_\{\\theta\} Algorithm 3FullOSISalgorithm with𝒟^\\hat\{\\mathcal\{D\}\}1:Prior distribution qq, generator GθG\_\{\\theta\}, π\\pi\-invariant kernel KK, empirical divergence 𝒟^\\hat\{\\mathcal\{D\}\}, number of outer steps NN, number of inner steps MM, number of chain steps NMCMCN\_\{\\rm MCMC\}, buffer size BB, batch size B~\\tilde\{B\}, 2:for outer=1,…,N\\texttt\{outer\}=1,\\dots,Ndo 3: \{zi\}i=1B∼𝒩\(0,I\),xi←Gsg\[θ\]\(zi\)\\\{z\_\{i\}\\\}\_\{i=1\}^\{B\}\\sim\\mathcal\{N\}\(0,I\),\\;x\_\{i\}\\leftarrow G\_\{\\text\{sg\}\\left\[\\theta\\right\]\}\(z\_\{i\}\) 4: ⊳\\trianglerightRun a chain forNMCMCN\_\{\\rm MCMC\}steps 5: ⊳\\trianglerightwith inverse temperatureβ\\beta 6: β←get\_inverse\_temp\(outer\)\\beta\\leftarrow\\operatorname\{get\\\_inverse\\\_temp\}\(\\texttt\{outer\}\) 7: yi←run\_mcmc\(xi,K,NMCMC,β\)y\_\{i\}\\leftarrow\\operatorname\{run\\\_mcmc\}\(x\_\{i\},\\,K,\\,N\_\{\\rm MCMC\},\\beta\) 8: ⊳\\trianglerightInstantiate buffer and save samples 9:Instantiate buffer ℬ\\mathcal\{B\} 10:with \(\{zi\}i=1B,\{yi\}i=1B\)\\left\(\\\{z\_\{i\}\\\}\_\{i=1\}^\{B\},\\,\\\{y\_\{i\}\\\}\_\{i=1\}^\{B\}\\right\) 11: ⊳\\trianglerightRun several inner iterations 12: ⊳\\trianglerightfor computational stability 13:for inner=1,…,M\\texttt\{inner\}=1,\\dots,Mdo 14:ifresampleisTruethen 15: \(\{zi\}i=1B~,\{yi\}i=1B~\)∼ℬ\\left\(\\\{z\_\{i\}\\\}\_\{i=1\}^\{\\tilde\{B\}\},\\,\\\{y\_\{i\}\\\}\_\{i=1\}^\{\\tilde\{B\}\}\\right\)\\sim\\mathcal\{B\} 16:else 17:Get \(\{zi\}i=1B~,\{yi\}i=1B~\)\\left\(\\\{z\_\{i\}\\\}\_\{i=1\}^\{\\tilde\{B\}\},\\,\\\{y\_\{i\}\\\}\_\{i=1\}^\{\\tilde\{B\}\}\\right\)from ℬ\\mathcal\{B\} 18:endif 19: xiθ←Gθ\(zi\)x^\{\\theta\}\_\{i\}\\leftarrow G\_\{\\theta\}\(z\_\{i\}\) 20: p^xθ←∑j=1B~δxjθ\\hat\{p\}^\{\\theta\}\_\{x\}\\leftarrow\\sum\_\{j=1\}^\{\\tilde\{B\}\}\\delta\_\{x^\{\\theta\}\_\{j\}\}, p^y←∑j=1B~δyj\\hat\{p\}\_\{y\}\\leftarrow\\sum\_\{j=1\}^\{\\tilde\{B\}\}\\delta\_\{y\_\{j\}\} 21: ℒ\(θ\)←𝒟^\(p^xθ,p^y\)\\mathcal\{L\}\(\\theta\)\\leftarrow\\hat\{\\mathcal\{D\}\}\(\\hat\{p\}^\{\\theta\}\_\{x\},\\hat\{p\}\_\{y\}\) 22:Gradient update on θ\\thetawith ∇θℒ\\nabla\_\{\\theta\}\\mathcal\{L\} 23:endfor 24:endfor 25:return GθG\_\{\\theta\} Algorithm 4FullOSISalgorithm withVV1:Prior distribution qq, generator GθG\_\{\\theta\}, π\\pi\-invariant kernel KK, vector field V\(⋅,⋅\)V\_\{\(\\cdot,\\cdot\)\}, number of outer steps NN, number of inner steps MM, number of chain steps NMCMCN\_\{\\rm MCMC\}, buffer size BB, batch size B~\\tilde\{B\}, step size η\\eta\. 2:for outer=1,…,N\\texttt\{outer\}=1,\\dots,Ndo 3: \{zi\}i=1B∼𝒩\(0,I\),xi←Gθ\(zi\)\\\{z\_\{i\}\\\}\_\{i=1\}^\{B\}\\sim\\mathcal\{N\}\(0,I\),\\;x\_\{i\}\\leftarrow G\_\{\\theta\}\(z\_\{i\}\) 4: ⊳\\trianglerightRun a chain forNMCMCN\_\{\\rm MCMC\}steps 5: ⊳\\trianglerightwith inverse temperatureβ\\beta 6: β←get\_inverse\_temp\(outer\)\\beta\\leftarrow\\operatorname\{get\\\_inverse\\\_temp\}\(\\texttt\{outer\}\) 7: yi←run\_mcmc\(sg\[xi\],K,NMCMC,β\)y\_\{i\}\\leftarrow\\operatorname\{run\\\_mcmc\}\(\\text\{sg\}\\left\[x\_\{i\}\\right\],\\,K,\\,N\_\{\\rm MCMC\},\\beta\) 8: ⊳\\trianglerightInstantiate buffer and save samples 9:Instantiate buffer ℬ\\mathcal\{B\}with \(\{zi\}i=1B,\{yi\}i=1B\)\\left\(\\\{z\_\{i\}\\\}\_\{i=1\}^\{B\},\\,\\\{y\_\{i\}\\\}\_\{i=1\}^\{B\}\\right\) 10: ⊳\\trianglerightRun several inner iterations 11: ⊳\\trianglerightfor computational stability 12:for inner=1,…,M\\texttt\{inner\}=1,\\dots,Mdo 13:ifresampleisTruethen 14: \(\{zi\}i=1B~,\{yi\}i=1B~\)∼ℬ\\left\(\\\{z\_\{i\}\\\}\_\{i=1\}^\{\\tilde\{B\}\},\\,\\\{y\_\{i\}\\\}\_\{i=1\}^\{\\tilde\{B\}\}\\right\)\\sim\\mathcal\{B\} 15:else 16:Get \(\{zi\}i=1B~,\{yi\}i=1B~\)\\left\(\\\{z\_\{i\}\\\}\_\{i=1\}^\{\\tilde\{B\}\},\\,\\\{y\_\{i\}\\\}\_\{i=1\}^\{\\tilde\{B\}\}\\right\)from ℬ\\mathcal\{B\} 17:endif 18: xiθ←Gθ\(zi\)x^\{\\theta\}\_\{i\}\\leftarrow G\_\{\\theta\}\(z\_\{i\}\) 19: ⊳\\trianglerightz∼qz\\sim q,x=Gθ\(z\)x=G\_\{\\theta\}\(z\) 20: ℒ\(θ\)←𝔼z,x‖Gθ\(z\)−sg\[x\+ηVρsg\[θ\],ρh\(x\)\]‖2\\begin\{aligned\} \\mathcal\{L\}\(\\theta\)&\\leftarrow\\\\ &\\mathbb\{E\}\_\{z,x\}\\\|G\_\{\\theta\}\(z\)\-\\text\{sg\}\\left\[x\+\\eta V\_\{\\rho\_\{\\text\{sg\}\\left\[\\theta\\right\]\},\\,\\rho\_\{h\}\}\(x\)\\right\]\\\|^\{2\}\\end\{aligned\} 21:Gradient update on θ\\thetawith ∇θℒ\\nabla\_\{\\theta\}\\mathcal\{L\} 22:endfor 23:endfor 24:return GθG\_\{\\theta\} ## Appendix BExtended theoretical analysis ### B\.1Additional notation and temperature scaling ##### KL divergence and Fisher information\. For distributionswρ,μ∈𝒫2\(ℝd\)\\rho,\\mu\\in\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\), we define KL divergence and Fisher information as follows: KL\(ρ∥μ\)≔∫ρ\(x\)logρ\(x\)μ\(x\)dx,\\operatorname\{KL\}\\left\(\\rho\\\!\\parallel\\\!\\mu\\right\)\\coloneqq\\int\\rho\(x\)\\log\\frac\{\\rho\(x\)\}\{\\mu\(x\)\}\\,\\mathrm\{d\}x,\(13\) FI\(ρ∥μ\)≔∫\|∇logρ\(x\)μ\(x\)\|2ρ\(x\)dx\.\\operatorname\{FI\}\\left\(\\rho\\\!\\parallel\\\!\\mu\\right\)\\coloneqq\\int\\left\|\\nabla\\log\\frac\{\\rho\(x\)\}\{\\mu\(x\)\}\\right\|^\{2\}\\rho\(x\)\\,\\mathrm\{d\}x\.\(14\) We assume that for all cases we consider both KL divergence and Fisher information are finite\. ##### 2\-Wasserstein distance\. For distributionsρ,μ∈𝒫2\(ℝd\)\\rho,\\mu\\in\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\), we define 2\-Wasserstein distance as 𝒲22\(ρ,μ\)=infγ∈Γ\(ρ,μ\)∫ℝd×ℝd\|x−y\|2𝑑γ\(x,y\),\\mathcal\{W\}\_\{2\}^\{2\}\(\\rho,\\mu\)=\\inf\_\{\\gamma\\in\\Gamma\(\\rho,\\mu\)\}\\int\_\{\\mathbb\{R\}^\{d\}\\times\\mathbb\{R\}^\{d\}\}\|x\-y\|^\{2\}\\,\\mathrm\{d\}\\gamma\(x,y\),\(15\)whereΓ\(ρ,μ\)\\Gamma\(\\rho,\\mu\)is the set of probability measures onℝd×ℝd\\mathbb\{R\}^\{d\}\\times\\mathbb\{R\}^\{d\}whose marginals areρ\\rhoandμ\\mu\. \([15](https://arxiv.org/html/2610.12679#A2.E15)\) defines the Kantorovich version of OT\. For measures in𝒫2\(ℝd\)\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\), the infimum is achieved \(*i\.e\.*, the minimal coupling exists\) and the minimiser is unique\. Moreover, for𝒫2\(ℝd\)\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\)the optimal coupling has the formγ∗=\(Id,T\)\#ρ\\gamma^\{\*\}=\(\\operatorname\{Id\},T\)\_\{\\\#\}\\rhofor someT:ℝd→ℝdT:\\mathbb\{R\}^\{d\}\\to\\mathbb\{R\}^\{d\}\. ThisTTis the solution to the Monge OT problem, which appears in another \(in this case, equivalent\) definition of 2\-Wasserstein distance: 𝒲22\(ρ,μ\)=infT:T\#ρ=μ∫\|T\(x\)−x\|2ρ\(x\)dx\.\\mathcal\{W\}^\{2\}\_\{2\}\(\\rho,\\mu\)=\\inf\_\{T:T\_\{\\\#\}\\rho=\\mu\}\\int\|T\(x\)\-x\|^\{2\}\\,\\rho\(x\)\\,\\mathrm\{d\}x\.\(16\)Yet a third useful definition of 2\-Wasserstein distance, which we use to prove, is due to Benamou\-Brenier\([Ambrosio et al\., 2005](https://arxiv.org/html/2610.12679#bib.bib21), Equation 8\.0\.3\)and employs the dynamic formulation of optimal transport: 𝒲22\(ρ,μ\)=inf∫01∫\|vt\|2ρt\(x\)𝑑x𝑑t,\\mathcal\{W\}^\{2\}\_\{2\}\(\\rho,\\mu\)=\\inf\\int\_\{0\}^\{1\}\\int\|v\_\{t\}\|^\{2\}\\rho\_\{t\}\(x\)\\,\\mathrm\{d\}x\\,\\mathrm\{d\}t,\(17\)where the infimum is taken over time\-dependent vector fieldsvtv\_\{t\}and densitiesρt\\rho\_\{t\}satisfying the continuity equation∂∂tρt\+∇⋅\(ρtvt\)=0\\frac\{\\partial\}\{\\partial t\}\\rho\_\{t\}\+\\nabla\\cdot\(\\rho\_\{t\}\\,v\_\{t\}\)=0and boundary conditionsρ0=ρ,ρ1=μ\\rho\_\{0\}=\\rho,\\rho\_\{1\}=\\mu\. In any of these definitions,𝒲2\\mathcal\{W\}\_\{2\}satisfies the triangle inequality, meaning that 2\-Wasserstein is a distance metric in𝒫2\(ℝd\)\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\)\. ### B\.2Proofs We first recall some standard definitions and results in Wasserstein geometry \(see,*e\.g\.*,[Santambrogio \(2015\)](https://arxiv.org/html/2610.12679#bib.bib20)\), then prove the results from the main text\. ###### Definition B\.1\(Geodesic convexity\)\. A functional𝒥\(ρ\):𝒫2\(ℝd\)→ℝ\\mathcal\{J\}\(\\rho\):\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\)\\to\\mathbb\{R\}is said to be - •geodesically convexif, for everyρ0,ρ1∈𝒫2\(ℝd\)\\rho\_\{0\},\\rho\_\{1\}\\in\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\), whenρt\\rho\_\{t\}is the geodesic betweenρ0\\rho\_\{0\}andρ1\\rho\_\{1\}, ∀t∈\[0,1\]𝒥\(ρt\)≤t𝒥\(ρ1\)\+\(1−t\)𝒥\(ρ0\);\\forall t\\in\[0,1\]\\quad\\mathcal\{J\}\(\\rho\_\{t\}\)\\leq t\\mathcal\{J\}\(\\rho\_\{1\}\)\+\(1\-t\)\\mathcal\{J\}\(\\rho\_\{0\}\);\(18\) - •λ\\lambda\-strongly geodesically convex\(for someλ\>0\\lambda\>0\) if, for everyρ0,ρ1∈𝒫2\(ℝd\)\\rho\_\{0\},\\rho\_\{1\}\\in\\mathcal\{P\}\_\{2\}\(\\mathbb\{R\}^\{d\}\), whenρt\\rho\_\{t\}is the geodesic betweenρ0\\rho\_\{0\}andρ1\\rho\_\{1\}, ∀t∈\[0,1\]𝒥\(ρt\)≤t𝒥\(ρ1\)\+\(1−t\)𝒥\(ρ0\)−λ2t\(1−t\)𝒲22\(ρ0,ρ1\)\.\\forall t\\in\[0,1\]\\quad\\mathcal\{J\}\(\\rho\_\{t\}\)\\leq t\\mathcal\{J\}\(\\rho\_\{1\}\)\+\(1\-t\)\\mathcal\{J\}\(\\rho\_\{0\}\)\-\\frac\{\\lambda\}\{2\}t\(1\-t\)\\mathcal\{W\}^\{2\}\_\{2\}\(\\rho\_\{0\},\\rho\_\{1\}\)\.\(19\) ## Appendix CExtended practical considerations This appendix details the main design choices ofOSIS\. It expands on the kernel, loss, tempering, and inference choices in[§5](https://arxiv.org/html/2610.12679#S5)\. This section aims to provide a general overview\. Specific details about the experiments are provided in[AppendixD](https://arxiv.org/html/2610.12679#A4)\. ### C\.1Choice of MCMC kernel We present here various methods to update the generated samples towards the targetπ\\pi\. They often consist of aπ\\pi\-invariant MCMC kernel, but other methods, such as Stein variational gradient descent\([Liu and Wang, 2016](https://arxiv.org/html/2610.12679#bib.bib26)\), can also be investigated\. In this subsection,q\(y∣x\)q\(y\\mid x\)denotes a proposal density\. We describe kernels for a fixed targetπ\\pi\. Under an annealing schedule, the same constructions useπβt\\pi\_\{\\beta\_\{t\}\}instead\. ##### Metropolis\-Hastings\. The Metropolis\-Hastings\([Metropolis et al\., 1953](https://arxiv.org/html/2610.12679#bib.bib36);[Hastings, 1970](https://arxiv.org/html/2610.12679#bib.bib37), MH\)kernel modifies a continuous full\-support proposal to beπ\\pi\-invariant by mixing it with a rejection kernel\. To be precise, given a proposal kernelqq, define K\(y∣x\)=α\(x,y\)q\(y∣x\)\+\(1−A\(x\)\)δx\(y\),K\(y\\mid x\)=\\alpha\(x,y\)q\(y\\mid x\)\+\(1\-A\(x\)\)\\delta\_\{x\}\(y\),\(20\)whereδx\\delta\_\{x\}is a Dirac delta atxxand α\(x,y\)=min\(1,π\(y\)q\(x∣y\)π\(x\)q\(y∣x\)\),A\(x\)=∫α\(x,y\)q\(y∣x\)𝑑y\.\\alpha\(x,y\)=\\min\\left\(1,\\frac\{\\pi\(y\)q\(x\\mid y\)\}\{\\pi\(x\)q\(y\\mid x\)\}\\right\),\\quad A\(x\)=\\int\\alpha\(x,y\)q\(y\\mid x\)\\,\\mathrm\{d\}y\.\(21\)By construction,KKisπ\\pi\-invariant \(in particular, ifqqis reversible with respect toπ\\pi, thenA\(x\)=1A\(x\)=1forπ\\pi\-almost everyxxandK=qK=q\)\. However, the choice ofqqgreatly affects the rate of convergence toπ\\piof a Markov chain with transition operatorKK\. ##### Random\-walk proposals\. Random\-walk Metropolis \(RWM\) is a commonly used form of MH\. A standard variant takesqqto be the kenrel of a Gaussian randmo walk:q\(y∣x\)=𝒩\(y,x,σ2I\)q\(y\\mid x\)=\\mathcal\{N\}\(y;x,\\sigma^\{2\}I\)for some chosen varianceσ2\\sigma^\{2\}\. The proposal is symmetric, so its acceptance probability simplifies tomin\(1,π\(y\)/π\(x\)\)\\min\(1,\\pi\(y\)/\\pi\(x\)\)\. RWM requires target energy evaluations but no target energy gradients\. ##### Langevin\-based proposals\. One can make use of target gradient information when defining a proposal distribution\. Metropolis\-Adjusted Langevin\([Rossky et al\., 1978](https://arxiv.org/html/2610.12679#bib.bib88);[Besag, 1994](https://arxiv.org/html/2610.12679#bib.bib87);[Roberts and Tweedie, 1996](https://arxiv.org/html/2610.12679#bib.bib14), MALA;\)is MH based on a proposalqqthat is an Euler\-Maruyama integration step of a Langevin SDE: q\(y∣x\)=𝒩\(y;x−δ∇ℰ\(x\),2δId\)q\(y\\mid x\)=\\mathcal\{N\}\\left\(y;x\-\\delta\\nabla\\mathcal\{E\}\(x\),2\\delta I\_\{d\}\\right\)for some chosen step sizeδ\>0\\delta\>0\. At inverse temperatureβt\\beta\_\{t\}, the drift becomes−δβt∇ℰ\(x\)\-\\delta\\beta\_\{t\}\\nabla\\mathcal\{E\}\(x\)\. Unlike RWM, this proposal is generally asymmetric, so both proposal\-density terms in \([21](https://arxiv.org/html/2610.12679#A3.E21)\) must be retained\. Omitting the accept/reject step gives the unadjusted Langevin algorithm \(ULA\), which does not in general preserve the target at a finite step size\([Roberts and Tweedie, 1996](https://arxiv.org/html/2610.12679#bib.bib14)\)\. Neither a MALA step nor a ULA step is the exact Langevin semigroup transitionKhK^\{h\}used in[§4](https://arxiv.org/html/2610.12679#S4)\. ##### Higher\-order splitting schemes\. Splitting integrators provide further discretisations of Langevin dynamics and can reduce finite\-step sampling bias\([Leimkuhler and Matthews, 2013a](https://arxiv.org/html/2610.12679#bib.bib62)\)\. The idea is to divide the dynamics into simpler parts that can each be integrated exactly, then compose their updates\. We consider the matrix formulation of the underdamped Langevin targetingπ\(x\)𝒩\(v,0,Id\)\\pi\(x\)\\mathcal\{N\}\(v;0,I\_\{d\}\), whereπ∝e−ℰ\\pi\\propto e^\{\-\\mathcal\{E\}\}: \[dXsdVs\]=\[Vs0\]ds⏟A\+\[0−∇ℰ\(Xs\)\]ds⏟B\+\[0−ζVsds\+2ζdBs\]⏟O,\\displaystyle\\begin\{bmatrix\}\\mathrm\{d\}X\_\{s\}\\\\ \\mathrm\{d\}V\_\{s\}\\end\{bmatrix\}=\\underbrace\{\\begin\{bmatrix\}V\_\{s\}\\\\ 0\\end\{bmatrix\}\\mathrm\{d\}s\}\_\{\\mathrm\{A\}\}\+\\underbrace\{\\begin\{bmatrix\}0\\\\ \-\\nabla\\mathcal\{E\}\(X\_\{s\}\)\\end\{bmatrix\}\\mathrm\{d\}s\}\_\{\\mathrm\{B\}\}\+\\underbrace\{\\begin\{bmatrix\}0\\\\ \-\\zeta V\_\{s\}\\,\\mathrm\{d\}s\+\\sqrt\{2\\zeta\}\\,\\mathrm\{d\}B\_\{s\}\\end\{bmatrix\}\}\_\{\\mathrm\{O\}\},\(22\)whereζ\>0\\zeta\>0is the friction coefficient andv∈ℝdv\\in\\mathbb\{R\}^\{d\}is the auxiliary velocity variable\. The underdamped Langevin dynamics is split into three parts: free transport of the position at fixed velocity \(A\), force\-driven evolution of the velocity at fixed position \(B\), and Ornstein\-Uhlenbeck evolution of the velocity, combining friction and Gaussian noise at fixed position \(O\)\. Their exact updates over a durationt≥0t\\geq 0are At:\\displaystyle\\mathrm\{A\}\_\{t\}:\\quad\(x,v\)⟼\(x\+tv,v\),\\displaystyle\(x,v\)\\longmapsto\(x\+tv,v\),\(23\)Bt:\\displaystyle\\mathrm\{B\}\_\{t\}:\\quad\(x,v\)⟼\(x,v−t∇ℰ\(x\)\),\\displaystyle\(x,v\)\\longmapsto\(x,v\-t\\nabla\\mathcal\{E\}\(x\)\),\(24\)Ot:\\displaystyle\\mathrm\{O\}\_\{t\}:\\quad\(x,v\)⟼\(x,e−ζtv\+1−e−2ζtξ\),ξ∼𝒩\(0,Id\)\.\\displaystyle\(x,v\)\\longmapsto\\left\(x,e^\{\-\\zeta t\}v\+\\sqrt\{1\-e^\{\-2\\zeta t\}\}\\,\\xi\\right\),\\qquad\\xi\\sim\\mathcal\{N\}\(0,I\_\{d\}\)\.\(25\)These exact substeps generally do not commute, so their composition provides a specific approximation of the full dynamics\([Leimkuhler and Matthews, 2013b](https://arxiv.org/html/2610.12679#bib.bib63), Section IV\)\. For example, the chronological order of a BAOAB step of sizeδ\\deltais Bδ/2⟶Aδ/2⟶Oδ⟶Aδ/2⟶Bδ/2\.\\mathrm\{B\}\_\{\\delta/2\}\\ \\longrightarrow\\ \\mathrm\{A\}\_\{\\delta/2\}\\ \\longrightarrow\\ \\mathrm\{O\}\_\{\\delta\}\\ \\longrightarrow\\ \\mathrm\{A\}\_\{\\delta/2\}\\ \\longrightarrow\\ \\mathrm\{B\}\_\{\\delta/2\}\.\(26\)The letters specify the order of the updates; each repeated letter here uses a half step\. Each update acts on the current state, with independent noise for each O substep\. Symmetric Strang composition \(*i\.e\.*, letters forming a palindrome\) give second\-order weak accuracy under appropriate conditions\([Strang, 1968](https://arxiv.org/html/2610.12679#bib.bib65)\)\. ##### HMC\. Hamiltonian Monte Carlo introduces auxiliary momentum variables and uses Hamiltonian trajectories, with a correction for numerical integration error, to construct a target\-invariant transition\([Neal, 2011](https://arxiv.org/html/2610.12679#bib.bib38)\)\. ##### Stein variational gradient descent\. Stein variational gradient descent \(SVGD\) transports a set of interacting particles using a deterministic, kernel\-based field\([Liu and Wang, 2016](https://arxiv.org/html/2610.12679#bib.bib26)\)\. For a differentiable positive\-definite kernelkkand step sizeδ\>0\\delta\>0, its update for particles\{xi\}i=1B\\\{x\_\{i\}\\\}\_\{i=1\}^\{B\}is yi=xi\+δB∑j=1B\[−k\(xj,xi\)∇ℰ\(xj\)\+∇xjk\(xj,xi\)\],i=1,…,B\.y\_\{i\}=x\_\{i\}\+\\frac\{\\delta\}\{B\}\\sum\_\{j=1\}^\{B\}\\left\[\-k\(x\_\{j\},x\_\{i\}\)\\nabla\\mathcal\{E\}\(x\_\{j\}\)\+\\nabla\_\{x\_\{j\}\}k\(x\_\{j\},x\_\{i\}\)\\right\],\\qquad i=1,\\ldots,B\.\(27\) ##### Multiple\-step kernels\. IfKKis aπ\\pi\-invariant kernel, so isKnK^\{n\}forn≥1n\\geq 1\. Thus one can construct several\-step kernels by iterating the primitive ones defined above, as we do in the experiments\. ##### Acceptance\-targeting step adaptation\. We adapt the step sizeδ\\deltaof random\-walk or Langevin MCMC during training to maintain a target acceptance rateα⋆\\alpha^\{\\star\}\([Andrieu and Thoms, 2008](https://arxiv.org/html/2610.12679#bib.bib15)\)\. The update is log\-space stochastic approximation, loght←loght\+ηadapt\(α^−α⋆\),\\log h\_\{t\}\\leftarrow\\log h\_\{t\}\+\\eta\_\{\\mathrm\{adapt\}\}\(\\hat\{\\alpha\}\-\\alpha^\{\\star\}\),whereα^\\hat\{\\alpha\}is the batch mean acceptance andht≔δβth\_\{t\}\\coloneqq\\frac\{\\delta\}\{\\beta\_\{t\}\}is the temperature\-rescaled step size\. We then set the sampler’s step size toδt=htβt\\delta\_\{t\}=h\_\{t\}\\beta\_\{t\}and clip it to\[δmin,δmax\]\[\\delta\_\{\\min\},\\delta\_\{\\max\}\]\. ### C\.2Choice of divergence To compute the distance between two distributions, we rely on sample\-based distribution metrics\. In practice, these distances are computed using samples from both distributions\. ##### Maximum mean discrepancy\. One possible way to compute the distance between two distributions is use maximum mean discrepancy\([Gretton et al\., 2012](https://arxiv.org/html/2610.12679#bib.bib35), MMD;\)\. Given a positive definite kernel function, such ask\(x,y\)=exp\(−‖x−y‖h\)k\(x,y\)=\\exp\\left\(\-\\frac\{\\\|x\-y\\\|\}\{h\}\\right\), wherehhis a positive bandwidth and∥⋅∥\\\|\\cdot\\\|isL2L\_\{2\}distance, the MMD between two distributionsμ1\\mu\_\{1\}andμ2\\mu\_\{2\}is defined as: MMD2\(μ1,μ2\)\\displaystyle\\text\{MMD\}^\{2\}\(\\mu\_\{1\},\\mu\_\{2\}\)≔𝔼x,x′∼μ1\[k\(x,x′\)\]\+𝔼y,y′∼μ2\[k\(y,y′\)\]\\displaystyle\\coloneqq\\mathbb\{E\}\_\{x,x^\{\\prime\}\\sim\\mu\_\{1\}\}\[k\(x,x^\{\\prime\}\)\]\+\\mathbb\{E\}\_\{y,y^\{\\prime\}\\sim\\mu\_\{2\}\}\[k\(y,y^\{\\prime\}\)\]−2𝔼x∼μ1,y∼μ2\[k\(x,y\)\],\\displaystyle\\phantom\{=\\;\}\-2\\mathbb\{E\}\_\{x\\sim\\mu\_\{1\},y\\sim\\mu\_\{2\}\}\[k\(x,y\)\],\(28\)which can be estimated by Monte Carlo sampling\. ##### Squared 2\-Wasserstein distance\. For probability distributionsμ1\\mu\_\{1\}andμ2\\mu\_\{2\}onℝd\\mathbb\{R\}^\{d\}with finite second moments, recall from \([16](https://arxiv.org/html/2610.12679#A2.E16)\) that the squared 2\-Wasserstein distance is an optimal transport cost 𝒲22\(μ1,μ2\)≔infγ∈Γ\(μ1,μ2\)∫ℝd×ℝd\|x−y\|2𝑑γ\(x,y\),\\displaystyle\\mathcal\{W\}^\{2\}\_\{2\}\(\\mu\_\{1\},\\mu\_\{2\}\)\\coloneqq\\inf\_\{\\gamma\\in\\Gamma\(\\mu\_\{1\},\\mu\_\{2\}\)\}\\int\_\{\\mathbb\{R\}^\{d\}\\times\\mathbb\{R\}^\{d\}\}\|x\-y\|^\{2\}\\,\\mathrm\{d\}\\gamma\(x,y\),\(29\)whereΓ\(μ1,μ2\)\\Gamma\(\\mu\_\{1\},\\mu\_\{2\}\)contains joint distributions overℝd×ℝd\\mathbb\{R\}^\{d\}\\times\\mathbb\{R\}^\{d\}whose marginals areμ1\\mu\_\{1\}andμ2\\mu\_\{2\}\.The𝒲2\\mathcal\{W\}\_\{2\}distance between empirical approximations toμ1\\mu\_\{1\}andμ2\\mu\_\{2\}is known as the minibatch 2\-Wasserstein, and it converges to \([29](https://arxiv.org/html/2610.12679#A3.E29)\) as the number of samples increases\([Fatras et al\., 2021](https://arxiv.org/html/2610.12679#bib.bib47)\)\. ##### Sliced Wasserstein distance\. The sliced Wasserstein distance\([Bonneel et al\., 2015](https://arxiv.org/html/2610.12679#bib.bib17);[Kolouri et al\., 2019](https://arxiv.org/html/2610.12679#bib.bib18);[Nadjahi et al\., 2020](https://arxiv.org/html/2610.12679#bib.bib19), SWD;\)is a tractable surrogate that computes OT distances using11\-dimensional projections\. It is defined by 𝕊22\(μ1,μ2\)≔𝔼ω\[𝒲22\(Pω\#μ1,Pω\#μ2\)\],\\mathbb\{S\}\_\{2\}^\{2\}\(\\mu\_\{1\},\\mu\_\{2\}\)\\coloneqq\\mathbb\{E\}\_\{\\omega\}\\left\[\\mathcal\{W\}^\{2\}\_\{2\}\(P\_\{\\omega\\,\\\#\}\\mu\_\{1\},P\_\{\\omega\\,\\\#\}\\mu\_\{2\}\)\\right\],\(30\)where the expectation is taken overω\\omegauniform on the unit sphere inℝd\\mathbb\{R\}^\{d\}\. In practice, the expectation over the projection directionω\\omegais estimated with random projections\. The quantity inside the expectation is a 1D OT distance between empirical distributions and is thus easily computed\. ##### Sinkhorn divergence\. Entropic regularisation is introduced to the transport problem \([29](https://arxiv.org/html/2610.12679#A3.E29)\) and provides a smoother optimisation problem that can be solved using Sinkhorn iterations\. For regularisation strengthε\>0\\varepsilon\>0, define\([Cuturi, 2013](https://arxiv.org/html/2610.12679#bib.bib24)\) OTε\(μ1,μ2\)=infγ∈Γ\(μ1,μ2\)\[∫\|x−y\|2dγ\(x,y\)\+εKL\(γ∥μ1⊗μ2\)\]\.\\operatorname\{OT\}\_\{\\varepsilon\}\(\\mu\_\{1\},\\mu\_\{2\}\)=\\inf\_\{\\gamma\\in\\Gamma\(\\mu\_\{1\},\\mu\_\{2\}\)\}\\left\[\\int\|x\-y\|^\{2\}\\,\\mathrm\{d\}\\gamma\(x,y\)\+\\varepsilon\\operatorname\{KL\}\\left\(\\gamma\\\!\\parallel\\\!\\mu\_\{1\}\\otimes\\mu\_\{2\}\\right\)\\right\]\.\(31\)The corresponding debiased Sinkhorn divergence subtracts the two self\-interaction terms, Sε\(μ1,μ2\)=OTε\(μ1,μ2\)−12OTε\(μ1,μ1\)−12OTε\(μ2,μ2\),\\operatorname\{S\}\_\{\\varepsilon\}\(\\mu\_\{1\},\\mu\_\{2\}\)=\\operatorname\{OT\}\_\{\\varepsilon\}\(\\mu\_\{1\},\\mu\_\{2\}\)\-\\tfrac\{1\}\{2\}\\operatorname\{OT\}\_\{\\varepsilon\}\(\\mu\_\{1\},\\mu\_\{1\}\)\-\\tfrac\{1\}\{2\}\\operatorname\{OT\}\_\{\\varepsilon\}\(\\mu\_\{2\},\\mu\_\{2\}\),\(32\)so thatSε\(μ,μ\)=0\\operatorname\{S\}\_\{\\varepsilon\}\(\\mu,\\mu\)=0\. For empirical measures and the squared\-Euclidean cost used here,Sε\(μ1,μ2\)\\operatorname\{S\}\_\{\\varepsilon\}\(\\mu\_\{1\},\\mu\_\{2\}\)approaches𝒲22\(μ1,μ2\)\\mathcal\{W\}^\{2\}\_\{2\}\(\\mu\_\{1\},\\mu\_\{2\}\)asε↓0\\varepsilon\\downarrow 0\. Asε→∞\\varepsilon\\to\\infty, it approaches the squared distance between the means ofμ1\\mu\_\{1\}andμ2\\mu\_\{2\}, equivalently, the squared MMD associated with the linear kernelk\(x,y\)=⟨x,y⟩k\(x,y\)=\\langle x,y\\rangle\([Feydy et al\., 2019](https://arxiv.org/html/2610.12679#bib.bib25), Section 1\.2\)\. ## Appendix DDetailed experiment setup ### D\.1Synthetic targets ##### Setup\. We trainOSISand the baselines for five random seeds\. After training, we draw five independent batches of10,00010\{,\}000samples with different seeds at the selected checkpoint: the final training iteration forOSISand the checkpoint with the highest ELBO for baselines, following[Blessing et al\. \(2024\)](https://arxiv.org/html/2610.12679#bib.bib5)\. Each sampler’s evaluation batch is compared with the same batch of10,00010\{,\}000reference samples for all samplers\. We average the results within each training seed, then report the mean and sample standard deviation across the five training\-seed means\. SMC uses five repetition groups instead of training seeds\. The reference rows provide a floor on each metric by comparing two independent target batches of samples\. We report empirical𝒲22\\mathcal\{W\}\_\{2\}^\{2\}with unnormalised squared Euclidean cost, computed using POT\([Flamary et al\., 2021](https://arxiv.org/html/2610.12679#bib.bib23)\)with convergence checks\. MMD corresponds to the square root of the biased Gaussian kernel MMD\-squared estimate, with bandwidth fixed per target to the median pairwise distance in an independent10,00010\{,\}000\-sample batch from the target distribution\. ##### Baselines\. We retrain FAB\([Midgley et al\., 2023](https://arxiv.org/html/2610.12679#bib.bib7)\), PIS\([Zhang and Chen, 2022](https://arxiv.org/html/2610.12679#bib.bib39)\), and DIS\([Berner et al\., 2024](https://arxiv.org/html/2610.12679#bib.bib9)\), and run SMC\([Del Moral et al\., 2006](https://arxiv.org/html/2610.12679#bib.bib6)\)using the reference implementation of[Blessing et al\. \(2024\)](https://arxiv.org/html/2610.12679#bib.bib5)\. FAB uses an 8\-layer Real\-NVP flow with\[64,64\]\[64,64\]conditioner MLPs, 12 intermediate distributions, 10 HMC transitions per intermediate distribution, a replay buffer, andα\\alpha\-divergence loss withα=2\\alpha=2\. It trains for 3,000 iterations with four replay updates per iteration, at batch size 2,000 \(1,000 for both NICE targets\)\. PIS and DIS use the PISGRADNet with width\-64 MLPs and Fourier time features, trained for 40,000 updates with a batch size of 2,000 and 128 integration steps\. For other hyperparameters, we use those reported in[Blessing et al\. \(2024\)](https://arxiv.org/html/2610.12679#bib.bib5)\. For ManyWelld=32,64d=32,64and 40GMMd=10d=10, since the paper does not provide default hyperparameters, we performed a grid search over the key hyperparameters, as provided in[Table5](https://arxiv.org/html/2610.12679#A4.T5), and selected the one with the best ELBO\. SMC uses 2,000 particles, 10 HMC transitions per annealing interval, and resampling threshold of0\.30\.3\. It uses 256 temperatures for GMM\-10D and ManyWell \(dimension3232and6464\) targets, and 128 otherwise\. For TB\+\+MALA\([Sendera et al\., 2024](https://arxiv.org/html/2610.12679#bib.bib10)\)and SubTB\+\+SMC\([Choi et al\., 2026](https://arxiv.org/html/2610.12679#bib.bib52)\), we use the original implementation of[Choi et al\. \(2026\)](https://arxiv.org/html/2610.12679#bib.bib52)and reimplement TB\+\+MALA within the same codebase to simplify benchmarking\. While the reimplemented TB\+\+MALA is not identical to the original from[Sendera et al\. \(2024\)](https://arxiv.org/html/2610.12679#bib.bib10), we ensure that the algorithm performs at least as well as, or better than, the original on synthetic targets\. For these two baselines, we set the network architectures and common hyperparameters \(learning rate and initial std\.\) to match those of DIS, while fixing algorithm\-specific hyperparameters as reported in the original papers\. Table 5:Baseline hyperparameter search grids\. ##### OSIS\. For the main evaluation study[Tables1](https://arxiv.org/html/2610.12679#S6.T1)and[2](https://arxiv.org/html/2610.12679#S6.T2),OSISwas trained using a squared SWD\([Bonneel et al\., 2015](https://arxiv.org/html/2610.12679#bib.bib17)\)with20482048projections, AdamW with weight decay10−210^\{\-2\}, batch size of40964096,T=10,000T=10\{,\}000training updates,2020steps of MCMC refinement andM=50M=50inner gradient iterations[Algorithm3](https://arxiv.org/html/2610.12679#alg3)\. All the other experiments were run forM=1M=1\. We parameterise the functionGθG\_\{\\theta\}with a ResNet\([He et al\., 2016](https://arxiv.org/html/2610.12679#bib.bib28)\)architecture for synthetic targets and a small UNet\([Ronneberger et al\., 2015](https://arxiv.org/html/2610.12679#bib.bib29)\)for the NICE targets\. The temperature follows the schedule of[Fogliani et al\. \(2026\)](https://arxiv.org/html/2610.12679#bib.bib89): at timett, the inverse temperature is β\(t\)=β0max\(0,1−t/\(0\.8T\)\),\\displaystyle\\beta\(t\)=\\beta\_\{0\}^\{\\max\(0,\\,1\-t/\(0\.8T\)\)\},so the last20%20\\%of training usesβ=1\\beta=1\. Funnel keepsβ=1\\beta=1during its entire training\. Additional hyperparameters are reported in[Table6](https://arxiv.org/html/2610.12679#A4.T6)\. Table 6:OSIShyperparameters\.ResNet\(hidden, blocks\) andUNet\(base channels, attention heads\)\. The temperature schedule is as defined in[AppendixD](https://arxiv.org/html/2610.12679#A4)\. ##### Kernel comparison\. For this study the model is kept fixed across kernels as described above\. Each kernel uses2020refinement steps\. Other parameters are tuned per kernel and presented in Table[7](https://arxiv.org/html/2610.12679#A4.T7)\. MALA in this study usesM=1M=1inner iterations \(see[Algorithm3](https://arxiv.org/html/2610.12679#alg3)\), so its parameters differ from[Table6](https://arxiv.org/html/2610.12679#A4.T6)\. Table 7:Selected hyperparameters for ManyWell \(d=32d=32\)\.KernelLearningrateInitialtemperatureInitial proposalstandarddeviationTargetacceptancerateULAdrift stepIntegrationstep sizeFrictionLeapfrogstepsSVGDupdate stepRBF kernelvarianceRWMH10−310^\{\-3\}410\.5——————MALA10−310^\{\-3\}16—0\.90——————ULA10−310^\{\-3\}16——0\.008—————BAOAB1\.9⋅10−51\.9\\cdot 10^\{\-5\}2———0\.00252———OBABO10−310^\{\-3\}16———0\.0062———HMC1\.7⋅10−31\.7\\cdot 10^\{\-3\}16———0\.019—10——SVGD10−310^\{\-3\}16——————1\.6⋅10−41\.6\\cdot 10^\{\-4\}2 ##### Parameter counts\. Trainable parameter counts are reported in Table[8](https://arxiv.org/html/2610.12679#A4.T8)\. SMC is non\-amortised and has no learned parameters\. PIS, DIS, and TB\+\+MALA share the same backbone and parameter count\.OSISuses aResNeton the synthetic targets and aUNeton the NICE targets\. Table 8:Number of trainable parameters in the benchmark study \(×103\\times 10^\{3\}\)\. ##### Sampling time\. Figure[3](https://arxiv.org/html/2610.12679#S6.F3)provides the median time over 250 runs for a batch of 2,000 samples on one NVIDIA GH200 GPU, excluding compilation and warm\-up\. The mean𝒲22\\mathcal\{W\}^\{2\}\_\{2\}metric is computed independently following the previously described 10,000\-sample protocol\.OSIS, TB\+\+MALA, and SubTB\+\+SMC use compiled PyTorch, while FAB, PIS, DIS and SMC use native JAX JIT\. ### D\.2Outsourced sampling ##### Target densities\. For MNIST we use a VAE\([Kingma and Welling, 2014](https://arxiv.org/html/2610.12679#bib.bib80)\)which we train ourselves\. Then, we use VAE decoder with a1616\-dimensional latent, and a Bernoulli likelihood, trained for100100epochs\. For the reward function, we train a CNN classifierp\(c∣x\)p\(c\\mid x\)\. The reward is defined aslogr\(z\)=log∑cw~c,p\(c∣x\(z\)\)\\log r\(z\)=\\log\\sum\_\{c\}\\tilde\{w\}\_\{c\},p\(c\\mid x\(z\)\), and the target log\-density islogπ\(z\)=logr\(z\)−12\|z\|2\\log\\pi\(z\)=\\log r\(z\)\-\\tfrac\{1\}\{2\}\|z\|^\{2\}\. Both networks remain frozen during the training ofOSIS\. For the experiments on FFHQ\-64, we use a pretrained frozen R3GAN\([Huang et al\., 2024](https://arxiv.org/html/2610.12679#bib.bib48)\)generator\. The reward is ImageReward\([Xu et al\., 2023](https://arxiv.org/html/2610.12679#bib.bib81)\)\. The target log\-density is defined aslogπ\(z\)=3\.0⋅IR\(G\(z\),caption\)−12\|z\|2\\log\\pi\(z\)=3\.0\\cdot\\mathrm\{IR\}\(G\(z\),\\text\{caption\}\)\-\\tfrac\{1\}\{2\}\|z\|^\{2\}\. Reward and gradient passes runbfloat16\. ##### Model architecture forGθG\_\{\\theta\}\. In all our experiments we use a ResNet architecture with MLP layers\. The inputs are first projected to a256256\-dimensional latent, then follow 3 residual blocks of the formx↦x\+SiLU\(Wx\+b\),W∈ℝ256×256x\\mapsto x\+\\mathrm\{SiLU\}\(Wx\+b\),\\,W\\in\\mathbb\{R\}^\{256\\times 256\}, then the outputs are projected to the output dimension\. ##### Training\. We trainOSISwith12001200outer steps and5050inner steps, following the procedure presented in[Algorithm3](https://arxiv.org/html/2610.12679#alg3)\. We use an empirical approximation of OT loss \([SectionC\.2](https://arxiv.org/html/2610.12679#A3.SS2)\) and use MALA with adaptive step size \([SectionC\.1](https://arxiv.org/html/2610.12679#A3.SS1)\) as the MCMC algorithm, which we run for2020steps\. For the initial10001000outer iterations we train the algorithm on a tempered target, the temperature follows a cosine schedule with initial temperature2\.02\.0and final temperature1\.01\.0\. Next, we trainOSISfor200200more iterations keeping the temperature fixed at1\.01\.0\. We use the AdamW\([Loshchilov and Hutter, 2019](https://arxiv.org/html/2610.12679#bib.bib27)\)optimiser with learning rate0\.0010\.001and no weight decay\. ##### Evaluations\. Each method in[Table4](https://arxiv.org/html/2610.12679#S6.T4)is evaluated usingN=10,000N=10\{,\}000\. The reference is rejection sampling from the same tilted density\. Since ImagerReward is unbounded by default, we bound it with a constant to perform rejection sampling\.𝔼\[logr\(x\)\]\\mathbb\{E\}\[\\log r\(x\)\]corresponds to the averaged ImageReward score computed with10,00010,000samples\. FID is computed against the rejection\-sampling reference\. We compute FID score usingtorchmetrics\. ## Appendix EAdditional experiment results ### E\.1Extended evaluation ofOSIS In this section, we report results forOSISwith OBABO kernels on the main synthetic targets\. This provides further evidence on whether splitting schemes can improveOSISperformance[Tables9](https://arxiv.org/html/2610.12679#A5.T9)and[10](https://arxiv.org/html/2610.12679#A5.T10)\. Overall, the benefits depend on the training objective: with OT, OBABO achieves lower mean squared 2\-Wasserstein distance and MMD than MALA on all evaluated targets, whereas with SWD, the results are mixed\. These experiments compare equal numbers ofOSIStraining steps, which does not necessarily imply equal computational cost\. Each OBABO step requires two energy\-gradient evaluations, whereas each MALA step requires two energy\-gradient evaluations and two additional energy\-value evaluations for the Metropolis\-Hastings correction\. Table 9:Sampler comparison on GMM, Funnel, and ManyWell targets\. OSIS rows specify the training kernel and loss\. This table was computed usingM=1M=1inner iteration \(see[Algorithm3](https://arxiv.org/html/2610.12679#alg3)\)\.Table 10:Sampler comparison on high\-dimensional targets\. OSIS rows specify the training kernel and loss\. This table was computed usingM=1M=1inner iteration \(see[Algorithm3](https://arxiv.org/html/2610.12679#alg3)\)\. ### E\.2OSISagainst its long run refinement In this section we compareOSISto its long run MCMC kernel\. Both methods are run with the same MALA settings, number of MCMC steps and temperature annealing schedule\. This table assesses how much of the refinement signalOSIScaptures relative to running the equally long rollout of its refinement kernel\. We note that the MALA parameters used here come fromOSISruns meaning, that they were not tuned forOSIStraining and may not be optimal for direct MCMC sampling\. We observe thatOSIShas higher reported discrepancies than MALA on 40GMM, motivating further investigation of whether additional inner optimisation steps helpOSIStrack its refinement targets[Table1](https://arxiv.org/html/2610.12679#S6.T1)\. On ManyWell, however, SWD\-trainedOSISoutperforms its refinement kernel run for the same number of steps\. These results show effective amortisation of the refinement signal in our sampler on specific targets\. Table 11:Comparison ofOSIS\+MALA with long\-run MALA only\.OSISrows specify the training kernel and loss\. MALA\-only runs start from Gaussian noise\. We compare one set of MALA samples against five reference sets and report the average and standard deviation\. This table was computed usingM=1M=1inner iteration \(see[Algorithm3](https://arxiv.org/html/2610.12679#alg3)\)\. ### E\.3Drift parameterisation ofOSIS We evaluate whetherOSIScan benefit from a drift\-based regression formulation motivated by Wasserstein gradient flows\. This formulation expresses the generator update as a regression task towards samples displaced by an estimated drift field, with the regression targets held fixed during differentiation\. For OT\-based objectives, computing these targets still requires the corresponding transport information\. The drift parameterisation yields lower mean values for both metrics on ManyWell \(d=32d=32\)\. On 40GMM \(d=2d=2\), both mean Sinkhorn divergence and mean MMD are higher, although the differences are within the scale of the reported variability\. Table 12:Comparison of discrepancy and drift objectives forOSIStrained with the OBABO refinement kernel\. Results are mean±\\pmstandard deviation over five training seeds, after averaging five sampling\-seed evaluations per checkpoint\. This table was computed usingM=1M=1inner iteration \(see[Algorithm3](https://arxiv.org/html/2610.12679#alg3)\)\. ### E\.4Ablations We provide in this section detailed ablations of the different design choices ofOSIS\. Table 13:Ablation on ManyWell \(d=32d=32\)\. The defaultOSIStraining configuration uses 20\-step MALA, batch size40964096, SWD\-20482048, and50,00050,000training updates\. Results are mean±\\pmstandard deviation over five training seeds, after averaging five sampling\-seed evaluations per checkpoint\. This table was computed usingM=1M=1inner iteration \(see[Algorithm3](https://arxiv.org/html/2610.12679#alg3)\)\. ### E\.5Additional samples from Bayesian posteriors on MNIST Figure 6:Samples obtained using outsourced sampling\. Samples correspond to the classifier\-based target with the rewardr\(x\)=0\.5p\(0\|x\)\+0\.5p\(1\|x\)r\(x\)=0\.5p\(0\|x\)\+0\.5p\(1\|x\)\. From left to right, samples are obtained withOSIS,OSISwith 20\-step MALA refinement, and rejection sampling\.Figure 7:Samples obtained using outsourced sampling\. Samples correspond to the classifier\-based target with the rewardr\(x\)=1p\(2∣x\)\+0\.5p\(5∣x\)\+2p\(7∣x\)r\(x\)=1p\(2\\mid x\)\+0\.5p\(5\\mid x\)\+2p\(7\\mid x\)\. From left to right, samples are obtained withOSIS,OSISwith 20\-step MALA refinement, and rejection sampling\. ### E\.6Additional samples from Bayesian posteriors on FFHQ\-64 Figure 8:Comparison of samples obtained withOSIS,OSISwith 20\-step MALA refinement, and rejection sampling on outsourced sampling task with ImageReward\-based target\. The following 4 prompts are considered \(from to to bottom\):a young blond girl,smiling toddler,a man with a beard,person wearing glasses\.OSISRejection samplingTB\+\+MALA\([Sendera et al\., 2024](https://arxiv.org/html/2610.12679#bib.bib10)\)MALAFigure 9:Visualisation of samples fromOSIS, rejection sampling, TB\+\+MALA\([Sendera et al\., 2024](https://arxiv.org/html/2610.12679#bib.bib10)\)and MALA\. All algorithms, except rejection sampling, are matched in terms of the number of log\-density evaluations\.
相似文章
用于模型校准的摊销局部后验推断及其在海洋湍流闭合模型中的应用
提出了一种基于条件归一化流的摊销变分推断方法，构建加权似然的局部后验来处理模型校准中的模型偏差与上下文不确定性，并在Branin-Hoo孪生实验及单柱海洋模型的k-ε湍流闭合校准中验证了其相较于全局后验的误差降低效果。
用于引导式推测推理的多重重要性采样
该论文将多重重要性采样(MIS)引入引导式推测推理(GSI)框架，使用多个小型大语言模型作为草稿提案模型并由大型模型与过程奖励模型验证，证明了平衡启发式的最优性，并给出KL散度与奖励退化的新界。
通过熵增益采样与答案混合实现递归离线蒸馏
该论文提出一种递归离线蒸馏方法，结合基于熵差的采样与答案混合，将大型模型的推理链蒸馏到 3-4B 小模型中，在相同数据预算下显著优于朴素蒸馏（平均 57.15 vs 51.61），并开源了代码与数据。
Generator-Guided Inverse Sampling for L\'evy-Driven Generative Models
This paper studies inverse sampling for Lévy-driven generative models, proposing a structured reverse sampler that decomposes dynamics into diffusion, small jump, and large jump components, with neural networks amortizing jump rates. The method is applied to OFDM-SISO channel estimation under mixed Gaussian and impulsive noise.
模型能力主导：AIMO 3推理时优化的经验启示
本论文分析了AIMO 3的推理时优化技术，发现模型能力优于提示工程和多样化采样策略。研究表明高温度采样已经能够最大程度地去相关化误差，为基于提示的改进留下了很少余地，并识别出单个模型pass@20与多数投票共识之间存在6分的选择损失差距。