Memory-Efficient Deep end-to-end Posterior Network (DEEPEN) for inverse problems

Current EBM methods pre-learn $p({\boldsymbol{x}})$ from the training data and then use it in (3), similar to PnP methods. Motivated by the improved performance of E2E methods, we learn the posterior distribution $p_{\boldsymbol{\theta}}({\boldsymbol{x}}|{\boldsymbol{b}})$ using the training data and for a specific forward model $\mathbf{A}$ in an E2E fashion using maximum likelihood optimization. In particular, we determine the optimal weights of the energy $E_{\boldsymbol{\theta}}({\boldsymbol{x}})$ by minimizing the negative log-likelihood of the training data set:

where $q({\boldsymbol{x}})$ is the probability distribution of the data and the posterior $p_{\boldsymbol{\theta}}({\boldsymbol{x}}|{\boldsymbol{b}})$ is defined as in (3) whose prior is defined in (4). We thus have:

where $\tilde{{\boldsymbol{Z}}}_{\boldsymbol{\theta}}=\int\exp\left(-\dfrac{1}{2}\|{\mathbf{A}}{\boldsymbol{x}}-{\boldsymbol{b}}\|^{2}-E_{\boldsymbol{\theta}}({\boldsymbol{x}})\right)d{\boldsymbol{x}}$ is a normalizing constant. For simplicity, we absorb $\eta^{2}$ parameter into the definition of the energy. Consequently we have:

The second term can be computed using the chain rule [5]:

where ${\boldsymbol{x}}\sim p_{\boldsymbol{\theta}}({\boldsymbol{x}}|{\boldsymbol{b}})$ are samples from the learned posterior distribution $p_{\boldsymbol{\theta}}({\boldsymbol{x}}|{\boldsymbol{b}})$. In the EBM literature, these are the generated or fake samples denoted by ${\boldsymbol{x}}^{-}$, while the training samples are referred to as true samples, denoted by ${\boldsymbol{x}}^{+}\sim q({\boldsymbol{x}})$. Thus, the ML estimation of $\boldsymbol{\theta}$ amounts to minimization of the loss:

Intuitively, the training strategy (II-A) will seek to decrease the energy of the true samples ($E_{\boldsymbol{\theta}}({\boldsymbol{x}}^{+})$) and increase the energy of the fake samples ($E_{\boldsymbol{\theta}}({\boldsymbol{x}}^{-})$). This may be seen as an adversarial training scheme similar to [7], where the classifier involves the energy itself as shown in Fig. 2. The algorithm converges when the the fake samples are identical in distribution to the training samples; i.e, $E_{\boldsymbol{\theta}}({\boldsymbol{x}}^{+})\approx E_{\boldsymbol{\theta}}({\boldsymbol{x}}^{-})$.

We generate the fake samples ${\boldsymbol{x}}^{-}\sim p_{\boldsymbol{\theta}}({\boldsymbol{x}}|{\boldsymbol{b}})$ using Langevin MCMC method, which only requires the gradient of $\log p_{\boldsymbol{\theta}}({\boldsymbol{x}}|{\boldsymbol{b}})$ w.r.t. ${\boldsymbol{x}}$ :

where $\epsilon>0$ is the step-size, ${\boldsymbol{z}}_{n}\sim\mathcal{N}(0,{\mathbf{I}})$ and ${\boldsymbol{x}}_{0}$ is drawn from a simple posterior distribution. The entire training process is summarized in Fig. 2.

Once the posterior is learned, we obtain the MAP estimate by minimizing (6) w.r.t. ${\boldsymbol{x}}$ using the classical gradient descent algorithm:

where $\alpha_{k}$ is the step-size and is found using the backtracking line search method [8]. This line search method starts with a large step-size value, i.e., $\alpha_{k}=1$, and keeps decreasing it by a factor of $\beta\in(0,1)$ until the chosen step-size sufficiently decreases the cost function. The decrease in the cost function is typically measured using the Armijo–Goldstein rule, which is $\mathcal{L}_{\theta}\left({\boldsymbol{x}}_{k}-\alpha_{k}\nabla_{x}\mathcal{L}_{\theta}({\boldsymbol{x}}_{k})\right)>\mathcal{L}_{\theta}({\boldsymbol{x}}_{k})-\beta\alpha_{k}\|\nabla_{x}\mathcal{L}_{\theta}({\boldsymbol{x}}_{k})\|_{2}^{2}$, thereby ensuring that the chosen step-size monotonically decreases the cost function at every iteration. The following result shows that the proposed algorithm converges monotonically to a local minimum of the cost function (6).

Note that the above theorem is applicable irrespective of the Lipschitz constant of the CNN, and when $\nabla_{\boldsymbol{x}}\mathcal{L}_{\theta}({\boldsymbol{x}})$ is a conservative vector field i.e., it is the gradient of $\mathcal{L}_{\theta}({\boldsymbol{x}})$.

We evaluated the performance of DEEPEN on the publicly available parallel fastmri brain data set which contains FLAIR and T2-weighted images [9]. It is a twelve-channel brain data set and consists of complex images of size $320\times 320$. The matrix ${\mathbf{A}}$ in (1) in this case is ${\mathbf{A}}={\boldsymbol{S}}{\mathbf{F}}{\boldsymbol{C}}$, where ${\boldsymbol{S}}$ is the sampling matrix, ${\mathbf{F}}$ is the Fourier transform and ${\boldsymbol{C}}$ is the coil sensitivity map which is estimated using ESPIRiT algorithm [10]. The data set for each contrast was split into $45$ training, $5$ validation, and $50$ test subjects. DEEPEN was trained and evaluated on both contrasts separately for a four-fold retrospective undersampled measurement, which was obtained by sampling the data along the phase-encoding direction using a 1D non-uniform variable density mask.

The network $E_{\boldsymbol{\theta}}({\boldsymbol{x}})$ defined in (4) was built using five 3x3 convolutional layers with 64 channels each, followed by a linear layer. A ReLU was used between each convolutional layer and at the end of the linear layer. We evaluated the gradient $\nabla_{\boldsymbol{x}}E_{\boldsymbol{\theta}}({\boldsymbol{x}})$ using the chain rule.

The MCMC sampling was performed for $30$ iterations and was initialized with ${\boldsymbol{x}}_{0}=({\mathbf{A}}^{H}{\mathbf{A}}+\tilde{\lambda}{\mathbf{I}})^{-1}{\mathbf{A}}^{H}{\boldsymbol{b}}$ which is obtained by minimizing the negative logarithmic posterior distribution $-\log p_{0}({\boldsymbol{x}}|{\boldsymbol{b}})=\|{\mathbf{A}}{\boldsymbol{x}}-{\boldsymbol{b}}\|_{2}^{2}+\tilde{\lambda}\|{\boldsymbol{x}}\|_{2}^{2}$ w.r.t. ${\boldsymbol{x}}$. Next, similar to [11], for stable training we found it beneficial to: a) add Gaussian noise of standard deviation $2\epsilon$ to the training data $\{{\boldsymbol{x}}_{i}^{+}\}$ that effectively made the training data smooth, and b) scale $\nabla_{x}\log p_{\boldsymbol{\theta}}({\boldsymbol{x}}|{\boldsymbol{b}})$ by $\dfrac{\epsilon^{2}}{2}$ which is equivalent to using the following Langevin MCMC update:

DEEPEN was trained with $\epsilon=0.001$ and the optimal $\boldsymbol{\theta}$ was found using the Adam optimizer. The gradient descent algorithm was run until $\dfrac{|\mathcal{L}_{\boldsymbol{\theta}}({\boldsymbol{x}}_{k+1})-\mathcal{L}_{\boldsymbol{\theta}}({\boldsymbol{x}}_{k})|}{|\mathcal{L}_{\boldsymbol{\theta}}({\boldsymbol{x}}_{k})|}\leq 10^{-6}$.
We compared the performance of the proposed algorithm with SENSE [12], MoDL [2], and MoL [4]. The unrolled algorithms were trained in an E2E fashion for $10$ iterations. A five-layer CNN was used for both the unrolled algorithms. The Lipschitz constraint in case of MoL was implemented using the log-barrier approach.

Table I compares the performance of the reconstruction algorithms on the test data. Fig. 3(a) and Fig. 3(b) show the reconstructed images of T2-weighted and FLAIR images, respectively. PSNR and SSIM are used as evaluation metrics. Table I shows that DEEPEN and MoDL perform better than SENSE and MoL, and perform comparably with each other. While SENSE is a CS-based algorithm and does not have a neural network module, the lower performance of the unrolled MoL algorithm is due to the Lipschitz constraint on its CNN. This constraint is required to ensure convergence of MoL to the fixed point. Note that, according to Theorem 2.1, the DEEPEN algorithm can ensure convergence to a stationary point without any constraint on $\nabla_{x}E_{\boldsymbol{\theta}}({\boldsymbol{x}})$ - which results in improved performance. We would like to remind the reader that although DEEPEN and MoDL have similar performance, DEEPEN is memory-efficient and hence paves the way to reconstruct images of higher dimensions (for example, 3D brain volumes) which is not possible with MoDL.

Given the posterior distribution $p_{\boldsymbol{\theta}}({\boldsymbol{x}}|{\boldsymbol{b}})$, we can now draw samples from it. This aids in estimating the Minimum Mean Square Error (MMSE) and the uncertainty map of the reconstructed image. Fig. 4 shows the MMSE and the uncertainty map provided by DEEPEN on a four-fold undersampled FLAIR-contrast MR brain image. The MMSE and the uncertainty map are estimated by computing the mean and variance of $100$ samples from the posterior distribution. The samples are obtained using the Langevin MCMC method initialized randomly from a Gaussian distribution.