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: This paper proposes a reparameterization of non-linear non-convex optimization problems. This reparameterization amounts to a linear map (i.e., "optimization params = linear operation of a different set of parameters). These linear maps are interpreted as a graph convolution network. The experimental results are validated on "Kuramoto models" and "persistent homology models".