Optimal Eye Surgeon: Finding Image Priors through Sparse Generators at Initialization

TL;DR

The paper presents Optimal Eye Surgeon (OES), a pruning-at-initialization framework that extracts Sparse-DIP subnetworks from a densely parameterized image generator to serve as underparameterized image priors for unsupervised reconstruction. By learning a Bernoulli mask at random initialization with KL regularization toward a target sparsity and using the Gumbel-Softmax reparameterization, OES yields Sparse-DIP that denoises effectively without overfitting and can be transferred across images and datasets. Empirically, OES outperforms Lottery Ticket-based and other pruning methods for DIP tasks, demonstrates robust denoising and reconstruction across benchmarks, and extends to MRI reconstruction scenarios, highlighting the practical impact of sparse generators as priors in image restoration. The work also clarifies the limitations and trade-offs of pruning strategies, emphasizing the encoder-decoder structure and the importance of the initial masking for forming reliable priors.

Abstract

We introduce Optimal Eye Surgeon (OES), a framework for pruning and training deep image generator networks. Typically, untrained deep convolutional networks, which include image sampling operations, serve as effective image priors (Ulyanov et al., 2018). However, they tend to overfit to noise in image restoration tasks due to being overparameterized. OES addresses this by adaptively pruning networks at random initialization to a level of underparameterization. This process effectively captures low-frequency image components even without training, by just masking. When trained to fit noisy images, these pruned subnetworks, which we term Sparse-DIP, resist overfitting to noise. This benefit arises from underparameterization and the regularization effect of masking, constraining them in the manifold of image priors. We demonstrate that subnetworks pruned through OES surpass other leading pruning methods, such as the Lottery Ticket Hypothesis, which is known to be suboptimal for image recovery tasks (Wu et al., 2023). Our extensive experiments demonstrate the transferability of OES-masks and the characteristics of sparse-subnetworks for image generation. Code is available at https://github.com/Avra98/Optimal-Eye-Surgeon.git.

Figures (62)

• Figure 1: Sparse-DIP lessens overfitting
• Figure 2: Sparse-DIPs are transferrable
• Figure 3: Subnetworks learned by OES are image generators with good image priors. On the contrary, the range of subnetworks learned by LTH is close to the overfitted noisy image, which is far from being an image prior. Image adapted from ulyanov2018deep
• Figure 4: Masking network at initialization induces image prior. Figures show the images after masking image generator parameters at random-initialization $G(\boldsymbol{\theta}_{in} \circ {\mathbf{m}}^{*}({\mathbf{y}}),{\mathbf{z}})$. The mask ${\mathbf{m}}^{*}$ was learned using the OES algorithm. Images corresponding to several sparsity levels are shown. We show that Strong Lottery Ticket Hypothesis also holds for image reconstruction.
• Figure 5: Comparative Analysis of Denoising Performance on 'Baboon' and 'Pepper' Images at $\sigma=25$ dB.