System-Embedded Diffusion Bridge Models

TL;DR

System-Embedded Diffusion Bridge Models (SDBs) integrate the known linear measurement system directly into the coefficients of a matrix-valued SDE to solve linear inverse problems via diffusion bridges. By embedding the measurement operator through $ oldsymbol{H}_t$ and $ oldsymbol{ abla}_t$ with scalar schedules $ oldsymbol{b1}_t, oldsymbol{eta}_t, oldsymbol{b3}_t$, SDB creates coupled range and null-space dynamics and enables principled posterior sampling from $ p(oldsymbol{x}|oldsymbol{y})$. Empirically, SDB (SB) outperforms supervised bridges and unsupervised baselines across inpainting, super-resolution, CT, and MRI tasks, and shows robust generalization under system misspecification. The work advocates for more expressive diffusion processes and demonstrates feasibility for nonlinear extensions, offering practical improvements for real-world inverse problems such as medical imaging and remote sensing.

Abstract

Solving inverse problems -- recovering signals from incomplete or noisy measurements -- is fundamental in science and engineering. Score-based generative models (SGMs) have recently emerged as a powerful framework for this task. Two main paradigms have formed: unsupervised approaches that adapt pretrained generative models to inverse problems, and supervised bridge methods that train stochastic processes conditioned on paired clean and corrupted data. While the former typically assume knowledge of the measurement model, the latter have largely overlooked this structural information. We introduce System embedded Diffusion Bridge Models (SDBs), a new class of supervised bridge methods that explicitly embed the known linear measurement system into the coefficients of a matrix-valued SDE. This principled integration yields consistent improvements across diverse linear inverse problems and demonstrates robust generalization under system misspecification between training and deployment, offering a promising solution to real-world applications.

Figures (9)

• Figure 1: SDB learns a diffusion bridge between pseudoinverse reconstructions and clean samples by embedding the measurement system into the coefficients of the linear SDE, allowing the score model to distinguish between the range and null spaces of the task.
• Figure 2: Qualitative comparison of SDB (SB) with the best-performing baselines (bridge methods). Rows depict the results for inpainting, superresolution, CT and MRI reconstruction respectively.
• Figure 3: Quantitative comparison of SDB (SB) with other bridge methods in a misspecified MRI reconstruction setting. The top row evaluates bridges trained with $\lambda_1 = 16$, $\sigma_2^2 = 5$ on data generated from systems with decreasing $\lambda_1$. The bottom row evaluates performance on data with $\lambda_1 = 14$ and increasing $\sigma_2^2$. Perturbing the original system makes the problem harder in both cases.
• Figure 4: Quantitative comparison of SDB (SB) with other bridge methods in a misspecified CT reconstruction setting. The top row evaluates bridges trained with $\tau = 3.2$, $\sigma_1^2 = 0.0001$ on data generated from systems with increasing $\tau$. The bottom row evaluates performance on data with $\tau=3.6$ and increasing $\sigma_1^2$. Perturbing the original system makes the problem harder in both cases.
• Figure 5: Top: Performance heatmaps for MRI reconstruction under varying $(\epsilon_1, \epsilon_2)$, using the optimal $(b_0, b_1)$ configuration. The red box highlights the LPIPS-optimal setting. Bottom: Boxplots showing the best achievable performance across 64 $(b_0, b_1)$ pairs on the superresolution task, where $(\epsilon_1, \epsilon_2)$ are tuned separately for each $(b_0, b_1)$ pair.

Theorems & Definitions (9)

• Theorem 1
• Theorem 2
• Theorem 3
• Proposition 1
• Proposition 2
• Lemma 1
• Lemma 2
• Lemma 3
• Lemma 4