Poisson flow consistency models for low-dose CT image denoising

Our objective is to obtain a high quality reconstruction $\bm{\hat{x}}\in\mathbb{R}^{N}$ of $\bm{x}\in\mathbb{R}^{N}$ based on noisy observations $\bm{y}=\mathcal{F}(\bm{x})\in\mathbb{R}^{N}$, where $\mathcal{F}:\mathbb{R}^{N}\rightarrow\mathbb{R}^{N}$ denotes the noise degradation operator, including factors such as quantum noise [20], and $N:=n\times n.$ We treat this as a statistical inverse problem and, given the prior $\bm{x}\sim p(\bm{x})$, the solution is given by a sample $\hat{\bm{x}}$ from the posterior $p(\bm{x}|\bm{y})$. This approach to inverse problem solving is commonly referred to as posterior sampling.

Consistency models [52, 53] build on the continuous-time formulation of diffusion models [40, 42], which realize generative trajectories by progressively perturbing the data distribution, $p_{\text{data}}$, into a noise distribution via a stochastic differential equation (SDE) and then learning to reverse the process, generating data from noise. Remarkably, this SDE has a corresponding ordinary differential equation (ODE) called the probability flow (PF) ODE [40]. Using the settings in [42] the PF ODE takes the form

for $\sigma\in[\sigma_{\text{min}},\sigma_{\text{{max}}}]$, where $\nabla_{\bm{{x}}}\log p_{\sigma}(\bm{x})$ is the score function of the perturbed data distribution $p_{\sigma}(\bm{x})$, and $\sigma(t)$ is the predefined noise scales. $\sigma_{\text{min}}$ is chosen sufficiently small, but not zero for numerical stability, such that $p_{\sigma_{\text{min}}}(\bm{x})\approx p_{\text{data}}(\bm{x})$ and $\sigma_{\text{max}}$ sufficiently large such that $p_{\sigma_{\text{max}}}(\bm{x})\approx\pi(\bm{x})$, some unstructured, easy-to-sample, noise distribution. In this case $\pi(\bm{x}):=\mathcal{N}(\bm{0},\sigma_{\text{max}}\bm{I}).$

In particular, for the solution trajectory $\{\bm{x}_{\sigma}\}_{\sigma\in[\sigma_{\text{min}},\sigma_{\text{max}}]}$ consistency models set out to learn a consistency function

by enforcing self-consistency

Let $F_{\bm{\theta}}(\bm{x},t)$ denote a free-form deep neural network, the desired solution will then be enforced via the boundary condition

where $c_{\text{skip}}(\sigma)$ and $c_{\text{out}}(\sigma)$ are differentiable functions such that $c_{\text{skip}}(\epsilon)=1$ and $c_{\text{out}}(\epsilon)=0$ and $f_{\bm{\theta}}(\bm{x},\sigma)$ denotes the deep neural network used to approximate the consistency function. To learn the consistency function, the domain $[\sigma_{\text{min}},\sigma_{\text{max}}]$ is discretized using a sequence of noise scales $\sigma_{\text{min}}=\sigma_{1}<\sigma_{2}<...<\sigma_{N}=\sigma_{\text{max}}.$ Following [42], reference [52] sets $\sigma_{i}=(\sigma_{\text{min}}^{1/\rho}+\frac{i-1}{N-1}(\sigma_{\text{max}}^{1/\rho}-\sigma_{\text{min}}^{1/\rho}))^{\rho}$ with $\rho=7$ for $i=1,...,N.$ Consistency models are then trained using the consistency matching loss

where

is a single-step in the reverse direction solving the PF ODE in Eq. (1), $\lambda(\cdot)\in\mathbb{R}_{+}$ is a positive weighting function, and $d(\cdot,\cdot)$ is a metric function such that $d(\bm{x},\bm{y})\geq 0$ and $d(\bm{x},\bm{y})=0\iff\bm{x}=\bm{y}.$ Reference [52] proposes to use the Learned Perceptual Image Patch Similarity (LPIPS) [58] as metric function. The expectation in Eq. (5) is taken with respect to $i\sim\mathcal{U}[1,N-1]$, a uniform distribution over integers $i=1,...,N-1$, and $\bm{x}_{\sigma_{i+1}}\sim p_{\sigma_{i+1}}(\bm{x}).$ The objective in Eq. (5) is minimized via stochastic gradient descent on parameters $\bm{\theta}$ whilst $\bm{\theta}^{-}$ are updated via the exponential moving average (EMA)

where $0\leq\mu\leq 1$ is the decay rate. $f_{\bm{\theta}}$ and $f_{\bm{\theta}^{-}}$ are called the student and teacher network, respectively. Directly minimzing the objective in Eq. (6) is not feasible in practice since $\nabla_{\bm{x}}\log p_{\sigma_{i+1}}(\bm{x})$ is unknown. Reference [52] presents two algorithms to circumvent this issue: consistency distillation and consistency training. Consistency distillation depends on a pre-trained score network, $s_{\bm{\phi}}(\bm{x},\sigma),$ to estimate the score function, $\nabla_{\bm{x}}\log p_{\sigma_{i+1}}(\bm{x}).$ The pre-trained score network is used instead of $\nabla_{\bm{x}}\log p_{\sigma_{i+1}}(\bm{x})$ in Eq. (6) to generate a pair of adjacent points $(\breve{\bm{x}}_{\sigma_{i}},\bm{x}_{\sigma_{i+1}})$ on the solution trajectory. Consistency training, on the other hand, is a stand-alone model. A perturbed sample $\bm{x}_{\sigma}$ from ground truth $\bm{x}$ can be obtained via the Gaussian perturbation kernel $p_{\sigma}(\bm{x}_{\sigma}|\bm{x})=\mathcal{N}(\bm{x},\sigma\bm{I})$ as $\bm{x}+\sigma\bm{z}$, where $\bm{x}\sim p(\bm{x})$ and $\bm{z}\sim\mathcal{N}(\bm{0},\bm{I}).$ Hence, consistency training obtains the pair of adjacent points on the same solution trajectory $(\bm{x}_{\sigma_{i+1}},\breve{\bm{x}}_{\sigma_{i}})$ via $\bm{x}_{\sigma_{i+1}}=\bm{x}+\sigma_{i+1}\bm{z}$ and $\breve{\bm{x}}_{\sigma_{i}}=\bm{x}+\sigma_{i}\bm{z}$ using the same $\bm{x}$ and $\bm{z}.$ Consistency training then employs the loss

which is equivalent to the consistency matching loss in Eq. (5) in the $N\rightarrow\infty$ limit [52]. Once we have obtained an approximation of the consistency function, a sample can simply be generated by drawing an initial noise vector from the noise distribution and passing through the consistency model $\hat{\bm{x}}=f_{\bm{\theta}}(\bm{z},\sigma_{\text{max}}).$ Importantly, consistency models do not sacrifice the ability to trade-off image quality for compute and zero-shot editing of data as they still support multi-step sampling [52].

PFGM++[44] is closely related to continuous-time diffusion models [40, 42] in theory and in practice. The training and sampling algorithms from [42] can be applied directly via a simple change-of-variables and an updated noise distribution [44]. It is also possible to show that the training and sampling processes of PFGM++ will converge to that of diffusion models [42] is the $D\rightarrow\infty,r=\sigma\sqrt{D},$ limit. Inspired by electrostatics, PFGM++ treats the $N$-dimensional data as an electric charge in a $N+D$-dimensional augmented space. Let $\tilde{\bm{x}}:=(\bm{x},\bm{0})\in\mathbb{R}^{N+D}$ and $\tilde{\bm{x}}_{\sigma}:=(\bm{x}_{\sigma},\bm{z})\in\mathbb{R}^{N+D}$ denote the augmented ground truth and perturbed data, respectively. The objective of interest is then the high dimensional electric field

where $S_{N+D-1}(1)$ is the surface area of the unit $(N+D-1)$-sphere, $p(\bm{x})$ is the distribution of the ground truth data. Notably, this electric field is rotationally symmetric on the $D$-dimensional cylinder $\sum_{i=1}^{D}z_{i}^{2}=r^{2},\forall r>0$ and thus it suffices to track the norm of the augmented variables $r=r(\bm{\tilde{x}}_{\sigma}):=||\bm{z}||_{2}.$ For notational brevity, redefine $\tilde{\bm{x}}:=(\bm{x},0)\in\mathbb{R}^{N+1}$ and $\tilde{\bm{x}}_{\sigma}:=(\bm{x}_{\sigma},r)\in\mathbb{R}^{N+1}.$ This dimensionality reduction, afforded by symmetry, yields the following ODE

where $\bm{E}(\bm{\tilde{x}}_{\sigma})_{\bm{x}_{\sigma}}=\frac{1}{S_{N+D-1}(1)}\int\frac{\bm{x}_{\sigma}-\bm{x}}{||\bm{\tilde{x}}_{\sigma}-\bm{\tilde{x}}||^{N+D}}p(\bm{x})d\bm{x}$ and $E(\bm{\tilde{x}}_{\sigma})_{r}=\frac{1}{S_{N+D-1}(1)}\int\frac{r}{||\bm{\tilde{x}}_{\sigma}-\bm{\tilde{x}}||^{N+D}}p(\bm{x})d\bm{x}$ are the $\bm{x}_{\sigma}$ and $r$ components of Eq. (9), respectively. Eq. (10) defines a bijective mapping between the ground truth data on the $r=0$ $(\bm{z}=\bm{0})$ hyperplane and an easy-to-sample prior noise distribution on the $r=r_{\max}$ hyper-cylinder [44]. As is the case for diffusion models [40, 42], PFGM++ uses a perturbation based objective function. Let $p_{r}(\bm{x}_{\sigma}|\bm{x})$ denote the perturbation kernel, which samples perturbed $\bm{x}_{\sigma}$ from ground truth $\bm{x}$, and $p(r)$ the training distribution over $r.$ Then the objective of interest is

We can ensure that the minimizer of Eq. (11) is $s^{*}_{\phi}(\bm{\tilde{x}}_{\sigma})=\sqrt{D}\bm{E}(\bm{\tilde{x}}_{\sigma})_{\bm{x}_{\sigma}}\cdot E(\bm{\tilde{x}}_{\sigma})_{r}^{-1}$ by choosing the perturbation kernel to be $p_{r}(\bm{x}_{\sigma}|\bm{x})\propto 1/(||\bm{x}_{\sigma}-\bm{x}||_{2}^{2}+r^{2})^{\frac{N+D}{2}}.$ As with diffusion models, a sample can then be generated by solving $d\bm{x}_{\sigma}/dr=\bm{E}(\bm{\tilde{x}}_{\sigma})_{\bm{x}_{\sigma}}/E(\bm{\tilde{x}}_{\sigma})_{r}=s^{*}_{\phi}(\bm{\tilde{x}}_{\sigma})/\sqrt{D}.$

We extend the idea of consistency models to where the underlying model is a PFGM++ instead of a diffusion model. PFGM++ accepts diffusion models as a special case but at the same time allows for trading-off robustness for rigidity by tuning $D$, the dimensionality of the augmentation variables. This yields a novel family of deep generative models which we call Poisson flow consistency models (PFCM). In this paper, we only treat the case of consistency distillation and leave consistency training, and improved consistency training [53], to future work. For consistency distillation, the PF ODE solver from the pre-trained diffusion model is applied as is, including for the second-order Heun solver proposed in [42]. Using the noise distribution from PFGM++ we can sample $\bm{x}_{\sigma_{i+1}}$ from the PFGM++ solution trajectory and, as with consistency models, we can subsequently obtain the correspond $\breve{\bm{x}}_{\sigma_{i}}$ by taking a single reverse step of the PF ODE solver using the pre-trained network from PFGM++, which estimates $\sqrt{D}\bm{E}(\bm{\tilde{x}}_{\sigma})_{\bm{x}_{\sigma}}\cdot E(\bm{\tilde{x}}_{\sigma})_{r}^{-1}.$ Hence, exactly as in PFGM++[44], we can re-purpose the consistency distillation training algorithm to train a PFCM via updating the noise distribution and applying a change-of-variables.

In particular, for our discretized domain $[\sigma_{\text{min}},\sigma_{\text{{max}}}]$, let $\Phi(\cdot,\cdot,\bm{\phi})$ denote the update function of a one-step ODE solver applied to the PF ODE. We can then write Eq. (6) as

where $\Phi(\bm{x}_{\sigma_{i+1}},\sigma_{i+1},\bm{\phi}):=-\sigma_{i+1}s_{\bm{\phi}}(\bm{x}_{\sigma_{i+1}},\sigma_{i+1})$, using the first order Euler step for simplicity. More generally, Eq. (12) is taking a single-step in the reverse direction solving a PF ODE on the form

Using the fact that $\sigma(t)=t,$ $\bm{\tilde{x}}_{\sigma}:=(\bm{x}_{\sigma},r),$ and the alignment formula $r=\sigma\sqrt{D}$ [44], we can do the change-of-variables $dr=d\sigma\sqrt{D}=dt\sqrt{D}$ and rewrite $d\bm{x}/dr$ in Eq. (10) as

Hence, we can write the PF ODE in Eq. (10) on the form in Eq. (13) with $\Phi_{\text{PFGM++}}(\bm{x},\sigma,\bm{\phi}):=s_{\bm{\phi}}(\tilde{\bm{x}})=\sqrt{D}\bm{E}(\bm{\tilde{x}}_{\sigma})_{\bm{x}_{\sigma}}\cdot E(\bm{\tilde{x}}_{\sigma})_{r}^{-1}.$ A PFCM can then be trained via consistency distillation by using $\Phi_{\text{PFGM++}}(\cdot,\cdot,\bm{\phi})$, from a pre-trained PFGM++, and an updated noise distribution.

As for diffusion models, $p_{r}(\cdot|\bm{x})$ enables sampling of perturbed $\bm{x}_{\sigma}$ from ground truth $\bm{x}.$ However, sampling from $p_{r}(\cdot|\bm{x})$ is slightly more involved than from a simple Gaussian. In particular, $p_{r}(\cdot|\bm{x})$ is first split into hyperspherical coordinates to $\mathcal{U}_{\psi}(\psi)p_{r}(R)$, where $\mathcal{U}_{\psi}$ denotes the uniform distribution of the angle component and the distribution of $R=||\bm{x}_{\sigma}-\bm{x}||$, the perturbed radius, is

In practice, reference [44] proposes to sample from $p_{r}(\cdot|\bm{x})$ through the following steps: 1) get $r_{i}$ from $\sigma_{i}$ via the alignment formula $r_{i}=\sigma_{i}\sqrt{D}$. 2) sample radii $R_{i}$ from $p_{r_{i}}(R)$, 3) sample uniform angles $\bm{v}_{i}\leftarrow\bm{u}_{i}/||\bm{u}_{i}||_{2},$ where $\bm{u}_{i}\sim\mathcal{N}(\bm{0};\bm{I})$, 4) get perturbed data $\bm{x}_{\sigma_{i}}=\bm{x}+R_{i}\bm{v}_{i}\sim p_{r}(\bm{x}_{\sigma_{i}}|\bm{x}).$

Finally, we are interested in inverse problem solving, in particular image denoising. As with diffusion models and PFGM++, one can move from an unconditional generator to a conditional one in two main ways: via zero-shot data editing using a network trained in an unsupervised fashion and training the network to directly learn a trajectory from the prior noise distribution to the posterior distribution of interest. The main advantages for the former is that it imposes a much laxer data requirement and allows for using one single pre-trained network for a range of downstream tasks (e.g., denoising, inpainting, colorization, super-resolution). The main advantage for the latter is that much more information is available at training time. Hence, we expect that, all else equal, this will yield superior performance. Since we have paired data available in the Mayo low-dose CT data, we opt for the supervised approach in this paper and leave the unsupervised approach for future work. Hence, we will do one more update to the framework in [52]: the pre-trained PFGM++ and PFCM are both trained in a supervised fashion to directly learn a mapping between the noise distribution and the posterior distribution of interest. Feeding the condition image, $\bm{y}$, as additional input to the network to directly learn a given inverse problem has been demonstrated to work well empirically for diffusion models [45, 48, 49], PFGM++ [34], and consistency models [34]. Our results indicate that this approach generalizes well to PS-PFCM. Our proposed training of PS-PFCM is available in Algorithm 1 with updates to consistency distillation in [52] highlighted in blue. The proposed sampling is available in Algorithm 2.

All networks are trained on randomly extracted $256\times 256$ patches using rectified Adam [56] with learning rate $1\times 10^{-4}$ for PFGM++ and $1\times 10^{-5}$ for PS-PFCM. In addition to yielding more efficient training in terms of lower graphics memory requirements, training on randomly extracted patches also serves as data augmentation. To further reduce graphics memory requirements the networks are trained using mixed precision. Hyperparameters for training and sampling are set as in [52] for the LSUN $256\times 256$ experiments¹¹1As specified in https://github.com/openai/consistency_models/blob/main/scripts/launch.sh. with the exception of batch size which we set to 4 and $\sigma_{\text{max}}$ which we set to $\sigma_{\text{max}}=380$ as in PFGM++[44]. The network architecture is as in [57] and we use channel multiplier 256, channels per resolution [1, 1, 2, 2, 4, 4], attention layers at resolutions [32,16,8], and two res-blocks per resolution. We use a dropout probability of $10\%$ when training PFGM++. As for CD in [52], we set $\lambda(\sigma):=1$ and using LPIPS[58] as metric function. To get conditional models, trained in a supervised fashion, the network is updated to accept one additional input channel and the training and sampling scripts are updated accordingly. The networks are trained using $4\times$ NVIDIA A6000 48 GPUs. We train PS-PFCM for $D\in\{4096,8192,131072\}$, where we have followed the recommendation to re-scale $D$ with $N.$²²2https://github.com/Newbeeer/pfgmpp/issues/10. The PFGM++ are first trained for 300k iterations. We subsequently train the PS-PFCM for another 300k iterations.

In addition to PFGM++[44], which is a state-of-the-art diffusion-style model, we also compare PS-PFCM to the corresponding method using consistency models [52, 53], the current state-of-the-art diffusion-style models with NFE=1, as well EDM [42]. All models are trained in a supervised fashion, using the “trick” of feeding the condition image, $\bm{y}$, as additional input to the network to directly learn a trajectory between the prior noise distribution and the posterior distribution of interest. EDM and CD are trained with identical hyperparameters as for PFGM++ and PS-PFCM. The improved version of consistency training, iCT [53], demonstrates impressive performance, outperforming even consistency distillation [53]. However, we still choose to compare to consistency distillation for two reasons. First, to the best of our knowledge, consistency training or iCT has never be trained in a supervised fashion to directly learn a trajectory to the desired posterior. Reference [34], on the other hand, trained a consistency model using consistency distillation in a supervised fashion and obtained good results. Second, since the proposed method is based on distillation, not consistency training, comparing to consistency distillation, trained for the same task with identical hyperparameters, to PS-PFCM is tantamount to an ablation study of the impact of choosing $D$ finite versus $D\rightarrow\infty.$ As in reference [52], we refer to the consistency model trained using consistency distillation as CD.

The results are evaluated qualitatively, via visual inspection, and quantitatively using average peak signal-to-noise ratio (PSNR), structural similarity index (SSIM [59]), and LPIPS over the low-dose CT validation dataset. Extensive experiments conducted in [58] illustrate that LPIPS corresponds much better to human perception than more traditional metrics such as SSIM and PSNR. However, one important caveat is that we now apply it to gray-scale CT images whereas the experiments, and training of the underlying networks, were on natural RGB images. The results correlate well with human perception; however, the interpretation of LPIPS is a bit unclear since we are essentially applying it to an out-of-distribution case. We employ the same version of LPIPS as used in the score-matching loss.

Quantitative results are available in Table I. We include these metrics for the low-dose CT (LDCT) images as a baseline. The overall top performer is EDM. However, this top performance incurs the computational cost of $\text{NFE}=79.$ For the $D$ considered here, the best performance for PFGM++ is obtained for $D=4096$, which is marginally worse than for EDM. CD achieves $\text{NFE}=1$ but it comes with a quite significant drop in performance compared to EDM. Moving from PFGM++ to single-step PS-PFCM also results in the expected drop in performance. However, PS-PFCM $D=4096$ outperforms the corresponding CD according to each metric considered, despite distilling from a marginally worse performing underlying model. Hence, we have demonstrated that there exists a finite $D\in(0,\infty)$ for which the penalty incurred for moving to a single-step sampler is lower for PS-PFCM than for consistency models. This is likely due to the additional robustness afforded by setting $D$ relatively low in PFGM++. In this setup, the single-step sampler yields $79\times$ faster sampling. To put this in perspective, processing the 1136 slices in the low-dose CT validation set with NFE=1 takes about 3 minutes whereas it takes around 4 hours for NFE=79 on a single NVIDIA A6000 48GB GPU.

Qualitative results for a representative sample from the low-dose CT validation set is available in Fig. 1 and 2. This patient has a metastasis in the liver and we show a magnified version of the region-of-interest (ROI) in Fig. 1, containing this lesion, in Fig. 2. SSIM and PSNR for this particular slice, in addition to NFE for easy comparison, are also included in Fig. 1. The normal-dose CT (NDCT) image, shown in a), is treated as ground truth. We can see that the metastasis is clearly visible in all images, including for the NDCT image. EDM, PFGM++, CD, and PS-PFCM all produce realistic noise characteristics that closely matches that of the NDCT image. How they differ is the accuracy in which certain details are reproduced. Since all models are 2D denoising networks, it is feasible that there is simply insufficient information in the LDCT image to fully reproduce all features visible in the NDCT image. We have marked a feature that seem particularly difficult to reproduce with a yellow arrow. There is a large variation in how this feature is reproduced both when comparing EDM to PFGM++, with $D\in\{4096,8192,131072\}$, but also when comparing EDM or PFGM++ to the corresponding distilled consistency model or PS-PFCM. Qualitatively, it is difficult to state which model does the best job reproducing this feature.