Lottery Image Prior

1 INTRODUCTION . Background Deep neural networks ( DNNs ) , in particular convolutional neural networks ( CNNs ) , have been powerful tools for solving various image inverse problems such as denoising ( Zhang et al. , 2017 ; Guo et al. , 2019 ; Lehtinen et al. , 2018 ) , inpainting ( Pathak et al. , 2016 ; Yu et al. , 2018 ; 2019b ) , and super resolution ( Ledig et al. , 2017 ; Lim et al. , 2017 ; Zhang et al. , 2018 ) . Conventional wisdom believes that is owing to DNNs ’ universal approximation ability and learning from massive training data . Yet , recent studies have revealed the specific architectures of CNNs have the inductive bias to represent and generate natural images well , and such favorable architecture inductive bias can work independently from fitting specific training sets ( Ulyanov et al. , 2018 ; Cheng et al. , 2019 ; Bora et al. , 2017 ; Heckel & Hand , 2019 ; Jalal et al. , 2020 ) . For example , deep image prior ( DIP ) ( Ulyanov et al. , 2018 ) shows an untrained neural network can be used as a handcrafted prior that transfers well across multiple inverse problems . The authors attributed the success to the CNN architecture itself , that appeared to possess high noise impedance even with only random initializations . As another example , in compressive sensing , Bora et al . ( 2017 ) ; Jalal et al . ( 2020 ) replaced the common structural assumptions such as sparsity with a pretrained generative adversarial networks ( GAN ) . The underlying rationale lies in that a pre-trained generator should ( approximately ) represent the notion of a vector being or more likely in the target domain such as natural images ; in other words , a sample more like a natural image will be closer to the output range of the pre-trained generator . In this paper , we refer to such general-purpose image prior parameterized by ( either untrained or pre-trained ) DNNs as DNN-based image priors . Recall that , classical image regularizers in the spatial or frequency domains are often not learningbased ( Tomasi & Manduchi , 1998 ; Sardy et al. , 2001 ; Dabov et al. , 2007 ) , or rely on compact learning models ( Cao et al. , 2008 ; Elad & Aharon , 2006 ; He et al. , 2015 ) . On contrary , DNN-based image priors have a massive number of parameters ( we compare the parameter numbers between the full model and the sparse subnetwork in Table . 2 ) , typically magnitudes more than the image ( even the image size ) size . The two extremes invite the natural question : Can we identify highly compact DNN-based image priors , that are same effective ? We note that , diving into this question has twofold appeals . On the algorithmic side , that could help us understand further how the topology and connectivity of CNN architecture itself will affect the effectiveness of those priors , and to what extent sparsity could be relevant . On the practical side , if we could provide an affirmative answer to this question , that would potentially lead to more computationally savings when applying those DNN-based priors in practice , leading to faster restoration or computational imaging with them . Towards the above question , the tool we refer to in this paper is the recently emerged Lottery Ticket Hypothesis ( LTH ) ( Frankle & Carbin , 2018 ; Frankle et al. , 2020a ) . LTH suggests that every dense DNN has an extremely sparse “ matching subnetwork ” , that can be trained in isolation to match the original dense DNN ’ s accuracy . While the vanilla LTH studies training from random scratch , the latest works also extend similar findings to fine-tuning the pre-trained models ( Chen et al. , 2020a ; 2021a ) . LTH has widespread success in image classification , language modeling , reinforcement learning and multi-modal learning , e.g. , ( Yu et al. , 2019a ; Renda et al. , 2020 ; Chen et al. , 2020a ; Gan et al. , 2021 ) . Our Contributions Drawing inspirations from the LTH literature , we conjecture and empirically study a novel “ lottery image prior ” ( LIP ) , stated as : Given an ( untrained or trained ) DNN-based image prior , it will have a sparse subnetwork that can be trained in isolation , to match the original DNN ’ s performance when being applied as a prior to regularizing various image inverse problems . Studying this new problem is , however , NOT a naive extension from the existing LTH methods , owing to several technical barriers : ( a ) till now , LTH has not been demonstrated for image inverse problem or DNN-based priors , to our best knowledge . Most LTH works studied discriminative tasks , with one exception ( Chen et al. , 2021c ) . It is therefore uncertain whether high-sparsity DNN is still viable for reconstruction-oriented tasks ; ( b ) existing LTH works typically require a full training set to locate the sparse subnetwork mask , whereas our LIP settings are only not data-rich . For example , DIP needs the DNN to be trained to overfit one specific image , making it drastically different from previous problems ; ( c ) the objectives between finding the sparse mask ( e.g. , learning the prior ) and fitting the sparse subnetwork ( e.g. , using the prior ) are often unaligned in LIP problems . For example , DIP will overfit a corrupted input image ( during which the sparse mask will be found ) in order to reconstruct a clean output image ( when the found sparse subnetwork will be used ) ; the pre-trained generator will also be used towards a different goal ( regularizing compressive sensing ) from their original pre-training task ( generating realistic images ) . Our extensive experimental study confirms the existence of LIP in two representative settings : ( i ) image restoration with the deep image prior , using an untrained DNN ( as shown in Fig . 2 ) ; and ( ii ) compressive sensing image reconstruction , using a pre-trained GAN generator . Using iterative magnitude pruning ( IMP ) with surrogate tasks ( the overview of our work paradigm is in Fig . 1 ) , we can successfully locate the LIP subnetworks at the sparsity range of 20 % -86.58 % in setting i ; and those at the sparsity range of 5 % -36 % in setting ii . Those LIP subnetworks also possess high transferability . For example , the LIP ticket found in the setting i transfer well across not only different images , but also different tasks such as denoising , inpainting and super-resolution . Our contributions are summarized below : • The first comprehensive study on LTH in DNN-based image priors and inverse problems , establishing the “ lottery image prior ” ( LIP ) and demonstrating the prevailing relevance of LTH more broadly than previously typical settings . • The investigation of both untrained and pre-trained DNNs as image priors , verifying that LIP subnetworks can be found in both settings , by overcoming several impediments such as severely limited training data and task objective mismatch . • High transferability of LIP subnetworks across data/datasets , and various inverse problem tasks . Our finding reflects the underlying common image prior that is agnostic to specific data or task , through the lens of the CNN architecture together with sparsity . 2 BACKGROUND WORK . Lottery Ticket Hypothesis LTH ( Frankle & Carbin , 2018 ) states that the dense , randomly initialized DNN contains a sparse matching subnetwork , which could reach the comparable or even better performance by independently being trained for the same epoch number as the full network do . Since then , the statement has been verified in a variety of fields , such as image classification ( Frankle & Carbin , 2018 ; Liu et al. , 2018 ; Wang et al. , 2020 ; Evci et al. , 2019 ; Frankle et al. , 2020b ; Savarese et al. , 2019 ; Yin et al. , 2019 ; You et al. , 2019 ; Ma et al. , 2021 ; Chen et al. , 2021a ) , natural language processing ( Gale et al. , 2019 ; Chen et al. , 2020a ) , reinforcement learning ( Yu et al. , 2019a ) , lifelong learning ( Chen et al. , 2020b ) , graph neural networks ( Chen et al. , 2021b ) , and adversarial robustness ( Cosentino et al. , 2019 ) . Rewinding was proposed by ( Frankle et al. , 2019 ) to scale up the LTH to large models and datasets . The found matching subnetworks also demonstrate transferability across datasets and tasks ( Morcos et al. , 2019 ; Desai et al. , 2019 ) . Deep Image Prior and Its Variants Despite CNN ’ s tremendous success on various imaging tasks , their outstanding performance is often attributed to massive data-driven learning . DIP ( Ulyanov et al. , 2018 ) pioneered to show that CNN architecture alone has captured important natural image priors : by over-fitting a randomly initialized untrained CNN to a single degraded image ( plus some early stopping ) , it can restore the clean output without accessing ground truth . Follow-up work ( Mataev et al. , 2019 ) strengths DIP performance by incorporating it with the regularization by denoising ( RED ) framework and a series of works ( Mastan & Raman , 2020 ; 2021 ) use the contextual feature learning method to achieve the same goal of DIP . Besides natural image restoration , DIP was successfully applied to PET image reconstruction ( Gong et al. , 2018 ) , dynamic magnetic resonance imaging ( Jin et al. , 2019 ) , unsupervised image decomposition ( Gandelsman et al. , 2019 ) and quantitative phase imaging ( Yang et al. , 2021 ) . Heckel & Hand ( 2018 ) further demonstrated that even an under-parameterized non-convolutional model , named “ Deep Decoder ” , can over-fit a sin- gle degraded image like DIP did , which does not critically rely on early stopping or so . Chen et al . ( 2020c ) was the first to study the possibility to optimize CNN architectures for capturing stronger image priors in DIP , via leveraging Neural Architecture Search ( NAS ) . Compressive Sensing using Generative Models Compressive Sensing ( CS ) reconstructs an unknown vector from under-sampled linear measurements of its entries ( Foucart & Rauhut , 2013 ) , by assuming the unknown vector to admit certain structural priors . The most common structural prior is to assume that the vector is k-sparse in some known bases ( Candes et al. , 2006 ; Donoho , 2006 ) , and more sophisticated statistical assumptions were also considered ( Baraniuk et al. , 2010 ) . However , those priors are inevitably oversimplified to depict the high-dimensional manifold of natural images . Bora et al . ( 2017 ) presented the first algorithm that used pre-trained generative models such as GANs , as the prior for compressed sensing . As a prior , the pre-trained generator encourages CS to produce vectors close to its output distribution , which approximates its training image distribution . Significant research has since followed to better understand the behaviours and theoretical limits of CS using generative priors , e.g. , ( Hand & Voroninski , 2018 ; Bora et al. , 2018 ; Hand et al. , 2018 ; Kamath et al. , 2019 ; Liu & Scarlett , 2020 ; Jalal et al. , 2020 ) .

Summary: Summary This submission studies lottery ticket hypothesis for deep neural network based inverse imaging. It considers two scenarios: 1-inference for compressed sensing based on pre-trained deep generative models, 2-deep image prior where all network parameters are fit to a single image. Both scenarios are dealing with a single image training task, or, an overfitting task, which makes this work different from the past work using LTH. The empirical results suggest that with a high level of sparsity not only the model size can be made much smaller, but also the generalization can improve which suggests pruning as a regularization. They also suggest good transferability from one image restoration task to another.