Sequential Posterior Sampling with Diffusion Models

TL;DR

This work addresses the slow runtime of diffusion-model-based posterior sampling in real-time sequential imaging. It introduces SeqDiff and SeqDiff+, autoregressive initializations that reuse past posterior estimates, with SeqDiff+ using a ViViT to model frame transitions for improved starts. On high-frame-rate cardiac ultrasound, the methods achieve comparable performance to full diffusion with up to a $25×$ reduction in diffusion steps and up to an $8\%$ PSNR improvement under motion, enabling real-time posterior sampling. The approach broadens the applicability of diffusion models to dynamic inverse problems and time-series imaging by exploiting temporal structure to dramatically accelerate inference.

Abstract

Diffusion models have quickly risen in popularity for their ability to model complex distributions and perform effective posterior sampling. Unfortunately, the iterative nature of these generative models makes them computationally expensive and unsuitable for real-time sequential inverse problems such as ultrasound imaging. Considering the strong temporal structure across sequences of frames, we propose a novel approach that models the transition dynamics to improve the efficiency of sequential diffusion posterior sampling in conditional image synthesis. Through modeling sequence data using a video vision transformer (ViViT) transition model based on previous diffusion outputs, we can initialize the reverse diffusion trajectory at a lower noise scale, greatly reducing the number of iterations required for convergence. We demonstrate the effectiveness of our approach on a real-world dataset of high frame rate cardiac ultrasound images and show that it achieves the same performance as a full diffusion trajectory while accelerating inference 25$\times$, enabling real-time posterior sampling. Furthermore, we show that the addition of a transition model improves the PSNR up to 8\% in cases with severe motion. Our method opens up new possibilities for real-time applications of diffusion models in imaging and other domains requiring real-time inference.

Figures (4)

• Figure 1: Geometric representation of the reverse diffusion process and corresponding manifolds $\mathcal{M_\tau}$ for each diffusion timestep $\tau$. In (a) a standard conditional reverse diffusion trajectory starting from a Gaussian sample ${\mathbf x}_\mathcal{T}\sim\mathcal{N}$ is shown with DPS as guidance rule chung2022diffusion. For initialization of the next frame $t+1$, we propose two different methods SeqDiff and SeqDiff+, depicted in (b) and (c) respectively. In the first option we initialize the trajectory from a noised version of the Tweedie estimate of the previous frame, $p({\mathbf x}_{\tau^\prime}^{t+1}|{\mathbf x}_0^{t})$ with $\tau^\prime \ll \mathcal{T}$. The second option improves upon this by predicting the next frame with $\tilde{{\mathbf x}}_0^{t+1} \approx f(\cdot)$, accounting for any motion between frames. This leads to the initialization $p({\mathbf x}_{\tau^\prime}^{t+1}|\tilde{{\mathbf x}}_0^{t+1})$, with $\tau^\prime_{\text{SeqDiff+}} < \tau^\prime_{\text{SeqDiff}}$.
• Figure 2: Qualitative comparison of Vanilla DPS (for $N=4$ and $N=100$ steps), and the two proposed initialization methods SeqDiff and SeqDiff+ for only $N^\prime=4$ diffusion steps. Target images ${\mathbf x}^t$ are 80% masked by ${\mathbf A}^t$ to produce observation ${\mathbf y}^t$. Initialization with SeqDiff(+) is able to improve on full diffusion trajectories with $25\times$ speedup.
• Figure 3: Comparison of SeqDiff(+) performance in PSNR against various motion conditions. (a) For every sample in the test set. The advantage of using a transition model (SeqDiff+) is most advantageous with high motion (see linear fit $m$). (b) For a single sequence of frames. SeqDiff+ is less correlated with the motion, whereas the error of SeqDiff increases with more movement, emphasizing the importance of the transition model. (c) Best performing $N^\prime$ for each initialization method against motion. SeqDiff+ outperforms the other methods for all motion levels. For lower motion levels, SeqDiff is a valid option.
• Figure 4: PSNR against number of diffusion steps on sequences of frames from the test split of EchoNet-Dynamic dataset. Confidence Interval (CI) is taken over 3 splits with different masks and seeds. SeqDiff+ shows a notable improvement, particularly with fewer diffusion steps $N^\prime$.