Year: 2021

Abstract:
We tackle the problem of accelerating non-linear non-convex optimization problems. We discover that reparametrizing the optimization variables as the output of a neural network can lead to significant speedup. We examine the dynamics of gradient flow of such neural reparametrization. We find that to obtain the maximum speed up, the neural network architecture needs to be a specially designed graph convolutional network (GCN). The aggregation function of the GCN is constructed from the gradients of the loss function and reduces to the Hessian in early stages of the optimization. We show the utility of our method on two optimization problems: network synchronization and persistent homology optimization, and find an impressive speedup, with our method being 4 ∼ 80× faster.

heading: 1 INTRODUCTION
text: 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.

heading: 2 METHODOLOGY
text: 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.

heading: 2.1 FAST AND SLOW DYNAMICS OF GRADIENT FLOW
text: 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.

heading: 2.2 NEURAL PARAMETRIZATION OF OPTIMIZATION PROBLEMS
text: 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.

heading: 2.3 PRACTICAL IMPLEMENTATION
text: To find a neural network model that satisfies equation 14, we need to consider the fact thatM(w(t)) changes during optimization. To minimize the overhead from constructing w(θ) (and hence K) as a function of M , the architecture has to be as simple as possible and involve steps which do not increase the time complexity of each iteration significantly. We will now discuss when such neural reparametrization is feasible and how to implement equation 14 in practice.
Hybrid optimization To use equation 14 at points other than the initial steps, we need to dynamically estimate and updateM . Though estimatingM(t) efficiently with ideas from Adam is possible, in the current work we focus on the speeding up at early stages using equation 4. Thus, our main proposal is to use a two stage hybrid optimization: 1) use the Hessian to reparametrize w(θ) in early stages for some number of iterations and then 2) switch to the original optimization over w.
Simple implementation Because M may depend on w(θ), as in equation 10, we need a K that is independent of θ to ensure E[MK] = E[M ]E[K]. The simplest architecture is a linear model wi = ∑ a κiaθa, which yields K = κκ
T ≈ (M/mmax)−1/2. We need K to be a full rank matrix, meaning that κ needs to be at least n×n. Since we need the least computationally expensive K, we choose θ ∈ Rn, and let κ be a symmetric n × n matrix. This way, κ = √ K = (M/mmax)
−1/4. For the initial stages, where M ≈ H2/n we have κ = (H/hmax)−1/2, withH being the Hessian.
Per-step time complexity Since w ∈ Rn, the complexity of computing L(w) is O(np) where p depends on the sparsity structure of L. M is an n × n matrix and computing M−1/2 is O(n3) , which is very expensive when n 1. The eigenvaluemmax can be estimated iteratively inO(qn2). Hence, computing K by exactly following equation 14 is not feasible, unless M is fixed and can be calculated offline. Even then, evaluating M −1/2 only makes sense if the expected GD steps for convergence is more than O(n), otherwise we don’t gain any speedup. We may instead try to approximate M−1/2 and mmax iteratively with complexity O(qn2) where q n is the number of iterations. Hence we note that using equation 14 to speedup the optimization may generally be useful if the the per step complexity of GD on L has a time complexity of least O(qn2).
Implementation constraints In summary, we have the following constraints on the structure of the neural network w(θ)
1. Computing w(θ) ∈ Rn must be at most O(qn2) with q n
2. To have K independent of M , ∂w/∂θ must be independent of w and hence, only linear maps w(θ) = κθ are allowed, with κ independent of w. 3. θ ∈ Rm with m ≥ n. We choose m = n, which yields κ = √ K = (M/mmax) −1/4
4. Instead of κ = (M/mmax)−1/4, which is O(n3), we should use O(qn2) approximations.
Relation to GCN Since both θ ∈ Rn and w ∈ Rn, our linear w = κθ architecture is essentially a Graph Convolutional Network (GCN) (Kipf & Welling, 2016) with the aggregation rule κ = (M/mmax)
−1/4, or a weighted adjacency matrix. In fact, all of our derivations remain unchanged if w ∈ Rn×d and θ ∈ Rn×m (i.e. if there are d and m features per node in w and θ). Thus, our claim is that we can speed up the optimization process by reparametrizing w using a GCN with linear activation, as w = κθ, where θ are trainable and κ is the aggregation function derived from the loss gradients. However, evaluating this κ during GD can be quite expensive (O(n3)) for large n (GD steps ∼ O(n2)). Ideally we want an approximation for κ which is O(qn2) with q ∼ O(1).
Computationally efficient implementation We will focus on speeding up the early stages, in which we use the HessianH in M ≈ H2/n. Define H ≡ (1− ξ)H/hmax, where hmax = √ mmax and ξ 1. To get an O(qn2) approximation for κ = (M/mmax)−1/4 = H−1/2 we can take the first q terms in the binomial expansion (see Appendix A.2). A rough estimate would be
H−1/2 ≈ I − 1 2 (I −H)− 3 4 (I −H)2. (15)
Since H is positive semi-definite and its largest eigenvalue is 1 − ξ < 1, the sum in equation 15 can be truncated after q ∼ O(1) terms. The first q terms of equation 25 can be implemented as a q layer GCN with aggregation function f(M) = I −M0 and residual connections. We choose q to be as small as 1 or 2, as larger q may actually slow down the optimization. Note that computing f(M0)
qθ is O(qn2) because we do not need to first compute f(M)q (which is O(qn3)) and instead use vi+1f(M)vi (O(n2)) q times with v1 = θ.
Lastly, to evaluate M0 we need to estimate the leading eigenvalue mmax. Given any vector v ∈ Rn and using spectral expansion, we have
Hqv = ∑ i m q/2 i (ψ T i v)ψi ⇒ hmax ≈ 1TH21 1TH1 = ∑ ij H2ij∑ ij Hij
(16)
where, sinceH is positive semi-definite, for q > 1 the leading eigenvector ψmax quickly dominates the sum and we have Hqv ≈ hqmax(ψmaxv)ψmax. Here we chose q = 2 to get a crude approximation for hmax. The generalized Perron-Frobenius theorem (Berman & Plemmons, 1994), state the components of the leading eigenvector ψmax should be mostly positive. Therefore, we chose v = 1/ √ n, where 1 = (1, . . . , 1), which should be close to the actual ψmax.
If we think of M as the weighted adjacency matrix of a graph with n vertices, the vector component Di = [H1]i = ∑ j Hij is the weighted degree of node i. Hence equation 16 states H/mmax ≈(∑
iDi/‖D‖2 ) H. This is similar to the graph diffusion operatorD−1H andD−1/2HD−1/2 which
are used as the aggregation functions in GCN (here Dij = Diδij is the degree matrix).
When using modern optimizers, the adaptove learning rates is a term within K. This will interfere with the optimal choice for κ and hence our GCN aggregation function. Therefore, in practice, we do not fix the coefficients for H and H2 in our GCN. Instead, we use single or two layer GCN with aggregation function f(H) = D−1/2HD−1/2.

heading: 3 EXPERIMENTS
text: We showcase the acceleration effect of neural reparametrization on two non-convex optimization problems: network of Kuramoto oscillators which are widely-used for modeling synchronization phenomena; and persistent homology, a mathematical tool that computes the topology features of data that persist across scales . We use Adam Kingma & Ba (2014) optimizer and compare different reparameterization models. We implemented all models with Pytorch.

heading: 3.1 NETWORK OF KURAMOTO OSCILLATORS
text: Kuramoto model Kuramoto (1975; 1984) are widely used for synchronization problems, which have profound impact on engineering, physics and social systems (Pikovsky et al., 2003). As shown in Fig. 3(right), Kuramoto model describes the behavior of a large set of coupled oscillators. Each
oscillator is defined by an angle θi = ωit + φi, where ωi is the frequency and φi is the phase. We consider the case where ωi = 0. The coupling strength is represented by a graph adjacency matrix with elements Aij ∈ R. Defining ∆ij ≡ φi− φj , the dynamics of the phases φi(t) in the Kuramoto model follows
dφi dt = −ε n∑ j=1 Aji sin ∆ij , L(φ) = n∑ i,j=1 Aji cos ∆ij . (17)
Hopf-Kuramoto (HK) model Lauter et al. (2015) is a more general version of the Kuramoto model. It includes second neighbor interactions. HK has a rich phase space with qualitatively different solutions. The phase space includes regions where simulations becomes slow and difficult. We use the HK model to showcase our method in more complex and nonlinear scenarios. The HK model’s dynamics follows
dφi dt = c ∑ j Aji [cos ∆ij − s1 sin ∆ij ]
+ s2 ∑ k,j (AijAjk [sin (∆ji + ∆jk)− sin (∆ji −∆jk)] +AijAik sin (∆ji + ∆ki)) . (18)
with the following loss function
L = c ε ∑ i,j Aji [sin ∆ij + s1 cos ∆ij ] + s2 2ε ∑ i,k,j AijAjk [cos (∆ji + ∆jk) + cos (∆ji −∆jk)]
(19) Our goal is to minimize the phase drift dφi/dt to synchronize the oscillators. Existing numerical methods directly optimize the loss L(φ) with gradient-based algorithms, which we refer to as linear. We apply our method to reparametrize the phase variables φ and speed up convergence towards synchronization.
Implementation. For early stages, we use a GCN with the aggregation function derived from the Hessian which for the Kuramoto model simply becomes Hij(0) = ∂2L/∂φi∂φj |φ→0 = Aij −∑ k Aikδij = −Lij , where L = D −A is the graph Laplacian of A. We found that NR in the early stages of the optimization gives more speed up. We implemented the hybrid optimization described earlier, where we reparemtrize the problem in the first 100 iterations and then switch to the original linear optimization for the rest of the optimization.
We experimented with three Kuramoto oscillator systems with different coupling structures: square lattice, circle graph, and tree graph. For each system, the phases are randomly initialized between 0 and 2π from uniform distribution. We let the different models run until the loss converges (10 patience steps for early stopping, 10−15 loss fluctuation limit).
Results of Kuramoto Model. Figure 2 shows the results of Kuramoto model on a square lattice. Additional results on circle graph, and tree graph can be found in Appendix B. Figure 2 (a) shows that NR with one-layer GCN (GCN-1) and GCN with residual connection (RGCN-1) achieves significant speedup. In particular, we found 3.6± .5 speed improvement for the lattice, 6.1± .1 for the
circle graph and 2.7± .3 for tree graphs. We also experimented with two layer (GCN/RGCN-2) and three layer (GCN/RGCN-3) GCNs. As expected, the overhead of deeper GCN models slows down optimization and offsets the NR speedup gains. Figure 2 (b) visualizes the evolution of the oscillator phase φi on a square lattice over iterations. We can see that even though different reparametrization models reach the same loss value, the final solutions are quite different. The linear version (without NR) arrives at the final solution smoothly, while GCN models form dense clusters at the initial steps and reach an organized state before 100 steps. To quantify the level of synchronization, we measure a quantity ρ known as the “global order parameter” (Sarkar & Gupte (2021)): ρ = 1N
∣∣∣∑j eiφj ∣∣∣. Figure 2 (c) shows the convergence of the global order parameter over time. We can see that one-layer GCN and RGCN gives the highest amount of acceleration, driving the system to synchronization.
Results of Hopf-Kuramoto Model. We report the comparison of NR models on synchronizing Hopf-Kuramoto dynamics. According to the Lauter et al. (2015) paper, we identify two different main patterns on the phase diagram Fig. 9(b): ordered (small s2/c, smooth patterns) and disordered (large s2/c, noisy) phases. In all experiments we use the same lattice size 50×50, with the the same stopping criteria (10 patience steps and 10−10 loss error limit) and switch between the Linear and GCN model after 100 iteration steps. Fig. 3 (a) shows the loss at convergence versus the speedup. We compare different GCN models and observe that GCN with A2 as the propagation rule achieves the highest speedup. Also, we can see that we have different speedups in each region, especially in the disordered phases. Furthermore, we observed that the Linear and GCN models converge into a different minimum state in a few cases. However, the patterns remain the same. (b) shows how the level of ordering changes region by region. If the global order parameter is closer to 0, we have more of a disordered phase while the parameter is closer to 1, meaning it is a more organized pattern.
We know that running dynamics in a large lattice are computationally expensive, especially close to the phase transition points where difficult to separate the slow and fast modes.

heading: 3.2 PERSISTENT HOMOLOGY
text: Persistent homology Edelsbrunner et al. (2008) is an algebraic tool for measuring topological features of shapes and functions. It provides precise definition of qualitative features, is computable with linear algebra and robust to perturbation of input data (Otter et al., 2017), see more details in Appendix B.3. An example of persistent homology is point cloud optimization Gabrielsson et al. (2020); Carriere et al. (2021). Given a random point cloud w that lacks any observable characteristics, we aim to produce persistent homological features by optimizing the position of data points.
L(w) = − ∑ p∈D ‖p− π∆(p)‖2∞ + n∑ i=1 ‖wi − r‖2 (20)
where p ∈ D = {(bi, di)}i∈Ik denotes the homological features in the the persistence diagram D, consisting of all the pairs of birth bi and death di filtration values of the set of k-dimensional homological features Ik. π∆ is the projection onto the diagonal ∆ and ∑ i d(φi, S) constrains the points within the a square centered at the origin with length 2r (denote r as range of the point cloud). Carriere et al. (2021) optimizes the point cloud positions directly with gradient-based optimization.
We will refer theirs as the baseline linear model. In contrast, we integrate GCN and reparameterize the point cloud positions, and will refer ours as the GCN model.
Implementation. To measure the speed-up of GCN model, we used the same Gudhi library for computing persistence diagram as Gabrielsson et al. (2020); Carriere et al. (2021). The runtime of learning persistent homology is dominated by computing persistence diagram in every iteration, which has the time complexity of O(n3). Thus, the runtime per iteration for GCN model and linear model are very similar, and we demonstrate that the GCN model can reduce convergence time by a factor of ∼ 4. We ran the experiments for point cloud of 100,200,300 points, with ranges of 0.5,1.0,2.0,4.0. The hyperparameters of the GCN model are kept constant, including network dimensions. The result for each setting is averaged from 5 consecutive runs.
Results. Fig. 4 and 5(a) show that the speedup of the GCN model is related to point cloud density. Training converges faster as the point cloud becomes more sparse, but the speedup gain saturates as point cloud density decreases. On the other hand, time required for initial point cloud fitting increases significantly with the range of point cloud. Consequently, the overall speedup peaks when the range of point cloud is around 4 times larger than what is used be in Gabrielsson et al. (2020); Carriere et al. (2021), which spans over an area 16 times larger. Further increase in point cloud range causes the speedup to drop as the extra time of initial point cloud fitting outweighs the reduced training time. The loss curve plot in Fig. 4 shows the convergence of training loss of the GCN model and the baseline model in one of the settings when GCN is performing well. Fig. 5 shows the initial random point cloud and the output from the GCN model. In the Appendix B.3, we included the results of GCN model hyperparameter search and a runtime comparison of the GCN model under all experiment settings.
Conclusion We propose a neural reparametrization scheme to accelerate a large class of optimization problems. Our method is grounded on analysis that the dynamics of gradient flow are related to the condition number of the system. By reparametrizing the optimization problem with a graph convolutional network, we can modify the condition number and obtain the maximum speed up. The aggregation function of the GCN is constructed from the gradients of the loss function and reduces to the Hessian in early stages of the optimization. We demonstrate our method on optimizing synchronization problems and persistent homology of pointclouds. Depending on the experiment, we obtain a best case speedup that ranges from 4 to 80.

Review summaries:
This work studies the question to what extend a reparametrization of an optimization problem, i.e. representing the original parameters w to optimize for as a function of some other parameters theta, can accelerate the convergences of the gradient flow / gradient descent for nonconvex optimization problems. It studies the dynamics of the flow via eigenvectors of a matrix M formed as the expectation over the outer product of the gradient of the loss with itself to reveal 'slow' and 'fast' modes of the evolution. It subsequently derives sufficient conditions for the reparametrization (which is chosen to be linear but time varying) to balance the decay on all modes. After discussing an efficient approximation of the theoretically derived scheme, numerical results demonstrate the effectiveness of the proposed reparametrization in two exemplary applications.  

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".

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. 

This work proposed a neural reparametrization scheme to accelerate a large class of nonconvex nonlinear optimization problems. The proposed method is grounded on analysis that the dynamics of gradient ﬂow are related to the condition number of the system. More specifically, by reparametrizing the optimization problem with a graph convolutional network (GNN), the proposed method can modify the condition number and obtain convergence speed up, the acceleration is demonstrated on optimizing synchronization problems and persistent homology of point-clouds.