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: In order to eliminate the need for large training sets, one can consider a transition from (1) fully-supervised to (2) self-supervised methods, and then from (2) self-supervised methods to (3) single-instance reconstruction methods.   In the context of accelerated MRI reconstruction which is considered in this work, the above categories translate into models that are trained based on (1) having access to a fully-sampled dataset, (2) having access to a dataset of under-sampled measurements, and (3) having access to only one under-sampled measurement. The paper targets (3), that is proposing a zero-shot learning approach for accelerated MRI reconstruction.  The algorithm is based on the idea proposed in paper [1] combined with building a dataset for the given sample in order to eventually perform self-supervised training on the synthesized dataset. In order to prevent overfitting, the authors propose a way to do automatic early stopping.  [1] Yaman, Burhaneddin, et al. "Self-supervised physics-based deep learning MRI reconstruction without fully-sampled data." 2020 IEEE 17th International Symposium on Biomedical Imaging (ISBI). IEEE, 2020.