Accelerating Optimization using Neural Reparametrization

1 INTRODUCTION . Gradient-based optimization lies at the heart of machine learning . For non-linear , non-convex largescale optimization problems , gradient-based optimization can suffer from slow convergence and significant computational bottleneck . In linear systems , preconditioning Axelsson ( 1996 ) ; Saad & Van Der Vorst ( 2000 ) can accelerating convergence by multiplying a symmetric positive-definite preconditioner matrix to the original problem . In effect , preconditioning re-scales the loss landscape to have more circle-like level sets , hence changing the condition number of the system . Adaptive gradient optimizers such as AdaGrad Duchi et al . ( 2011 ) and Adam Kingma & Ba ( 2014 ) uses the same idea but updates the preconditioner during iterative optimization with adaptive step size , which leads to great improvement in convergence speed . Our goal is to improve the convergence of non-convex optimization further . First , we examine the dynamics of the gradient flow and identify the condition number that dictates the speed of convergence . We discover that changing the condition number naturally requires the reparametrization to be a type of graph neural networks ( GNN ) . Based on the analysis , we introduce a novel neural reparametrization scheme , which generalizes preconditioning to non-linear systems . Note that our method is complementary to existing gradient-based optimization algorithms and can be used in conjunction with any adaptive gradient optimizer ( e.g . Adam ) . We test our hypothesis on two highly nonlinear optimization problems and obtain impressive speedups . A few earlier work Sosnovik & Oseledets ( 2019 ) ; Hoyer et al . ( 2019 ) have demonstrated that convolutional neural networks can serve as priors to improve the parameterization of a class of structural optimization problems . But the theoretical foundation behind these techniques are not well understood . Concurrently , the implicit acceleration of over-parametrization in linear neural networks have been analyzed in Arora et al . ( 2018 ) ; Tarmoun et al . ( 2021 ) . Specifically , Arora et al . ( 2018 ) shows that reparametrizing the linear weights with deep linear networks impose a pre-conditioning scheme on gradient descent . Tarmoun et al . ( 2021 ) further exploits the conservation laws in gradient flow and derive convergence rate for two-layer linear networks . In contrast , our method applies to highly nonlinear problems , guided by theoretical insights . Our contributions can be summarized as follows 1 . We analyze the dynamics of the gradient flow and derive how a neural reparametrization ( NR ) can lead to a speedup . 2 . We prove that maximum speedup is achieved if the NR is a special GCN which uses gradient of the loss in its aggregation function . 3 . For early steps , we show that the NR should use the Hessian of the loss , leading to an efficient hybrid implementation . 4 . We showcase NP on two optimization tasks : network synchronization and persistent homology and observe consistent speedup . Our method also demonstrates a novel usage and unique perspective for GNN , specifically GCN ( Kipf & Welling , 2016 ) . The majority of literature treat GNNs as a representation learning tool for graph-structure data ( see surveys and the references in Bronstein et al . ( 2017 ) ; Zhang et al . ( 2018 ) ; Wu et al . ( 2019 ) ; Goyal & Ferrara ( 2018 ) ) . However , our analysis shows that GNN can be used to modify optimization dynamics in problems without explicit graph data . 2 METHODOLOGY . We consider the problem of minimizing a smooth , lower-bounded loss function L ( w ) ∈ R on abounded variable w ∈ Rn : argminwL ( w ) When the optimization variables are matrices or tensors , w represents a flattened vector containing all variables . w can also represent all trainable parameters in a deep neural network . For supervised learning with a dataset Z = { ( xi , yi ) Ni=1 } , the loss function L ( w ) becomes L ( w ; Z ) . The dataset and the model class ( e.g . architecture of a neural network ) together define a loss landscape as a function of the optimization variables w. Gradient-based optimization updates the variables w by taking repeated steps in the direction of the steepest descent wt+1 = wt−ε ∂L∂wt for a learning rate ε . With very small time steps δt , we examine the dynamics of w ( t ) ∈ Rn during the optimization process , treating t as a continuous time . Using steepest descent for the optimization , dw/dt follows the gradient flow : dwi dt = −εii ∂L ∂wi , ( 1 ) where denote the learning rate ε as a diagonal matrix of learning rates , because modern optimizers such as RMSProp ( Tieleman et al. , 2012 ) or Adam ( Kingma & Ba , 2014 ) choose ε adaptively for each wi . Specifically , in Adam and RMSProp , which we use in experiments , we have εij ( ∇L , t ) ≈ η√ Et [ ( ∇iL ) 2 ] + ξ δij ( 2 ) where δij is the Kronecker delta , and η and ξ are small numbers . The expected value E is calculated using stochastic samples for ∇iL over multiple iteration steps , with more weights given to more recent steps . Equation 1 is also an ordinary differential equation ( ODE ) . The local minima of L ( w ) correspond to the stationary solutions of the ODE , i.e . dw/dt = 0 . Convergence rate We are interested in the expected convergence rate E [ dL/dt ] of the loss function . This expectation can be either approximated in a fashion similar to adaptive gradient optimizers ( i.e . exponential time average ) or be over random initializations under a distribution P ( w ( 0 ) ) . We will assume the expectation value for ε and E [ dL/dt ] are evaluated with the same method , meaning we assume E [ ε ] = ε . The convergence rate is then given by E [ dL ( w ) dt ] = E [ ∑ i ∂L ∂wi dwi dt ] = −Tr [ εM ] , M ij ≡ E [ ∂L ∂wi ∂L ∂wj ] ( 3 ) where we used our assumption to get E [ εM ] = εE [ M ] , with Mij = ∇iL∇jL being the “ square gradients ” matrix and M = E [ M ] . With this definition , equation 2 becomes ε = η/ √ diag [ M ] + ξ . Note that M is time-dependent and can change during optimization . Estimating M If we know the distribution P ( w ( t ) ) for some t , we can approximate M . One way to do this is similar to Adam as a weighted average over the past few steps . Also , if we know the distribution of the initialization P ( w ( 0 ) ) , we can evaluate M for early stages , which is what we will use in the current paper . When using Xavier initialization ( Glorot & Bengio , 2010 ) , we have w ( 0 ) ∼ N ( 0 , n−1/2 ) . In this case we can show ( see appendix A.1 ) M ij ( t→ 0 ) = ∑ k , l E [ wkwl ] ∂2L ∂wk∂wi ∂2L ∂wl∂wj +O ( n−2 ) = 1 n [ H2 ] ij ∣∣∣ w→0 +O ( n−2 ) , ( 4 ) whereHij ( w ) ≡ ∂2L ( w ) /∂wi∂wj is the Hessian of the loss at t = 0 . The eigenvectors of M with zero eigenvalue are modes that do not evolve during optimization . As gradients are zero in these directions , GD can never find and explore those directions . Therefore , we may exclude them from the beginning and claim that M is full rank . The rate dwi/dt can be different along different directions of the parameter space , with fast and slow modes . Next , we quantify this difference in rates and suggest ways for acceleration . 2.1 FAST AND SLOW DYNAMICS OF GRADIENT FLOW . M is positive semi-definite because for any vector v ∈ Rn we have vTMv = E [ ( vT∇L ) 2 ] ≥ 0 . Thus , M is Hermitian with a spectral expansion M = ∑ imiψiψ T i . Using equation 1 , we can show that the dynamics of w along different ψi are orthogonal to each other E [ ( ψTi dw dt ) ( ψTj dw dt ) ] = ε2ψTi Mψj = εiεjmiδij . ( 5 ) If ε is constant and not adaptive , w evolves faster along modes ψi with larger eigenvalues mi , since if mi > mj we have ε2mi = E [ ( ψTi dw dt ) 2 ] > E [ ( ψTj dw dt ) 2 ] = ε2mj ( 6 ) This is precisely what optimizers such as Adam try to address . Ideally , having εi = η/ √ mi would make up for the rate difference , resulting in a uniform GD where no mode ψi is evolving more slowly than others . However , calculating mi is O ( n3 ) and can be quite expensive during optimization , hence the approximate version equation 2 is used . We show here that in many cases it is worthwhile to do a full correction to GD using M , at least in early stages of the optimization . The eigenvalues of M determine the rate of evolution of the overlaps along each of its eigenvectors . In particular , if some eigenvalue mslow meani [ mi ] is much smaller than the mean , the evolution of w along ψslow will be much slower than other directions . Therefore , we will refer to all ψslow as the slow modes of the optimization problem . Conversely , the ψfast where mfast meani [ mi ] will be referred to as the fast modes . When running GD , the maximum change in w is bounded to ensure numerical stability . Because of the orthogonality in equation 5 , we can enforce a numerical bound η via max i E [ ( ψTi dw dt ) 2 ] ≤ η2 ⇒ εmax ≤ η √ mmax , E [ ( ψTi dw dt ) 2 ] ≤ η2 mi mmax . ( 7 ) where mmax = maximi is the largest eigenvalue of M . Therefore , the learning rate ε is bounded by the largest eigenvalues ( fastest modes ) . Thus , in order to speed up convergence , we need to focus on the slow modes , which can be achieve with a reparametrization of w . 2.2 NEURAL PARAMETRIZATION OF OPTIMIZATION PROBLEMS . Using neural networks to reparametrize optimization problems have been suggested in ( Hoyer et al. , 2019 ; Sosnovik & Oseledets , 2019 ) . The idea is that the optimization variable w can be defined as the output of a deep neural network w ( θ ) ≡ f ( θ ) with trainable parameters θ . Then , rather than optimizing over the original variables w directly , we can instead optimize over the neural network parameters θ . We will take this idea further and analyze the type of neural network architectures that can accelerate the optimization . Modified convergence rate The neural reparametrization of the variables w ( θ ) leads the following updating rule for θ : dθa dt = −ε̂a ∂L ∂θa = −ε̂a ∑ i ∂L ∂wi ∂wi ∂θa ( 8 ) The convergence rate is then given by dL dt = ∑ i , a ∂L ∂wi ∂wi ∂θa dθa dt = − ∑ i , j ∂L ∂wi ∂L ∂wj ∑ a εa ∂wi ∂θa ∂wj ∂θa = − ∑ i , j MijKij = −Tr [ MKT ] , Kij ≡ ∑ a ε̂a ∂wi ∂θa ∂wj ∂θa . ( 9 ) Here K is the neural tangent kernel ( NTK ) ( Jacot et al. , 2018 ) of w ( θ ) , and Mij is the squared gradients . Note , we have absorbed ε̂a into K for convenience of notation . Expected convergence rates The expected convergence rate E [ dL/dt ] equation 3 depended on M = E [ M ] , the reparametrized one from equation 9 depends on E [ MK ] . In order to quantify any improvements in E [ dL/dt ] from equation 9 , we need to be able to express E [ MK ] in terms of E [ M ] , which is not possible for a general K. If we assume M and K to be independent from each other , i.e . E [ MK ] becomes E [ M ] E [ K ] . Since M is a function of w , K must not be a function of w. One way is to have w = σ ( Aθ + b ) or any other similar neural network . We will discuss explicit choices for w below in sec . 2.3 . In summary , an important constraint that we will put on w ( θ ) is that it yields a K independent of w choose w ( θ ) s.t . : E [ MK ] = E [ M ] E [ K ] ( 10 ) Now we can compute the expected rate of change of w ( θ ) along the previous fast and slow mode eigenvectors ψi of M , similar to equation 6 . Using equation 9 and equation 8 we have ψTi dw dt = ψTi ∑ a ∂w ∂θa dθa dt = −ψTi ∑ a ε̂a ∂w ∂θa ∑ j ∂wj ∂θa ∂L ∂wj = −ψTi K ∂L ∂w ( ∈ R ) ( 11 ) where ∂L/∂w = ∇wL is the gradient vector . Multiplying equation 11 by its transpose , denoting K ≡ E [ K ] , and using equation 10 the expected value becomes E [ ( ψTi dw dt ) 2 ] = ψTi K M K T ψi = ∑ j mj ( ψTi Kψj ) 2 . ( 12 ) which shows how the neural reparametrization may speed up the convergence . First , note that NR changes the rate of evolution of modes ψi in equation 7 . We can quantify this change when [ K , M ] = 0 . In this case , K and M can be diagonalized simultaneously with shared eigenvectors ψi . Plugging the spectral expansion K = ∑ i kiψiψ T i into equation 12 and using ε < η/ √ mmax ( equation 7 ) we have if : [ K , M ] = 0 ⇒ E [ ∥∥∥∥ψTi dwdt ∥∥∥∥2 ] = k2imi < η 2 mik 2 i mmax . ( 13 ) Thus , one way to speed up the convergence , is to have [ K , M ] = 0 and have K increase the slowest rates ‖ψslowdw/dt‖ , while keeping the fastest rate mostly unchanged . As a result , we have : Theorem 1 . With a neural reparametrization yielding [ K , M ] = 0 , to make all modes ψi reach the ideal maximum rate maxi ‖ψTi dw/dt‖ < η , we need to have K ≈ ( M mmax ) −1/2 ( 14 ) Proof . This is because plugging k2i ≈ mmax/ ( mi + ξ ) ( ξ 1 ) into equation 13 makes the r.h.s . equal to η2 . Next , we discuss how we can achieve this in practice . Specifically , we derive what architecture for the neural network w ( θ ) = f ( θ ) will result in a K satisfying the condition in equation 14 .

Summary: The authors derive a neural reparameterization of non-convex optimization problems in order to accelerate their convergence. They do this by deriving how the slowest components of the optimization variables can have their convergence rate improved by preconditioning with a NTK-based matrix. They make connections between this approach and Group Convolutional Networks. Experimentally, they show this approach improves upon baseline gradient-based optimization on a two datasets.