The source is a design variable
Along the straight path from a source sample $x_0$ to a posterior sample $x_1$, the conditional velocity is
$x_1 - x_0$. SNaP keeps the path and the endpoint coupling fixed and designs the source, choosing a
family with a linear, measurement-dependent center $x_0 = Ky + \xi$, where $\xi$ has covariance $\Sigma$.
The expected displacement splits into a center term and a spread term:
$$ \mathbb{E}\big[\lVert x_1 - x_0 \rVert_2^2\big] \;=\; \operatorname{tr} R(K) \;+\; \operatorname{tr}\Sigma,
\qquad R(K) := \mathbb{E}\big[(x_1 - Ky)(x_1 - Ky)^\top\big]. $$
Proposition 1 shows the center term is uniquely minimized by the linear MMSE estimator
$K_C = CA^\top(ACA^\top + \sigma_n^2 I)^{-1}$, for any image distribution with covariance $C$. With the
isotropic working covariance $C = \tau^2 I$ this is the Tikhonov estimate, and matching the spread to its
error covariance gives the source
$$ q_\tau(x \mid y) = \mathcal{N}\big(x;\, m_y,\, \tau^2 W\big), \qquad \lambda = \sigma_n^2/\tau^2, $$
$$ m_y = A^\top(AA^\top + \lambda I)^{-1} y, \qquad
W = I - A^\top(AA^\top + \lambda I)^{-1} A. $$
In the right singular basis of $A$, the source variance along $v_i$ is $\tau^2 w_i$ with
$w_i = \sigma_n^2 / (\sigma_n^2 + \tau^2 s_i^2)$: small in well-measured directions, approaching
$\sigma_n^2/s_i^2$ when the signal dominates, and equal to $\tau^2$ on $\mathrm{null}(A)$. The Tikhonov center
also avoids amplifying noise in weakly measured directions, unlike the least-squares anchor $A^\dagger y$.
The working model is used only to build the source; it does not assume the images are Gaussian. As
$\sigma_n \to 0$ the source tends to $\mathcal{N}(A^\dagger y, \tau^2 P)$, which recovers the noiseless
NullFlow source.
The target is still the true posterior
Proposition 2. Let $x_0 \sim q_\tau(\cdot \mid y)$ and $x_1 \sim p(\cdot \mid y)$ be
conditionally independent and joined by straight lines. If the posterior has finite second moments and the
marginal flow has a terminal limit, the exact one-step map
$$ T_y(x) = x + u(x, 0, 1 \mid y) \quad\text{satisfies}\quad
(T_y)_\# \, q_\tau(\cdot \mid y) = p(\cdot \mid y) \quad \text{for every } \tau > 0. $$
The source changes the transport, not its destination; the trained network $u_\theta$ approximates $u$, so
SNaP samples inherit the guarantee up to that approximation error. A companion result (Lemma 2) shows
that the irreducible velocity-regression error—the part no network can remove—is smaller
under $q_\tau$ than under any isotropic source $\mathcal{N}(0, \sigma^2 I)$ with $\sigma \ge \tau$.
Training and sampling
Sampling $q_\tau$ directly would need $W^{1/2}$. Perturb-and-solve avoids it: with independent
$\epsilon \sim \mathcal{N}(0, \tau^2 I)$ and $n' \sim \mathcal{N}(0, \sigma_n^2 I)$,
$$ x_0 = \epsilon + A^\top (AA^\top + \lambda I)^{-1} \big(y - A\epsilon - n'\big)
\;\sim\; \mathcal{N}(m_y, \tau^2 W) \text{ exactly.} $$
The solve is elementwise for inpainting and super-resolution, two FFTs for deblurring, and 10–20
conjugate-gradient iterations for multi-coil MRI. The network is trained with the conditional MeanFlow
identity (one Jacobian–vector product per step) and the adaptive-weight loss of MeanFlow. At test time
each sample is one source draw and one network evaluation; redrawing $(\epsilon, n')$ with
$y$ fixed produces conditionally independent samples.
SampleSource—perturb-and-solve
require: y, A, σₙ, τ
λ ← σₙ² / τ²
ε ~ N(0, τ²I), n′ ~ N(0, σₙ²I)
r ← y − Aε − n′
x₀ ← ε + Aᵀ(AAᵀ + λI)⁻¹ r
return x₀
x₀ | y ~ N(mᵧ, τ²W) exactly (Lemma 1);
no square root of W is formed
Training step
require: A, σₙ, τ, pratio = 0.5
x₁ ~ p, y ← A x₁ + n
t ~ logit-normal
r ← t w.p. pratio, else r ~ U(0, t)
x₀ ← SampleSource(y, A, σₙ, τ)
zᵣ ← (1−r) x₀ + r x₁, v ← x₁ − x₀
u, dᵣu ← JVP(uθ; (zᵣ, r, t), (v, 1, 0))
utgt ← v + (t − r) sg[dᵣu]
ℓ ← ‖u − utgt‖² / n
θ ← AdamW(θ, ∇θ sg[ω(ℓ)] ℓ)
Sampling—M posterior draws
require: y, A, σₙ, trained uθ, M
for j = 1 … M do
x₀⁽ʲ⁾ ← SampleSource(y, A, σₙ, τ)
u ← uθ(x₀⁽ʲ⁾, 0, 1 | y, σₙ)
x̂₁⁽ʲ⁾ ← x₀⁽ʲ⁾ + u
return {x̂₁⁽ʲ⁾}
M draws = M NFEs, one batched pass;
mean → distortion, spread → uncertainty
