Invertible Learned Primal-Dual
Abstract
We propose invertible Learned
Primal-Dual as a method for tomographic image reconstruction.
This is a learned iterative method based on the Learned PrimalDual neural network architecture, which incorporates ideas from invertible neural networks.
The invertibility signiﬁcantly reduces the GPU memory footprint of the
Learned Primal-Dual architecture, thus making it applicable to 3D tomographic reconstruction as demonstrated in the experiments
Introduction
Combining model-based and data driven methods has demonstrated promising performance in medical image reconstruction [4].
A key idea is to deﬁne domain adapted neural networks that integrate components from regularization theory.
In learned iterative schemes this is done by unrolling a suitable optimization scheme and replacing selected terms with trainable neural networks [23].
Several unrolled reconstruction methods have been proposed for various imaging modalities: computed tomography (CT) [3], photoacoustic tomography [14], magnetic resonance (MR) [12], and lensless imaging [22].
These are based on unrolling different schemes, such as gradient descent [2], [29], primal-dual [3], alternating direction method of multiplier (ADMM)
Here we focus on CT reconstruction, where unrolled methods have shown an outstanding performance in the low dose 2D setting.
However, training and deploying these methods for 3D CT reconstruction is still a challenge due to the signiﬁcantly increased memory requirements.
Recent progress in direction of decreasing the GPU memory footprint of such methods is presented in [13], where a multi-scale approach is combined with learned gradient-descent (LGD), as well as in [29] that relies on greedy training of LGD on image patches.
The aim here is to explore the idea of invertibility as an alternative path for reducing GPU memory requirements.
Invertible neural networks (INNs) were introduced in NICE [7] for the purpose of modelling complex probability densities.
Here invertibility was achieved by additive coupling, which was later extended to afﬁne coupling in [8].
Since then many works rely on additive/afﬁne coupling to construct invertible architectures [11, 15, 18, 9] for classiﬁcation, density modelling and Algorithm ( x , u 0 ) ∈ ( X N , Y N
Note
[Table] [Table] denoising. As outlined in [11], one of the key beneﬁts of INN is that depth of the network can be increased while maintaining a constant memory footprint.
This allows to increase the number of unrolled iterations in a learned iterative method i-RIM [26, 25] and to apply deep learning to 3D
MR image reconstruction.
Similarly to [26, 25], we combine the idea of INN with the learned iterative method in [3] that introduces the Learned Primal-Dual (LPD) architecture for CT reconstruction.
The key difference in our work is that we do not rely on afﬁne coupling, since a slight change in the LPD architecture is sufﬁcient to make it invertible.
We demonstrate that our architecture signiﬁcantly reduces the
GPU memory requirements in comparison to the original LPD in the 2D setting, while achieving comparable image reconstruction quality.
Furthermore, we are able to apply our method in the 3D setting, where it achieves superior results comparing to other methods.
[-u i − 
-]
[-Primal update: (cid:40) x i +
-]
Background 2.1 Learned Primal-Dual
The LPD architecture is a domain adapted neural network especially suited for solving an inverse problem.
It is typically trained in a supervised manner and it has been used for CT reconstruction.
The latter can be formulated as the task of recovering a 2D/3D image x ∈ X from noisy tomographic data/sinogram y ∈ Y , where
[Equation]
In the above, the (linear) forward operator T : X → Y models how data is generated in the absence of noise, and δ ( x ) ∈ Y is observation noise.
We phrase a machine learning based approach applied to (1) as the problem of ﬁnding a (non-linear) parametrized mapping T † θ : Y → X that satisﬁes the pseudo-inverse property: T † θ ( y )
≈ x .
A key element in this approach is to parametrize the set of pseudo-inverse mappings with parameter θ ∈ Θ .
Supervised training is performed by minimizing the loss functional
[Equation]
for training dataset D .
Here (cid:96) measures how close the obtained reconstruction is to the ground truth.
The LPD architecture is inspired by the iterative scheme in the primal-dual hybrid gradient (PDHG) algorithm
This architecture incorporates a forward operator into a deep neural network by unrolling a proximal primal-dual optimization scheme and replacing proximal operators with convolutional neural networks (CNNs).
More precisely, the LPD architecture is given in Algorithm 1, where M is the number of unrolling iterates, N is the total number of memory channels and superscripts (1) and (2) denote the 1st and the 2nd channels of assigned variables.
The functions Γ i and Λ i correspond to CNNs with different learned parameters but with the same architecture for each unrolled iteration i .
To simplify notation, we suppress the explicit dependence of Γ i and Λ i on the learned parameters θ ∈ Θ .
2.2 Invertible Neural Networks via Coupling Layers
Additive coupling [7] is a simple method to construct invertible layers.
Let ˜ X
= X C 1 + C 2 represent a space of feature maps with C 1 + C 2 channels.
Let f : X C 1 → X C 2 be an arbitrary function (meant to be a neural network).
Then, an invertible mapping Φ : ( x i − 1 , u i − 1 ) (cid:55)→ ( x i , u i ) and its inverse Φ − 1 are deﬁned by x i = x i − 1 + f ( u i − 1 ) u i =
u
i −and
u i − 1 = u
i x i − 1 = x
i − f ( u i ) .
A INN of arbitrary depth can then be constructed by a simple composition of these invertible mappings.
If activations are restored by successively applying the layers’ inverses (instead of storing them), the memory demand is independent of the network’s depth, as described in [11].
Due to the need for re-computing activations, one additional forward pass has to be performed per training step.
Invertible Learned Primal-Dual
Inspired by additive coupling layers, we deﬁne the invertible dual update mapping G iy and primal update mapping F i .
While standard coupling layers typically map between image spaces of the same resolution, we rely on the primal and dual operators T and T ∗ to convert between the primal and dual spaces, extended to multiple channels ( X N and Y N ).
The dual update mapping G iy : ( x i − 1 , u i − 1 ) (cid:55)→ ( x i , u i ) is given by
[Equation]
where the primal-to-dual mapping is deﬁned by g iy : = g iy, post ◦ T ◦
g
i pre : X N → Y N .
Here, g i pre ( pre-operator network ) and g iy, post ( post-operator network ) are CNNs, where the pre-operator network operates on the primal space and the post-operator network operates on the dual space.
In this context, tomographic data y is a parameter of the post-operator network g iy, post , even though in practice it is incorporated via concatenation with the input variable x i − 1 .
Note that these networks allow for arbitrarily changing the number of channels if desired.
The primal update mapping F i : ( x i − 1 , u i − 1 ) (cid:55)→ ( x i , u i ) is deﬁned as
[Equation]
where the dual-to-primal mapping is given by f i :
= f i post ◦ T ∗ ◦
f i pre : Y N → X N , similarly to the g iy .
We can now deﬁne an unrolled scheme through the composition
[Equation]
F y : X N × Y N → X N × Y N
[Equation]
Since F i and G iy are invertible, the whole unrolled scheme F y is invertible, and we can use memoryefﬁcient backpropagation schemes.
Combined with projection operator P : X N × Y N → X (e.g. retaining only one channel) and an initialization scheme init , the unrolled scheme F y deﬁnes the iLPD reconstruction method:
[Equation]
The above parametrization of T † θ is general and results in a deep neural network that can be adapted for solving a particular inverse problem.
We simplify this general approach to the form described in
Algorithm 2, by making several speciﬁc choices.
First, we choose the number of memory channels for dual and primal variables to be N = 1 .
Next, we substitute the pre-operator networks for both dual and primal updates with identity functions, g i pre = f i pre = Id , and chose the post-operator networks to be 3-layer CNNs, g iy, post = Γ i and f iy, post = Λ i (see Appendix A.1).
Note that M steps in LPD are equivalently represented as 2 M steps in iLPD.
Therefore, the main difference between
LPD and iLPD is that Γ i and Λ i do not take u i − 1 and x i as inputs respectively, i.e. we achieve invertibility of LPD by removing skip connections.
Results 4.1 2D CT reconstruction
In the ﬁrst set of experiments we simulate tomographic data from CT images of human abdomen provided by Mayo Clinic for the AAPM Low Dose CT Grand Challenge [21].
The dataset contains [Table] 3D scans of 10 patients, of which 9 are used for training and one for testing.
We split the reconstructed 3D volumes with 3 mm slice thickness into axial slices.
This results in 2 378 images of size of size 512 × 512 , of which 2168 are used for training and 210 are used exclusively for testing.
In addition, 1% of training images is reserved for model validation during training.
Since the images are given in Hounsﬁeld scale, we re-scale those by a factor 10 − 3 , so that pixel intensity values are approximately in the interval [0; 2] , which is convenient for applying deep learning models.
We use the same procedure for simulating the tomographic data, as described in the original LPD paper [3].
In particular, the forward operator T is a fan-beam ray transform with 1 000 projection angles and 1 000 detector elements and the data is corrupted by Poisson noise (see Appendix A.2).
We compare the performance of the iLPD with the original LPD architecture [3] in terms of PSNR,
SSIM and GPU memory footprint.
The results are presented in Table 1.
The original LPD architecture consists of 10 unrolled iterations ( M = 10 ) and has 5 memory channels ( N = 5 ).
First, we can see that removing memory channels from LPD (as in LPD-NoMem with N = 1 ) results only in a slight drop in performance.
Furthermore, enforcing invertibility in iLPD-10 ( N = 1 , M = 10 ) leads to additional drop in performance comparing to LPD-NoMem.
However, this drop can be partially compensated by increasing the number of primal-dual iterations as in iLPD-20 ( M = 20 ) leading to performance more close to original LPD.
Furthermore, both invertible networks require much less
GPU memory during training.
In addition, the memory footprint stays almost constant while depth of the network (which depends on the number of unrolled iterations) is increased.
To summarize, we can see that enforcing invertibility leads to reduced memory consumption which comes at the cost of slight drop in performance.
4.2 3D CT reconstruction
We evaluate the applicability of our method to 3D data, using a CT database of 42 walnuts scanned with cone beam CT (CBCT), provided by the FleX-ray lab at Centrum Wiskunde & Informatica [6].
Data acquisition was performed with 0 .
3 ◦ uniform angle increment resulting in 1 200 projections.
All walnuts were scanned with source positioned at 3 different heights and the projection data were combined to create artefact-free ground-truth reconstructions.
In our experiments, we select 120 uniformly spaced angles out of 1200.
Additionally we down-sample both, measurements and ground-truth reconstructions, by a factor of 2.
This results in reconstructions of size 251 3 and projection data of size 384 × 486 .
This was the maximum reconstruction size possible under the GPU memory constraints of this study.
The maximum reconstruction size from previously published results in 3D reconstruction of this dataset performed by a learned iterative method was 168 [-
 [13],so-]3 [13], so with iLPD we managed to reconstruct 3 .
375 larger volumes than [13] thanks to the reduction of memory consumption achieved by utilizing invertibility.
Additionally, we perform reconstruction with Feldkamp-Davis-Kress (FDK) [10], learned FDK (NN-FDK)
[19], post-processing with U-Net [16] and greedily learned gradient descent on patches (gLGDp) [29] (see Appendix A.1).
For training of all methods we have chosen 40 walnuts, which leaves one for validation (#1) and one for testing (#2).
The obtained reconstructions for the test walnut are shown in Figure 1 (Appendix A.3) while performance metrics are presented in Table 2.
Both qualitative and quantitative results show that the proposed iLPD out-performs all competitors.
Concerning execution times, clearly FDK is the fastest to execute, this is simply because execution time largely depends on the number of evaluations of T and T ∗ (FDK has one evaluation of T ∗ , and iLPD has 20 evaluations of T and T ∗ ).
However, for all methods execution time is within acceptable limits for practical applications. [Table]
Conclusion
We have proposed a new algorithm in the family of deep learning based iterative reconstruction schemes.
The algorithm is a modiﬁed LPD method, where by changing its original architecture we achieved invertibility.
The proposed iLPD method requires signiﬁcantly less GPU memory for training than the original LPD and therefore it is applicable for 3D CT reconstruction.
We demonstrated that iLPD algorithm gives state of the art results for both 2D and 3D CT
A Appendix A.1 Implementation details
The methods described above were implemented in Python using Operator Discretization Library (ODL) [1] and Pytorch [24].
All operator-related components, such as the forward operator T , were implemented in ODL with ASTRA [28] as back-end.
These were then converted into Pytorch layers utilizing the OperatorModule functionality of ODL.
The neural network layers and training were implemented using PyTorch.
The invertibility is incorporated employing MemCNN [20] library.
We re-implemented LPD network in Pytorch keeping the structure and hyper-parameters the same as in the original work [3].
In particular, we used M = 10 unrolled iterations and N = 5 memory channels for both primal and dual variables.
A three-layer CNN is applied in both image and projection domain at each unrolled iteration for both LPD and iLPD, i.e., Λ i and Γ i have two hidden convolutional layers with 32 ﬁlters and one convolutional output layer.
Parametric Rectiﬁed Linear Units (PReLU) are used as activation functions after hidden layers.
In 3D experiments 2-dimensional convolutional layers are replaced with 3-dimensional layers, but in both cases the kernel size is 3 along each dimension.
Both LPD and iLPD were initialized with zeros in both 2D and 3D experiments.
An implementation of NN-FDK has been provided by the authors of the method [19].
The U-Net post-processor has the same architecture as in [16].
The only difference is that instead of 64 number of ﬁlters in convolutional layers, we used 40 due to memory constraints.
For the same reason, training was performed on patches of size 128 3 instead of on entire volumes.
Here an initial reconstruction is ﬁrst performed using FDK and a neural network is trained on pairs of noisy FDK images and noiseless/low noise ground truth images, learning a mapping between them.
We implemented gLGDp method as suggested in [29].
We used only 3 unrolled iterations (since no beneﬁt was observed from using more), patch size of 160 3 and same U-net modiﬁcations.
Although we did not split the projection angles to subsets as in the original work, we did not observe additional improvements from including more than one gradient step within a learned iteration.
The algorithm was initialized with FDK reconstructions.
All learned iterative methods as well as U-net post-processor were trained using Adam as the optimizer [17] using the (cid:96) 2 -distance to the ground-truth as loss.
Networks are trained for 100 ,iterations in 2D and for 10 , 000 iterations in 3D. In all cases, the batch size was chosen to be 1.
The initial learning rate is set to 10 − 3 for LPD and U-net post-processor, 5 × 10 − 4 for iLPD and 10 −for gLGDp with cosine annealing.
All experiments in 2D and in 3D are executed on single GeForce RTX 3090 GPU with available memory of [-
-]24234 MiB. Finally, we provide an implementation of proposed method 1 .
A.2 Data simulation
The data used in 2D experiments (Sec.4.1) are corrupted by noise such that the number of photons absorbed by the detector is a Poisson distributed random variable
[Equation]
where H 0 the number of photons per pixel before attenuation and µ
= 0 .
mm − 1 is X-ray attenuation of water at the mean X-ray energy of 70 keV.
Then, we linearize the tomographic data
[Equation]
so that y noisy ≈ T ( x ) .
We use value H
= 10 4 , which corresponds to a low dose CT scan.
A 3 Images yz-plane xz-plane xy-plane
P h a n t o m
F DK
NN - F
DK yz-plane xz-plane xy-plane
U - N e t g L
GD p i L P
D