Computer Science Engineering
Computer Science-2021.gz:143364
Computer Science
ba3e2aa923f5eca5602aa9168062b646382d4498
Regularization-Agnostic Compressed Sensing MRI Reconstruction with Hypernetworks

Reconstructing under-sampled k-space measurements in Compressed Sensing MRI (CS-MRI) is classically solved with regularized least-squares. Recently, deep learning has been used to amortize this optimization by training reconstruction networks on a dataset of under-sampled measurements. Here, a crucial design choice is the regularization function(s) and corresponding weight(s). In this paper, we explore a novel strategy of using a hypernetwork to generate the parameters of a separate reconstruction network as a function of the regularization weight(s), resulting in a regularization-agnostic reconstruction model. At test time, for a given under-sampled image, our model can rapidly compute reconstructions with different amounts of regularization. We analyze the variability of these reconstructions, especially in situations when the overall quality is similar. Finally, we propose and empirically demonstrate an efficient and data-driven way of maximizing reconstruction performance given limited hypernetwork capacity. Our code is publicly available at https://github.com/alanqrwang/RegAgnosticCSMRI.


Introduction
Magnetic resonance imaging (MRI) can be accelerated by under-sampling k-space -a technique known as Compressed Sensing MRI (CS-MRI) [21]. This yields a well-studied ill-posed inverse problem. Classically, this problem is reduced to iteratively minimizing a regularized regression cost function on each collected measurement set [2,3,5,8,11,31]. In contrast, learning-based amortized optimization trains a neural network to minimize the cost function over a dataset of under-sampled measurements [28,30]. This method has been shown to significantly speed up computation time while outperforming traditional instance-based methods.
However, both the classical and amortized methods suffer from a strong dependence on the regularization function, the weighting of which can vastly affect reconstructions. Thus, this hyperparameter must be carefully tuned to achieve optimal reconstruction quality. As a result, considerable time and resources can be spent setting the weight(s) of regularization, using methods such as grid-search, random search, or Bayesian optimization.
An emerging solution to the problem of hyperparameter tuning involves the use of hypernetworks which enable the user to jointly learn the model weights and optimal hyperparameters concurrently. Such methods have been applied broadly in learning models [4,20] and recently in medical image registration [16]. In this paper, we build on this line of thinking, applying it to the unsupervised reconstruction problem in CS-MRI and proposing a regularization-agnostic strategy for performing amortized optimization. Specifically, our model uses a hypernetwork that takes as input a value (or values) for the regularization weight(s) and outputs the parameters of a reconstruction network. At test-time, arbitrary weight values can be provided and corresponding reconstructions will be efficiently computed. Thus, our method is capable of producing a range of reconstructions over the entire landscape of hyperparameter values, instead of a single reconstruction associated with a single hyperparameter value setting. We envision a workflow in which a user will select from many diverse reconstructions for further use based on visual inspection.
Additionally, we investigate different methods of sampling from the hyperparameter space during training. In particular, we examine two sampling procedures: uniform hyperparameter sampling (UHS) and a novel data-driven non-uniform hyperparameter sampling (DHS) strategy. We visualize and quantify the impact of sampling on the quality and range of reconstructions, using empirical results to demonstrate significant improvement when data-driven hyperparameter sampling is used during training.

Amortized Optimization of CS-MRI
In the classical CS-MRI formulation, a reconstruction of the unobserved fully-sampled MRI scan 4 x ∈ C N is obtained by solving an ill-posed inverse problem via an optimization of the form arg min for each under-sampled measurement y ∈ C M , where N is the number of pixels of the full-resolution grid and M < N is the number of measurements. The first term, called the data-consistency loss, quantifies the agreement between the measurement vector y and reconstruction and can be defined as: where F u denotes the under-sampled Fourier operator. The second term is composed of a sum of p regularization functions, R i (x), each weighted by a hyperparameter α i ∈ R + , that determines the degree of trade-off between the p + 1 competing terms. The regularization term is often carefully engineered to restrict the solutions to the space of desirable images. Common choices include sparsityinducing norms of wavelet coefficients [12], total variation (TV) [17,26], and their combinations [21,24]. A recently popular approach involves solving the instance-based optimization of Eq. (1) with a neural network [30]. For a neural network G θ parameterized with θ and a training set D of undersampled measurements y, the problem reduces to: This formulation can be viewed as an amortization of the instance-specific optimization of Eq. (1), via a neural network G θ [1,9,14,22,28].
Amortized optimization provides several advantages over classical solutions. First, at test-time, it replaces an expensive iterative optimization procedure with a simple forward pass of a neural network. Second, since the function G θ is tasked with estimating the reconstruction for any viable input measurement vector y and not just a single instance, amortized optimization has been shown to act as a natural regularizer for the optimization problem [1,28].

Hypernetworks
Originally introduced for achieving weight-sharing and model compression [15], hypernetworks take as input a set of hyperparameters that are in turn converted to the weights of the network that solves the main task, such as classification. In this framework, the parameters of the hypernetwork, and not the main-task network, are learned. This idea has found numerous applications including neural architecture search [4,32], Bayesian neural networks [19,29], multi-task learning [18,23,27], and hyperparameter optimization [20]. In an application related to ours, for example, hypernetworks have been used as a means of replacing cross-validation for finding the optimal level of weight decay on MNIST classification [20]. The authors propose a training scheme which allows the hypernetwork to output the optimal weights of a main network as a function of the weight decay parameter. Similarly, in parallel work hypernetworks were used to build a hyperparameter-invariant medical image registration framework [16].

Proposed Method
Inspired by the prior applications of hypernetworks, we propose to use them for the CS-MRI reconstruction problem, where our goal is to produce a model that can efficiently compute a reconstruction that approximates the solution to Eq. (1) for arbitrary regularization weight values. We call this model regularization-agnostic, and illustrate it in Fig. 1.

Regularization-Agnostic Reconstruction Network
Suppose all regularization weights are gathered in a p-vector As before, let G θ denote a main network which maps under-sampled measurements y to reconstructionsx, whose parameters are denoted as θ ∈ Θ. We define a hypernetwork H φ : R p + → Θ that maps a regularization weight vector α to the parameters θ of the main network G θ . During training we only learn φ, the parameters of the hypernetwork.
Given a dataset D of under-sampled measurements y, the objective is to minimize the following loss with respect to φ: In the hypothetical scenario of infinite hypernetwork capacity, the hypernetwork can capture a mapping of any input α to the optimal θ * = H φ * (α) that minimizes Eq. (3) with the corresponding α. The training of the hypernetwork is not sample limited, since one can draw as many hyperparameter samples as needed. Thus, overfitting for the hypernetwork is not a practical concern. However, in practice, we will have a limited capacity hypernetwork and its learned parameters will in general depend on the assumed distribution for α. In a sense, while G θ amortizes over the space of images, H φ can be thought of as amortizing over the space of regularization weights.
In general, the expression inside the braces of Eq. (4) can be manipulated such that the hyperparameter support is bounded to α ∈ [0, 1] p . For example, for one and two regularization weights, we can use: wherex = G H φ (α) (y) denotes the model reconstruction.

Training
Uniform Hyperparameter Sampling (UHS): A straightforward strategy for training the hypernetwork involves sampling the regularization weights from a uniform distribution p(α) = U [0, 1] p and an under-sampled measurement vector y for each forward pass during training. The gradients are then computed with respect to the loss evaluated at the sampled α via a backward pass. This corresponds to minimizing Eq. (4) with a uniform distribution for α. Algorithm 1 details the pseudo-code for the resulting training procedure. However, the finite model capacity of the hypernetwork constrains the ability to achieve optimal loss for every hyperparameter value, particularly without resorting to training large and computationally expensive hypernetworks. In addition, sampling hyperparameters from the entire support [0, 1] p might "waste" model capacity on regularization weight values that have no hope of producing acceptable reconstructions, even if solved optimally. These two observations suggest that modifying the hyperparameter sampling distribution could lead to better performance by making better use of limited hypernetwork model capacity.
Unfortunately, in most real-world reconstruction scenarios, we don't have a good prior distribution model for sampling the regularization weights. In the remainder of this section, we propose an unsupervised, data-driven sampling scheme which learns the optimal sampling region during training, and show in the Experiments section how this significantly improves reconstruction performance.

Data-driven Hyperparameter Sampling (DHS):
In the unsupervised setting, we propose to use the data-consistency loss induced by a setting of the regularization weights to assess whether the reconstruction will be useful or not. Intuitively, values of α which lead to high data-consistency loss J(G H φ (α) , y) will produce reconstructions that deviate too much from the underlying anatomy, and which therefore can be ignored during training.
We propose the following strategy to exploit this idea. During training, one can compute gradients only on values of α which induce a data-consistency loss below a pre-specified threshold. However, this presents two problems. First, calibrating and tuning this threshold can be difficult in practice. Second, at the beginning of training, since the main network will not produce good reconstructions,
Algorithm 2 outlines our proposed DHS strategy, which uses the best K samples with the lowest data-consistency loss within a mini-batch of size B to calculate the gradients. In effect, this dedicates model capacity for the subset of the hyperparameter landscape that can produce reconstructions that are most consistent with the data, while ignoring those that do not. The percentage of the landscape which is optimized is K/B.

Experiments
In all experiments, we restrict our attention to the case of p = 2 regularization functions, although we emphasize that our method works for any number of regularization loss terms. We choose layerwise total ℓ 1 -penalty on the weights θ of the main reconstruction network and the total variation of the reconstruction image as two regularization loss terms: whereθ = H φ (α),θ i denotes the weights of the reconstruction network for the ith layer, L is the total number of layers of the reconstruction network, and ∇ is the spatial gradient operator.
Implementation The reconstruction network consists of a residual U-Net architecture [25] with 64 hidden channels per encoder layer, yielding a total of n = 592, 002 (non-trainable) parameters. The hypernetwork consists of fully-connected layers with leaky ReLU activations and batch normalizations between intermediate layers. Hypernetwork weights were initialized to approximate standard Kaiming initialization in the main network [6]. We experiment with three hypernetwork architectures: "small" consisting of 1-2-4-n fully-connected layers, "medium" consisting of 1-8-32-n fully-connected layers, and "large" consisting of 1-8-32-32-32-n fully-connected layers. For all models and baselines, we set the mini-batch size to B = 32 and used ADAM optimizer with learning rate 1e − 5. For DHS, we took the top 25% of samples within a mini-batch, i.e. K = 8, chosen using a grid-search, although we find that results are not sensitive to this optimal value. All training and testing experiments in this paper were performed on a machine equipped with an Intel Xeon Gold 6126 processor and an NVIDIA Titan Xp GPU, and all models were implemented in Pytorch. Data. We trained and evaluated networks on a dataset of T1-weighted axial brain images [10]. All images were intensity-normalized to the range [0, 1] and cropped and re-sampled to a pixel grid of size 256 × 256. Dataset splits were 2000, 500, and 1000 slices for training, validation, and testing. No subjects and scans were overlapping between the splits. To generate the under-sampled measurements, we applied retrospective down-sampling with 4-fold acceleration sub-sampling masks generated using a 2nd-order polynomial Poisson-disk variable-density sampling strategy [7,13,21].
Evaluation Metrics. Reconstructions were evaluated against ground-truth full-resolution images on peak signal-to-noise ratio (PSNR). We constructed the relative PSNR (RPSNR) value for a given reconstruction by subtracting the corresponding PSNR value for the zero-filled reconstruction, computed by applying the inverse Fourier transform directly to the k-space data where missing values are filled with zeros. We refer to these models as benchmarks and emphasize that they demand significant computational cost, since each of these models need to be trained separately (see Table 1). We chose hyperparameter values non-uniformly to more densely sample in high RPSNR regions.

Hypernetwork Capacity and Hyperparameter Sampling
We evaluate the performance of the proposed models using RPSNR landscapes over the space of permissible hyperparameter values (α 1 , α 2 ) ∈ [0, 1] 2 . We generated landscapes for visualization by densely sampling the support [0, 1] 2 to create a grid of size 100×100. For each grid point, we computed the value by passing the corresponding hyperparameter values to the model along with each undersampled measurement y in the test set and taking the average RPSNR value. For baselines, the 18 × 18 grid was linearly interpolated to 100 × 100 to match the hypernetwork landscapes. Fig. 3 shows RPSNR landscapes for different reconstruction models. The top left map corresponds to baselines. The remaining images in the top row correspond to UHS models with varying hypernetwork capacities. The first image in the second row shows an example histogram of the regularization weight samples used for gradient computation over one DHS training epoch. The rest of the second row shows RPSNR landscapes for DHS models at the same capacities. We observe that higher capacity hypernetworks approach the baseline models' performance, at the cost of computational resources and training time (see Table 1). We also observe significant improvement in performance using DHS as compared to UHS, given a fixed hypernetwork capacity.
We find that the performance improvement achieved by DHS is less for the large hypernetwork, validating the expectation that the sampling distribution plays a more important role when the hypernetwork capacity is restricted. In Fig. 4a, we quantify these observations by plotting the percent area (y-axis) that exceed the corresponding RPSNR threshold (x-axis), for all models and baselines. Better-performing models have a larger area under the curve.

Range of Reconstructions
Since our aforementioned workflow involved providing many different reconstructions for a user, we were interested in examining the range of reconstructions that were computed by the hypernetwork models. Fig. 4b shows the reconstruction range of different models. We plot the total pixel-wise range (y-axis) in the set of reconstructions that exceed an RPSNR threshold (x-axis). A larger area under these curves is desirable, particularly for relatively high RPSNR thresholds. We observe that higher capacity models and DHS yield a consistent increase in the range of reconstructions, compared to their lower capacity and UHS counterparts. Fig. 5 shows some example reconstructions computed by the DHS large model. Along each row, we selected two reconstructions for the corresponding ground truth by collecting the subset of images within RPSNR range [4, 4.5] and finding the two which had maximum distance in terms of ℓ 2 . In particular, we observe the varying levels of smoothness and sharp detail for a single slice despite the fact that the overall quality of the reconstructions (as measured by RPSNR) is virtually indistinguishable.

Conclusion
We introduce a novel unsupervised MR reconstruction method that is agnostic to regularizationweight hyperparameters in the amortized optimization formulation. We propose and demonstrate two hyperparameter sampling methods to maximize the performance of our proposed model, and  show that it performs comparably to non-agnostic and computationally-expensive baselines. We furthermore highlight and quantify the range of reconstructions capable of being produced by our models. While our experiments focused on MRI reconstruction, this method can be applied broadly in order to render agnosticism towards any hyperparameter that may arise in many machine learning models. Furthermore, the reconstruction setting can be improved with more expressive networks and better regularization functions, particularly data-driven variants.