We assume an unobserved (vectorized) image $\bm{x}\in\mathbb{C}^{N}$ is transformed by a forward model $A$ into a measurement vector $\bm{y}$:

Here, $N$ is the number of pixels of the full-resolution grid and $\epsilon$ encapsulates unmodeled effects such as noise. In image denoising, $A=I$. In image superresolution, $A=D$, where $D$ is a down-sampling operator. In single-coil CS-MRI, $A=\mathcal{F}_{u}$, where $\mathcal{F}_{u}\in\mathbb{C}^{M\times N}$ is the under-sampled discrete Fourier operator and $M<N$.

Classically, inverse problems for image reconstruction are reduced to iteratively minimizing a regularized regression cost function on each collected measurement set (Beck and Teboulle, 2009; Boyd et al., 2011; Chambolle and Pock, 2011; Combettes and Pesquet, 2009; Daubechies et al., 2003; Ye et al., 2019).

The first term, called the data-consistency loss, quantifies the agreement between the measurement vector $\bm{y}$ and reconstruction. One common choice for $J$ is:

The $K-1$ regularization terms $\mathcal{R}_{i}(\cdot)$ are hand-crafted priors which restrict the space of permissible solutions. A set of hyperparameters $\{\lambda_{1},...,\lambda_{K-1}\}$ weights the competing contributions of the $K$ terms. Common choices for these priors include sparsity-inducing norms of wavelet coefficients (Figueiredo and Nowak, 2003) and total variation (TV) (Liu et al., 2018; Hu and Jacob, 2012; Rudin et al., 1992). While considerable research has been dedicated to designing suitable priors, such methods are invariably limited by lack of flexibility and adaptivity to the data. In addition, solving the minimization involves expensive iterative procedures that can take minutes per instance.

Recent advances in data-driven algorithms and deep learning seek to address both aforementioned shortcomings of instance-based algorithms. Using large datasets and highly-parametrized and nonlinear neural network models, these algorithms learn priors directly from data and enable efficient inference via nearly-instantaneous forward passes.

Given a training dataset $\mathcal{D}_{tr}$ and a neural network model $M_{\theta}$ with parameters $\theta$, the objective to minimize is:

where $\mathcal{L}$ is a loss function parameterized by $\theta$, $\mathbb{E}_{\mathcal{D}_{tr}}$ denotes the empirical average, and $\bm{\lambda}$ encapsulates hyperparameters of interest.

In supervised learning (SL), $\mathcal{D}_{tr}$ is assumed to contain pairs of observations and corresponding ground-truth images $(\bm{y},\bm{x})\sim\mathcal{D}_{tr}$, and $\mathcal{L}$ is some combination of $K$ supervised losses $L(\cdot,\cdot)$:

Common choices for $L(\cdot,\cdot)$ include mean-squared error (MSE), mean-absolute error (MAE), and structural similarity index measure (SSIM).

In addition to the SL set-up, we can consider the amortized optimization (AO) scenario, where $\mathcal{D}_{tr}$ is assumed to only contain observations $\bm{y}\sim\mathcal{D}_{tr}$. In AO, $\mathcal{L}$ takes a similar form to Eq. (2):

Thus, $M_{\theta}$ is trained to minimize Eq. (2) for all observations in the dataset. Although it minimizes the same cost function, AO 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 model $M_{\theta}$ estimates the reconstruction for any viable input measurement vector $\bm{y}$ and not just a single instance, AO acts as a natural regularizer for the optimization problem (Balakrishnan et al., 2019; Shu et al., 2018; Wang et al., 2020). Third, recent work has shown that inductive biases of convolutional neural networks provide favorable implicit priors for image reconstruction (Heckel and Hand, 2019; Liu et al., 2018; Ulyanov et al., 2020). Since no ground-truth images are required, AO is sometimes referred to as unsupervised learning.

Importantly, similar to the classical cost function, both the SL and AO class of loss functions involve the tuning of a set of hyperparameters $\{\lambda_{1},...,\lambda_{K-1}\}$, which can significantly affect the end reconstructions. Therefore, hyperparameter tuning is typically carefully done using expensive methods like cross-validation. In this work, we propose a strategy that can replace the expensive hyperparameter tuning step, such that reconstructions with different $\bm{\lambda}$ values can be efficiently generated at inference time. To achieve this, we employ hypernetworks.

A hypernetwork is a neural network that generates the weights of a network which solves the main task (e.g. classification, segmentation, reconstruction). In this framework, the parameters of the hypernetwork, and not the main network, are learned. While originally introduced for achieving weight-sharing and model compression (Ha et al., 2016), this idea has found numerous applications including neural architecture search (Brock et al., 2018; Zhang et al., 2019), Bayesian neural networks (Krueger et al., 2018; Ukai et al., 2018), multi-task learning (Lin et al., 2021; Mahabadi et al., 2021; Pan et al., 2018; Shen et al., 2017; Klocek et al., 2019), and hyperparameter optimization (Lorraine and Duvenaud, 2018; Hoopes et al., 2021).

Notably, hypernetworks have been used to replace expensive hyperparameter tuning procedures like cross-validation (Lorraine and Duvenaud, 2018). Hyperparameter tuning can be formulated as a nested optimization. The inner optimization minimizes some loss with respect to the weights of the main network $\theta\in\Theta$ over a training dataset $\mathcal{D}_{tr}$, while the outer optimization minimizes that loss with respect to the hyperparameters $\bm{\lambda}\in\Lambda$ over a held-out validation dataset $\mathcal{D}_{val}$:

Hypernetworks can be trained to approximate the solution to the inner optimization $\theta^{*}(\bm{\lambda})$. In this setting, a hypernetwork maps from hyperparameters to main network weights $H_{\phi}:\Lambda\rightarrow\Theta$. The model is trained via a stochastic optimization scheme whereby hyperparameters are sampled from a pre-defined distribution over $\Lambda$ during training:

The approximate solution to the inner optimization is then $\theta^{*}(\bm{\lambda})=H_{\phi^{*}}(\bm{\lambda})$. Note that only the parameters of the hypernetwork $\phi$ are learned.

The upshot of this method is that rapid hyperparameter tuning is possible, since the inner optimization can be performed with a forward pass of the hypernetwork. While using this model enables us to address the issue of costly hyperparameter tuning in the context of image reconstruction (as discussed in the previous section), we further build on this idea in this work by leveraging the vast set of reconstructions that are capable of being generated for a given measurement as a result of changing $\bm{\lambda}$. That is, we propose this method as a means of creating an interactive and controllable image reconstruction tool.

Let $M_{\theta}$ denote a main network with parameters $\theta\in\Theta$ which maps an observation $\bm{y}$ to a reconstruction $\hat{\bm{x}}$. We define a hypernetwork $H_{\phi}:\mathbb{R}_{+}^{K-1}\rightarrow\Theta$ that maps a hyperparameter weight vector $\bm{\lambda}$ to the parameters $\theta$ of the main network $M_{\theta}$. A reconstruction for a given observation $\bm{y}$ and hyperparameter vector $\bm{\lambda}$ is then $\hat{\bm{x}}=M_{H_{\phi}(\bm{\lambda})}(\bm{y})$. Effectively, this makes $\bm{\lambda}$ an input to the model.

The objective is given in Eq. (8), where $\Lambda=\mathbb{R}_{+}^{K-1}$.

We restrict our attention to the case of two and three loss terms (one and two hyperparameters, respectively), although we emphasize that our method is applicable to an arbitrary number of loss terms. In general, Eqs. (5) and (6) can be manipulated such that the hyperparameter support is bounded to $\bm{\lambda}\in\Lambda=[0,1]^{K-1}$. For example, for one and two hyperparameter weights and arbitrary loss terms $L_{1},L_{2}$, and $L_{3}$ (omitting $\theta$, $\bm{x}$, and $\bm{y}$ for brevity):

A straightforward strategy for training the hypernetwork involves sampling the coefficients from a uniform distribution $p(\bm{\lambda})=U[0,1]^{K-1}$ and sampling an under-sampled measurement vector $\bm{y}$ for each forward pass during training. The gradients are then computed with respect to the loss evaluated at the sampled $\bm{\lambda}$ via a backward pass. This corresponds to minimizing Eq. (8) with a uniform distribution for $\bm{\lambda}$. We denote this sampling strategy as uniform hyperparameter sampling (UHS).

In the hypothetical scenario of infinite hypernetwork capacity, the hypernetwork can capture a mapping of any input $\bm{\lambda}$ to the optimal $\theta^{*}=H_{\phi^{*}}(\bm{\lambda})$ that minimizes Eqs. (5) and (6) with the corresponding $\bm{\lambda}$. However, in practice, the finite model capacity of the hypernetwork constrains the ability to achieve optimal loss for every hyperparameter value. In this case, the model performance will depend on the adopted distribution for $\bm{\lambda}$.

In SL loss functions, we are typically interested in the reconstructions associated with all possible settings of the coefficients. Thus, UHS is a necessary sampling strategy. In contrast, in AO loss functions, some coefficients will not produce acceptable reconstructions, even if solved optimally. For example, a reconstruction with a large coefficient for TV tends to be overly smooth with little to no structural content. Thus, sampling hyperparameters from the entire support $[0,1]^{K-1}$ “wastes” model capacity on undesirable regions of hyperparameter space.

In most real-world reconstruction scenarios, we don’t have a good prior distribution from which to sample desirable regions. Instead, we propose a data-driven sampling scheme (DHS) which learns the prior over the course of training. We leverage the data-consistency loss induced by a setting of the coefficients to assess whether the reconstruction will be useful or not. Intuitively, values of $\bm{\lambda}$ which lead to high data-consistency loss $J\left(M_{H_{\phi}(\bm{\lambda})}(\bm{y}),\bm{y}\right)$ will produce reconstructions that deviate too much from the underlying anatomy, and which therefore can be ignored during training.

One can enforce this idea by computing gradients only on values of $\bm{\lambda}$ which induce a data-consistency loss below a pre-determined threshold during training. 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, this threshold will likely not be satisfied.

In lieu of a complex training/threshold scheduler, we adapt the threshold with the quality of reconstructions directly within the training loop. The proposed DHS strategy works by using the best $b$ samples with the lowest data-consistency loss within a mini-batch of size $B$ to compute the gradients. In effect, this induces a variable threshold $b/B$ of the landscape percentage which is optimized, where this threshold adapts dynamically over training. Algorithm 1 details the training loop for both UHS and DHS.

We demonstrate in the Experiments section that DHS leads to improved model performance on the most promising hyperparameter space regions compared to UHS, given a fixed hypernetwork capacity.

In general, a given loss function $L_{i}$ can be arbitrarily scaled by a constant $\alpha_{i}L_{i}$ and still yield the same reconstruction results (ignoring the influence on learning rate, for example). In multi-term loss functions, the relative scales of the $K$ loss outputs should be approximately matched during training to allow for equal contribution of gradients during backpropagation.

In the context of our proposed method, these loss output scales have an added effect: different scales will change the resulting reconstructions as a function of $\bm{\lambda}$. If one loss dominates another, then most of the landscape will be dedicated to minimizing the dominant loss and very little variation will exist for different $\bm{\lambda}$. Indeed, we ideally would like a landscape which varies maximally over $\bm{\lambda}$, which would (hopefully) capture as wide a range of diverse reconstructions as possible.

For SL, we account for this by normalizing the loss functions by its “best-case” loss averaged over a validation set. Let $s_{i}$ denote the average validation loss on $L_{i}$ for a model only optimized for this specific loss function, and we set $\alpha_{i}=1/s_{i}$. During training, validation losses associated with all loss terms will (roughly) converge to a value of $1$, and thus the loss output scales will be approximately matched. Further details are given in Section 5.4.

For AO, we cannot take the same approach since the “best-case” losses for $J(\cdot,\cdot)$ and $\mathcal{R}(\cdot)$ are $0$. When Eq. (2) is derived from the maximum a posteriori (MAP) estimate, the scaling factors can be computed with respect to the parameters of the likelihood and/or prior distributions (Chambolle et al., 2010; Hoopes et al., 2021). However, in cases where the regularization term cannot be expressed as a valid probability or when multiple terms are used, such an approximation cannot be obtained. In this work, we tuned the scaling factors to align the magnitudes of the separate terms. Although this works well in practice, this issue is still an open question in the general case and may be an important direction for further research.

We conducted our experiments using two large, publicly-available datasets. The first is the NYU fastMRI dataset (Zbontar et al., 2018) composed of proton density (PD) and proton density fat-suppressed (PDFS) weighted knee MRI scans. The second is the ABIDE dataset, composed of T1-weighted brain MRI scans (Di Martino et al., 2014).

For both datasets, we used 100, 25, and 50 subjects for training, validation, and testing, respectively. Volumes were separated into 2D slices and slices with only background were removed, resulting in a final train/val/test split of 3500/875/1750 for knee and 11400/2850/5700 for brain slices. All images were intensity-normalized to the range $[0,1]$ and cropped and re-sampled to a pixel grid of size $256\times 256$.

This work presents a general strategy for computing hyperparameter-agnostic reconstructions with a single model, which is applicable to any main network architecture. In this work, the main network $M_{\theta}$ is a commonly-used Unet architecture (Ronneberger et al., 2015) adapted for reconstruction, where the network input is a 2-channel, complex-valued input and the network output is a single-channel, real-valued output. A hidden channel dimension of $32$ is used for all layers, resulting in a total main network parameter count of $149,409$.

Since $M_{\theta}$ receives noisy images as input, this network can be viewed as a post-processing network preceded by an optional model-based inversion step (e.g. applying the inverse Fourier transform to under-sampled measurements), which is a common reconstruction pipeline in the literature (Jin et al., 2017; Kang et al., 2017; Wang et al., 2021b). We expect the conclusions drawn in this work to hold for any deep learning-based main network, particularly state-of-the-art unrolled networks (Monga et al., 2019).

The hypernetwork $H_{\phi}$ consists of a 5-layer fully-connected (FC) architecture with intermediate LeakyReLUs. The input to the network is the number of hyperparameters, and the hidden dimension and output dimension are of size $d$. Each convolution layer in the main network takes as input the $d$-dimensional output embedding of the hypernetwork and linearly projects it to the size of the kernel and bias, where the parameters of the projection are learned. We treat the hypernetwork and the main network as a single, large network, and refer to the overall model as HyperRecon. We experiment with hypernetwork hidden dimensions $d\in\{4,32,128\}$, and refer to them as HyperRecon-{S, M, L}, respectively. We used a batch size of 32 and the Adam optimizer for all models (Kingma and Ba, 2017).

All training and testing experiments were performed on a machine equipped with an Intel Xeon Gold 6126 processor and an NVIDIA Titan Xp GPU. All models were implemented in Pytorch. Table 1 outlines the training time, inference time, and number of parameters for all models used in experiments. We train each model until the loss converges on the validation set, typically for 1000 epochs.

For comparison, we trained separate Unet reconstruction networks for each fixed hyperparameter value. We refer to these models as baselines and emphasize that they demand significant computational resources, since each of these models needs to be trained and saved separately (see Table 1). For SL experiments, we trained five baseline models for $\lambda\in\{0.0,0.25,0.5,0.75,1.0\}$. For AO experiments, we trained $361$ baseline models with hyperparameters chosen non-uniformly on a $19\times 19$ grid over the space $[0,1]\times[0,1]$, in order to more densely sample in high-performing regions²²2 The 19 values were {0.0,0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.85,0.9,0.93,0.95,0.98, 0.99,0.995,0.999,1.0}..

For all tasks, noisy inputs to our model were obtained retrospectively and performed in-silico. Although this is a slight simplification, we believe that these experiments are consistent with the main methodological message of this work, and furthermore that these results will translate to real-world scenarios.

We evaluate the performance of HyperRecon-L using metric curves over the space of permissible hyperparameter values $\lambda\in[0,1]$. We generated the curves for visualization by densely sampling the support $[0,1]$ to create $100$ discrete samples. For each grid point, we computed the value by passing the corresponding hyperparameter values to the model along with each under-sampled measurement $\bm{y}$ in the test set and taking the average PSNR value.

Fig. 3 shows metric performance for brain data. Each panel shows Unet and HyperRecon-L performance for the specified metric across $\lambda$ values. Similarly, Fig. 4 shows metric performance for knee data on three different reconstruction tasks. We observe very close matching of performance across all metrics between Unet and HyperRecon-L models, indicating that the hypernetwork is able to generate the optimal weights for the entire hyperparameter space. For downstream brain segmentation, the optimal $\lambda$ for maximizing Dice is around $0.25$. With HyperRecon-L, this optimal point can be easily obtained with a single model.

Fig. 5(a) shows representative brain and knee slices for the CS-MRI task, and Fig. 5(b) shows zoomed-in versions of the red arrow regions for the brain slices. We notice a significant visual similarity between the same hyperparameters for Unet and HyperRecon-L models. In general, a higher weight on MAE tends to lead to more blurry and smoother reconstructions, while a higher weight on SSIM tends to lead to more noisy reconstructions with more high frequency content. In addition, MAE tends toward less contrast between hypo-intense and hyper-intense regions, whereas SSIM accentuates this contrast more. With baseline models where $\bm{\lambda}$ must be chosen before training, these variations would be completely missed and end users would be stuck with one reconstruction. Additional slices for CS-MRI, denoising, and superresolution tasks are presented in the Appendix.

We evaluate the performance of HyperRecon-{S, M, L} using rPSNR landscapes over the space of permissible hyperparameter values $(\lambda_{1},\lambda_{2})\in[0,1]\times[0,1]$. We generated landscapes for visualization by densely sampling the support $[0,1]\times[0,1]$ to create a grid of size $100\times 100$. For baselines, the $19\times 19$ grid was linearly interpolated to $100\times 100$ to match the hypernetwork landscapes.

Fig. 6 shows rPSNR landscapes for different reconstruction models. The top left map corresponds to baselines. 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 remaining images in the top row correspond to UHS models with varying hypernetwork capacities. Similarly, the bottom row shows rPSNR landscapes for DHS models at the same capacities. We find 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.

Fig. 7 shows representative brain and knee slices from the DHS HyperRecon-L model. The two corresponding reconstructions are selected as follows. First, we densely sample uniformly from $[0,1]\times[0,1]$ to generate $100\times 100$ reconstructions. Then, we filter out the reconstructions which are below some threshold PSNR value (we choose the 90th percentile). Finally, we choose the two reconstructions from this filtered set which are maximally separated by $\ell_{2}$ distance.

We notice that the two reconstructions are significantly dissimilar despite the similarity in PSNR value. In the baseline setting, only one of these reconstructions would be available and finding another reconstruction would require training from scratch. With our model, users can search over the entire space of possible reconstructions and choose the one(s) they prefer. This highlights the value of this model as a tool for the interactive selection of many diverse reconstructions for further use based on visual inspection.