PLUGIn: A simple algorithm for inverting generative models with recovery guarantees

p 4dn0/ndk✏k . Thus , the algorithm can denoise when the expansion ratio nd/n0 is large . Numerical experiments on synthetic data and real data are provided to validate our theoretical results and to illustrate that the algorithm can effectively remove artifacts in an image . 1 Introduction . We consider the inverse problem of recovering an unknown vector x⇤ 2 Rn0 from a noisy observation y 2 Rnd of the form y = G ( x⇤ ) + ✏ , ( 1 ) where ✏ 2 Rnd is noise and G : Rn0 ! Rnd is a known d-layer feed-forward neural network with ReLU activation functions ( in fact , it is straightforward to generalize to positively homogeneous activation functions , which we will elaborate later ) . Precisely , G has the form G ( x ) = ( Ad ( Ad 1 . . . ( A1x ) . . . ) ) , ( 2 ) where ( · ) = max ( · , 0 ) is the ReLU activation function and Ai 2 Rni⇥ni 1 is the weight matrix in the i-th layer . The inverse problem ( 1 ) has applications in , for example , signal denoising and signal compression [ 1 , 2 , 3 ] . In the case of denoising the typical goal is to recover a clean signal y⇤ from its noisy observation y⇤ + ✏ and , in the case of compression , the goal is to find an efficient low dimensional representation of y⇤ . Traditional approaches for these problems often use priors on signals , for example , a sparsity prior with respect to a fixed basis or dictionary [ 4 , 5 , 6 ] . An emerging viewpoint is to use a generative prior that assumes the unknown signal y⇤ is in the range of a deep generative model G , i.e. , y⇤ = G ( x⇤ ) , and develop algorithms that can estimate the latent code vector x⇤ from y ⇤ + ✏ , thus recovering y⇤ . Recent advancements in training deep neural networks have shown that priors in the form of generative models can effectively map low dimensional vectors to the space of natural image classes 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . [ 7 , 8 , 9 ] . Learned generative models can then be used as priors to solve various inverse problems including denoising [ 1 , 10 ] , compressive sensing [ 11 , 12 , 13 , 14 , 15 , 16 ] , phase retrieval [ 17 ] , blind deconvolution [ 18 , 19 ] , low-rank matrix recovery [ 20 ] and have been shown to perform on par or outperform classical sparsity based approaches for these inverse problems . For example , in [ 1 ] the authors empirically showed that an end-to-end approach for denoising using a neural network that maps noisy patches in an image to noise-free ones achieves state-of-the-art performance and is on par with BM3D . Similarly , in [ 11 ] the authors empirically showed that for compressive sensing using generative prior , optimization of the empirical risk objective over the latent code space ( of the generative prior ) can recover a vector that effectively estimates the uncompressed signal with 5-10 times less measurements compared to Lasso in some cases . Given that y equals G ( x⇤ ) , with possibly some additive noise , a standard way to estimate x⇤ , i.e. , to invert G ( x⇤ ) , would be to look for a minimizer of the program min x2Rn0 ky G ( x ) k2 . ( 3 ) Unfortunately , this program is non-convex and to our knowledge there is no known efficient method that can achieve its global minimum in general . On the other hand , for generative networks with random ( Gaussian ) weight matrices , a line of papers showed that gradient-based algorithms can provably avoid local minima with high probability [ 10 , 14 , 17 , 20 ] . In particular , [ 10 ] considers a random Gaussian noise model , and random Gaussian weight matrices Ai which are highly expansive at each layer . Under these conditions , the authors show that the latent code vector x⇤ can be accurately estimated using a gradient-based method that uses the ( sub- ) gradient updates given by x k+1 = xk ⌘ ( D1A1 ) | ( D2A2 ) | · · · ( DdAd ) | G ( xk ) y , ( 4 ) where xk is the k-th estimate , ⌘ 2 R is step size , and Dj is a diagonal matrix with entries that are either zero or one . Each Dj zeros out the inactive rows of Aj with respect to the estimate xk and so it is a function of xk ( and Ap for p < j ) . Thus , at each iteration all Dj need to be updated . In this paper , we show that we can drop the Dj and still recover an accurate estimate of x⇤ . We propose the following iterative algorithm , Partially Linearized Update for Generative Inversion ( PLUGIn ) , to estimate x⇤ : PLUGIn : x k+1 = xk ⌘A|1A | 2 · · ·A | d G ( xk ) y . ( 5 ) This algorithm was inspired by previous work showing that latent vectors for non-linear single-index function can be approximately estimated by treating the function as linear [ 21 , 22 ] . Similar to [ 10 , 11 , 12 , 14 , 17 , 19 , 23 ] , for theoretical analysis , we assume the weight matrices are Gaussian . To show that the algorithm works more broadly , we conduct real data simulations . Applying the ideas in [ 21 , 22 ] one can show that for any fixed x0 , the first iteration of ( 5 ) provides an unbiased estimate of x⇤ with ⌘ = 2d , which is generally not the case for the gradient descent estimates given by ( 4 ) , see Section 3 for further details . Additionally , each iteration of PLUGIn maps the difference G ( xk ) y to the low dimensional latent code space using a static matrix A|1A | 2 . . . A | d , which can be pre-multiplied and reused in subsequent iterations . In addition to the randomness assumption on the weight matrices , when the width of the each layer satisfies ni & 5in0 ( up to log factors ) , noise ✏ is independent of weight matrices , and the step size ⌘ is fixed , we show that the estimates provided by PLUGIn converge to a neighbourhood of x⇤ geometrically with high probability . This result allows ni > ni+1 for i 1 and thus can provide theoretical guarantees for estimating x⇤ from ( 1 ) even when G has some contractive layers . To our knowledge , this is the first such result . Additionally , we show that the k-th recovery and reconstruction errors satisfy kxk x⇤k . s 4dn0 nd k✏k and kG ( xk ) G ( x⇤ ) k . s 4dn0 nd k✏k with high probability when k is large . Thus , PLUGIn can effectively denoise noisy observations satisfying ( 1 ) provided that the expansion ratio nd/n0 is large . These contributions are highlighted below : 1 . We introduce a simple novel algorithm ( 5 ) , called Partially Linearized Update for Generative Inversion ( PLUGIn ) , for inverting a generative model . Note that each update of PLUGIn can be implemented with two blackboxes . The first blackbox provides the G ( xk ) for each iterate and the second blackbox provides the product of the weight matrices A|1A | 2 . . . A | d . Thus , given these two blackboxes the individual matrices are not needed . 2 . Under relatively mild conditions , we show that for generative priors with possibly contractive layers , PLUGIn can effectively estimate the unknown latent code x⇤ and the unknown signal G ( x⇤ ) from observation y satisfying ( 1 ) . The convergence is geometric and holds for a range of step sizes . 3 . In the presence of arbitrary noise ✏ , we show that the recovery and reconstruction error corresponding to the estimate xk , for sufficiently large k , is approximately p 4dn0/ndk✏k . In contrast to typical random noise assumptions , we only require the noise to be independent of the weight matrices . This allows stochastic noise ( independent of weight matrices ) as a special case if one can bound the norm of the stochastic noise with high probability . Organization of the paper : In Section 1.1 , we introduce notations used through the paper . In Section 2 , we state the main results . In Section 3 , we provide an outline of the proof . In Section 4 , we provide numerics on synthetic and real data . In Section 5 , we briefly discuss a generalization of activation functions , as well as limitations in our results . 1.1 Notations . For a positive integer n , let [ n ] = { 1 , 2 , . . . , n } . For a vector x , let kxk be its Euclidean norm ; for a matrix A , let kAk be its operator norm ; for a matrix A and a set T , let kAkT : = supx2T \ { 0 } kAxk kxk . Let N ( 0 , t2 ) be the normal distribution with mean zero and standard deviation t. Let B ( x , r ) be the Euclidean ball of radius r centered at x and let Bn ( 0 , r ) be the Euclidean ball in Rn with radius r , centered at origin . We use C and c to denote absolute constants ( often c for small ones and C for large ones ) which may vary from line to line . We also use C0 , C1 , etc. , to denote particular absolute constants , which do not change throughout the paper . 2 Main Results . Let the generative model G : Rn0 ! Rnd be defined as in Equation ( 2 ) . We consider the inverse problem : Let : x⇤ 2 Rn0 , ✏ 2 Rnd , y = G ( x⇤ ) + ✏ , Given : generative model G and observation y , Estimate : latent code vector x⇤ and G ( x⇤ ) . We assume : A1 . All activation functions are ReLU . A2 . Each Ai 2 Rni⇥ni 1 has i.i.d . N ( 0 , 1/ni ) entries and { Ai } i2 [ d ] are independent . A3 . Layer widths ( number of nodes in each layer ) satisfy ni C05in0 log 0 @ i 1Y j=0 enj n0 1 A , i 2 [ d ] ( 6 ) for some ( sufficiently large ) absolute constant C0 . A4 . The noise ✏ does not depend on { Ai } i2 [ d ] . ( The noise may be deterministic or random . ) Let x0 2 Rn0 and let xk be the result of applying k iterations of PLUGIn ( 5 ) . Under assumptions A1-A4 , xk geometrically converges to a neighbourhood of x⇤ ( and also G ( xk ) to a neighbourhood of G ( x⇤ ) ) for a range of step sizes near 2d . Precisely , we have the following theorem . Theorem 1 . Let ✓ 2 ( 0 , 43 ) and let ↵ = |1 ✓|+ 1 2✓ 2 ( 0 , 1 ) . Let R be a positive number such that kx0 x⇤k R. Under assumptions A1-A4 , the k-th estimate xk given by PLUGIn algorithm ( 5 ) with constant step size ⌘ = ✓2d satisfies kxk x⇤k ↵kR+ 15✓ 1 ↵2 d p n0/ndk✏k , and kG ( xk ) G ( x⇤ ) k 3↵kR+ 45✓ 1 ↵2 d p n0/ndk✏k with probability at least 1 2 ( k + 3 ) e 10n0 . When ✓ = 1 , Theorem 1 reduces to the following corollary . Corollary 1 . Let R be a positive number such that kx0 x⇤k R. Under assumptions A1-A4 , the k-th estimate x k given by PLUGIn algorithm ( 5 ) with constant step size ⌘ = 2d satisfies kxk x⇤k 2 kR+ 30 · 2d p n0/ndk✏k , and kG ( xk ) G ( x⇤ ) k 2 k ( 3R ) + 90 · 2d p n0/ndk✏k with probability at least 1 2 ( k + 3 ) e 10n0 . Remark 1 ( Allowed number of iterations and comparison of error bound summands ) . The probability in Theorem 1 ( and also Corollary 1 ) decreases as the number of iterations increases . However , this is generally not a problem because we can take up to exponentially many iterations before this probability becomes trivial . In fact , if we require a probability of at least 99 % on the error bounds , then we can take any k k ? ( n0 ) : = 1200e 10n0 3 . Note that k ? grows exponentially in n0 and it is already large when n0 is relatively small1 . The noise term in the error bounds ( proportional to k✏k ) will dominate the other term ( long ) before k ? many iterations are reached , unless the noise is virtually zero . Indeed , in Corollary 1 , if 2 k ? R p 4dn0/ndk✏k , then the noise must satisfy k✏k < R q nd 4dn0 1 2 k ? = 8R q nd 4dn0 1 2 exp ( 10n0 ) /200 . Note that 1 2 exp ( 10n0 ) /200 is a double exponential and extremely small2 , so we will consider the noise as being virtually zero in this case . On the other hand , if 2 k ? R < p 4dn0/ndk✏k , then after at most k ? iterations , with probability at least 99 % , the estimation errors for x⇤ and G ( x⇤ ) will be at most C p 4dn0/ndk✏k . The same argument can be made for Theorem 1 when ↵ is close to 12 ( i.e. , ✓ is close to 1 ) so that ↵k decays fast and 11 ↵ is not close to being singular . Remark 2 ( Contractive layers ) . In A3 , ( 6 ) states a lower bound on ni with respect to the latent code dimension n0 ( up to log factors ) . While this bound strictly increases with layer depth i , it is not necessary for ni to always increase with i ( except in the first layer ) . For example , consider ni = C05dn0d ( 2d i ) where is any fixed number such that C0 2 N and 4 + logC0 . It is easy to see n1 > n2 > · · · > nd , and we can also verify ( see Appendix D ) that such ni satisfy ( 6 ) . In this case , the network is contractive in each layer after the first , and Theorem 1 still applies . Remark 3 ( Initialization may depend on random weight matrices ) . The results of the theorem can still hold when x0 is chosen randomly , dependent on the weight matrices Ai . In this case , suppose that kx0 x⇤k R with probability at least 1 . Then , the error bounds hold with probability at least 1 2 ( k + 3 ) e 10n0 . This does not follow directly from the theorem as stated ( which fixes x 0 , then takes random weight matrices ) , but follows from the proof . Remark 4 ( Comparison to guarantees for gradient-based method ) . Here we compare our results to the ones in [ 10 ] , which uses ( 4 ) for iterations and considers a random noise model . When the noise ✏ ⇠ N ( 0 , 1nd 2 Ind ) , the error bounds on the k-th PLUGIn iterate in Corollary 1 , with high probability , reduces to max { kxk x⇤k , kG ( xk ) G ( x⇤ ) k } C ⇣ 2 kR+ 2d p n0/nd ⌘ . Thus , we get geometric convergence with a rate of 1/2 and after sufficiently many iterations , the errors are at most C2d p n0/nd . A similar result was shown in [ 10 ] , which studies a gradient-based method for solving the inverse problem ( 1 ) . They show that when the weight matrices are random Gaussian and sufficiently 1For example , k ? ( 5 ) ⇡ 2.59⇥ 1019 and k ? ( 10 ) ⇡ 1.34⇥ 1041 . 2For example , 2 exp ( 10 ) /200 ⇡ 7⇥ 10 34. expansive at each layer , the iterates of the gradient-based method converge to a neighborhood of the target signal x⇤ . The result holds for a fixed step size , dependent on d. After sufficiently many iterations , the iterates converge geometrically to a neighborhood of x⇤ of radius at most on the order of p n0/nd , up to log factors . This rate of convergence takes the form ( 1 C/d2 ) , thus giving slower convergence for deeper nets . On the other hand , we note that dependence on d is of relatively minor concern . Generative models usually have small depth in practice , our MNIST experiments ( below ) work well with depth 3 , and typical applications use depth less than 8 .

Summary: This paper introduces an algorithm for recovering an unknown latent code vector under a specific generative model: a deep generative network with ReLU activation functions. Innovation of paper compared with [10] is to show that one drops the $D_j$ and still recover an accurate estimate of $x^*$. Theoretical results on geometrical convergence to a neighborhood of $x^{*}$ are confirmed by experiments on synthetic random data and MNIST images.