Sampling, Diffusions, and Stochastic Localization

1.1 Sampling

• The sampling framework targets both explicitly specified distributions and distributions known only through independent samples, with implementation differences concentrated in the denoising oracle.The paper focuses on structural properties of the generating process rather than on one particular data-access setting.Source: §1.1 · pp. 1–2

1.2 Diffusions

• Time reversal converts a forward diffusion into a sampling process, and suitable discretization yields approximate samples from the target distribution.For the Ornstein–Uhlenbeck process, the reverse process can be initialized from a Gaussian law and reaches the target distribution at its endpoint; earlier stopping gives an approximation.Source: §1.2 · pp. 3–4

1.3 A special stochastic localization process

• Stochastic localization represents progressively informative observations through conditional measures that localize to a random target point, whose barycenter approaches a sample from the original distribution.The associated SDE provides an algorithmic route: discretize it to a large time T and output Y_T/T or m(Y_T;T).Source: §1.3 · pp. 4–5
• The paper shows that the diffusion and stochastic-localization constructions coincide after a change of variables, motivating a generalization to broader stochastic localization processes.The observation process need not share the state space or dimension of the sampled variable.Source: §1.3 · pp. 4–5

• The paper’s viewpoint unifies recent deep-learning developments and suggests a different palette of sampling schemes through stochastic localization.Source: §2 · p. 6
• Non-Gaussian noise extends the discrete-time Ornstein–Uhlenbeck construction beyond its Gaussian autoregressive form [NRW21].Source: §2 · p. 6
• Discrete diffusion variants use coordinate flipping, arbitrary time-dependent transition rates, absorbing states, or general continuous-time Markov processes [SYD+22].Source: §2 · p. 6
• An alternative represents discrete variables as binary strings and obtains discrete samples by rounding without modifying the forward or reverse diffusion [CZH22].Source: §2 · p. 6
• The time-reversal approach now covers a broad array of sampling schemes, since sampling can be viewed as reversing a noising process.Source: §2 · p. 6
• Continuous-time probability flows provide a complementary generative-modeling perspective, including stochastic flows that encompass diffusions [AVE22] [ABVE23].Source: §2 · p. 6

• An observation process requires reverse-time Markov structure and completeness: the full observation path determines the latent variable x.Source: §3 · p. 7
• The stochastic localization scheme is the sequence of posterior measures µ_t induced by observations Y_t of x.Source: §3 · p. 7
• Observations become less noisy as time or signal strength increases, with their ordering corresponding to Blackwell ordering or physical degradation.Source: §3 · p. 7
• The framework encompasses earlier machine-learning constructions and permits observations to occupy different spaces from x, or different spaces at different times.Source: §3 · p. 7
• Measure-valued stochastic localization is equivalent to the posterior formulation because its posteriors form a Doob martingale and localize to δ_x.Source: §3 · pp. 7–8
• Completeness can be ensured when a function of the observations converges to x or recovers x at the terminal time.Source: §3 · p. 8
• Sampling proceeds by discretizing time, approximating forward transition kernels, and reversing those kernels to generate samples.Source: §3 · p. 8
• The construction is algorithmic only when transitions between nearby times can be approximated efficiently.Source: §3 · p. 8

4.1 The isotropic Gaussian process; 4.2 The anisotropic Gaussian process

• The isotropic Gaussian process is sampled by Euler-discretizing its Markov SDE, while anisotropic generalization replaces identity covariance with a time-dependent positive semidefinite matrix Q(t).Improved discretizations are cited in [SME20] [KAAL22].Source: §4.1 · p. 9, §4.2 · pp. 9–10

4.3 The erasure process

• The erasure process independently reveals coordinates over t ∈ [0,1], producing standard sequential sampling or any prescribed coordinate order through deterministic revealing times.Each coordinate is erased with probability 1−t, and the associated sampler proceeds sequentially.Source: §4.3 · p. 10

4.4 The binary symmetric process

• The binary symmetric process evolves from uniform random signs at t=0 to x at t=1 through coordinatewise resampling, forming a continuous-time Markov chain and reformulating binary sampling from.Its coordinates are independently resampled backward in time, with uniformly distributed first resampling times.Source: §4.4 · pp. 10–11

4.5 The symmetric process

• The symmetric process extends the binary construction to q-ary alphabets using coordinatewise Poisson resampling with rate t^-1, yielding conditionally independent observations and explicit transition rates.This process is equivalent to the construction introduced in.Source: §4.5 · pp. 11–12

4.6 The linear observation process

• Linear observation processes reveal Ax through Brownian or noisy linear measurements; discretization samples an approximation of Ax, after which pseudoinversion or compressed sensing can recover structured x.Full column rank permits x=A†Y∞, while sparse signals can be recovered with compressed sensing.Source: §4.6 · p. 12

4.7 The information percolation process

• The information percolation process reveals grid-edge differences sequentially, reducing sampling to computing each next conditional difference given the information graph already exposed.It is a change of variables of the erasure process.Source: §4.7 · p. 13

4.8 The Poisson observation process; 4.9 The half-space process

• The Poisson observation process records independent coordinatewise Poisson counts with Y_t,k∼Poisson(tx_k), while the half-space process records successive membership bits for a sequence of half-spaces.The Poisson transition rules use the conditional means m_k(t;y)=E[X_k|Y_t=y].Source: §4.8 · p. 13, §4.9 · p. 14

4.10 Revisiting reverse diffusions; 4.11 The Eucldean invariant process; 4.12 All of the above

• Any reverse diffusion can be represented as a stochastic localization observation process; the equivalence follows by viewing diffusion trajectories as progressively revealed paths, and broader increasing observation sets extend beyond time reversal.The construction also gives a symmetry-preserving Gram-matrix sampler for Euclidean-invariant particle distributions and supports combining observation processes into a new one.Source: §4.10 · p. 14, §4.11 · p. 15, §4.12 · p. 15

4.12 All of the above

• The SDE drift is modified to incorporate observations from the second component, including restrictions to coordinates that are not erased.The modified dynamics track which coordinates remain un-erased in y^(2) and use the corresponding restricted vectors.Source: §4.12 · p. 15
• The second component is initialized with all coordinates erased and then evolves over small time intervals coordinate by coordinate.The initialization is Y^(2)_0 = (*, ..., *), with subsequent updates specified for each i in [n].Source: §4.12 · p. 16
• Implementing the modified process requires expectations and marginal probabilities under the time-dependent distribution µ_t, or conditional densities relative to a reference measure.For finite-alphabet variables, the conditional is expressed using probability mass functions; more generally, component densities P1,t and P2,t are taken with respect to ν.Source: §4.12 · p. 16

Problem setting

• The mixture consists of two well-separated Gaussians with identity covariance, fixed mixture weights, and separation norm comparable to the Gaussian radius.The centered model is µ = p·N((1−p)a; I_n) + (1−p)·N(−pa; I_n), with p independent of n.Source: §5 · p. 16

Isotropic diffusion

• Isotropic diffusion generates a distribution significantly different from the target along the one-dimensional projection, even across increasingly large training configurations.The compared models use the same two-layer ReLU architecture but vary sample size, training epochs, depth, and width.Source: §5 · p. 17

Modified sampling process

• The modified generation first samples a component label from the empirical half-space fraction, runs isotropic diffusion with the corresponding denoiser, and returns the terminal estimate.Separate denoisers are learned from samples on either side of the empirical principal direction.Source: §5 · p. 17

Experimental comparison

• Figure 1 compares projected empirical distributions with the correct distribution across posterior-mean models and generation times, while Figure 2 uses the same architecture with a different generating process.Figure 2 also reveals the projection onto a fixed vector alongside Gaussian-perturbed data.Source: Figure 1 · p. 18, §5 · p. 18, Figure 2 · p. 19

Modified sampling process

• The modified process produces a distribution that better matches the target along the mixture-separation direction.The process separates samples using a learned principal direction, then applies a denoiser selected by the inferred half-space.Source: §5 · pp. 17–18

Why isotropic diffusion struggles

• The isotropic posterior mean changes sharply across a midpoint window of width Θ(1/n), making the evolution highly sensitive to denoiser error.When ||a|| = Θ(√n), the derivative of the relevant posterior function in this window is positive and of order n.Source: §5 · pp. 18–19

6 The role of architecture: Images with long range correlations

• The synthetic distribution consists of 512×512 RGB images that are mostly blue or mostly red, each with probability 1/2, plus smooth variations.Samples from this distribution are shown in Figure 3.Source: §6 · pp. 19–20, Figure 3 · p. 20
• Both approaches learn the distribution with a simple two-layer convolutional denoiser, differing in whether they use isotropic diffusion or a linear observation process.The linear process appends channel averages, scaled by b, to the observations.Source: §6 · p. 20, §6.2 · p. 22
• The linear observation process captures long-range image correlations while the isotropic process remains locally correlated, despite using the same short-ranged denoiser architecture.Figure 4 compares the two generating processes on synthetic RGB images.Source: §6 · p. 20, Figure 4 · p. 21

6.1 Isotropic diffusion

• In the isotropic setup, 5×5 convolutions make each output pixel depend on a 9×9 input patch, yielding short-range correlations.The network uses one hidden layer and encodes time through the angle variable α.Source: §6.1 · pp. 20, 22

6.2 Linear observation process

• The linear-observation denoiser minimally modifies the isotropic network by using six input channels and fifteen middle-layer channels.Channels 4–6 contain spatially constant channel averages, with the same noise added to all entries within each channel.Source: §6.2 · pp. 22–23

7.1 Sampling via isotropic Gaussian diffusions

• The isotropic sampler can be accurate only after t ≳ 1/λ_min(Σ), while accurate discretization requires δλ_max(Σ) ≲ 1, yielding iteration counts scaling with κ(Σ).The condition number is κ(Σ) = λ_max(Σ)/λ_min(Σ).Source: §7.1 · p. 23
• The convolutional restriction is imposed through a shift-invariant matrix whose kernel vanishes beyond radius r, preventing the denoiser from directly representing global interactions.This finite-window architecture is analyzed for shift-invariant Gaussian targets.Source: §7.1 · p. 23, §7 · p. 23
• For Σ_n = I_n + α1_n1_n^T, isotropic diffusion with window size 2r + 1 generates a distribution whose correlations extend only to distances of order r.The target has E_µ(x_i x_j) = α for every i ≠ j, so its correlations are long range.Source: §7.1 · pp. 24–25

7.2 Sampling via the linear observation process

• The linear-observation construction uses an augmented observation with one global coordinate and n coordinate-specific components, with a linear, shift-invariant conditional mean.The resulting estimator has a circulant matrix component and scalar coefficients.Source: §7.2 · pp. 25–26
• The linear observation process appends scaled channel averages, revealing the global component faster than other coordinates and enabling estimators adapted to that latent structure.The proposed estimators are Bayes optimal up to O(1/n) for the global observation and use it to estimate the remaining coordinates.Source: §7.2 · pp. 25–26

8 The role of architecture: Images with long range correlations

• Sampling algorithms based on diffusions or stochastic localization reduce sampling from µ to approximating transition probabilities.For known µ, this requires analytic approximation; for sample-defined µ, the transition probabilities must be learned from data.Source: §8 · p. 27

8.1 Approximating transition probabilities for explicitly known µ

• For explicitly known µ, transition probabilities can be approximated by computing conditional means or marginals using variational methods or correlation decay.Correlation decay applies when the target energy decomposes over a graph structure.Source: §8.1 · pp. 27–28

8.2 Estimating transition probabilities from samples

• For targets known through samples, one can generate observation-process realizations and fit conditional transition densities by minimizing empirical negative likelihood.In the isotropic Gaussian case, this objective is essentially equivalent to standard score matching.Source: §8.2 · p. 28

8.3 On the choice of the stochastic localization scheme

• Observation processes should respect target symmetries because this keeps conditional measures invariant and focuses estimation on non-trivial degrees of freedom.The paper also recommends adding observations that explicitly depend on important latent variables in strongly multimodal distributions.Source: §8.3 · p. 29
• An information-computation gap can make the denoiser intractable over an interval of t, causing polynomial-time approximations and sampling to fail in some cases.The paper therefore recommends avoiding hard regions where conditional means cannot be approximated efficiently.Source: §8.3 · p. 29
• Sharp algorithmic MMSE drops signal phase transitions, so localization schemes should avoid observation paths with such threshold behavior.For complex distributions known only through samples, polynomial-time approximability of denoisers or transitions may remain unclear.Source: §8.3 · p. 29

• The appendix provides explicit formulas for the general KL divergence and covers the examples developed in the preceding sections.Source: A Loss functions · p. 34

A.1 Gaussian observation process

• The estimated drift defines a Gaussian observation process stopped at time T, with diffusion noise scaled by Q(t).Source: §A.1 · p. 34
• Girsanov’s theorem provides the change-of-measure step, while standard diffusions and linear information processes appear as special cases.Source: §A.1 · p. 34

A.2 Discrete time Markov chains

• Discrete-time chains are described through transition probabilities P_t and estimated transitions ˆP_t, assuming the latter have densities relative to the former.Source: §A.2 · p. 34

A.3 Continuos time Markov chains

• Continuous-time discrete-state chains use rates r_t(y,y′) for jumps and q_t(y) for remaining in the current state.Source: §A.3 · p. 34
• The framework includes symmetric and Poisson processes, using estimated conditional expectations or transition quantities to construct their probabilities and divergences.Source: §A.3 · p. 35

C.1 The distribution over images

• The image experiments generate RGB pixels by applying tanh to a sinusoidal field with randomly drawn offset, amplitude, and frequency parameters.Source: §C.1 · pp. 35–36

• Training uses stochastic gradient descent with batch size 4 on 300 generated images for 100 epochs.The noise vector G and signal-to-noise ratio t are resampled independently at each SGD sample.Source: §C.2 · p. 36

D.1 General formulas for isotropic diffusions

• The optimal convolution with window size 2r + 1 is obtained by solving Eq. (7.7).Source: §D.1 · p. 36
• The covariance matrix is circulant and is analyzed through an eigenvalue decomposition involving c_0 and ν(q).Source: §D.1 · pp. 36–37
• The function F(ν;c) is positive, continuously differentiable, and convex for c in (0,1), with F(1;c)=1/c and F(0;c)=1.Source: §D.1 · p. 37
• The logarithmic factor is bounded between 1 and log(1/c) for s ≥ 0 and c in (0,1).Source: §D.1 · p. 37
• The derived correlation-length formulas characterize covariance through coefficients ˆc_k and specialize to Σ = I + α11^T.Source: §D.1 · pp. 36–37

D.2 Proof of Proposition 7.2

• The drift estimate is expressed explicitly as m*(Y_0,t,Y*_t;t), combining the common and coordinate-specific process components.Source: §D.2 · pp. 37–38
• The process has an almost-sure limiting representation involving a standard normal variable and the generated Gaussian component.Source: §D.2 · p. 38
• Each generated coordinate process satisfies a linear stochastic differential equation driven by Y_0,t and Brownian noise.Source: §D.2 · p. 38
• For fixed t, the vector Y*_t is centered Gaussian with covariance c_1(t)I + c_2(t)11^T.Source: §D.2 · p. 38
• The limiting right-hand side has target distribution μ_n, whereas the left-hand side remains Gaussian for every t.Source: §D.2 · p. 38
• Almost-sure convergence of Gaussian vectors to a non-degenerate Gaussian limit implies convergence in W2 and total variation.Source: §D.2 · p. 38