Zero-Shot Self-Supervised Learning for MRI Reconstruction

1 INTRODUCTION . Magnetic resonance imaging ( MRI ) is a non-invasive , radiation-free medical imaging modality that provides excellent soft tissue contrast for diagnostic purposes . However , lengthy acquisition times in MRI remain a limitation . Accelerated MRI techniques acquire fewer measurements at a subNyquist rate , and use redundancies in the acquisition system or the images to remove the resulting aliasing artifacts during reconstruction . In clinical MRI systems , multi-coil receivers are used during data acquisition . Parallel imaging ( PI ) is the most clinically used method for accelerated MRI , and exploits the redundancies between these coils for reconstruction ( Pruessmann et al. , 1999 ; Griswold et al. , 2002 ) . Compressed sensing ( CS ) is another conventional accelerated MRI technique that exploits the compressibility of images in sparsifying transform domains ( Lustig et al. , 2007 ) , and is commonly used in combination with PI . However , PI and CS may suffer from noise and residual artifacts at high acceleration rates ( Robson et al. , 2008 ; Sandino et al. , 2020 ) . Recently , deep learning ( DL ) methods have emerged as an alternative for accelerated MRI due to their improved reconstruction quality compared to conventional approaches ( Hammernik et al. , 2018 ; Knoll et al. , 2020b ) . Particularly , physics-guided deep learning reconstruction ( PG-DLR ) approaches have gained interest due to their robustness and improved performance ( Hammernik et al. , 2018 ; Aggarwal et al. , 2019 ; Schlemper et al. , 2018 ; Hosseini et al. , 2020b ) . PG-DLR explicitly incorporates the physics of the data acquisition system into the neural network via a procedure known as algorithm unrolling ( Monga et al. , 2021 ) . This is done by unrolling iterative optimization algorithms that alternate between data consistency ( DC ) and regularization steps for a fixed number of iterations . Subsequently , PG-DLR methods are trained in a supervised manner using large databases of fully-sampled measurements ( Hammernik et al. , 2018 ; Aggarwal et al. , 2019 ) . More recently , self-supervised learning has shown that reconstruction quality similar to supervised PG-DLR can be achieved while training on a database of only undersampled measurements ( Yaman et al. , 2020 ) . While such database learning strategies offer improved reconstruction quality , acquisition of large datasets may often be infeasible . In some MRI applications involving time-varying physiological processes , dynamic information such as time courses of signal changes , contrast uptake or breathing patterns may differ substantially between subjects , making it difficult to generate high-quality databases of sufficient size for the aforementioned strategies . Furthermore , database training , in general , brings along concerns about robustness and generalization ( Eldar et al. , 2017 ; Knoll et al. , 2020c ) . In MRI reconstruction , this may exhibit itself when there are mismatches between training and test datasets in terms of image contrast , sampling pattern , SNR , vendor , and anatomy . While it is imperative to have high-quality reconstructions that can be used to correctly identify lesions/disease for every individual , the fastMRI transfer track challenge shows that pretrained models fail to generalize when applied to patients/scans with different distribution or acquisition parameters , with potential for misdiagnosis ( Muckley et al. , 2021 ) . Finally , training datasets may lack examples of rare and/or subtle pathologies , increasing the risk of generalization failure ( Knoll et al. , 2019 ; 2020c ) . In this work , we tackle these challenges associated with database training , and propose a zero-shot self-supervised learning ( ZS-SSL ) approach , which performs subject-specific training of PG-DLR without any external training database . Succinctly , ZS-SSL partitions the acquired measurements into three types of disjoint sets , which are respectively used only in the PG-DLR neural network , in defining the training loss , and in establishing a stopping strategy to avoid overfitting . Thus , our training is both self-supervised and self-validated . In cases where a database-pretrained network is available , ZS-SSL leverages transfer learning ( TL ) for improved reconstruction quality and reduced computational complexity . Our contributions can be summarized as follows : • We propose a zero-shot self-supervised method for learning subject-specific DL MRI reconstruction from a single undersampled dataset without any external training database . • We provide a well-defined methodology for determining stopping criterion to avoid overfitting in contrast to other single-image training approaches ( Ulyanov et al. , 2018 ) . • We apply the proposed zero-shot learning approach to knee and brain MRI datasets , and show its efficacy in removing residual aliasing and banding artifacts compared to supervised database learning . • We show our ZS-SSL can be combined with with TL in cases when a database-pretrained model is available to reduce computational costs . • We show that our zero-shot learning strategies address robustness and generalizability issues of trained supervised models in terms of changes in sampling pattern , acceleration rate , contrast , SNR , and anatomy at inference time . 2 BACKGROUND AND RELATED WORK . 2.1 ACCELERATED MRI ACQUISITION MODEL . In MRI , raw measurement data is acquired in the frequency domain , also known as k-space . In current clinical MRI systems , multiple receiver coils are used , where each is sensitive to different parts of the volume . In practice , MRI is accelerated by taking fewer measurements , which are characterized by an undersampling mask that specifies the acquired locations in k-space . For a multi-coil MRI acquisition , the forward model is given as yi = PΩFCix + ni , i ∈ { 1 , . . . , nc } , ( 1 ) where x is the underlying image , yi is the acquired data for the ith coil , PΩ is the masking operator for undersampling pattern Ω , F is the Fourier transform , Ci is a diagonal matrix characterizing the ith coil sensitivity , ni is measurement noise for ith coil , and nc is the number of coils ( Pruessmann et al. , 1999 ) . This system can be concatenated across the coil dimension for a compact representation yΩ = EΩx + n , ( 2 ) where yΩ is the acquired undersampled measurements across all coils , EΩ is the forward encoding operator that concatenates PΩFCi across i ∈ { 1 , . . . , nc } . The general inverse problem for accelerated MRI is given as arg min x ‖yΩ −EΩx‖22 +R ( x ) , ( 3 ) where the ‖yΩ−EΩx‖22 term enforces consistency with acquired data ( DC ) andR ( · ) is a regularizer . 2.2 PG-DLR WITH ALGORITHM UNROLLING . Several optimization methods are available for solving the inverse problem in ( 3 ) ( Fessler , 2020 ) . Variable-splitting via quadratic penalty is one such approach that decouples the DC and regularizer units . It introduces an auxiliary variable z that is constrained to be equal to x , and ( 3 ) is reformulated as an unconstrained problem with a quadratic penalty arg min x , z ‖yΩ −EΩx‖22 + µ‖x− z‖22 +R ( z ) , ( 4 ) where µ is the penalty parameter . The optimization problem in ( 4 ) is then solved via z ( i ) = arg min z µ‖x ( i−1 ) − z‖22 +R ( z ) , ( 5a ) x ( i ) = arg min x ‖yΩ −EΩx‖22 + µ‖x− z ( i ) ‖22 , ( 5b ) where z ( i ) is an intermediate variable and x ( i ) is the desired image at iteration i . In PG-DLR , an iterative algorithm , as in ( 5a ) and ( 5b ) is unrolled for a fixed number of iterations ( Liang et al. , 2020 ) . The regularizer sub-problem in Eq . ( 5a ) is implicitly solved with neural networks and the DC sub-problem in Eq . ( 5b ) is solved via linear methods such as gradient descent ( Hammernik et al. , 2018 ) or conjugate gradient ( CG ) ( Aggarwal et al. , 2019 ) . There have been numerous works on PG-DLR for accelerated MRI ( Schlemper et al. , 2018 ; Hammernik et al. , 2018 ; Aggarwal et al. , 2019 ; Liang et al. , 2020 ; Yaman et al. , 2020 ) . Most of these works vary from each other on the algorithms used for DC and neural networks employed in the regularizer units . However , all these works require a large database of training samples . 2.3 SUPERVISED LEARNING FOR PG-DLR . In supervised PG-DLR , training is performed using a database of fully-sampled reference data . Let ynref be the fully-sampled k-space for subject n and f ( y n Ω , E n Ω ; θ ) be the output of the unrolled network for under-sampled k-space ynΩ , where the network is parameterized by θ. End-to-end training minimizes ( Knoll et al. , 2020b ; Yaman et al. , 2020 ) min θ 1 N N∑ n=1 L ( ynref , Enfullf ( ynΩ , EnΩ ; θ ) ) , ( 6 ) where N is the number of samples in the training database , Enfull is the fully-sampled encoding operator that transform network output to k-space and L ( · , · ) is a loss function . 2.4 SELF-SUPERVISED LEARNING FOR PG-DLR . Unlike supervised learning , self-supervised learning enables training without fully-sampled data by only utilizing acquired undersampled measurements ( Yaman et al. , 2020 ) . A masking approach is used for self-supervision in this setting , where a subset Λ ⊂ Ω is set aside for checking prediction performance/loss calculation , while the remainder of points Θ = Ω\Λ are used in the DC units of the PG-DLR network . End-to-end training is performed using the loss function min θ 1 N N∑ n=1 L ( ynΛ , E n Λ ( f ( ynΘ , E n Θ ; θ ) ) ) . ( 7 ) 3 ZERO-SHOT SELF-SUPERVISED LEARNING FOR PG-DLR . As discussed in Section 1 , lack of large datasets in numerous MRI applications , as well as robustness and generalizability issues of pretrained models pose a challenge for the clinical translation of DL reconstruction methods . Hence , subject-specific reconstruction is desirable in clinical practice , since it is critical to achieve a reconstruction quality that can be used for correctly diagnosing every patient . While the conventional self-supervised masking strategy , as in ( Yaman et al. , 2020 ) can be applied for subject-specific learning , it leads to overfitting unless the training is stopped early ( Hosseini et al. , 2020a ) . This is similar to other single-image learning strategies , such as the deep image prior ( DIP ) or zero-shot super-resolution ( Ulyanov et al. , 2018 ; Shocher et al. , 2018 ) . DIP-type approaches shows that an untrained neural network can successfully perform instance-specific image restoration tasks such as denoising , super-resolution , inpainting without any training data . However , such DIPtype techniques requires an early stopping for avoiding over-fitting , which is typically done with a manual heuristic selection ( Ulyanov et al. , 2018 ; Hosseini et al. , 2020a ; Darestani et al. , 2021 ) . While this may work in a research setting , having a well-defined automated early stopping criterion is critical to fully harness the potential of subject-specific DL MRI reconstruction in practice . Early stopping regularization in database-trained setting is conventionally motivated through the bias-variance trade-off , in which a validation set is used as a proxy for the generalization error to identify the stopping criterion . Using the same bias-variance trade-off motivation , having a validation set can aid in devising a stopping criterion , but this has not been feasible in existing zero-shot learning approaches , which either use all acquired measurements ( Ulyanov et al. , 2018 ; Senouf et al. , 2019 ) or partition them into two sets for training and defining loss ( Hosseini et al. , 2020a ) . Hence , existing zero-shot learning techniques lack a validation set to identify the stopping criterion . ZS-SSL Formulation and Training : We propose a new ZS-SSL partitioning framework to enable subject-specific self-supervised training and validation with a well-defined stopping criterion . We define the following partition for the available measurement locations from a single scan , Ω : Ω = Θ t Λ t Γ , ( 8 ) where t denotes a disjoint union , i.e . Θ , Λ and Γ are pairwise disjoint ( Figure 1 ) . Similar to Section 2.4 , Θ is used in the DC units of the unrolled network , and Λ is used to define a k-space loss for the self-supervision of the network . The third partition Γ is a set of acquired k-space indices set aside for defining a k-space validation loss . Thus , ZS-SSL training is both self-supervised and self-validated . In general , since zero-shot learning approaches perform training using a single dataset , generation of multiple data pairs from this single dataset is necessary to self-supervise the neural network ( Quan et al. , 2020 ) . Hence , we generate multiple ( Θ , Λ ) pairs from the acquired locations Ω of the single scan . In ZS-SSL , this is achieved by fixing the k-space validation partition Γ ⊂ Ω , and performing the retrospective masking on Ω\Γ multiple times . Formally , Ω\Γ is partitioned K times such that Ω\Γ = Θk t Λk , k ∈ { 1 , . . . , K } , ( 9 ) where Λk , Θk and Γ are pairwise disjoint , i.e . Ω = Γ tΘk t Λk , ∀ k. ZS-SSL training minimizes min θ 1 K K∑ k=1 L ( yΛk , EΛk ( f ( yΘk , EΘk ; θ ) ) ) In the proposed ZS-SSL , this is supplemented by a k-space self-validation loss , which tests the generalization performance of the trained network on the k-space validation partition Γ . For the lth epoch , where the learned network weights are specified by θ ( l ) , this validation loss is given by : L ( yΓ , EΓ ( f ( yΩ\Γ , EΩ\Γ ; θ ( l ) ) ) ) . ( 10 ) Note that in ( 10 ) , the network output is calculated by applying the DC units on Ω\Γ = Θ t Λ , i.e . all acquired points outside of Γ , to better assess its generalizability performance . Our key motivation is that while the training loss will decrease over epochs , the k-space validation loss will start increasing once overfitting is observed . Thus , we monitor the loss in ( 10 ) during training to define an early stopping criterion to avoid overfitting . Let L be the epoch in which training needs to be stopped . Then at inference time , the network output is calculated as f ( yΩ , EΩ ; θ ( L ) ) , i.e . all acquired points are used to calculate the network output . ZS-SSL with Transfer Learning ( TL ) : While pretrained models are very efficient in reconstructing new unseen measurements from similar MRI scan protocols , their performance degrades significantly when acquisition parameters vary ( Muckley et al. , 2021 ) . Moreover , retraining a new model on a large database for each acquisition parameter , sampling/contrast/anatomy/acceleration , may be very computationally expensive ( Knoll et al. , 2019 ) . Hence , TL has been used for re-training DL models pre-trained on large databases to reconstruct MRI data with different characteristics ( Knoll et al. , 2019 ) . However , such transfer still requires another , often smaller , database for re-training . In contrast , in the presence of pre-trained models , ZS-SSL can be combined with TL , referred to as ZSSSL-TL , to reconstruct a single slice/instance with different characteristics by using weights of the pre-trained model for initialization . Thus , ZS-SSL-TL ensures that the pretrained model is adapted for each patient/subject , while facilitating faster convergence time and reduced reconstruction time .

Summary: This article propose a "zero-shot" method for MRI reconstruction, which is a well-studied inverse problem. The method is based on the ideas of deep image prior, i.e. the ability of correctly sized neural networks to learn about the structure of a single image, sufficiently well for denoising tasks. This self-supervised learned network can then be used in a plug-and-play architecture to solve inverse problem with a variational approach, i.e. as if the learned denoiser were a Total Variation (TV) minimiser. Their denoiser can be improved in a transfer learning approach to benefit from a more complex network trained on similar images than those at hand, and fined-tuned on the image to be reconstructed  The authors go on to apply their plug-and-play architecture to solve the MRI reconstruction problem. They provide experiments and comparisons.