Stochastic Prior Injection (SPI) Explained

• Stochastic Prior Injection (SPI) is a methodology that injects controlled Gaussian noise into intermediate representations to regularize optimization in inverse problems and generative models.
• It is applied in plug-and-play denoising and flow matching, leveraging score-based priors to improve convergence and mitigate issues like distribution mismatch.
• SPI provides theoretical guarantees by performing stochastic gradient descent on a Gaussian-smoothed objective, enabling efficient escape from saddle points and enhancing overall model robustness.

Stochastic Prior Injection (SPI) refers to a family of methodologies for introducing controlled stochasticity into generative modeling and inverse problem-solving frameworks by perturbing intermediate representations with random noise, chiefly to enhance generalization, robustness, and convergence properties. Two principal lines of development are prominent: (i) SPI within plug-and-play (PnP) optimization for ill-posed inverse problems, particularly leveraging pre-trained score-based diffusion model priors (Park et al., 4 Apr 2026), and (ii) SPI as source or interpolant noise in flow-matching frameworks for distribution-to-distribution generative modeling (Su et al., 8 Oct 2025). Both approaches exploit stochastic injection to address limitations of deterministic inference and learning, including distribution mismatch and brittle supervision, by smoothing the underlying optimization or learning objectives.

1. Mathematical Formalizations of SPI

SPI is instantiated differently depending on the base framework:

1.1. Plug-and-Play Inverse Problems

Given observations $y \in \mathbb{R}^m$ modeled as $y = \mathcal{A}(x) + e$, with $e \sim \mathcal{N}(0, \eta^2 I)$, the MAP estimate solves

$\hat{x} = \arg\min_x \{ f(x) := g(x) + h(x) \}$

where $g(x) = - \log p(y|x)$ and $h(x) = -\log p(x)$. Classical proximal methods depend on $\mathrm{prox}_{\gamma h}(z)$. PnP methods replace $\mathrm{prox}_{\gamma h}$ by a trained denoiser $D_\theta(\cdot;\sigma)$:

• $$x_k \leftarrow \mathrm{prox}_{\gamma g}(z_{k-1} - s_{k-1})$$
• $y = \mathcal{A}(x) + e$0
• $y = \mathcal{A}(x) + e$1

SPI modifies this by injecting Gaussian noise prior to denoising:

• $y = \mathcal{A}(x) + e$2, $y = \mathcal{A}(x) + e$3

In gradient-based variants:

• $y = \mathcal{A}(x) + e$4, where $y = \mathcal{A}(x) + e$5.

The resulting procedure is a stochastic gradient descent on a Gaussian-smoothed MAP objective,

$y = \mathcal{A}(x) + e$6

with $y = \mathcal{A}(x) + e$7 (Park et al., 4 Apr 2026).

1.2. Flow Matching with Prior Perturbation

For distribution-to-distribution generative modeling via flow matching, the velocity field $y = \mathcal{A}(x) + e$8 transports $y = \mathcal{A}(x) + e$9 and $e \sim \mathcal{N}(0, \eta^2 I)$0 via

$e \sim \mathcal{N}(0, \eta^2 I)$1

Standard (deterministic) flow matching interpolates:

$e \sim \mathcal{N}(0, \eta^2 I)$2

SPI in this context ("prior perturbation") injects Gaussian noise to source samples:

$e \sim \mathcal{N}(0, \eta^2 I)$3

and uses

$e \sim \mathcal{N}(0, \eta^2 I)$4

with modified loss:

$e \sim \mathcal{N}(0, \eta^2 I)$5

(Su et al., 8 Oct 2025).

2. Score-Based and Smoothed Objective Interpretations

SPI reinterprets classical denoising or velocity regression steps as approximate descent on a smoothed objective—one that is convolutionally regularized by a Gaussian kernel.

• In PnP, Tweedie’s formula links the MMSE denoiser $e \sim \mathcal{N}(0, \eta^2 I)$6 to the score of a Gaussian-smoothed prior:

$e \sim \mathcal{N}(0, \eta^2 I)$7

Thus, traditional PnP denoising corresponds to a gradient step on the smoothed log-prior. SPI formalizes this relationship and introduces explicit sampling from perturbed inputs to ensure out-of-distribution iterates are regularized (Park et al., 4 Apr 2026).

• For flow matching, Lemma 1 demonstrates that without prior noise, source distributions remain as mixtures of delta masses under deterministic flows, precluding recovery of truly continuous targets. The loss with SPI trains the velocity field over a local stochastic neighborhood, reducing overfitting and mitigating mode-collapse due to discrete empirical priors (Su et al., 8 Oct 2025).

3. Theoretical Guarantees and Optimization Behaviour

SPI admits concrete theoretical benefits in both optimization and generalization regimes.

• Convergence: Under mild regularity and unbiasedness assumptions, the SPI update implements unbiased stochastic gradient descent on the Gaussian-smoothed objective $e \sim \mathcal{N}(0, \eta^2 I)$8, converging (with $e \sim \mathcal{N}(0, \eta^2 I)$9 annealing) to stationary points of the true MAP $\hat{x} = \arg\min_x \{ f(x) := g(x) + h(x) \}$0 (Park et al., 4 Apr 2026).
• Strict Saddle Point Escape: The stochastic gradient has nondegenerate variance along directions of negative curvature, ensuring that iterates escape strict saddle points with high probability (by leveraging results from stochastic nonconvex optimization such as the CNC condition). This holds provided the denoiser exhibits sufficient variance when perturbed, as quantified by the projection of $\hat{x} = \arg\min_x \{ f(x) := g(x) + h(x) \}$1 along negative eigenvectors (Park et al., 4 Apr 2026).
• Generalization for Flow Matching: The generalization gap of flow-matching using SPI is bounded above by that on the deterministic interpolant, as the $\hat{x} = \arg\min_x \{ f(x) := g(x) + h(x) \}$2 distance between population and empirical distribution is reduced by prior perturbation. The stochastic injection "fills in" the sparse tubes of deterministic supervision in high dimensions, densifying paths and improving coverage (Su et al., 8 Oct 2025).

4. Algorithmic Schemes and Practical Implementation

4.1. SPI in SGPnP (Stochastic Generative Plug-and-Play)

The core SGPnP-PGM algorithm proceeds as:

$h(x) = -\log p(x)$6 [(Park et al., 4 Apr 2026), Algorithm 2].

Parameters involve log-linear noise schedules, with typical initialization (from Table 5):

• Inpainting: $\hat{x} = \arg\min_x \{ f(x) := g(x) + h(x) \}$3, $\hat{x} = \arg\min_x \{ f(x) := g(x) + h(x) \}$4, $\hat{x} = \arg\min_x \{ f(x) := g(x) + h(x) \}$5
• Deblurring: $\hat{x} = \arg\min_x \{ f(x) := g(x) + h(x) \}$6–$\hat{x} = \arg\min_x \{ f(x) := g(x) + h(x) \}$7, $\hat{x} = \arg\min_x \{ f(x) := g(x) + h(x) \}$8, $\hat{x} = \arg\min_x \{ f(x) := g(x) + h(x) \}$9
• CS-MRI: $g(x) = - \log p(y|x)$0–$g(x) = - \log p(y|x)$1, $g(x) = - \log p(y|x)$2, $g(x) = - \log p(y|x)$3

Ablation studies reveal that performance is maximized when the injected noise and denoiser conditioning ($g(x) = - \log p(y|x)$4) evolve on decoupled schedules, rather than keeping them equal [(Park et al., 4 Apr 2026), Table 4].

4.2. SPI in Flow Matching

At each training step:

1. Sample $g(x) = - \log p(y|x)$5 and $g(x) = - \log p(y|x)$6
2. Sample $g(x) = - \log p(y|x)$7
3. Compute $g(x) = - \log p(y|x)$8
4. Interpolate $g(x) = - \log p(y|x)$9
5. Compute and regress to target velocity $h(x) = -\log p(x)$0
6. Minimize $h(x) = -\log p(x)$1

SPI introduces two Gaussian draws and negligible compute overhead per step. Typically, $h(x) = -\log p(x)$2 in a VAE latent space (Su et al., 8 Oct 2025).

5. Empirical Performance and Evaluation

5.1. Plug-and-Play Inverse Problems

SGPnP is evaluated on FFHQ (256×256) for large-mask inpainting, super-resolution, deblurring, and on fastMRI multi-coil data (256×256) for CS-MRI reconstruction. Baselines include deterministic PnP, SNORE, and diffusion solvers such as DPS and DiffPIR. Quantitative metrics are PSNR, SSIM, and LPIPS.

Task	Baseline	SGPnP-PGM (PSNR↑)	SSIM↑, LPIPS↓	Notable Outcome
Inpainting, MRI	DPIR	+0.2–3dB	Top	SGPnP yields best or competitive
All tasks	SDPnP vs PnP	+~0.7dB	Improved	Score priors consistently better
fastMRI (same prior)	Det → SGPnP	+0.2–2dB	Improved	Stochastic always improves over det.

Qualitatively, SGPnP reconstructs plausible, artifact-free content where deterministic or prior variants hallucinate holes or produce blurred reconstructions [(Park et al., 4 Apr 2026), Figs. 1,2].

5.2. Flow-Matching Generative Modeling

SPI is tested on BBBC, SeasoNet, Yosemite, MIMIC-CXR, and GalaxiesML. It is compared against GAN-based bridges, DDIB, and SDEdit.

Dataset	Standard Flow FID	Flow + SPI FID	FID Improvement	Pixel MSE	Assignment Accuracy
BBBC	33.6	19.9	-13.7	-22%	+15pp
SeasoNet	80.0	60.5	-19.5	-	-

Ablations confirm that eliminating prior noise or other stochastic forms reliably degrades results (Su et al., 8 Oct 2025).

6. Practical Considerations and Extensions

• Noise Schedule Tuning: For both PnP and flow matching, excessive noise blurs identity, while too little under-utilizes stochastic smoothing. Empirical validation or small-scale grid search is recommended.
• Decoupled Schedules: Separating injected noise and inference (conditioning) noise often yields superior performance. This decoupling is empirically validated in ablation studies (Park et al., 4 Apr 2026).
• Inference and Sampling: For flow matching, models trained with SPI can often be sampled deterministically by jittering $h(x) = -\log p(x)$3 at test time, sidestepping stochastic differential equation solvers and reducing sampling complexity (Su et al., 8 Oct 2025).
• Extensions: Adaptive or learned noise scales, domain-specific perturbations (e.g., non-Gaussian), and hybridization with conditional bridge architectures remain open research avenues.

7. Limitations and Future Directions

• Distribution Support: SPI does not resolve difficulties arising from highly disjoint support between $h(x) = -\log p(x)$4 and $h(x) = -\log p(x)$5; additional mechanisms, such as conditional bridges, may be necessary (Su et al., 8 Oct 2025).
• Nature of Noise: Current practice is strictly Gaussian; alternative or data-adaptive noise processes remain largely unexplored.
• Optimization Pathology: In PnP, underlying assumptions (unbiasedness, smoothness) may not hold universally, limiting theoretical guarantees in practice, especially under severe data or model misspecification (Park et al., 4 Apr 2026).

Stochastic Prior Injection systematically leverages noise to regularize training and inference within both inverse problem and generative modeling settings. Its principled integration with score-based priors and flow matching substantially improves robustness, sample quality, and optimization tractability, with rigorous theoretical and empirical support across diverse high-dimensional imaging benchmarks (Park et al., 4 Apr 2026, Su et al., 8 Oct 2025).