summaries: In this manuscript, the authors proposes a Stochastic Anderson Mixing (SAM) scheme to solve nonconvex optimization problems. Under the smoothness, unbiased gradient estimate with bounded variance assumptions, the authors show different variants of the scheme (SAM, SAM with variance reduction, Ada-SAM) have asymptotic and non-asymptotic convergence. Experiments on MNIST, CIFAR and Penn Treebank seem to show the superior performance over other baselines.

summaries: This paper proposes a stochastic Anderson mixing method to solve non-convex stochastic optimization problems. The main contribution is a  convergence theory and the applications of this method for deep learning problems. Adaptive, variance reduced, and preconditioned versions of this methods are also studied.

summaries: The paper applies Anderson mixing to stochastic optimization problems and shows its effectiveness in training neural networks. To make Anderson mixing work in stochastic settings, the paper proposes a series of strategies to handle noise and uncertainty, including damped projection, adaptive regularization, variance reduction, and preconditioned mixing. Proof of convergence is provided, together with numerical evaluations.

summaries: The paper under review consider the stochastic Anderson Mixing algorithm and  shows its convergence analysis. An enhanced variance reduction of the algorithms is also derived. The algorithms are tested on extensive data sets with other existing method. This paper seems to be the first one to address the theoretical convergence analysis in this direction.

summaries: In this paper, the author propose a stochastic version of Anderson acceleration (widely used in fixed point iterations)  to solve non convex optimization problem. Some important modifications of the deterministic  version are introduced in order to stabilize and improve the algorithm. The paper theoretically shows that the new algorithm can achieve O(epsilon^{-2}) sample complexity. Numerical results are provided to verify the performance of SAM.

abstractText: Anderson mixing (AM) is an acceleration method for fixed-point iterations. Despite its success and wide usage in scientific computing, the convergence theory of AM remains unclear, and its applications to machine learning problems are not well explored. In this paper, by introducing damped projection and adaptive regularization to the classical AM, we propose a Stochastic Anderson Mixing (SAM) scheme to solve nonconvex stochastic optimization problems. Under mild assumptions, we establish the convergence theory of SAM, including the almost sure convergence to stationary points and the worst-case iteration complexity. Moreover, the complexity bound can be improved when randomly choosing an iterate as the output. To further accelerate the convergence, we incorporate a variance reduction technique into the proposed SAM. We also propose a preconditioned mixing strategy for SAM which can empirically achieve faster convergence or better generalization ability. Finally, we apply the SAM method to train various neural networks including the vanilla CNN, ResNets, WideResNet, ResNeXt, DenseNet and LSTM. Experimental results on image classification and language model demonstrate the advantages of our method.

authors:
Zeyuan Allen-Zhu
Elad Hazan

title: Variance reduction for faster non-convex optimization

venue: In International conference on machine learning,

title: Iterative procedures for nonlinear integral equations

venue: Journal of the ACM (JACM),

authors:
Wei Bian
Xiaojun Chen
CT Kelley

title: Anderson acceleration for a class of nonsmooth fixed-point problems

venue: SIAM Journal on Scientific Computing,

authors:
Léon Bottou
Frank E. Curtis
Jorge Nocedal

title: Optimization methods for large-scale machine learning

venue: SIAM Review,

authors:
Claude Brezinski
Michela Redivo-Zaglia
Yousef Saad

title: Shanks sequence transformations and Anderson acceleration

authors:
Claude Brezinski
Stefano Cipolla
Michela Redivo-Zaglia
Yousef Saad

title: Shanks and Anderson-type acceleration techniques for systems of nonlinear equations

venue: arXiv preprint arXiv:2007.05716,

authors:
Richard H Byrd
Samantha L Hansen
Jorge Nocedal
Yoram Singer

title: A stochastic quasi- Newton method for large-scale optimization

venue: SIAM Journal on Optimization,

authors:
Eric Cancès
Claude Le Bris

title: Can we outperform the DIIS approach for electronic structure calculations

venue: International Journal of Quantum Chemistry,

authors:
Xiangyi Chen
Sijia Liu
Ruoyu Sun
Mingyi Hong

title: On the convergence of a class of Adam-type algorithms for non-convex optimization

venue: In International Conference on Learning Representations,

title: On a stochastic approximation method

venue: The Annals of Mathematical Statistics,

authors:
Jing Dong
Xin T Tong

title: Replica exchange for non-convex optimization

venue: Journal of Machine Learning Research,

authors:
John Duchi
Elad Hazan
Yoram Singer

title: Adaptive subgradient methods for online learning and stochastic optimization

title: Probability: Theory and examples, volume 49

venue: Cambridge university press,

authors:
Haw-ren Fang
Yousef Saad

title: Two classes of multisecant methods for nonlinear acceleration

venue: Numerical Linear Algebra with Applications,

authors:
Jerome Friedman
Trevor Hastie
Robert Tibshirani

title: The elements of statistical learning, volume 1. Springer series in statistics

venue: New York,

authors:
Alejandro J Garza
Gustavo E Scuseria

title: Comparison of self-consistent field convergence acceleration techniques

venue: The Journal of chemical physics,

authors:
Saeed Ghadimi
Guanghui Lan

title: Stochastic first- and zeroth-order methods for nonconvex stochastic programming

authors:
Robert Gower
Dmitry Kovalev
Felix Lieder
Peter Richtárik

title: RSN: Randomized subspace Newton

venue: In Conference on Neural Information Processing Systems,

authors:
Robert M Gower
Peter Richtárik

title: Randomized quasi-Newton updates are linearly convergent matrix inversion algorithms

venue: SIAM Journal on Matrix Analysis and Applications,

authors:
Kaiming He
Xiangyu Zhang
Shaoqing Ren
Jian Sun

title: Deep residual learning for image recognition

venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,

authors:
Sepp Hochreiter
Jürgen Schmidhuber

title: Long short-term memory

venue: Neural computation,

authors:
Gao Huang
Zhuang Liu
Laurens Van Der Maaten
Kilian Q Weinberger

title: Densely connected convolutional networks

authors:
Shaojun Huang
Yuanzhang Sun
Qiuwei Wu

title: Stochastic economic dispatch with wind using versatile probability distribution and L-BFGS-B based dual decomposition

venue: IEEE Transactions on Power Systems,

authors:
Xunpeng Huang
Xianfeng Liang
Zhengyang Liu
Lei Li
Yue Yu
Yitan Li

title: SPAN: A stochastic projected approximate Newton method

venue: In Proceedings of the AAAI Conference on Artificial Intelligence,

authors:
Brian Irwin
Eldad Haber

title: Secant penalized BFGS: A noise robust quasi-Newton method via penalizing the secant condition

venue: arXiv preprint arXiv:2010.01275,

authors:
Rie Johnson
Tong Zhang

title: Accelerating stochastic gradient descent using predictive variance reduction

venue: Advances in neural information processing systems,

authors:
Diederik P Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

venue: arXiv preprint arXiv:1412.6980,

title: Korpelevich. The extragradient method for finding saddle points and other problems

authors:
Alex Krizhevsky
Geoffrey Hinton

title: Learning multiple layers of features from tiny images

authors:
Yann LeCun
Léon Bottou
Yoshua Bengio
Patrick Haffner

title: Gradient-based learning applied to document recognition

venue: Proceedings of the IEEE,

authors:
Jian Li
Xuanyuan Luo
Mingda Qiao

title: On generalization error bounds of noisy gradient methods for non-convex learning

authors:
Yuanzhi Li
Colin Wei
Tengyu Ma

title: Towards explaining the regularization effect of initial large learning rate in training neural networks

venue: Advances in Neural Information Processing Systems,

authors:
Mitchell P Marcus
Mary Ann Marcinkiewicz
Beatrice Santorini

title: Building a large annotated corpus of English: The penn treebank

venue: Computational Linguistics,

title: Deep learning via Hessian-free optimization

venue: In Proceedings of the 27th International Conference on Machine Learning,

authors:
Aryan Mokhtari
Alejandro Ribeiro

title: Stochastic quasi-Newton methods

authors:
Arkadi Nemirovski
Anatoli Juditsky
Guanghui Lan
Alexander Shapiro

title: Robust stochastic approximation approach to stochastic programming

venue: SIAM Journal on optimization,

title: New stochastic approximation type procedures

venue: Automat. i Telemekh,

authors:
Boris T Polyak
Anatoli B Juditsky

title: Acceleration of stochastic approximation by averaging

venue: SIAM journal on control and optimization,

authors:
Florian A Potra
Hans Engler

title: A characterization of the behavior of the Anderson acceleration on linear problems

venue: Linear Algebra and its Applications,

authors:
Phanisri P Pratapa
Phanish Suryanarayana
John E Pask

title: Anderson acceleration of the Jacobi iterative method: An efficient alternative to Krylov methods for large, sparse linear systems

venue: Journal of Computational Physics,

title: On the momentum term in gradient descent learning algorithms

venue: Neural networks,

authors:
Sashank J Reddi
Ahmed Hefny
Suvrit Sra
Barnabas Poczos
Alex Smola

title: Stochastic variance reduction for nonconvex optimization

authors:
Herbert Robbins
Sutton Monro

title: A stochastic approximation method

venue: The annals of mathematical statistics,

authors:
Youcef Saad
Martin H Schultz

title: GMRES: A generalized minimal residual algorithm for solving nonsymmetric linear systems

venue: SIAM Journal on scientific and statistical computing,

title: Iterative methods for sparse linear systems

title: Asymptotic distribution of stochastic approximation procedures

authors:
Nicol N Schraudolph
Jin Yu
Simon Günter

title: A stochastic quasi-Newton method for online convex optimization

venue: In Artificial intelligence and statistics,

authors:
Damien Scieur
Alexandre d’Aspremont
Francis Bach

title: Regularized nonlinear acceleration

venue: In Proceedings of the 30th International Conference on Neural Information Processing Systems,

authors:
Damien Scieur
Francis Bach
Alexandre d’Aspremont

title: Nonlinear acceleration of stochastic algorithms

venue: In Proceedings of the 31st International Conference on Neural Information Processing Systems,

authors:
Damien Scieur
Edouard Oyallon
Alexandre d’Aspremont
Francis Bach

title: Nonlinear acceleration of CNNs

venue: In ICLR Workshop track,

title: Regularized nonlinear acceleration

venue: Mathematical Programming,

authors:
Shai Shalev-Shwartz
Shai Ben-David

title: Understanding machine learning: From theory to algorithms

authors:
Tijmen Tieleman
Geoffrey Hinton

title: Lecture 6.5-RMSprop: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural networks for machine learning

authors:
Alex Toth
CT Kelley

title: Convergence analysis for Anderson acceleration

venue: SIAM Journal on Numerical Analysis,

authors:
Alex Toth
J Austin Ellis
Tom Evans
Steven Hamilton
CT Kelley
Roger Pawlowski
Stuart Slattery

title: Local improvement results for Anderson acceleration with inaccurate function evaluations

authors:
Henk A Van der Vorst
C Vuik

title: The superlinear convergence behaviour of GMRES

venue: Journal of computational and applied mathematics,

authors:
Homer F Walker
Peng Ni

title: Anderson acceleration for fixed-point iterations

authors:
Xiao Wang
Shiqian Ma
Donald Goldfarb
Wei Liu

title: Stochastic quasi-Newton methods for nonconvex stochastic optimization

authors:
Saining Xie
Ross Girshick
Piotr Dollár
Zhuowen Tu
Kaiming He

title: Aggregated residual transformations for deep neural networks

authors:
Pan Xu
Jinghui Chen
Difan Zou
Quanquan Gu

title: Global convergence of Langevin dynamics based algorithms for nonconvex optimization

venue: In Proceedings of the 32nd International Conference on Neural Information Processing Systems,

authors:
Zhewei Yao
Amir Gholami
Sheng Shen
Mustafa Mustafa
Kurt Keutzer
Michael Mahoney

title: Adahessian: An adaptive second order optimizer for machine learning

authors:
Sergey Zagoruyko
Nikos Komodakis

title: Wide residual networks. In British Machine Vision Conference 2016

venue: British Machine Vision Association,

authors:
Michael Zhang
James Lucas
Jimmy Ba
Geoffrey E Hinton

title: Lookahead optimizer: k steps forward, 1 step back

venue: In Advances in Neural Information Processing Systems,

authors:
Juntang Zhuang
Tommy Tang
Yifan Ding
Sekhar C Tatikonda
Nicha Dvornek
Xenophon Papademetris
James Duncan

title: Adabelief optimizer: Adapting stepsizes by the belief in observed gradients

heading: 1 Introduction

text:
Stochastic optimization is important in various areas such as statistics [15], machine learning [4, 54] and power systems [25]. In this paper, we consider the following stochastic optimization problem:
min x∈Rd
f(x) = Eξ [F (x; ξ)] , (1)
where F : Rd × Ξ → R is continuously differentiable and possibly nonconvex and the random variable ξ ∈ Ξ may follow an unknown probability distribution. It is assumed that only noisy information about the gradient of f is available through calls to some stochastic first-order oracle (SFO). One special case of (1) is the empirical risk minimization problem:
min x∈Rd
f(x) def =
1
T T∑ i=1 fξi(x), (2)
where fξi : Rd → R is the loss function corresponding to the i-th data sample and T denotes the number of data samples. T can be extremely large such that it prohibits the computation of the full gradient∇f . Thus designing efficient and effective numerical algorithms for solving problem (1) or (2) with rigorous theoretical analysis is a challenging task.
∗Corresponding author. clbao@mail.tsinghua.edu.cn.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
One classical approach for solving (1) is the stochastic gradient descent (SGD) method [45]. It mimics the gradient descent (GD) method with noisy gradients and exhibits optimal convergence rate for some strongly convex stochastic problems [10, 48]. Some early related works of SGD in convex optimization can be found in [38, 39, 40]. For nonconvex cases, Ghadimi and Lan [17] proposed a randomized stochastic gradient (RSG) method that randomly selects a solution x̄ from previous iterates. To ensure x̄ satisfying E [‖∇f(x̄)‖22] ≤ , the total number of SFO-calls2 needed by RSG is O ( −2 ) . For training neural networks, adaptive learning rate methods were proposed to accelerate SGD, e.g. Adagrad [12], RMSprop [55] and Adam [29], though their theoretical properties in nonconvex optimization were unknown until recently [9]. Regarding to second-order optimization methods [49, 36, 37], one notable work is the framework of stochastic quasi-Newton (SQN) methods proposed by Wang et al. [60], which covers a class of SQN methods and has theoretical guarantees in nonconvex stochastic optimization. However, these second-order methods usually demand more gradient evaluations in every iteration and less noisy gradient information to achieve actual acceleration [4]. There are also some recent works inspired from the dynamics perspective, such as the Langevin dynamics based algorithms [62, 33, 11], but many of them focus on the theory and their applications in deep learning are still under-explored.
In practice, the choices of optimizers can vary for different applications [63] and their practical performances are still unsatisfactory in terms of convergence rate or generalization ability [66]. For these reasons, we develop a novel second-order method based on Anderson mixing (AM) [2]. Our analysis and experiments show the proposed method is competent from both theoretical and practical perspectives and performs well in training various neural networks in different tasks.
AM is a sequence acceleration method in scientific computing [5]. It is widely used to accelerate the slow convergence of nonlinear fixed-point iterations arising in computational physics and quantum chemistry, e.g. the self-consistent field iteration in electronic structure calculations [16, 8], where the function evaluation is costly. It turns out that AM is closely related to multisecant quasi-Newton methods in nonlinear problems [14, 5] or the GMRES method [46] in linear problems [59, 41].
Inspired by the great success of AM in accelerating fixed-point iterations, it is natural to ask whether AM can be applied to accelerate nonlinear optimization since the gradient descent with constant stepsize is a fixed-point iteration. This idea has been explored in [50, 51], but the proposed Regularized Nonlinear Acceleration (RNA) method is built on the minimal polynomial extrapolation (MPE) approach [5], a sequence transformation method that has subtle difference from AM. Also, their methods rely heavily on the contraction assumption of the fixed-point map and the strong convexity. For AM, although current research has proved local linear convergence of AM for fixed-point iterations under some conditions [56, 57, 3], there exists no version of AM that guarantees convergence for nonconvex optimization, let alone stochastic optimization.
In this paper, we develop a stochastic version of AM. Due to the nonconvexity and noise inside the problems and lack of line search or trust-region, a straightforward migration of AM to nonconvex stochastic optimization is infeasible. As a result, we make several fundamental modifications to AM. We highlight the main contributions of our works as follows:
1. We develop a stochastic version of AM, namely Stochastic Anderson Mixing (SAM), by introducing damped projection and adaptive regularization. We prove its almost sure convergence to a stationary point and analyze its work complexity. When a randomly chosen iterate xR is returned as the output of SAM, we prove that the worst-case SFO-calls complexity to guarantee E [ ‖∇f(xR)‖22 ] ≤ is O ( −2 ) . (See Theorems 1-4.)
2. We give a variance reduced extension of SAM by borrowing the stochastic variance reduced gradient (SVRG) [28] technique and analyze its SFO-calls complexity. (See Theorem 5.) We also propose a preconditioned mixing strategy for AM and obtain the preconditioned SAM method which can empirically converge faster or generalize better. (See Section 2.3.)
3. Extensive experiments on training Convolutional Neural Network (CNN), ResNet, WideResNet, ResNeXt, DenseNet, and LSTM on different tasks and datasets show the faster convergence or better generalization ability of our method compared with the state-of-the-art methods. (See Section 4.)
2One call to the SFO in (1) or one evaluation of ∇fξi(x) in (2) is counted as one SFO-call.

heading: 2 Methodology

heading: 2.1 Anderson Mixing

text:
AM is proposed for acceleration of fixed-point iterations. We assume the fixed-point iteration is xk+1 = g(xk) def = xk + rk, where rk def = −∇f(xk). Then g(xk) − xk = rk. Here, we adopt the description of AM in [14, 59]. Let ∆ denote the forward difference operator, say, ∆xk = xk+1 − xk. Let Xk and Rk record the most recent m(m ≤ k) iterations:
Xk = [∆xk−m,∆xk−m+1, · · · ,∆xk−1], Rk = [∆rk−m,∆rk−m+1, · · · ,∆rk−1]. (3)
AM can be decoupled into two steps. We call them the projection step and the mixing step:
x̄k = xk −XkΓk, (Projection step) (4a) xk+1 = x̄k + βkr̄k, (Mixing step) (4b)
where βk is the mixing parameter, and r̄k def = rk −RkΓk is reminiscent of extragradient [30]. Γk is determined by solving Γk = arg min
Γ∈Rm ‖rk −RkΓ‖2. (5)
Combining (4a) and (4b), we obtain the full form of AM [14, 59, 5, 42]:
xk+1 = xk + βkrk − (Xk + βkRk) Γk. (6)
Remark 1. To see the rationality of AM, we assume f is twice continuously differentiable. Then a quadratic approximation of f implies∇2f(xk) (xj − xj−1) ≈ ∇f(xj)−∇f(xj−1) in a local small region around xk, so it is reasonable to assume Rk ≈ −∇2f(xk)Xk. Thus we see ‖rk −RkΓ‖2 ≈ ‖rk +∇2f(xk)XkΓ‖2. Hence, we can recognize (5) as solving ∇2f(xk)pk = ∇f(xk) in a leastsquares sense, where pk = XkΓk. When the quadratic approximation is exact, solving (5) is a minimal residual procedure, thus being verified as a residual projection method [47]. Moreover, let Hk be the solution to the constrained optimization problem [14]:
min Hk ‖Hk − βkI‖F subject to HkRk = −Xk, (7)
then iterate (6) is xk+1 = xk+Hkrk, which is indeed a multisecant quasi-Newton method. Note that a key simplification in AM is using differences of historical gradients Rk to approximate−∇2f(xk)Xk, which reduces the heavy cost to compute Hessian-vector products [7, 19, 26].

heading: 2.2 Stochastic Anderson Mixing

text:
We describe our method Stochastic Anderson Mixing (SAM) in this section. At the k-th iteration, let Sk ⊆ [T ] def = {1, 2, . . . , T} be the sampled mini-batch and the corresponding objective function
value is fSk (xk) = 1 |Sk| ∑ i∈Sk fξi(xk). Then rk def = −∇fSk (xk) and the noisy Rk is defined correspondingly (cf. (3)). Due to the instability and inaccurate estimation of Rk, we stabilize the projection step by proposing damped projection and adaptive regularization techniques. Algorithm 1 is a sketch of our method. Now, we elaborate the mechanism of this algorithm.
Damped projection. From Remark 1, we see the determination of Γk in (5) relies on the local quadratic approximation of (2), which can be rather inexact in general nonlinear optimization. To improve the stability, we propose a damped projection method for (4a). Let αk be the damping parameter, we obtain x̄k via
x̄k = (1− αk)xk + αk(xk −XkΓk) = xk − αkXkΓk. (8)
Combining (8) and (4b) and noting that r̄k = rk − αkRkΓk, the new iterate xk+1 is given by
xk+1 = xk + βkrk − αk (Xk + βkRk) Γk. (9)
It is worth noting that βk and αk in (9) behave like stepsize or learning rate in SGD, and the extra term (αkXk + αkβkRk) Γk can be viewed as a generalized momentum term.
Algorithm 1 Stochastic Anderson Mixing (SAM) Input: x0 ∈ Rd,m = 10, αk = 1, βk = 1, µ ∈ (0, 1),max_iter > 0. Output: x ∈ Rd
1: for k = 0, 1, . . . ,max_iter do 2: rk = −∇fSk (xk) 3: if k = 0 then 4: xk+1 = xk + βkrk 5: else 6: mk = min{m, k} 7: Xk = [∆xk−mk ,∆xk−mk+1, · · · ,∆xk−1] 8: Rk = [∆rk−mk ,∆rk−mk+1, · · · ,∆rk−1] 9: Check Condition (14) and use smaller αk if (14) is violated.
10: xk+1 = xk + βkrk − (αkXk + αkβkRk) ( RTkRk + δkX T k Xk )† RTk rk 11: end if 12: Apply learning rate schedule of αk, βk 13: end for 14: return xk
Adaptive regularization. Since Rk may be rank deficient and no safeguard method is used in AM, the least squares problem (5) can be unstable. A remedy is to add regularization [6] to (5). One well-known choice is the constant regularization introduced in [50, 53], which can be viewed as forcing ‖Γk‖2 to be small [53], leading to a penalty term to (5):
Γk = arg min Γ∈Rm
‖rk −RkΓ‖22 + δ‖Γ‖22, (10)
where δ ≥ 0 is the penalty constant. The solution of (10) is Γk = ( RTkRk + δI )† RTk rk, where “†" denotes the Penrose-Moore inverse. This regularized variant of AM is named as RAM and serves as a baseline in the experiments.
Here, we propose a new regularization, namely adaptive regularization, to better suit the stochastic optimization. Since −XkΓk = x̄k − xk denotes the update from xk to x̄k, a large magnitude of ‖XkΓk‖2 tends to make the intermediate step x̄k overshoot the trust region around xk. Thus it is more reasonable to force ‖XkΓk‖2 rather than ‖Γk‖2 to be small. We formulate this idea as
min Γ ‖rk −RkΓ‖22 + δk‖XkΓ‖22, (11)
where δk ≥ 0 is a variable determined in each iteration. Explicitly solving (11) leads to Γk = ( RTkRk + δkX T k Xk )† RTk rk. (12)
We call AM with this regularization and damped projection as SAM, i.e. the prototype algorithm given in Algorithm 1.
Positive definiteness. From (9) and (12), the SAM update is xk+1 = xk + Hkrk, where Hk = βkI − αkYkZ†kRTk , Yk = Xk + βkRk, Zk = RTkRk + δkXTk Xk. Hk is generally not symmetric. A critical condition for the convergence analysis of SAM is the positive definiteness of Hk, i.e.
pTkHkpk ≥ βkµ‖pk‖22, ∀pk ∈ Rd, (13) where µ ∈ (0, 1) is a constant. Next, we give an approach to guarantee it. Let λmin(·) denote the smallest eigenvalue, λmax(·) denote the largest eigenvalue. Since pTkHkpk = 1 2p T k (Hk +H T k )pk, Condition (13) is equivalent to λmin ( 1 2 ( Hk +H T k )) ≥ βkµ. With some simple
algebraic operations, we obtain λmin ( 1 2 ( Hk +H T k )) = βk − 12αkλmax(YkZ † kR T k + RkZ † kY T k ). Let λk def = λmax(YkZ † kR T k +RkZ † kY
T k ), then Condition (13) is equivalent to αkλk ≤ 2βk(1− µ). (14)
To check Condition (14), note that
λk = λmax ( (Yk Rk) ( 0 Z†k Z†k 0 )( Y Tk RTk )) = λmax (( Y Tk RTk ) (Yk Rk) ( 0 Z†k Z†k 0 )) . (15)
Algorithm 2 Stochastic Anderson Mixing with variance reduction (SAM-VR) Input: x̃0 ∈ Rd; βkt , αkt , δkt for SAM_update(xkt , gkt ); Batch size n ≥ 1 . Output: x ∈ Rd
1: for k = 0, . . . , N − 1 do 2: xk0 = x̃k 3: ∇f(x̃k) = 1T ∑T i=1∇fξi(x̃k) 4: for t = 0, . . . , q − 1 do 5: Sample a subset K ⊆ [T ] with |K| = n 6: gkt = ∇fK(xkt )−∇fK(x̃k) +∇f(x̃k) where∇fK(xkt ) = 1|K| ∑ i∈K∇fξi(xkt ) 7: xkt+1 = SAM_update(x k t , g k t ) 8: end for 9: Set x̃k+1 = xkq
10: end for 11: return Iterate x chosen uniformly random from {{xkt } q−1 t=0 } N−1 k=0
Since ( Y Tk RTk ) (Yk Rk) , ( 0 Z†k Z†k 0 ) ∈ R2m×2m, and m d, λk can be computed efficiently, say, using an eigenvalue decomposition algorithm with the time complexity of O(m3). This cost is negligible compared with those to form XTk Xk, R T kRk. After that, to guarantee the positive definiteness, we check if αk satisfies (14) and use a smaller αk if necessary.
AdaSAM method. We choose the δk in (11) as
δk = max
{ c1‖rk‖22
‖xk − xk−1‖22 + , c2β
−2 k
} , (16)
where c1, c2, are constants. Such choice reflects the curvature change in the vicinity of xk. A large ‖rk‖2 indicates a potential dramatic change in landscape, suggesting using a precautious tiny stepsize. The denominator in (16) behaves like annealing, which can measure the noise in gradients like that in secant penalized BFGS [27].

heading: 2.3 Enhancement of Stochastic Anderson Mixing

text:
We introduce the variance reduction and preconditioned mixing techniques to further enhance SAM.
Variance reduction. Variance reduction techniques turn out to be effective if a scan over the full dataset is feasible [1, 44]. Similar to SdLBFGS-VR proposed in [60], we also incorporate SVRG to SAM, which we call SAM-VR (Algorithm 2), for solving (2). To simplify the description, we denote one iteration of SAM in Algorithm 1 as SAM_update(xk, gk), i.e. one update of xk given the gradient estimate gk.
Preconditioned mixing. Motivated by the great success of preconditioning in solving linear systems and eigenvalue computation [18], we present a preconditioned version of SAM. The key modification is the mixing step (4b). We replace the simple mixing xk+1 = x̄k + βkr̄k with xk+1 = x̄k +M−1k r̄k where Mk approximates the Hessian. Combining it with (8) and (12), we obtain
xk+1 = xk + ( M−1k − αk ( Xk +M −1 k Rk ) ( RTkRk + δkX T k Xk )† RTk ) rk. (17)
Setting αk ≡ 1 and δk ≡ 0, (17) reduces to a preconditioned AM update, which can be recast as the solution to the constrained optimization problem: min
Hk ‖Hk −M−1k ‖F subject to HkRk = −Xk,
a direct extension of (7). This preconditioned version of AM is related to quasi-Newton updates [20]. We also point out that the action of M−1k can be implicitly done via an update of any optimizer at hand, i.e. xk+1 = optim(x̄k,−r̄k), where optim updates x̄k given the extragradient −r̄k. If Rk = −∇2f(xk)Xk, which is the case in deterministic quadratic optimization, the projection step in preconditioned AM is still a minimal residual procedure. Remark 2. The same as SdLBFGS [60], SAM needs another 2md space to store Xk and Rk. The extra main computational cost for SAM compared with SGD is O(m2d) +O(m3), which accounts for the matrix multiplications (Rm×d × Rd×m) and matrix decomposition of a small Rm×m matrix.
Since dense matrix multiplication can be ideally parallelized and the cost of gradient evaluations often dominates the computing, the benefit from SAM pays for this extra cost. Also, this cost can be further reduced to O(md) +O(m3) (see Appendix B.2). In our implementation, we incorporate sanity check of the positive definiteness, alternating iteration and moving average, and the details are given in the supplementary materials.

heading: 3 Theory

text:
In this section, we analyze the convergence and complexity of SAM. All the proofs are deferred to the Appendix A.1. We first give two assumptions about the objective function f that are widely used in stochastic programming [17, 60, 44]. Assumption 1. f : Rd → R is continuously differentiable. f(x) ≥ f low > −∞ for any x ∈ Rd. ∇f is globally L-Lipschitz continuous; namely ‖∇f(x)−∇f(y)‖2 ≤ L‖x−y‖2 for any x, y ∈ Rd. Assumption 2. For any iteration k, the stochastic gradient ∇fξk(xk) satisfies Eξk [∇fξk(xk)] = ∇f(xk), Eξk [‖∇fξk(xk)−∇f(xk)‖22] ≤ σ2, where σ > 0, and ξk, k = 0, 1, . . . , are independent samples that are independent of {xj}kj=0.
Throughout our analysis, we always assume both Assumptions 1 and 2 hold. Moreover, we state the following conditions about the three parameters αk, βk, δk of SAM.
+∞∑ k=0 βk = +∞, +∞∑ k=0 β2k < +∞, (18a)
δk ≥ Cβ−2k , αk ∈ [0,min{1, β 1 2
k }] and statisfies (14), (18b) where C > 0 is a constant.
Convergence and complexity. We give the convergence of SAM by analyzing the iterations {xk} generated by Algorithm 1. Theorem 1. Let {xk} be the sequence generated by Algorithm 1 with fixed batchsize n. Given a positive constant C, if βk ∈ (0, µ4L(1+C−1) ], then under the conditions (18a) and (18b), it has
lim inf k→∞
‖∇f(xk)‖2 = 0 with probability 1. (19)
Moreover, there exists a positive constant Mf such that
E[f(xk)] ≤Mf ∀k. (20) Theorem 2. Under the same conditions as Theorem 1, if we require that the noisy gradient is bounded, i.e., Eξk [‖∇fξk(xk)‖22] ≤Mg, (21) where Mg > 0 is a constant, we have
lim k→∞
‖∇f(xk)‖2 = 0 with probability 1. (22)
Now, we give the iteration complexity of SAM. Theorem 3. Let {xk} be the sequence generated by Algorithm 1 with batchsize n. Given constants C > 0 and r ∈ (0.5, 1), we choose βk = µ4L(1+C−1) (k + 1) −r. Under the condition (18b), it has
1
N N−1∑ k=0 E‖∇f(xk)‖22 ≤ 16L(1 + C−1)(Mf − f low) µ2 Nr−1 + (1 + L−1µ−1)σ2 (1− r)n (N−r − rN−1),
(23)
where n is the batchsize, andN is the iteration number. Moreover, for a given ∈ (0, 1), to guarantee that 1N ∑N−1 k=0 E‖∇f(xk)‖22 < , the number of iterations N needed is at most O( − 11−r ).
If the output xR is randomly chosen from previous iterates according to the certain probability distribution, SAM has the same complexity O( −2) as RSG [17] and SQN [60].
Theorem 4. Let {xk} be the sequence generated by Algorithm 1 with batchsize n and xR be the random output of the first N iterations where the random variable R follows the uniform distribution PR(k) def = Prob{R = k} = 1/N . Given constant C > 0, we choose βk = µ4L(1+C−1) and assume condition (18b) holds. For any > 0, there exist some n and N = O( −2)/n such that we have E [ ‖∇f(xR)‖22 ] ≤ , where the expectation is taken with respect to R and {Sj}N−1j=0 .
In fact, from the proof of Theorem 4, we can find the explicit choice of the batchsize. Define Df def = f(x0) − f low and D̃ is a problem-independent positive constant. The batchsize n can be
chosen as n = ⌈ min { N̄ ,max { 1, σL √ N̄ D̃ }}⌉ ,where N̄ is the total number of SFO-calls satisfying
N̄ ≥ max { C21 2 + 4C2 , σ2
L2D̃
} , and C1 = 32Df (1+C −1)σ
µ2 √ D̃
+ (L+ µ−1)σ √ D̃, C2 = 32DfL(1+C −1)
µ2 . Thus, to ensure E [ ‖∇f(xR)‖22 ] ≤ , the total number of SFO-calls is O( −2). Moreover, we give the SFO-calls complexity of the variance reduced version of SAM. Theorem 5. Let x be the output of Algorithm 2 with batchsize n. Given a positive constant C, suppose ν, µ0 ∈ (0, 1) to be two constants satisfying
µ0µ 2(1 + C−1)1/2 − µ
2 0(2L+ µ −1)(e− 1) L − 2µ20n− 2µ30(2L+ µ −1)(e− 1)n L ≥ ν,
where e is the Euler’s number. Set βkt = β := µ0n L(1+C−1)1/2T 2/3 , δkt ≥ Cβ−2, αkt ∈ [0,min{1, β 1 2 }] and satisfies (14), and the number of iteration q = ⌊
T µ0nd0 ⌋ , where d0 = 6 + 2L(1+C−1)1/2 . Then,
E [ ‖∇f(x)‖22 ] ≤ T 2/3L
( f(x0)− f low ) qNnν . (24)
To ensure E [ ‖∇f(x)‖22 ] ≤ , the total number of SFO-calls is O(T + (T 2/3/ )). Remark 3. It is worth noting that the approximated HessianHk in SAM depends on the data samples {ξi}i∈Sk of the current mini-batch, which violates the assumption AS.4 in [60].
Analysis of AdaSAM. All the results in Theorems 1-5 are applicable for AdaSAM as (16) satisfies the condition (18b). We further discuss the rationality of the first term in (16). Since rTkHkrk ≥ βkµ‖rk‖22, and ‖Hkrk‖22 ≤ 2 ( β2k + α 2 kδ −1 k ) ‖rk‖22 ≤ 2β2k ( 1 + C−1 ) ‖rk‖22 ifαk ≤ 1, δk ≥ Cβ−2k (see Corollary 1 in Appendix A.1), it is sensible to suppose ‖xk − xk−1‖2 ≈ ‖xk+1 − xk‖2 = ‖Hkrk‖2 = βkhk‖rk‖2, where hk ∈ (0, h) (h < +∞) is related to Hk. Therefore,
c1‖rk‖22 ‖xk − xk−1‖22 ≈ c1‖rk‖ 2 2 β2kh 2 k‖rk‖22 ≥ c1h−2β−2k , (25)
which coincides with the requirement that δk ≥ Cβ−2k . A further discussion is in Appendix D.3.

heading: 4 Experiments

text:
We used AdaSAM (AdaSAM-VR) for the experiments, which is a special case of SAM (SAM-VR) with δk specified as (16). The pseudocodes are given in Appendix B.1. To evaluate the effectiveness of our methods, we compared them with several first-order and second-order optimizers for mini-batch training of neural networks on different machine learning tasks. The datasets were MNIST [32], CIFAR-10/CIFAR-100 [31] for image classification and Penn Treebank [35] for language model. The experimental details and hyper-parameter tuning are referred to Appendix C.
Experiments on MNIST. We trained a simple convolutional neural network (CNN) 3 on MNIST, for which we were only concerned about the minimization of the empirical risk (2), i.e. the training loss, with large batch sizes. The training dataset was preprocessed by randomly selecting 12k images from the total 60k images to facilitate large mini-batch training. Neither weight-decay nor dropout was used. We compared AdaSAM with SGDM [43] (SGD with momentum), Adam [29], SdLBFGS [60], and RAM (cf. (10)). The learning rate was tuned and fixed for each optimizer. The historical lengths for SdLBFGS, RAM and AdaSAM were set as 20. δ = 10−6 for RAM and c1 = 10−4 for AdaSAM.
3Based on the official PyTorch implementation https://github.com/pytorch/examples/blob/master/mnist.
Figure 1 (a) and (b) show the curves of training loss when training 100 epochs with batch sizes of 6K and 3K, which indicate that AdaSAM can significantly minimize the empirical risk in large mini-batch training. The comparison with RAM verifies the benefit of adaptive regularization. We notice that there hardly exists any oscillation in AdaSAM during training except for the first few epochs, which suggests that AdaSAM is tolerant to noise. We also tested the effectiveness of variance reduction and preconditioning introduced in Section 2.3. The variance reduced extension AdaSAMVR was compared with AdaSAM for different batch sizes. The variants of AdaSAM preconditioned by Adagrad [12] and RMSprop [55] are denoted as Adagrad-AdaSAM and RMSprop-AdaSAM, respectively. Though AdaSAM-VR is more costly per iteration and the preconditioned variants seem to deteriorate the final training loss, we point out that AdaSAM-VR can achieve lower training loss (10−9) and the preconditioned variants converge faster to an acceptable training loss (e.g. 10−3).
Experiments on CIFAR. For CIFAR-10 and CIFAR100, both datasets have 50K images for training and 10K images for test. The test accuracy at the final epoch was reported as the evaluation metric. We trained ResNet18/20/32/44/56 [21] and WideResNet16-4 [64] on CIFAR-10, and ResNet18, ResNeXt50 [61] and DenseNet121 [24] on CIFAR-100. The baseline optimizers were SGDM, Adam, AdaBelief [66], Lookahead [65], AdaHessian [63] and RNA [52]. The hyperparameters were kept unchanged across different tests. We trained 160 epochs with batch size of 128 and decayed the learning rate at the 80th and 120th epoch. For AdaSAM/RAM, αk and βk were decayed at the 80th and 120th epoch.
Table 1 demonstrates the generalization ability of AdaSAM (AdaHessian ran out of memory for training ResNeXt50 and DenseNet121). It shows AdaSAM consistently outperforms other optimizers for various neural
networks on CIFAR10/CIFAR100. Compared with SGDM/Lookahead, AdaSAM is built on a noisy quadratic model to extrapolate historical iterates more elaborately, which may account for its effectiveness. We also conducted tests on training for 120 epochs and 80 epochs. Figure 2 shows AdaSAM
can achieve comparable or even lower test loss than SGDM/Lookahead when training with fewer epochs, thus saving a large number of iterations. More results can be found in Appendix C.3.
To compare the total training time, we set SGDM as the baseline and terminated the training of AdaSAM when the accuracy difference AdaSAM − SGDM > −0.05%. The memory, per-epoch computational cost and total training time of training ResNet20/WideResNet16-4 (abbr. WResNet) on CIFAR10 and ResNeXt50/DenseNet121 on CIFAR100 are reported in Table 2. It shows AdaSAM is more efficient than AdaHessian and saves more than 15% training time than SGDM in these models.
Experiments on Penn Treebank. We trained LSTM [22] on Penn Treebank and reported the perplexity on the validation set in Figure 3 and perplexity on the test set in Table 3, where pAdaSAM denotes the variant of AdaSAM preconditioned by Adam. The experimental setting was the same as that in AdaBelief [66]. In our practice, we found that the vanilla AdaSAM with default hyperparameter setting is not suitable for this task. Nevertheless, a suitable preconditioner (e.g. Adam) can largely
improve the behaviour of AdaSAM. Conversely, AdaSAM can also enhance a optimizer when the latter is used as a preconditioner.

heading: 5 Discussion

text:
In this section, we provide discussions that compare SAM with SQN [60] and SGD.
SAM versus SQN. Our theoretical results are similar to those of SQN [60], but there are some fundamental differences:
(i) The basic assumptions are different. (a) Recall that the update is xk+1 = xk − Hkgk, where gk = ∇fSk(xk). Let S[k−1] := (S0, . . . , Sk−1) denote the random samplings in the first k iterations. A key assumption in [60] is that Hk depends only on S[k−1], which guarantees E[Hkgk|S[k−1]] = Hk∇f(xk) that can simplify the proof. In contrast, in our methods SAM and SAM-VR, Hk depends on Sk since rk = −gk participates in constructing Hk, thus violating this assumption and complicating our analysis. (b) The proposed algorithm SdLBFGS in [60] needs twice continuously differentiable assumption of f and the boundedness of∇2f , which are not required in our algorithm.
(ii) The ways to construct the secant equations are different. In the deterministic case, a secant equation is Hksk = yk, where sk = xk − xk−1 and yk = gk − gk−1. In the stochastic case, various choices exist in forming yk. For SAM, the choice of yk is the same as that of the deterministic case, thus making Hk depend on Sk. For the SQN methods in [60], yk = ∇fSk−1(xk)−∇fSk−1(xk−1), which ensures Hk is independent of Sk. Such choice can reduce variance of noise in secant equations but demands evaluating gradients on the same mini-batch Sk−1 twice. Some other SQN methods [7] use subsampled Hessian-vector products to form yk = ∇2fS′k(xk − xk−1), where S ′ k denotes the data samples independent of Sk, which also ensures Hk is independent of Sk.
SAM versus SGD. We compare SAM with SGD from theoretical and practical perspectives.
(i) Theoretical perspective. (a) For general nonconvex stochastic optimization, SAM achieves the same worst-case complexity O(1/ 2) as SGD. (b) There are some specific cases that SAM can be theoretically better than SGD. For strongly convex quadratic minimization, we prove that SAM is essentially equivalent to GMRES (see Appendix A.2) that can exhibit superlinear convergence [58]. For nonlinear optimization, SAM can also have superlinear convergence rate in some cases as it is closely related to the multisecant quasi-Newton method [14], as also pointed out in Remark 1.
(ii) Practical perspective. (a) With the same number of epochs as SGD, it is shown in Table 1 and Table 3 that our method has better generalization performance than SGD in both image classification and language tasks. (b) Our method can largely reduce the number of epochs to achieve a comparable solution as SGD, as shown in Figure 2 in Section 4 and Figure 8 in Appendix C.3. The total running time can reduce correspondingly, as shown in Table 2. (c) SAM needs additional memory and computational cost, but it is still much more efficient and economical than other second-order optimizers, as shown in Table 2 in Section 4 and Figure 22 in Appendix D.8.

heading: 6 Conclusion

text:
In this paper, we develop an extension of Anderson mixing, namely Stochastic Anderson Mixing, for nonconvex stochastic optimization. By introducing damped projection and adaptive regularization, we establish the convergence theory of our new method. We also analyze its work complexity in terms of SFO-calls and show it can achieve the O(1/ 2) complexity for an -accurate solution. We also give a specific form of adaptive regularization. Then we propose two techniques to further enhance our method. One is the variance reduction technique, which can further improve the work complexity of our method theoretically and help achieve lower empirical risk in our experiments. The other one is the preconditioned mixing strategy that directly extends Anderson mixing. Experiments show encouraging results of our method and its enhanced versions in terms of convergence rate or generalization ability in training different neural networks in different machine learning tasks. These results confirm the suitability of Anderson mixing for nonconvex stochastic optimization.
Acknowledgments and Disclosure of Funding
This work was supported by the National Key R&D Program of China (No. 2018YFB1005103), Technology and Innovation Major Project (No. 2020AAA0108403), National Natural Science Foundation of China (No.61925601, No. 61772302, No. 11901338) and Huawei Noah’s Ark Lab. We thank all anonymous reviewers for their valuable comments and suggestions on this work.

title: Stochastic Anderson Mixing for Nonconvex Stochastic Optimization


summaries: This manuscript describes a novel method for sampling from the posterior in a linear inverse problem subjected to noise for an arbitrary measurement matrix. This is achieved by combining a black-box prior (such as that derived from a minimum MSE Gaussian denoiser) whose score function is easily computable with an annealed Langevin dynamics framework. The novelty is the construction of the Langevin noise term, which guarantees that the conditional score function is tractable. In addition, the authors propose a heuristic for step size calculation inspired by Newton's method in optimization. The method is evaluated on a variety of linear inverse problems: image deblurring, super resolution, and compressive sensing. For each task, realistic results are achieved and the residual is shown to conform to the imposed statistical model.

summaries: This paper proposes an algorithm called SNIPS that solves noisy linear inverse problems via an approximate MMSE estimator which makes use of pre-trained Gaussian denoisers. The estimator is based on MCMC sampling which is driven by a synthetically annealed version of Langevin dynamics. The motivations and background leading up to the algorithms derivation are well discussed. Experiments show interesting results with regard to the diversity and faithfulness of the images generated.   ---Edit after author responses---  Having read the other reviews and the responses posted by the authors, I am reducing my original scores from 6 --> 4.  My main issues are that the submission lacks (i) a comparison with an experimental baseline, and (ii) clarity in explaining the technical sections of the paper.

summaries: The authors address the problem of conditional sampling using Langevin dynamics. It is known that annealed Langevin dynamics speeds up convergence, and prior work that uses Langevin dynamics for conditional sampling does not correctly model the effects of annealing. Specifically, since there is some annealing noise added to the estimate, the likelihood distribution of the measurements given an estimate will change, and the authors derive approximations to the correct functional form. Experiments show that their new estimator can be implemented in practice.

summaries: This paper proposes a technique based on annealed Langevin dynamics to solve noisy linear inverse problems. The technique is built on top of [18, 19] and proposes to perform the Langevin dynamics in the spectral space of the linear degradation operator $H$. To improve the performance of the proposed method, the authors use coordinate wise step sizes obtained by approximating the hessian with its diagonal. Qualitative results from the proposed method are presented for compress sensing, inpainting, super-résolution and deblurring.

summaries: The paper describes an apparently new approach to solving noisy inverse problems.  It shows how the approach can be applied to three different noisy inverse problems: which it calls image deblurring, super-resolution, and compressive sensing.   It validates the approach on synthetic data which is synthesized (eg blurred and noise added) starting from publicly available imagery.  It is carefully written and I find it readable.

abstractText: In this work we introduce a novel stochastic algorithm dubbed SNIPS, which draws samples from the posterior distribution of any linear inverse problem, where the observation is assumed to be contaminated by additive white Gaussian noise. Our solution incorporates ideas from Langevin dynamics and Newton’s method, and exploits a pre-trained minimum mean squared error (MMSE) Gaussian denoiser. The proposed approach relies on an intricate derivation of the posterior score function that includes a singular value decomposition (SVD) of the degradation operator, in order to obtain a tractable iterative algorithm for the desired sampling. Due to its stochasticity, the algorithm can produce multiple high perceptual quality samples for the same noisy observation. We demonstrate the abilities of the proposed paradigm for image deblurring, super-resolution, and compressive sensing. We show that the samples produced are sharp, detailed and consistent with the given measurements, and their diversity exposes the inherent uncertainty in the inverse problem being solved.

authors:
S. Arridge
P. Maass
O. Öktem
C.-B. Schönlieb

title: Solving inverse problems using datadriven models

venue: Acta Numerica,

authors:
Y. Bahat
T. Michaeli

title: Explorable super resolution

venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,

title: Markov Chain Monte Carlo for statistical inference

venue: Center for Statistics and the Social Sciences,

authors:
S.A. Bigdeli
M. Jin
P. Favaro
M. Zwicker

title: Deep mean-shift priors for image restoration

authors:
Y. Blau
T. Michaeli

title: The perception-distortion tradeoff

venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,

authors:
A. Buades
B. Coll
J.-M. Morel

title: A non-local algorithm for image denoising

venue: IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR’05),

authors:
G.T. Buzzard
S.H. Chan
S. Sreehari
C.A. Bouman

title: Plug-and-play unplugged: Optimization-free reconstruction using consensus equilibrium

venue: SIAM Journal on Imaging Sciences,

authors:
R. D’Agostino
E.S. Pearson

title: Tests for departure from normality. Empirical results for the distributions of b2 and √

venue: b. Biometrika,

authors:
A. Danielyan
V. Katkovnik
K. Egiazarian

title: BM3D frames and variational image deblurring

venue: IEEE Transactions on Image Processing,

authors:
C. Dong
C.C. Loy
X. Tang

title: Accelerating the super-resolution convolutional neural network

venue: In European Conference on Computer Vision,

authors:
W. Dong
L. Zhang
G. Shi
X. Li

title: Nonlocally centralized sparse representation for image restoration

title: Tweedie’s formula and selection bias

venue: Journal of the American Statistical Association,

authors:
M. Elad
M. Aharon

title: Image denoising via sparse and redundant representations over learned dictionaries

authors:
I. Goodfellow
J. Pouget-Abadie
M. Mirza
B. Xu
D. Warde-Farley
S. Ozair
A. Courville
Y. Bengio

title: Generative adversarial nets

authors:
B. Guo
Y. Han
J. Wen

title: AGEM: solving linear inverse problems via deep priors and sampling

venue: Advances in Neural Information Processing Systems,

authors:
H. Gupta
K.H. Jin
H.Q. Nguyen
M.T. McCann
M. Unser

title: CNN-based projected gradient descent for consistent CT image reconstruction

venue: IEEE Transactions on Medical Imaging,

authors:
M. Haris
G. Shakhnarovich
N. Ukita

title: Deep back-projection networks for super-resolution

authors:
J. Ho
A. Jain
P. Abbeel

title: Denoising diffusion probabilistic models

authors:
C.M. Hyun
H.P. Kim
S.M. Lee
S. Lee
J.K. Seo

title: Deep learning for undersampled MRI reconstruction

venue: Physics in Medicine & Biology,

authors:
Z. Kadkhodaie
E.P. Simoncelli

title: Solving linear inverse problems using the prior implicit in a denoiser

venue: arXiv preprint arXiv:2007.13640,

authors:
B. Kawar
G. Vaksman
M. Elad

title: Stochastic image denoising by sampling from the posterior distribution

venue: arXiv preprint arXiv:2101.09552,

authors:
O. Kupyn
T. Martyniuk
J. Wu
Z. Wang

title: DeblurGAN-v2: deblurring (orders-of-magnitude) faster and better

venue: In Proceedings of the IEEE/CVF International Conference on Computer Vision,

authors:
R. Laumont
V. De Bortoli
A. Almansa
J. Delon
A. Durmus
M. Pereyra

title: Bayesian imaging using plug & play priors: When Langevin meets Tweedie

venue: arXiv preprint arXiv:2103.04715,

title: Probability: The Analysis of Data, volume 1, pages 104–107

venue: CreateSpace Independent Publishing Platform,

title: Non-local color image denoising with convolutional neural networks

authors:
H. Li
Y. Yang
M. Chang
H. Feng
Z. Xu
Q. Li
Y. Chen

title: SRDiff: single image superresolution with diffusion probabilistic models

venue: arXiv e-prints, pages arXiv–2104,

authors:
Z. Liu
P. Luo
X. Wang
X. Tang

title: Deep learning face attributes in the wild

venue: In Proceedings of the IEEE International Conference on Computer Vision,

authors:
A. Lucas
M. Iliadis
R. Molina
A.K. Katsaggelos

title: Using deep neural networks for inverse problems in imaging: beyond analytical methods

venue: IEEE Signal Processing Magazine,

authors:
M.T. McCann
K.H. Jin
M. Unser

title: Convolutional neural networks for inverse problems in imaging: A review

authors:
T. Meinhardt
M. Moller
C. Hazirbas
D. Cremers

title: Learning proximal operators: Using denoising networks for regularizing inverse imaging problems

authors:
S. Menon
A. Damian
S. Hu
N. Ravi
C. Rudin

title: PULSE: Self-supervised photo upsampling via latent space exploration of generative models

title: An empirical Bayes estimator of the mean of a normal population

venue: Bull. Inst. Internat. Statist.,

authors:
G. Ohayon
T. Adrai
G. Vaksman
M. Elad
P. Milanfar

title: High perceptual quality image denoising with a posterior sampling CGAN

venue: arXiv preprint arXiv:2103.04192,

authors:
S. Peng
K. Li

title: Generating unobserved alternatives: A case study through super-resolution and decompression

venue: arXiv preprint arXiv:2011.01926,

authors:
A. Radford
L. Metz
S. Chintala

title: Unsupervised representation learning with deep convolutional generative adversarial networks

venue: In 4th International Conference on Learning Representations,

authors:
I. Ram
M. Elad
I. Cohen

title: Image processing using smooth ordering of its patches

authors:
S. Ravishankar
J.C. Ye
J.A. Fessler

title: Image reconstruction: From sparsity to data-adaptive methods and machine learning

venue: Proceedings of the IEEE,

authors:
G.O. Roberts
R.L. Tweedie

title: Exponential convergence of Langevin distributions and their discrete approximations

authors:
Y. Romano
M. Elad
P. Milanfar

title: The little engine that could: Regularization by denoising (RED)

authors:
A. Rond
R. Giryes
M. Elad

title: Poisson inverse problems by the plug-and-play scheme

venue: Journal of Visual Communication and Image Representation,

authors:
E. Ryu
J. Liu
S. Wang
X. Chen
Z. Wang
W. Yin

title: Plug-and-play methods provably converge with properly trained denoisers

venue: In International Conference on Machine Learning,

authors:
C. Saharia
J. Ho
W. Chan
T. Salimans
D.J. Fleet
M. Norouzi

title: Image super-resolution via iterative refinement

venue: arXiv preprint arXiv:2104.07636,

authors:
U. Simsekli
R. Badeau
T. Cemgil
G. Richard

title: Stochastic Quasi-Newton Langevin Monte Carlo

authors:
Y. Song
S. Ermon

title: Generative modeling by estimating gradients of the data distribution

title: Improved techniques for training score-based generative models

authors:
Y. Song
J. Sohl-Dickstein
D.P. Kingma
A. Kumar
S. Ermon
B. Poole

title: Score-based generative modeling through stochastic differential equations

venue: In International Conference on Learning Representations,

title: Estimation of the mean of a multivariate normal distribution

venue: The annals of Statistics,

authors:
M. Suin
K. Purohit
A. Rajagopalan

title: Spatially-attentive patch-hierarchical network for adaptive motion deblurring

authors:
Y. Sun
B. Wohlberg
U.S. Kamilov

title: An online plug-and-play algorithm for regularized image reconstruction

venue: IEEE Transactions on Computational Imaging,

authors:
T. Tirer
R. Giryes

title: Image restoration by iterative denoising and backward projections

authors:
G. Vaksman
M. Zibulevsky
M. Elad

title: Patch ordering as a regularization for inverse problems in image processing

authors:
G. Vaksman
M. Elad
P. Milanfar

title: LIDIA: Lightweight learned image denoising with instance adaptation

venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops,

authors:
S.V. Venkatakrishnan
C.A. Bouman
B. Wohlberg

title: Plug-and-play priors for model based reconstruction

venue: IEEE Global Conference on Signal and Information Processing,

authors:
X. Wang
K. Yu
S. Wu
J. Gu
Y. Liu
C. Dong
Y. Qiao
C. Change Loy

title: ESRGAN: enhanced super-resolution generative adversarial networks

venue: In Proceedings of the European Conference on Computer Vision (ECCV) Workshops,

authors:
J. Yang
J. Wright
T.S. Huang
Y. Ma

title: Image super-resolution via sparse representation

venue: IEEE transactions on image processing,

authors:
F. Yu
A. Seff
Y. Zhang
S. Song
T. Funkhouser
J. Xiao

title: LSUN: construction of a large-scale image dataset using deep learning with humans in the loop

venue: arXiv preprint arXiv:1506.03365,

authors:
G. Yu
G. Sapiro
S. Mallat

title: Solving inverse problems with piecewise linear estimators: From Gaussian mixture models to structured sparsity

authors:
J. Zhang
B. Ghanem

title: ISTA-Net: interpretable optimization-inspired deep network for image compressive sensing

authors:
K. Zhang
W. Zuo
Y. Chen
D. Meng
L. Zhang

title: Beyond a Gaussian denoiser: Residual learning of deep CNN for image denoising

authors:
K. Zhang
W. Zuo
S. Gu
L. Zhang

title: Learning deep CNN denoiser prior for image restoration

authors:
K. Zhang
W. Zuo
L. Zhang

title: FFDNet: toward a fast and flexible solution for CNN-based image denoising

authors:
R. Zhang
P. Isola
A.A. Efros
E. Shechtman
O. Wang

title: The unreasonable effectiveness of deep features as a perceptual metric

authors:
D. Zoran
Y. Weiss

title: From learning models of natural image patches to whole image restoration

venue: In 2011 International Conference on Computer Vision,

heading: 1 Introduction

text:
Many problems in the field of image processing can be cast as noisy linear inverse problems. This family of tasks includes denoising, inpainting, deblurring, super resolution, compressive sensing, and many other image recovery problems. A general linear inverse problem is posed as
y = Hx + z, (1)
where we aim to recover a signal x from its measurement y, given through a linear degradation operator H and a contaminating noise, being additive, white and Gaussian, z ∼ N ( 0, σ20I ) . In this work we assume that both H and σ0 are known.
Over the years, many strategies, algorithms and underlying statistical models were developed for handling image restoration problems. A key ingredient in many of the classic attempts is the prior that aims to regularize the inversion process and lead to visually pleasing results. Among the various options explored, we mention sparsity-inspired techniques [13, 55, 11], local Gaussian-mixture modeling [57, 63], and methods relying on non-local self-similarity [6, 9, 36, 51]. More recently, and with the emergence of deep learning techniques, a direct design of the recovery path from y to an estimate of x took the lead, yielding state-of-the-art results in various linear inverse problems, such as denoising [25, 59, 61, 52], deblurring [22, 48], super resolution [10, 17, 54] and other tasks [29, 28, 19, 16, 37, 58].
Despite the evident success of the above techniques, many image restoration algorithms still have a critical shortcoming: In cases of severe degradation, most recovery algorithms tend to produce washed out reconstructions that lack details. Indeed, most image restoration techniques seek a reconstruction that minimizes the mean squared error between the restored image, x̂, and the unknown original one, x. When the degradation is acute and information is irreversibly lost, image reconstruction becomes a highly ill-posed problem, implying that many possible clean images could explain the given
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
measurements. The MMSE solution averages all these candidate solutions, being the conditional mean of the posterior of x given y, leading to an image with loss of fine details in the majority of practical cases. A recent work reported in [5] has shown that reconstruction algorithms necessarily suffer from a perception-distortion tradeoff, i.e., targeting a minimization of the error between x̂ and x (in any metric) is necessarily accompanied by a compromised perceptual quality. As a consequence, as long as we stick to the tendency to design recovery algorithms that aim for minimum MSE (or other distances), only a limited perceptual improvement can be expected.
When perceptual quality becomes our prime objective, the strategy for solving inverse problems must necessarily change. More specifically, the solution should concentrate on producing a sample (or many of them) from the posterior distribution p (x|y) instead of its conditional mean. Recently, two such approaches have been suggested – GAN-based and Langevin sampling. Generative Adversarial Networks (GANs) have shown impressive results in generating realistically looking images (e.g., [14, 35]). GANs can be utilized for solving inverse problems while producing high-quality images (see e.g. [2, 31, 34]). These solvers aim to produce a diverse set of output images that are consistent with the measurements, while also being aligned with the distribution of clean examples. A major disadvantage of GAN-based algorithms for inverse problems is their tendency (as practiced in [2, 31, 34]) to assume noiseless measurements, a condition seldom met in practice. An exception to this is the work reported in [33], which adapts a conditional GAN to become a stochastic denoiser.
The second approach for sampling from the posterior, and the one we shall be focusing on in this paper, is based on Langevin dynamics. This core iterative technique enables sampling from a given distribution by leveraging the availability of the score function – the gradient of the log of the probability density function [38, 3]. The work reported in [44, 20, 46] utilizes the annealed Langevin dynamics method, both for image synthesis and for solving noiseless inverse problems.1 Their synthesis algorithm relies on an MMSE Gaussian denoiser (given as a neural network) for approximating a gradually blurred score function. In their treatment of inverse problems, the conditional score remains tractable and manageable due to the noiseless measurements assumption.
The question addressed in this paper is the following: How can the above line of Langevin-based work be generalized for handling linear inverse problems, as in Equation 1, in which the measurements are noisy? A partial and limited answer to this question has already been given in [21] for the tasks of image denoising and inpainting. The present work generalizes these ([44, 20, 46, 21]) results, and introduces a systematic way for sampling from the posterior distribution of any given noisy linear inverse problem. As we carefully show, this extension is far from being trivial, due to two prime reasons: (i) The involvement of the degradation operator H, which poses a difficulty for establishing a relationship between the reconstructed image and the noisy observation; and (ii) The intricate connection between the measurements’ and the synthetic annealed Langevin noise. Our proposed remedy is a decorrelation of the measurements equation via a singular value decomposition (SVD) of the operator H, which decouples the dependencies between the measurements, enabling each to be addressed by an adapted iterative process. In addition, we define the annealing noise to be built as portions of the measurement noise itself, in a manner that facilitates a constructive derivation of the conditional score function.
Following earlier work [44, 20, 46, 21], our algorithm is initialized with a random noise image, gradually converging to the reconstructed result, while following the direction of the log-posterior gradient, estimated using an MMSE denoiser. Via a careful construction of the gradual annealing noise sequence, from very high values to low ones, the entries in the derived score switch mode. Those referring to non-zero singular values start by being purely dependent on the measurements, and then transition to incorporate prior information based on the denoiser. As for entries referring to zero singular values, their corresponding entries undergo a pure synthesis process based on the prior-only score function. Note that the denoiser blends values in the evolving sample, thus intermixing the influence of the gradient entries. Our derivations include an analytical expression for a positiondependent step size vector, drawing inspiration from Newton’s method in optimization. This stabilizes the algorithm and is shown to be essential for its success.
We refer hereafter to our algorithm as SNIPS (Solution of Noisy Inverse Problems Stochastically). Observe that as we target to sample from the posterior distribution p (x|y), different runs of SNIPS on the same input necessarily yield different results, all of which valid solutions to the given inverse
1The work reported in [18, 42, 26] and [15, 23] is also relevant to this discussion, but somewhat different. We shall specifically address these papers’ content and its relation to our work in section 2.
problem. This should not come as a surprise, as ill-posedness implies that there are multiple viable solutions for the same data, as has already been suggested in the context of super resolution [31, 2, 34]. We demonstrate SNIPS on image deblurring, single image super resolution, and compressive sensing, all of which contain non-negligible noise, and emphasize the high perceptual quality of the results, their diversity, and their relation to the MMSE estimate.
To summarize, this paper’s contributions are threefold:
• We present an intricate derivation of the blurred posterior score function for general noisy inverse problems, where both the measurement and the target image contain delicately inter-connected additive white Gaussian noise.
• We introduce a novel stochastic algorithm – SNIPS – that can sample from the posterior distribution of these problems. The algorithm relies on the availability of an MMSE denoiser.
• We demonstrate impressive results of SNIPS on image deblurring, single image super resolution, and compressive sensing, all of which are highly noisy and ill-posed.
Before diving into the details of this work, we should mention that using Gaussian denoisers iteratively for handling general linear inverse problems has been already proposed in the context of the Plugand-Play-Prior (PnP) method [53] and RED [39], and their many followup papers (e.g., [60, 30, 1, 49, 7, 50, 40, 4]). However, both PnP and RED are quite different from our work, as they do not target sampling from the posterior, but rather focus on MAP or MMSE estimation.

heading: 2 Background

text:
The Langevin dynamics algorithm [3, 38] suggests sampling from a probability distribution p (x) using the iterative transition rule
xt+1 = xt + α∇xt log p (xt) + √ 2αzt , (2)
where zt ∼ N (0, I) and α is an appropriately chosen small constant. The added zt allows for stochastic sampling, avoiding a collapse to a maximum of the distribution. Initialized randomly, after a sufficiently large number of iterations, and under some mild conditions, this process converges to a sample from the desired distribution p (x) [38].
The work reported in [44] extends the aforementioned algorithm into annealed Langevin dynamics. The annealing proposed replaces the score function in Equation 2 with a blurred version of it, ∇x̃t log p (x̃t), where x̃t = xt + n and n ∼ N ( 0, σ2I ) is a synthetically injected noise. The core idea is to start with a very high noise level σ and gradually drop it to near-zero, all while using a step size α dependent on the noise level. These changes allow the algorithm to converge much faster and perform better, because it widens the basin of attraction of the sampling process. The work in [20] further develops this line of work by leveraging a brilliant relation attributed to Miyasawa [32] (also known as Stein’s integration by parts trick [47] or Tweedie’s identity [12]). It is given as
∇x̃t log p (x̃t) = D (x̃t, σ)− x̃t
σ2 , (3)
where D (x̃t, σ) = E [x|x̃t] is the minimizer of the MSE measure E [ ‖x−D (x̃t, σ) ‖22 ] , which can be approximated using a denoising neural network. This facilitates the use of denoisers in Langevin dynamics as a replacement for the evasive score function.
When turning to solve inverse problems, previous work suggests sampling from the posterior distribution p (x|y) using annealed Langevin dynamics [20, 46, 21] or similar methods [15, 18, 42, 26], by replacing the score function used in the generation algorithm with a conditional one. As it turns out, if limiting assumptions can be posed on the measurements formation, the conditional score is tractable, and thus generalization of the annealed Langevin process to these problems is within reach. Indeed, in [44, 20, 46, 42, 26] the core assumption is y = Hx for specific and simplified choices of H and with no noise in the measurements. The works in [15, 23] avoid these difficulties altogether by returning to the original (non-annealed) Langevin method, with the unavoidable cost of becoming extremely slow. In addition, their algorithms are demonstrated on inverse problems in which the additive noise is restricted to be very weak. The work in [21] is broader, allowing for an arbitrary additive white Gaussian noise, but limits H to the problems of denoising or inpainting. While all these works demonstrate high quality results, there is currently no clear way for deriving the blurred score function of a general linear inverse problem as posed in Equation 1. In the following, we present such a derivation.

heading: 3 The Proposed Approach: Deriving the Conditional Score Function

heading: 3.1 Problem Setting

text:
We consider the problem of recovering a signal x ∈ RN (where x ∼ p (x) and p (x) is unknown) from the observation y = Hx + z, where y ∈ RM ,H ∈ RM×N ,M ≤ N, z ∼ N ( 0, σ20I ) , and H and σ0 are known.2 Our ultimate goal is to sample from the posterior p (x|y). However, since access to the score function∇x log p(x|y) is not available, we retarget our goal, as explained above, to sampling from blurred posterior distributions, p (x̃|y), where x̃ = x + n and n ∼ N ( 0, σ2I ) , with noise levels σ starting very high, and decreasing towards near-zero.
As explained in the supplemental material, the sampling should be performed in the SVD domain in order to get a tractable derivation of the blurred score function. Thus, we consider the singular value decomposition (SVD) of H, given as H = UΣVT , where U ∈ RM×M and V ∈ RN×N are orthogonal matrices, and Σ ∈ RM×N is a rectangular diagonal matrix containing the singular values of H, denoted as {sj}Mj=1 in descending order (s1 > s2 > · · · > sM−1 > sM ≥ 0). For convenience of notations, we also define sj = 0 for j =M + 1, . . . , N . To that end, we notice that
p (x̃|y) = p ( x̃|UTy ) = p ( VT x̃|UTy ) . (4)
The first equality holds because the multiplication of y by the orthogonal matrix UT does not add or remove information, and the second equality holds because the multiplication of x̃ by VT does not change its probability [24]. Therefore, sampling from p ( VT x̃|UTy ) and then multiplying the result by V will produce the desired sample from p (x̃|y). As we are using Langevin dynamics, we need to calculate the conditional score function∇VT x̃ log p ( VT x̃|UTy ) . For simplicity, we denote hereafter yT = UTy, zT = UT z, xT = VTx, nT = ΣVTn, and x̃T = VT x̃. Observe that with these notations, the measurements equation becomes
y = Hx + z = UΣVTx + z,
and thus
UTy = ΣVTx + UT z = ΣVT (x̃− n) + UT z = ΣVT x̃−ΣVTn + UT z, where we have relied on the relation x̃ = x + n. This leads to
yT = Σx̃T − nT + zT . (5)
In this formulation, which will aid in deriving the conditional score, our aim is to make design choices on nT such that zT − nT has uncorrelated entries and is independent of x̃T . This brings us to the formation of the synthetic annealed noise, which is an intricate ingredient in our derivations.
2We assume M ≤ N for ease of notations, and because this is the common case. However, the proposed approach and all our derivations work just as well for M > N .
We base this formation on the definition of a sequence of noise levels {σi}L+1i=1 such that σ1 > σ2 > · · · > σL > σL+1 = 0, where σ1 is high (possibly σ1 > ‖x‖∞) and σL is close to zero. We require that for every j such that sj 6= 0, there exists ij such that σijsj < σ0 and σij−1sj > σ0. This implies ∀i : σisj 6= σ0, which helps ease notations. SNIPS works just as well for σisj = σ0.
Using {σi}L+1i=1 , we would like to define {x̃i} L+1 i=1 , a sequence of noisy versions of x, where the noise level in x̃i is σi. One might be tempted to define these noise additions as independent of the measurement noise z. However, this option leads to a conditional score term that cannot be calculated analytically, as explained in the supplemental material. Therefore, we define these noise additions differently, as carved from z in a gradual fashion. To that end, we define x̃L+1 = x, and for every i = L,L− 1, . . . , 1: x̃i = x̃i+1 + ηi, where ηi ∼ N ( 0, ( σ2i − σ2i+1 ) I ) . This results in
x̃i = x + ni, where ni = ∑L k=i ηk ∼ N ( 0, σ2i I ) .
And now we turn to define the statistical dependencies between the measurements’ noise z and the artificial noise vectors ηi. Since ηi and z are each Gaussian with uncorrelated entries, so are the components of the vectors ΣVTηi, ΣV
Tni, and zT . In order to proceed while easing notations, let us focus on a single entry j in these three vectors, for which sj > 0, and omit this index. We denote these entries as ηT,i, nT,i and zT , respectively. We construct ηT,i such that
E [ηT,i · zT ] =  E [ η2T,i ] for i ≥ ij E [( zT − nT,ij )2] for i = ij − 1
0 otherwise.
This implies that the layers of noise ηT,L+1, . . . ηT,ij are all portions of zT itself, with an additional portion being contained in ηT,ij−1. Afterwards, ηT,i become independent of zT . In the case of sj = 0, the above relations simplify to be E[ηT,i · zT ] = 0 for all i, implying no statistical dependency between the given and the synthetic noises. Consequently, it can be shown that the overall noise in Equation 5 satisfies
( ΣVTni − zT ) j = nT,i − zT ∼
{ N ( 0, s2jσ 2 i − σ20 ) if σisj > σ0
N ( 0, σ20 − s2jσ2i ) otherwise.
(6)
The top option refers to high values of the annealed Langevin noise, in which, despite the possible decay caused by the singular value sj , this noise is stronger than zT . In this case, nT,i contains all zT and an additional independent portion of noise. The bottom part assumes that the annealed noise (with the influence of sj) is weaker than the measurements’ noise, and then it is fully immersed within zT , with the difference being Gaussian and independent.

heading: 3.2 Derivation of the Conditional Score Function

text:
The above derivations show that the noise in Equation 5 is zero-mean, Gaussian with uncorrelated entries and of known variance, and this noise is independent of x̃i. Thus Equation 5 can be used conveniently for deriving the measurements part of the conditional score function. We denote x̃T = V
T x̃i, x̃ = x̃i, n = ni for simplicity, and turn to calculate ∇x̃T log p (x̃T |yT ). We split x̃T into three parts: (i) x̃T,0 refers to the entries j for which sj = 0; (ii) x̃T,< corresponds to the entries j for which 0 < σisj < σ0; and (iii) x̃T,> includes the entries j for which σisj > σ0. Observe that this partition of the entries of x̃T is non-overlapping and fully covering. Similarly, we partition every vector v ∈ RN into v0,v<,v>, which are the entries of v corresponding to x̃T,0, x̃T,<, x̃T,>, respectively. Furthermore, we define v 60,v 6<,v6> as all the entries of v except v0,v<,v>, respectively. With these definitions in place, the complete derivation of the score function is detailed in the supplemental material, and here we bring the final outcome. For x̃T,0, the score is independent of the measurements and given by
∇x̃T,0 log p (x̃T |yT ) = ( VT∇x̃ log p (x̃) ) 0 . (7)
For the case of x̃T,>, the expression obtained is only measurements-dependent, ∇x̃T,> log p (x̃T |yT ) = ( ΣT ( σ2iΣΣ T − σ20I )† (yT −Σx̃T ) ) > . (8)
Lastly, for the case of x̃T,<, the conditional score includes two terms – one referring to the plain (blurred) score, and the other depending on the measurements,
∇x̃T,< log p (x̃T |yT ) = ( ΣT ( σ20I− σ2iΣΣ T )† (yT −Σx̃T ) ) < + ( VT∇x̃ log p (x̃) ) < . (9)
As already mentioned, the full derivations of equations 7, 8, and 9 are detailed in the supplemental material. Aggregating all these results together, we obtain the following conditional score function:
∇x̃T log p (x̃T |yT ) = ΣT ∣∣∣σ20I− σ2iΣΣT ∣∣∣† (yT −Σx̃T ) + (VT∇x̃ log p (x̃))∣∣ 6> , (10)
where (v)| 6> is the vector v, but with zeros in its entries that correspond to v>. Observe that the first term in Equation 10 contains zeros in the entries corresponding to x̃T,0, matching the above calculations. The vector ∇x̃ log p (x̃) can be estimated using a neural network as in [44], or using a pre-trained MMSE denoiser as in [20, 21]. All the other elements of this vector are given or can be easily obtained from H by calculating its SVD decomposition once at the beginning.

heading: 4 The Proposed Algorithm

text:
Armed with the conditional score function in Equation 10, the Langevin dynamics algorithm can be run with a constant step size or an annealed step size as in [44], and this should converge to a sample from p (x̃T |yT ). However, for this to perform well, one should use a very small step size, implying a devastatingly slow convergence behavior. This is mainly due to the fact that different entries of x̃T advance at different speeds, in accord with their corresponding singular values. As the added noise in each step has the same variance in every entry, this leads to an unbalanced signal-to-noise ratio, which considerably slows down the algorithm.
In order to mitigate this problem, we suggest using a step size vector αi ∈ RN . We denote Ai = diag (αi), and obtain the following update formula for a Langevin dynamics algorithm:
VT x̃i = V T x̃i−1 + c ·Ai · ∇VT x̃i log p ( VT x̃i|yT ) + √ 2 · cA 1 2 i · zi, (11)
where the conditional score function is estimated as described in subsection 3.2, and c is some constant. For the choice of the step sizes in the diagonal of Ai, we draw inspiration from Newton’s method in optimization, which is designed to speed up convergence to local maximum points. The update formula in Newton’s method is the same as Equation 11, but without the additional noise zi, and with Ai being the negative inverse Hessian of log p ( VT x̃i|yT ) . We calculate a diagonal approximation of the Hessian, and set Ai to be its negative inverse. We also estimate the conditional score function using Equation 10 and a neural network. Note that this mixture of Langevin dynamics and Newton’s method has been suggested in a slightly different context in [43], where the Hessian was approximated using a Quasi-Newton method. In our case, we analytically calculate a diagonal approximation of the negative inverse Hessian and obtain the following:
(αi)j =  σ2i , sj = 0 σ2i − σ20 s2j , σisj > σ0 σ2i · ( 1− s2j σ2i σ20 ) , 0 < σisj < σ0.
(12)
The full derivations for each of the three cases are detailed in the supplemental material. Using these step sizes, the update formula in Equation 11, the conditional score function in Equation 10, and a neural network s (x̃, σ) that estimates the score function ∇x̃ log p (x̃),3 we obtain a tractable iterative algorithm for sampling from p (x̃L | y), where the noise in x̃L is sufficiently negligible to be considered as a sampling from the ideal image manifold.
Algorithm 1: SNIPS
Input: {σi}Li=1, c, τ , y, H, σ0 1 U,Σ,V← svd (H) 2 Initialize x0 with random noise U [0, 1] 3 for i← 1 to L do 4 (Ai)0 ← σ2i I 5 (Ai)< ← σ2i · ( I− σ 2 i
σ20 Σ<Σ<
T )
6 (Ai)> ← σ2i I− σ20Σ † >Σ †T > 7 for t← 1 to τ do 8 Draw zt ∼ N (0, I)
9 dt ← ΣT · ∣∣∣σ20I− σ2iΣΣT ∣∣∣† · (UTy −ΣVTxt−1)+ (VT · s (xt−1, σi))∣∣ 6>
10 xt ← V · ( VTxt−1 + cAidt + √ 2cA 1 2 i zt ) 11 end 12 x0 ← xτ 13 end
Output: x0
Note that when we set H = 0 and σ0 = 0, implying no measurements, the above algorithm degenerates to an image synthesis, exactly as in [44]. Two other special cases of this algorithm are obtained for H = I or H = I with some rows removed, the first referring to denoising and the second to noisy inpainting, both cases shown in [21]. Lastly, for the choices of H as in [20] or [44, 46] and with σ0 = 0, the above algorithm collapses to a close variant of their proposed iterative methods.

heading: 5 Experimental Results

text:
In our experiments we use the NCSNv2 [45] network in order to estimate the score function of the prior distribution. Three different NCSNv2 models are used, each trained separately on the training sets of: (i) images of size 64× 64 pixels from the CelebA dataset [27]; (ii) images of size 128× 128 pixels from LSUN [56] bedrooms dataset; and (iii) LSUN 128× 128 images of towers. We demonstrate SNIPS’ capabilities on the respective test sets for image deblurring, super resolution, and compressive sensing. In each of the experiments, we run our algorithm 8 times, producing 8 samples for each input. We examine both the samples themselves and their mean, which serves as an approximation of the MMSE solution, E [x|y].
3Recall that s (x̃, σ) = (D (x̃, σ)− x̃) /σ2, being a denoising residual.
For image deblurring, we use a uniform 5 × 5 blur kernel, and an additive white Gaussian noise with σ0 = 0.1 (referring to pixel values in the range [0, 1]). Figure 1 demonstrates the obtained results for several images taken from the CelebA dataset. As can be seen, SNIPS produces visually pleasing, diverse samples.
For super resolution, the images are downscaled using a block averaging filter, i.e., each nonoverlapping block of pixels in the original image is averaged into one pixel in the low-resolution image. We use blocks of size 2× 2 or 4× 4 pixels, and assume the low-resolution image to include an additive white Gaussian noise. We showcase results on LSUN and CelebA in Figures 2, 4, and 5.
For compressive sensing, we use three random projection matrices with singular values of 1, that compress the image by 25%, 12.5%, and 6.25%. As can be seen in Figure 3 and as expected, the more aggressive the compression, the more significant are the variations in reconstruction.
We calculate the average PSNR (peak signal-to-noise ratio) of each of the 8 samples in our experiments, as well as the PSNR of their mean, as shown in Table 1. In all the experiments, the empirical conditional mean presents an improvement of around 2.4 dB in PSNR, even though it is less visually appealing compared to the samples. This is consistent with the theory in [5], which states that the difference in PSNR between posterior samples and the conditional mean (the MMSE estimator) should be 3 dB, with the MMSE estimator having poorer perceptual quality but better PSNR.
A comparison of our deblurring results to those obtained by RED [39] is detailed in the supplemental material. We show that SNIPS exhibits superior performance over RED, achieving more than 11% improvement in PSNR and more than 58% improvement in LPIPS [62], a perceptual quality metric.

heading: 5.1 Assessing Faithfulness to the Measurements

text:
A valid solution to an inverse problem should satisfy two conditions: (i) It should be visually pleasing, consistent with the underlying prior distribution of images, and (ii) It should be faithful to the given measurement, maintaining the relationship as given in the problem setting. Since the prior distribution
is unknown, we assess the first condition by visually observing the obtained solutions and their tendency to look realistic. As for the second condition, we perform the following computation: We degrade the obtained reconstruction x̂ by H, and calculate its difference from the given measurement y, obtaining y −Hx̂. According to the problem setting, this difference should be an additive white Gaussian noise vector with a standard deviation of σ0. We examine this difference by calculating its empirical standard deviation, and performing the Pearson-D’Agostino [8] test of normality on it, accepting it as a Gaussian vector if the obtained p-value is greater than 0.05. We also calculate the Pearson correlation coefficient (denoted as ρ) among neighboring entries, accepting them as uncorrelated for coefficients smaller than 0.1 in absolute value. In all of our tests, the standard deviation matches σ0 almost exactly, the Pearson correlation coefficient satisfies |ρ| < 0.1, and we obtain p-values greater than 0.05 in around 95% of the samples (across all experiments). These results empirically show that our algorithm produces valid solutions to the given inverse problems.

heading: 6 Conclusion and Future Work

text:
SNIPS, presented in this paper, is a novel stochastic algorithm for solving general noisy linear inverse problems. This method is based on annealed Langevin dynamics and Newton’s method, and relies on the availability of a pre-trained Gaussian MMSE denoiser. SNIPS produces a random variety of high quality samples from the posterior distribution of the unknown given the measurements, while guaranteeing their validity with respect to the given data. This algorithm’s derivation includes an intricate choice of the injected annealed noise in the Langevin update equations, and an SVD decomposition of the degradation operator for decoupling the measurements’ dependencies. We demonstrate SNIPS’ success on image deblurring, super resolution, and compressive sensing.
Extensions of this work should focus on SNIPS’ limitations: (i) The need to deploy SVD decomposition of the degradation matrix requires a considerable amount of memory and computations, and hinders the algorithm’s scalability; (ii) The current version of SNIPS does not handle general content images, a fact that is related to the properties of the denoiser being used [41]; and (iii) SNIPS, as any other Langevin based method, requires (too) many iterations (e.g., in our super-resolution tests on CelebA, 2 minutes are required for producing 8 sample images), and means for its acceleration should be explored.

heading: 7 Funding Transparency Statement

text: This research was partially supported by the Israel Science Foundation (ISF) under Grant 335/18 and the Technion Hiroshi Fujiwara Cyber Security Research Center and the Israel Cyber Bureau. Bahjat Kawar’s scholarship was partially provided by Li Ka Shing Fellowships and the Planning and Budgeting Committee of the Israel Council for Higher Education.

title: SNIPS: Solving Noisy Inverse Problems Stochastically