Particle Stochastic Dual Coordinate Ascent: Exponential convergent algorithm for mean field neural network optimization
Authors: Kazusato Oko
Taiji Suzuki
Atsushi Nitanda
Denny Wu
Venue: ICLR 2022 Poster (2022)
Keywords: Neural Network Optimization
Mean field Regime
Overparameterization

Abstract: We introduce Particle-SDCA, a gradient-based optimization algorithm for two-layer neural networks in the mean field regime that achieves exponential convergence rate in regularized empirical risk minimization. The proposed algorithm can be regarded as an infinite dimensional extension of Stochastic Dual Coordinate Ascent (SDCA) in the probability space: we exploit the convexity of the dual problem, for which the coordinate-wise proximal gradient method can be applied. Our proposed method inherits advantages of the original SDCA, including (i) exponential convergence (with respect to the outer iteration steps), and (ii) better dependency on the sample size and condition number than the full-batch gradient method. One technical challenge in implementing the SDCA update is the intractable integral over the entire parameter space at every step. To overcome this limitation, we propose a tractable \textit{particle method} that approximately solves the dual problem, and an importance re-weighted technique to reduce the computational cost. The convergence rate of our method is verified by numerical experiments.

Published as a conference paper at ICLR 2022
PARTICLE STOCHASTIC DUAL COORDINATE ASCENT: EXPONENTIAL CONVERGENT ALGORITHM FOR MEAN FIELD NEURAL NETWORK OPTIMIZATION
Kazusato Oko".t, Taiji Suzuki Atsushi Nitanda 3.1.2., Denny Suc.
University of Tokyo, 2RIKEN AIP, *Kyushu Institute of Technology. +JST PRESTO,
ò University of Toronto, Veector Institute
oko fkk--kazusatogg.o ecc. u-tokyo. .c. jp, ftaijimmi st .......... w.t.kko..a.. jp. ni mandaaafai Lattaaaaa...uutecc.eee. kww.ech.
"dennywuccs wttpy/uessttronto.... www.nto.
ABSTRACT
We introduce MettccleSDCA,, gradient- based optimization algorithm for two-layer neural networks in the mean field regime that achieves exponential convergence rate in regularized empirical risk minimization. The proposed algorithm be regarded as an infinite dimensional extension of Stochastic Dual Coordinate Ascent (SDCA) in the probability space: we exploit the convexity of the dual problem, for which the coordinate-wil method inherits advantages of the original SDCA, including (i) exponential con- vergence (with respect to the outer iteration steps), and (ii) better dependency on the sample size and condition number than the full-batch gradient method. One technical challenge in implementing the SDCA update is the intractable integral over the entire parameter space at every step. To overcome this limitation, we propose a tractable particle method that approximately solves the dual problem, and an importance re- re-weighting technique to reduce the computational cost. The convergence rate of our method is verified by numerical experiments.
proximal gradient method can be applied. Our proposed
INTRODUCTION
major challenge in developing an optimization theory for neural network to handle the non- convexity of loss landscape. Recent works showed that overparameterization is one key property that makes global optimization possible. particular, the effectiveness overparameterization has been extensively investigated in the theory of Neural Tangent Kernel (NTK) and the mean field analysis.
The NTK analysis considers a linear approximation of neural networks and casts the gradient descent dynamics to that in corresponding reproducing kemnel Hilbert space (Jacot et al., 2018). Since the whole argument can be carried out essentially a linear space, we can derive both global convergence and generalization guarantees (Du et al., 2019: Allen-Zhu et al., 2019; Zou et al., 2018; Nitanda & Suzuki, 2020). However, one of the biggest drawbacks of the NTK approach is that the fixed kernel fails capture the feature training aspect of deep learning. Indeed, this feature learning ability of neural network is an important ingredient that separates deep learning from shallow models such as kernel methods (Suzuki, 2019; Yehudai & Shamir, 2019; Ghorbani al., 2019).
In contrast, the mean field analysis considers a smaller scale of parameters than the NTK analysis, which allows the parameters to travel away from the initialization. However, the optimization problem is no longer convex and the analysis requires more involved mathematical tools. A typical approach to formulate the optimization of neural networks as a partial differential equation (PDE) of the distribution of parameters (Nitanda & Suzuki, 2017: Mei et al., 2018; Chizat & Bach, 2018; Rotskoff Vanden-Fin Fijnden, 2018: Sirignano & Spiliopoulos, 2020). A benefit of the PDE formulation is that the objective function becomes convex in the space of measures. This being said, since the dynamics is essentially infinite dimensional, most convergence results are shown in the continuous limit (infinite-width limit). Also, many existing guarantees only handle the continuous time dynamics, due to difficulty in analyzing the time discretization of the PDE. Indeed, Nitanda & Suzuki (2017); Chizat & Bach 22018); Rotskoff & Vanden Eijnnen (2018) proved convergence a global optimal solution for essentially infinite width networks under continuous time settings. On the other hand, global optimality may be shown under less restrictive settings if consider the noisy gradient
Vanden-Fijnden,

---

Published as a conference paper at ICLR 2022
descent algorithm; the corresponding stochastic dynamics can be formulated as McKean-Vlasse dynamics, and its convergence has been studied in Mei et al. (2018); Hu et al. (2019). However, most of these works establish convergence result for the discrete-time finite-width method by bounding the difference from its continuous-ting and infinite-widit limit counterpart, which usually incurs a large discretization error. Hence, it is important develop practical algorithm with optimization guarantees for finite width network and discrete time setting.
Recently Nitanda et al. (2021) proposed a completely time-space discretized algorithm termed Particle Dual Averaging (PDA) that attains polynomial time convergence guarantee to the (regularized) global optimal solution. This method combines the dual averaging technique (Nesterov, 2005; 2009; Xiao, 2009) and the gradient Langevin dynamics (Vempala & Wibisono, 2019). More specifically, the algorithm realized time discretization by extending the dual averaging technique for finite dimensional convex optimization to the infinite dimensional setting, and realized the space discretization by sampling finite number of particles via the gradient Langevin dynamics. Eventually, the method achieves ((1/t) convergence with respect to the outer-loop iteration t. Although (l(tt) is optimal in an online optimization setting, there is still room for improvement for a finite sample size setting. Motivated by this observation, propose method that improves the convergence rate from the polynomial order (O(1/t) to an exponential order (O(exp(-Ct))-
Our contribution. We propose Particle-SDCA, a novel optimization algorithm for two-layer neural networks in the mean field regime that achieves linear convergence (w.r.l the outer loop steps) for regularized empirical risk minimization' We adopt the stochastic dual coordinate ascent (SDCA) technique (Shalev-Shee -haltv--- & Zhang, 2013a; Takáč et al., 2013; Shalev-Shwartz & Zhang, 2013b) to achieve linear convergence rate and much better dependency on the sample size n. Importantly, the integration with respect to the probability measure on the parameters is approximated by finite number of particles. In contrast to existing methods, in which particles are sampled at each iteration, we employ a novel update rule of the weights on the particles so that the sampling procedure is only performed once in ii-iterations, which provides significant computational speedup.
Furthermore, unlike many existing analyses concerning the time-space discretized version of McK- can-Vlasov dynamics (Mei et al., 2018; Gou-Rabee& & Schuh, 2020; Bou-Rabea & Eberle, 2021), 1), we do not couple the discrete time update with the continuous time counterpart by taking the small step size limit. Instead, we directly analyze the convergence of the discrete time update; hence, our method converges to the optimal solution with a fixed (non-vanihhing) step size. Our contribution can be summarized as follows (for comparison with existing works see Table I):
We propose new optimization method for mean field representation of neural networks that is efficient in the finite sample setting. By utilizing the SDCA technique, the algorithm achieves exponential convergence with respect to the outer iteration. Moreover, the global convergence is guaranteed without letting the step size converge to O.
The integral required in each update is approximated by an average over finite number of particles. The number of particles can be of linear order with respect to a required precision, and the computational cost for re-sampling can be of poly-log order.
Thanks to a novel re- seighitttng scheme for particles, we only resample particles once updates, where ñ can be as large as the total sample size nı; this significantly reduces the computation cost.
Method (authors)
Outer-iteration
Inner-iteration
Total complexity
PDE
'p' (og(1/Ep)
1
M1²
Mr²ep sp log(1/ep)
(Bou-Rabee & Schuh, 2020)
PDA
ep'
M₁G
MGG'E7'
(Nitanda al., 2021)
P-SDCA
1 1 (n + >x, log(
n
MG(++G")
M5 R (n+ 72. log( (
G*
(ours)
Table I: Required computational complexity to achieve ac-accuracy (P(p) P(p") EP). ñ is the re- sampling interval. M. (i = 1,2,3) is the number of particles required to approximate the true distribution for each method: They are given as M the number of gradient evaluation to obtain the sampling distribution which has TV-distance ep from the true distribution: when MALA is used, this is given as G O log (og(p/))))/2)
0(E M; (epp log(n)) MM = O(e log(n)). G" is
We remark that our goal is to propose more efficient algorithm for optimizing mean field neural network with faster convergence guarantce, rather than gain better understanding on standard neural network training.
2

---

Published as a conference paper at ICLR 2022
2 PRELIMINARIES
We first introduce the problem setting and some notations used in the following sections. We consider the empirical risk minimization problem in the supervised learning setting with two- layer neural networks, where the input and output spaces are denoted by X' C E and y C R, respectively, Let hg: X x+ Y' be one neuron with the parameters 0 Rd For example, we may set he(x)= = tanh(r)o z) for f = (r, w) E IR x R and an activation function a (the tanh operation merely to ensure boundedness of the output). In the mean field regime, consider neural networks represented as an average of neurons: fa() = tV Em= ho., (z) where O (#m) C F is set of parameters. The continuous limit in the mean field regime is obtained by taking M - oo, and by an analogy to the law of large numbers, the function se converges to the following integral representation of a two-layer neural network:
fp(x) - Soo (ho
where p Rdd + R is a probability density function on Rd representing the weight of parameters. Here, we denote by P the set of probability density functions on Rd. As in the typical mean field analysis (Nitanda & Suzuki, 2017; Mei al., 2018; Chizat & Bach, 2018), we aim to optimize the density function p E P so that the neural network fn accurately predicts the output y y from the input æ E X. To define the empirical risk and predictive risk, we let (z,,y) : y x y + R be convex loss function, such as the squared loss {(2,y) (= y)²/2 for regression, and the logistic loss f(z. y) = log(1 + exp(-xe yz)) for classification. For each (z..85), we use the notation (0) and (z:)) to indicate (x;); and {(f(2;). (2:4).) y.) respectively. To estimate the density function p P. we minimize the regularized empirical risk defined by
min
Σ ((F,(T))) + AKL(P)AN(,,0) PEP
(1)
where A >> 0 is a regularization parameter and (L (0, (0, o²1)) the normal distribution NO0 2² T) with mean 0 and covariance o² I to the distribution with the density function p. This type of regularization naturally arises in training neural networks by the gradient Langevin dynamics (Hu et al., 2019; Chen et al., 2020b). In our framework, we consider directly optimizing the empirical risk with KL regularization, instead of running gradient Langevin dynamics (i.e., noisy gradient descent) on the unregularized objective. By decomposing the KL divergence in Eq. (1), we obtain the following equivalent representation of the objective:
represents the KL divergence from
1 ((fx(z)))) n
+ X1 / (000 X2 f log((((()))
min P(p) := PEP
where \ι, A2 > 0 are the regularization parameters. Importantly, the objective is convex with respect to the density p E P although it non-convex with respect each parameter 0. We make full use of this convexity in the following analysis.
The main difficulty to optimize (2) stems from the following two factors: (i) there an integral with respect to p that should be approximated by computationally tractable scheme, and (ii) there is summation over n data points, which could lead to costly gradient evaluations. To overcome the former difficulty, sample finite number of particles from the density p. As for the later, adopt the stochastic dual coordinate ascent (SDCA) technique, which we explain in the next section.
3 PROPOSED METHOD: PARTICLE-SDCA
We now explain our proposed Particle-Stochasee Dual Coordinate Ascent (P-SDCA) method.
3.1 DUAL PROBLEM AND ALGORITHM DESCRIPTION
The SDCA method relies the dual problem. The dual of the optimization problem (2) can be derived from the following Fenchel's duality theorem. Note that the convexity of P(p) with respect to the density function p is important in the derivation.
3

---

Published as a conference paper at ICLR 2022
Proposition 1 (Fenchel's duality theorem). Suppose that l R R is proper convex function, h4 I + R is bounded, and the primal problem satisfies pnf P(p) > -oc, then the following duality holds:
Σ (g.)) X2 log (f
inf P(p) = sup D(g)= == sup PEP
9ER"
BER"
b2 Ei (;(0) + x111113)] and E is the convex conjugate? of e.- Moreover, p² E P and 9¹ E RR" are both optimal solutions of primal and dual problems if and only if for (zi) E (gt))) and plg") = p" p' where plg](0) == ( sub-differential of f-
where g|9](0) exp
4)a((0) and df; denotes the
This proposition can be derived from Theorems 3 and 8 of Rockafellar (1967) (its finite dimensional version can be found Corollary 31.2.1 of Rockafellar (1970)). See Appendix A for the complete proof. By the optimality condition, we can recover the optimal primal solution p" from the optimal dual solution g" Hence, we aim to optimize the dual problem as p" updating one coordinate 9i can be carried out by picking up one data points (x4,84)- The main strategy of SDCA to randomly pick the sample index E [r]] 3 and update the corresponding coordinate gi by maximizing the dual (i.e., dual coordinate ascent). Before we state the algorithm, we assume the following condition.
". the dual problem,
Assumption 1.
(A1) 4 R PRRR is (1/7))smootit smooth for all i E [n].
(A2) s [0,( < forall
all i E [n]-
The first assumption (A1) is satisfied, for example, by the squared loss ((()) (y: =2)3/2 with = and the logistic loss (u) = log(1 + exp(-yu)) with y = 1. As for the second assumption (A2), our analysis valid for bounded neurons, and the upper-bound is chosen just for simplicity and can be replaced by another value. In fact, if we assume SUPøER [0;(0)] < R instead, the analysis can be reduced to the setting of R = 1 by rescaling h; and the loss function, i.e., R can be absorbed into the smoothness parameter, which then becomes R/7.
As stated above, the strategy optimize the dual problem is to (i) randomly pick E [n] and (ii) update 9: by approximately maximizing D(g) with respect to 9i with other coordinates fixed. For the update of each coordinate, we employ the proximal gradient-type update where we optimize the lower bound of D(g) given by the following lemma (the proof can be found in Appendix C).
Lemma 1. Under (A2) of Assumption 1, for g E R" and ög e R", the second term of D(g) can be bounded as
X2 log (ʃ alg + (0) < S
+ , |6 + >2 log (ʃ 4 [ (0)d0).
Using this inequality, we can derive the update rule as maximizer of an lower bound of the objective as follows. When we update only the i-th component (i.e., 8g; = (Vj + i)) for randomly chosen index i, the dual objective can be bounded as follows:
D + óg) 2 D + (4)) - E (ā: + &g:)) + f hy (0)p (0)d0gg
2n L ōg
Hence, by substituting ğ +- ((() and maximizing the right hand side with respect to 8g. gi 0(00 we obtain the following ideal version of the update rule of g(O:
(++ 1)
argmax
(.:) 1 (000))"")))0)0))0 -0(0) 2nA2
(g:
9
while the other components are not updated: -(((z)) = 9 (j + i). Note that this update requires only one data point while the full gradient descent requires the whole data points. Along with the updated dual variable g(f), we update the primal variable as pit) p[)( The whole optimization procedure based on this update rule is summarized in Algorithm
"hhe convex conjugate f of convex function f R R is defined by f'( == supeeafur f(x)). 3 WWe write [N] : (1...., N) for an integer N.
*f RD
R T-smooth if f is differentiable and [V) vf(y)) S TlI vll for all x, R..
4

---

Published as a conference paper at ICLR 2022
Algorithm Dual Coordinate Descent the continuous limit
Require: training data {IT, 8)]=== and number of iterations tend
1: Choose 9. s.t. e (g. (0)
S (i I.....n) and (0; - (0)
2: for tu 0, ........ tend do
3: Randomly choose it from 2,, :,; .n.
4: Pi
(((+1) argmax S., ER { {
f/ (0i,) hi, (0))00)) a!.")
2m1o
1 10s
-
5: Update the primal solution p((+1) p(g
6: end for
7: return g(fend)
3.2 OUR PROPOSAL: PARTICLE UPDATE METHOD
The update (3) requires integration with respect to plg("). We approximate the integral by weighted sum over M particles:
S' h, (1)p v EM ((0m)))
(4)
2 Tim
where (0., MM are particles that approximately cover the support of pla"'), and (rm) m the weights of each particle. Therefore, the key point of designing an efficient algorithm is to construct a "good" set of particles ( and their weights (rm) j.
The construction of particle approximation consists of two stages: (1) the sampling stage, and (2) the re-weighting stage. We perform the sampling stage once every ñ E [n]] iterations, and perform the re-weighting stage in the other iterations. The detailed description of each stage is given below.
(I) Sampling stage. In the sampling stage, we sample M i.i.d. particles (0m m= from the density
function (1] by using a Monte Carlo sampler, and place an even weight (tt)) = 1/M (m = 1...., .M) on the sampled particles (fm ) As for the sampler, we may employ the unadjusted gradient Langevin algorithm (ULA), or the Metropolis-adjuss
Langevin algorithm (MALA), both described the following subsection. Note that the distribution of particles obtained from the Monte Carlo samplers as listed above may be biased away from the target distribution. We tolerate this inaccuracy by only requiring the particles to obey an approximated density function (tt) that is in c distance from the target density (9] measured by the total variation distance Our convergence rate analysis will take this error into account.
(2) Re- Reightit stage. Since the sampling procedure requires relatively heavy computation, is desirable perform this operation once in while. We consider performing the sampling step once in ñ iterations. In between the sampling stages, we instead perform the re-weighting stage, in which we iteratively update the weight of each particle, rather than resampling them. Recall that (() (0) = exp (-
E h4 (0)g + A [0]1)) which yields ([g(++))] =
(f expl -h. (9) (t+1) - (1) (nd2)) for the index it chosen at the t-th iteration. We adapt this
(1+11
relation to update the weight (() of particles as
(t+1) r expl-hi, (0m (0m)) (t+1) -0""
during the re-weighting stage because r1 represents "importance" of 6m which is proportional so that the particle approximation (4) is (nearly) unbiased. We allow the number of particles depend the outer-loop iteration T, and write (6T) as the number of particles in the 7' outer-loop. The entire procedure is summarized Algorithm 2.
to (0
Finally, output the solution of the last iterate. We call this output Option (A) and write its solution
(A) as Dout Alternatively, there is another choice Option (B), which randomly selects the stopping time fnnd from the final nn-iterations, i.e., that is randomly chosen from {MTend
1, .,..., n. enTend }. and returns the solution (B) = g (ffmma) This technique improves the accuracy up to O(n) factor. Therefore,
reduce the number of particles by O(n) times.
'The total variation (TV) distance between two probability measures with densities p. q is defined as TV(pllg) := & lp g/d0.
5

---

Published as a conference paper at ICLR 2022
Algorithm 2 Dual Coordinate Descent with the particle method
Require: training data {(zi, u:) and numbers of inner-loop iterations ñ and outer-loop iterations Tend.
1: Choose 9t s.t. e (9 (9)
(0)
(0) S (i ,,...,,, n) and (7) (() S (0)
2: for T = 0, ..... Tend
do
3: Randomly (approximately) draw i.d. parameters 6m (m 1,... M (KT) from p ((TT) (0)c0
that satisfies TV ( (plgg1111ll( [p(g(()) S CC
(RT)
4: I' ((TT) f ( (m = , M(TT)
5: fort = nf. nT +1...., nT + h
do
6:
Randomly choose in from 2, (1,2,...,,
7:
94.
(1+1)
f- argmax
E, (9..)
m 11
I'm
hi (8m)
(gi. - (1)) ) 2nd2 -
(9i. - (
M (t)
ER
2m=1
8:
I'
(1+1)
rim exp
(1)
ndg hi, (9)
(1+1)
9 (1))) (m 1, l.....,.......... M(aT)
9:
end for
end for
(loaa for t'end that randomly chosen return Option (A): Jout (enaa) Option (B): Jout
(A)
(B)
11:
from ñ[eend 1 + +,,1,... 1..... }.
Sampling algorithms. Here, we present two examples that can be applied to sampling the particles (0m yd from (0) exp( (0)(-U(0) where (0) is the potential function defined U()) n i=1 (,)0) 1111011).
(1) ULA (unadjusted gradient Langevin algorithm): The algorithm generates samples from the following noisy gradient descent update:
g--1= 1 ( 22
where 77 > 0 the step-size and E ~ (0, I) is a random noise generated from the d-dimensional standard normal distribution. The sampling efficiency of ULA has been extensively studied in the literature (Dalalyan, 2017; Durmus & Moulines, 2017; Vempala & Wibisono, 2019).
(2) MALA (Metropolis-adjusted Langevin algorithm): The algorithm combines ULA with an accept-reject step according to the Metropolis Hastings algorithm. That is, we generate a proposal gh+1 by ULA, i.e., 0+++ = 8h 1) (Uk)) + V2TEA. which is accepted with probability a given by
a min { 11. UC )a(0" (Ā--1)
(0(0 1100)
where (0'10) x exp(-110' - - -V0))(())///))) If Ōk+1 is not accepted, we set @ = 0² The convergence of MALA has been studied by Bou-Rabee & Hairer (2013); Ma et al. (2019); Chen et al. (2020a), to name few. Due to the Metropolis- Haatopolis-Hastings step. MALA allows for larger step size than ULA, which leads an improved iteration complexity with respect the desired accuracy EC-
CONVERGENCE ANALYSIS
We present the convergence analysis of our method under the following assumptions. Assumption 2.
(A3) inf, (x(z) -0 holds for all i E [n]. There exists constant o' such that Vi E [n]. E is -mmmoohh in the interval {y dom(E)) (e) n 1-2, 1-2, *
0).
B =
sup (1),(x) - inf.r fx(x')) [x] S 1, e [r]]] and Bg == sup (1)()) [ are both finite.
S 1, i [m]
(A4) hy is L-smooth for all I,...,n.
It is straightforward to verify the finiteness of B1, B2 and inf (x))) in (A3) for several practical loss functions such as the squared loss and logistic loss. The smoothness assumption (7') ) on & is required to fill in the gap between (*) and g((), and is also satisfied by the loss functions listed above: the squared loss satisfies it with >'= and the logistic loss satisfies with 7 (e2 1)2 /ea. The

---

Sampling complexity. In Theorem 1, it is assumed that the density (At) from which the particles (AT' from the target density (8T( measured by the total variation distance. (RT) Published as a conference paper at ICLR 2022
second assumption (A4) is required to prove convergence of the sampling procedure, and can be omitted if we employ a particle sampler that does not require smoothness of the density function. This assumption is satisfied, for example, by a (truncated) two-layer neuron ha(x) tanh(r)o (w" x) for 0 = (r, w) E R x Id (which also satisfies the boundedness assumption (A2)). We are now ready to def R7d2 2(14))2)
present the convergence rate of our proposed method Algorithm 2. Here, we define 8
Theorem 1. Suppose that Assumptions I and 2 hold. Let tend = and. There exist constants Õ2
O2 and O3 which can depend on ā and X2 such that, if
cc S Oi exx(--11))) M
(RT)
Z
(ee Tp
C2
log log(AnTTona/O) (VT E [Tend]),
(5)
for 0 <8<< < 1 and X22 2 2B/n, then there exists event E satisfying P(E) 2 1 - 6 in which the
conditional expectation of the duality gap can be bounded by E P(plgout) ( ((guut) (A)) E S SpP
as long as the number of iterations tend satisfies the following conditions:
Lend 2 2 + 1
((1 n
1
C'
(Option A)
log
(6)
722
722
cp
C3s
(Option B)
Lend 2 2 n
7/22
log
( ny222
cp
+
h
Moreover, the constants have the following dependency on n, X2. 5: C2
exp(
4B1² Agm
C 1
O (~ (x=2 exp
(((20++1)) and C3 = Ommxx ,, 3/2 + 8 1
The proof and detailed descriptions of each constant are given in Appendix C. The theorem indicates that our proposed method achieves exponential convergence with respect to the number of outer iterations. This coincides with the convergence rate of the finite dimensional version of SDCA (Shalev-Shward) & Zhang, 20). In contrast, a (non-stochastic) full-batch gradient descent requires O((n/>27) log(1/ep)) total gradient evaluations to optimize convex objective with condition number 1/(22)). Importantly. SDCA (and our method) "decouples" the sample size n and the condition number 1/()27) yielding only O((n + 1/\27) log(L/ep)) iterations, which is beneficial in setting of big data and a small regularization parameter.
The difference between options (A) and (B) is only an (log(n)) factor in terms of the number of iterations. However, this improvement affects the number of particles given by Eq. (5). Indeed, the log(n) factor improvement of land allows us to take (RTama) 2 M(C2Os/(Cfen Either way, the number of
in the last
inner loop, whereas option (A) requires (āTuma)) 2 D(nC3Cs/(C7ep) particles can be of linear order in terms of the inverse of the target precision ep.
Remark 1. Depending on the regularization strength 22, the re-sampling interval ñ should be chosen differently. For constant order (non-vanishing) Xz we may take the remmmllingg interval as = n. so that all constants are bounded. On the other hand, in the small X2 regime, say X2 O(1/n), we have C2 (ñ))) and there also appears an exponential dependency C₁ w.E.t. â. In this case, ñ = O(1) is a better choice to avoid the large particle complexity. But for such choice of n, we should re-sample at almost every iteration, which is computationally demanding; this presents trade-off. In practice, re-sampling can be executed in parallel and thus efficient training (using GPU) possible. Importantly, to obtain the usual O(1/ (L/v/))-ge /no/-/n))geeee -eneeraaliiinnnnn error, x2 = 11//n is sufficient (Nitanda et al., 2021) and thus we may take ñ = ((vn) in that situation. We however conjecture that this trade-off off is unavoidable without additional assumptions.
are generated is within EC
The next proposition provides an upper bound of the number of iterations required to obtain the CC accurate distribution by the Monte-Carlo samplers such as ULA and MALA. The key property of the target density that determines the convergence rate of these methods is the log-Sobolev inequality with constant CLS (see Definition 2 for its precise definition).
'nntuitively, the event & represents an event in which "good" particles are sampled from the sampler to approximate the integral.
7

---

Published as a conference paper at ICLR 2022
Proposition 2. Under Assumptions I and 2, the target density p(g each outer iteration satisfies the log-Sobolev inequality with constant CLS = 2 iterations R obtain the CC
exp(-4B2/ (2))) Accordingly, the number of
X2
(8I) -accurate distribution
be evaluated as follows:
(ULA)
K!T)
AL Ls
max
{
(1,
4d 13
log 4B22 )
Bo(++ + log(1/<c (T 3/2
(MALA)
K. 0O
A2
The proof and a more detailed statement can be found in Appendix B. Roughly speaking, this proposition states that the number of iterations for ULA can be simplified as K(T) 2 ((CLSE )-2e)"")) (oolog()) On the other hand, MALA requires only poly-log order complexity with (RT) respect to CC as K{T) = R( 5/2 Lo ( which is therefore more preferable to obtain higher accuracy. Combining the required number of particles M and the number of iterations K!T) for the sampler, we see that the total computational complexity is still of polynomial order.
Remark 2. Note that the log Sobooee constant CLs also depends exponentially on 1/x2. While the exponential factor can be avoided the outer loop (Theorem 1) by appropriately setting û, such dependence is known to be unavoidable for cLs in the most general selting (Menz & Schlichting, 2014). An interesting direction is to explore additional assumptions that remove this dependence.
Total complexity. By summarizing the results we obtained above, the total computational com- plexity of our algorithm (Algorithm 2) with the MALA sampler can be roughly bounded
Etm M (RT) (1 nnk /m)n= = 0
log(n) hep
(n
>tt)) [(11+
log(e,")
5/27
CLS
where we assumed Cr,,Co O(1) (see Remark 1 on how this is achieved). Detailed evaluation of the computational complexity with the ULA sampler can be found in Appendix B.2.
5 ADDITIONAL RELATED WORKS
Hu al. (2019); Sirignano& & Spiliopoulos (2020) showed linear convergence rate of the McKean- Vlasov dynamics. However, the theory only handles the infinite width limit and continuous time dynamics, and it from trivial to covert their results to a discrete time and finite width setting. Mei et al. (2018) also analyzed space-time discretization of similar dynamics. However, quantitative convergence rate is not provided, and the time discretization error grows exponentially with the time horizon. Compared to these results, deal with general (1/7)-smoot) loss functions and made the dependency on the parameters λ2 and y explicit in our convergence rate. Bou-Rabee & Schuh (2020); Bou-Rabee & Eberle (2021) showed an exponential convergence of gradient Langevin dynamics with an interaction potential, which includes certain mean field models. When specialized to mean field neural networks, our formulation covers much wider range of objectives and loss functions. Furthermore, they considered only the Hamiltonian Monte Carlo sampler, and the analysis relied on the coupling between the discrete time dynamics and the continuous time counterpart. Consequently, their outer-loop iteration complexity is larger than O (o/gg/////////// whereas our bound requires only poly-log order cost due to the flexible choice of sampling algorithm. Jabir al. (2021) also considered spatial-time discretization of mean field dynamics of deep neural network. However, it contains the time discretization lag in their convergence rate, while our method is purely discrete time and is not affected by such time discretization eиror. Moreover, their analysis requires a strong regularization (which corresponds to A₃ in our setting) so that the objective becomes convex like with respect to each particle show the geometric ergodicity; this is not required in analysis.
6 NUMERICAL EXPERIMENTS
We verify our convergence rate analysis on regression problems with the squared loss. We consider teacher-student setup in which teacher and student are both two-layer neural networks of different width. The training dataset is generated by the teacher model: y; = aww" x; + b²) + 6,, where I;; and 6; independently drawn from Gaussian distributions. We optimize tanh neural network (he (I) = tanh(w' w æ + b) with 0 = (w, b)) of width M, as a student using Algorithm 2.
8

---

Published as a conference paper at ICLR 2022
Excess training loss
10-1
train
10-1
10a
SGD (NTK)
SCD (mean feeld
train fit test fit
PDA
Test loss
10
Excess
10-3
10"
200
0.25
0.50
0.75
1.00
iterations
Gradient evalueations
1e5
Linear convergence of P-SDCA.
Comparison of predictive accuracy.
Figure 1: (Left) Linear convergence of PSDCA to the optimal value of regularized objective, with test loss as for the supplementary. (Right) Comparison of predictive accuray of P-SDCA, PDA, and SGD. The error-bar indicates the standard error over 10 independent repetitions.
Convergence rate of P-SDCA. We verify the linear convergence rate of P-SDCA as suggested by our theoretical analysis in a simple setting. To verify the convergence rate, we run the method with 1000-steps of ULA. As for the teacher model, we use the ReI.U activation for and (note that even though this a simple setting, the teacher model "not" included the student model). The number of examples and dimensionality of the input space are 500 and d 50, respectively. The other hyperparameters are set to Å 10 X2 = 10-4, and M, = 200. The left side of Figure 1 depicts the convergence of the primal objective to the optimal value. The r-axis is the number of outer-iteration and the y-axis is the excess primal objective P(-) inf, P(p) including the regularization terms training set and test set in log-scale The integral with respect to the primal solution was computed by the particle average and the entropy is estimated by using the k-NN entropy estimator (Kozachenko & Leonenko, 1987; Brodersen, 2020). Since the y-axis is log-scaled, we can clearly observe the linear convergence of P-SDCA.
Comparison of predictive accuracy with SGD and PDA. We employ a sign activation for the teacher model (o) = sign(-)). (.)) and the student is an overparameterized tanh network of width M2 200. The number of examples and input dimensionality are n 1000 and d 50, and we set A = 10 3 X2 = 10- As for SGD, we cannot apply the entropic regularization and thus the term is omitted in the optimization. We also used the neural tangent kernel (NTK) scaling (fe(x) VM
Σm= he. (z)) for SGD optimization in addition to the mean field scaling. The test error with respect to the number gradient evaluations is displayed in the right side of Figure 1. We make the following observations: (i) P-SDCA achieves better test error than other methods; (ii) SGD achieves minimal test error at the early stage but does not perform well afterwards; (iii) the mean field parameterization outperforms the NTK parameterization in terms of test loss; one possible explanation of this observed advantage is the presence of feature learning, which we empirically demonstrate in Appendix D.
7 CONCLUSION
In this paper, we proposed new stochastic optimization technique for two-layer neural networks with the mean field representation. Our algorithm can be seen as an infinite dimensional extension of SDCA that maintains its advantage, namely, that converges linearly with respect to the outer- iteration, and is much more efficient in terms of the sample size than non-stochastic methods. In the algorithm, we approximate the integral required in each update by an average over finite number of particles. Remarkably, the number of particles can be of linear order with respect to the required precision, and its sampling cost can be of poly-log order. This is significantly different from existing numerical approximation approaches of McKeen-Vlasoo equations, which typically require either vanishing step size or exponentially large particle size. We validated the linear convergence of P-SDCA in numerical experiments, and demonstrated that it outperforms other existing methods.
In the convex optimization literature, it is known that the SDCA algorithm considered this work is not rate optimal - the optimal algorithm can be obtained by combining Nesterov's 's acceleration with stochastic methods (Allen- -Zu, 2017: Murata & Suzuki, 2017). Extending our algorithm to such an optimal accelerated method is interesting future direction. In addition, while the exponential dependence on 1 1/2₂ in the log-Sobolev constant may be unavoidable in the general setting, it is worth investigating whether such dependency can be improved under structural assumptions.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGMENT
TS and KO were partially supported by JSPS KAKENHI (18H03201). Japan Digital Design and JST CREST. AN was partially supported by JSPS Kakenhi (19K20337) and JST-PRESTO (JPMJPR 1928).
REFERENCES
Z. Allen-Zhu. Katyusha: The first direct acceleration of stochastic gradient methods. In Proceedings of Annual ACM SIGACT Symposium on Theory of Computing 49, pp. 1200-1205, ACM, 2017.
Z. Allen-Zhu, Y. Li, and Song. convergence theory for deep learning via over-vveernnne vver-aaramtterization.. In Proceedings of International Conference on Machine Learning 36, pp. 242-252, 2019.
N. Bou-Rabee and M. Hairer. Nonasymptotic mixing of the MALA algorithm. IMA Journal of Numerical Analysis, 33(1):80-115 2013.
N. Bou-Rabee and A. Eberle. Mixing time guarantees for unadjusted Hamiltonian Monte Carlo. arXiv preprint arXiv:2105.00883 2021.
N. Bou-Rabee and K. Schuh. Convergence of unadjusted Hamiltonian Monte Carlo for mean-field models. arXiv preprint wwXi:20099087355, 2020.
S. Boucheron, G. Lugosi, and P. Massart. Concentration Inqqualities: Nonasymptotic Theory of Independence. Oxford University Press, 2013.
P. Brodersen. Entropy estimators, 2020.
entropy_est imators.
Y. Chen, Dwivedi, M. Wainwright, and Yu. Fast mixing of Metropolized Hamiltonian Monte Carlo: Benefits of multi-step gradients. Journal of Machine Learning Research, 21(92):1-72, 2020a.
7. Chen, Y. Cao, Q. Gu, and Zhang. generalized neural tangent kernel analysis for two-layer neural networks. arXiv preprint arXiv:2002.04022 2020b.
L. Chizat and F. Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. In Advances in Neural Information Processing Systems 31, pp. 3040-3050, 2018.
N. Cristianini, J. Shyler A. Elisseeff, and J. Kandola. On kernel-target alignment. In Advances in Neural Information Processing Systems 14, pp. 367-373 2002.
A. S. Dalalyan. Theoretical guarantees for approximate sampling from smooth and log-concave densities. Journal of the Royal Statistical Society. Series B (Statistical Methodology). 79(3): 651-676, 2017.
S. S. Du, X. Zhai, B. Poczos, and A. Singh. Gradient descent provably optimizes over- peeemaarmmteee neural networks. International Conference on Learning Representations 7, 2019.
A. Durmus and E. Moulines. Nonasymptotic convergence analysis for the unadjusted Langevin 2017.
algorithm. The Annals of Applied Probability,
B. Ghorbani, S. Mei, T. Misiakiewicz, and A. Montanari. Limitations of lazy training of two-layers neural network. In Advances in Neural Information Processing Systems 32, pp. 9111-9121, 2019.
R. Holley and Stroock. Logarithmic Sobolev inequalities and stochastic ising models. Journal of statistical physics, 46(5-6):1159-1111 1987.
K. Hu, Z. Ren. D. Siska, and L. Szpruch. Mean-field Langevin dynamics and energy landscape of neural networks. arXiv preprint 19X.0:1905.07769, 2019.
J.FF. Jabir, D. Šiška, and Lukasz Szpruch. Mean-field neural ODEs via relaxed optimal control. arXiv preprint 19/2: 2021.
10

---

Published as a conference paper at ICLR 2022
A. Jacot, F. Gabriel, and C. Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems 31, pp. 3580-8589 2018.
L. F. Kozachenko and N. N. Leonenko. Sample estimate of the entropy of a random vector. Problems of Information Transmission, 23(2):95-101, 1987.
M. Ledoux and M. Talagrand. Probability in Banach Spaces. Isoperimetry and Processes. Springer, New York, 1991.
Y.-A. Ma, Y. Chen, Jin, Flammarion, and I. Jordan. Sampling can be faster than optimization. Proceedings of the National Academy of Sciences, 116(42)20881122888 2019.
S. Mei, A. Montanari, and P.-M. Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 15(33))E77665-E7777 2018.
G. Menz and A. Schlichting. Poincaré and logarithmic Sobolev inequalities by decomposition of the energy landscape. The Annals of Probability, 12(5):1809-1884, 2014.
S. Mischler. An introduction to evolution PDEs, Chapter 0: On the Gronwall lemma, 2019. URL hwtps:ttp///.. /tt///wwwww//.// cwwemade dauphine. htt.://eeeeeeeeee.ee///////eeeee MZevo12018/ Shaap..pp.
M. Mohri, A. Rostamizadeh, and A. Talwalkar. Foundations of Machine Learning. The MIT Press, 2012.
T. Murata and T. Suzuki. Doubly accelerated stochastic variance reduced dual averaging method for regularized empirical risk minimization. In Advances in Neural Information Processing Systems 30, pp. 608-617, 2017.
Y. Nesterov. Smooth minimization of non-smooth functions. Mathematical programming, 103(1): 127-152, 2005.
Y. Nesterov. Primal-dus subgradient methods for convex problems. Mathematical programming, 120(1):221-255) 2009.
A. Nitanda and T. Suzuki. Stochastic particle gradient descent for infinite ensembles. arXiv preprint arXiv:1712.054433 2017.
A. Nitanda and Suzuki. Optimal rates for averaged stochastic gradient descent under neural tangent kernel regime. arXiv preprint arXiv:2006. /2297, 2020.
A. Nitanda, D. Wu, and T. Suzuki. Partiele dual averaging: Optimization of mean field neural networks with global convergence rate analysis. In Advances in Neural Information Processing Systems 34, 2021.
E. Posner. Random coding strategies for minimum entropy. IEEE Transactions on Information Theory, 2144),388-39,, 1975.
R. Rockafellar. Convex Analysis. Princeton University Press, Princeton, 1970.
R. T. Rockafellar. Augmented Lagrangians and applications of the proximal point algorithm in convex programming. Mathematics of Operations Research, 1:97-116, 1976.
R. T. Rockafellar. Duality and stability in extremum problems involving convex functions. Pacific Journal of Mathematics, 21(1):167-187, 1967.
G. M. Rotskoff and E. Vanden-Eijndes Trainability and accuracy of neural networks: An interacting particle system approach. arXiv preprint 1rXiv:/855.009 2018.
S. Shalev-Shwartz and T. Zhang. Stochastic dual coordinate ascent methods for regularized loss minimization. Journal of Machine Learning Research, 14:567-599, 2013a.
S. and T. Zhang. Accelerated mini-bant bati- stochastic dual coordinate ascent. In Advances in Neural Information Processing Systems 26, 2013b.
11

---

Published as a conference paper at ICLR 2022
J. Sirignano and K. Spiliopoulos. Mean field analysis of neural networks: A central limit theorem. Stochastic Processes and their Applications, (000)):820-18222 2020.
T. Suzuki. Adaptivity of deep ReLU network for learning in Besov and mixed smooth Besov spaces: Optimal rate and curse of dimensionality. In International Conference on Learning Representations, 2019. URL https: / /pppreviee.....o/// net / forrmmid--eeeeeeeee s sctm.
M. Takáč, A. Bijral, P. Richtárik, and N. Srebro. Mini-batch primal and dual methods for SVMs. In Proceedings of the 30th International Conference on Machine Learning, volume 28 of Proceedings of Machine Learning Research, pp. 1022-1030. PMLR, 2013.
S. Vempala and A. Wibisono. Rapid convergence of the unadjusted Langevin algorithm: Isoperimetry suffices. In Advances in Neural Information Processing Systems 32, pp. 8094-8106, 2019.
L. Xiao. Dual averaging method for regularized stochastic learning and online optimization. In Advances in Neural Information Processing Systems 22, pp. 2116 2112-2124, 2009.
G. Yehudai and O. Shamir. On the power and limitations of random features for understanding neural networks. In Advances in Neural Information Processing Systems 32, pp. 6598-6608, 2019.
D. Zou, Y. Cao, D. Zhou, and Q. Gu. Stochastic gradient descent optimizes ver-pparameterized deep ReLU networks. arXiv preprint arXiv: 1008888, 2018.
12

---

Published as a conference paper at ICLR 2022
TABLE OF CONTENTS
Introduction
Preliminaries
3 Proposed method: Particle-SDCA
3.1 Dual problem and algorithm description
3.2 Our proposal: Particle update method
4
Convergence analysis
Additional related works
6
Numerical experiments
7
Conclusion
Fenchel's duality theorem (Proof Proposition
14
Computational complexity of sampling algorithms (Proof of Proposition 2)
15
B.1
Log-Sobolev inequality of the target function
15
B.2 Sampling complexity of ULA
16
B.3 Sampling complexity MALA
17
Proof of Theorem
19
C.I
Convergence of the ideal update j(() (Algorithm 1)
21
C.2 Convergence the particle- version update (0) (Algorithm 2)
................................
25
C.2.1
Convergence of auxiliary variables g(()) and g
29
C.2.2 Particle approximation, and analysis of (()
................................
30
C.2.3 Integration of two sources of discretization error
34
C.2.4 Discrete Gronwall's lemma
40
D Additional experiments and details of experimental settings
40
13

---

Published as a conference paper at ICLR 2022
Appendix
A FENCHEL'S S DUALITY THEOREM (PROOF OF PROPOSITION 1)
To prove Proposition 1, we utilize the infinite dimensional Fenchel's duality theorem by Rockafellar (1967). Two real vector spaces E and E' with locally convex Hausdorff topology called lopologically paired vector spaces if all the elements of each space can be identified with all the continuous linear functionals on the other. Therefore, for a continuous linear functional y X + R, there exists E such that we can write y(z) = (y. g , as a bilinear form E.
Definition (Fenchel- Legender convex conjugate). Let X, X be topologically paired vector spaces. For X R U {oo}, its Fenchel-Legendee convex conjugate f X - R U {oo) is defined by
f (v) = supl ().x)x-.x f(x) 1 EX.
Remark 3. When X = X H equipped with the usual inner product, the definition of the Fenchel-Legendre convex conjugate coincides with the definition introduced in the main text.
Lemma 2 (Theorems and of Rockafellar (1967)). Let (X, X*) and YYY be two pairs of topologically paired vector spaces. Let f. X - } and g Y + E U {oo} are proper lower semi-conee coniicontinuous convex functions, and let A X + Y be a bounded linear operator. Assume that Br E X such that (z) is finite and g is finite and continuous at Ar. Then,
inf (gxx) + f(x) | x E X} supl 9" (v) 1 f* IuEY")
where A is the adjoint operator of A such that (y. Ar)Y- = (A" x)xxx Moreover, it and ù the optimal solutions of the primal and dual problems respectively if and only if A& Ogg (ū) and A y E Of))).
Now, we are ready to prove Proposition
Proof of Proposition 1. We utilize Lemma 2 show the assertion. We let Y IR Then, know that Y is self-dual, that is, its topological dual Y is Y itself where the linear functional given by the standard inner product (u, u)r = E.= Maen for = (u,)²= E and (u:) Y We let X be the set of L -integra functions with respect the Lebesgue measure, i... X = ()d) = I (Ted 11), where is the Lebesgue measure on (Re ()(TT) It is known that its topological dual is the set of essentially bounded functions equipped with - i.e.,
X - - (R 0 (Rdd ()) = =f= Rd + + ʃ B(RR)) -measuuuee,,,,llee ||/ < <xx), where || = inf [ > 0 (((x [f(x) 2 M)) = 0} for the Lebesgue measure Jr. The bilinear functional (., xxxxxxxxx .x is given by (g. f)xx = fgdp. Then, X, X become topologically paired vector spaces by equipping X with its weak* topology with respect to the inner product (x,)xx xx-
We let A (RR) + R" be Af (f (() E R² for f L ((R)). Then, we can see that A is a bounded linear functional by the boundedness of h₁(-).
also known that the KL divergence is jointly lower semi-continuous with respect to the weak convergence topology (Posner, 1975, Theorem 1). Since the -nom between probability density functions is equivalent to the total variation norm between the corresponding probability distributions, and the total variation norm gives a stronger topology than the weak topology, KI. divergence lower semi-continuous with respect to the -.ommorm of the corresponding density functions. It is also well known that KI. divergence is jointly convex and it is easy to check that p KL,p,, N(O o3 T2) is proper, i.c., 3p E P such that (L(p,N(O. o²1)) < oc.
Under these settings, if we set g(u) E ((u,) for (w)) Y Rd and
f(p) { (1, (1222)) (pEP), (otherwise),
that P(p) g(Ap) f(p) for p P and infpEp P(p) inf (g(Ah) + fhh)]. We can also verify the condition in Lemma 2. Indeed, if we let po E P be the density function of N(0, A1 A27)), then (po) oo and y is continuous App because
for p E L' (Rd),, then
14

---

Published as a conference paper at ICLR 2022
convex function is continuous when its effective domain is the whole space I² (Rockafellar, 1976, Theorem 10.1), 1), and the domain of g is the whole space I" (remember that the range & is RR)
Therefore, by applying Lemma 2 and going through simple calculation to obtain f", we obtain the desired assertion.
B COMPUTATIONAL COMPLEXITY OF SAMPLING ALGORITHMS (PROOF OF PROPOSITION 2)
B. LOG-SOBOLEV INEQUALITY OF THE TARGET FUNCTION
In each outer iteration, we sample M particles from the density p[g(AT) so that the particles are sampled from an approximated density p p(it) with (ppl convergence of Monte-Carie methods such as ULA and MALA, the log- oog-Sooolev inequality associated with the target density plays the central role.
(+T)) I Eg To estimate the
Definition 2 (Log-Sobolev inequality). Let p(0) be smooth probability density function on IRd
Then, we say p(0) (or iss corresponding probability measure on (Tt, B(TR))) satisfies the log-Sobolev function
inequality (LSI) with a constant CLS > 0 if and only if. for any smooth function @ : RRP - with
Ep [x²]] oc, it holds that
2 log - Eplø" (Eg - Ep [
CLS
The following lemma shows that the target density (TT) satisfies the LSI.
Lemma 3. Suppose max, 19; ( S B (this is true by Lemma below with = B2) and assume (A2) of Assumption 1, then p g (mT) satisfies the LSI with the constant
2\1 exp(-4B/22).
CLS A2
Proof. Remember that (TT is given as
plg p (0) exp
1 (0)))
(RT)
2
110112
A2 nh i=1
A2
Therefore, (() can be written
(0)) exp(f(0))a(0) Ex(exp([())]
(8)
where f(0) = - E. (.)0)): (TT) and q0) is the probability density function of Gaussian
measure N(0, A2 I). Note that || S B/22 by the assumption (A2) and max; 1g; (RT) < B. 2X1
Then, is well known that satisfics the LSI with LSI constant c (g)(g) = 2 and Holley & Stroock
(1987) showed that the bounded perturbation as Eq. (8) preserves the LSI condition where the LSI constant of plg (RT) is bounded by cLS cLs(g)/ exp(4|]fll z concludes the proof.
2 exp(-4B/Aa) This
Xa
By Lemma 3, show that existing convergence rate analyses of ULA and MALA such as Vempala & Wibisono (2019); Ma al. (2019) be applied target density (TT) From now on, denote our target density (*T) by p", and the negative log density of (aT) by U, that is, (0()) = exp( (U -O)). In addition, let Pk the marginal distribution of 0 (sample the k-th iteration of ULA MALA).
15

---

Published as a conference paper at ICLR 2022
B.2 SAMPLING COMPLEXITY OF ULA
First, we apply the convergence rate analysis of ULA developed by Vempala & Wibisono (2019) to setting. We utilize the following proposition given by Vempala & Wibisono (2019).
Proposition 3 (Theorem 1 of Vempala & Wibisono (2019)). Assume that p" (0)d0 satisfies the log-Sobolev inequality with constant cts and U is L-smooth. Under these assunptions, for the step size yj satisfying S 2L we have that
SmdI2
KL (pis ll p²) S
(p1 p²)
CLS
Proposition implies that if the step size satisfies satisfies
She min I, S 11.2
and the number of iterations
id
2KL (po ll p²))
k
log
CLS7
then we obtain an e-accurate distribution рı in terms of the KL-divergence.
If we employ (f) exp(- A 21) ) (a Gaussian measure), then 2KL (po p²) is bounded by 4Bs from Lemma 6:
Xu
Po( do
KL (po |1 pI) = po) log p" (0)
= =- po()) [log p" (0) log po (0)] d0
A.10).. 1011
-- -- (0)
exp
22
log
log (0
de
f exp
>m
(,(0")) 19, 11110'11
do'
sup, lgal
A1
$ f (0)
exp
exp
Aa
log
log (
d0
exp
sup; Lovl
An
ʃ exp
Az
||8112
d0
f (O()
2 Bg
po) }
log
exp
log po (0)
d0
A2
2B₂
X2
Moreover, using Pinsker's inequality bounds total variation norm by KL-divergence:
l - lrv & /2KL (px p*).
Therefore, in order to have an ac-accurate distribution in total variation, suffices to let the step size and the number of iterations satisfy
satisfy S 4 min [,,
8d
2
4B2
k
log
CLS?
Agec
Hence, the total computational complexity can be bounded by O log
B Azec
Total computational complexity of sampling by ULA According to Theorem 1, sampling ac- curacy in total variation &c Therefore, the total computational complexity for sampling can be calculated as follows:
(RT) should be smaller than O(exp [T_T
(- 2
in each sampling step T.
Algorithm 2 with Option A: Suppose
we
run outer-loop
Teni
2
(Bi+Ba 1- Ba
log n
times. Indeed it sufficient to have an 6p-
approximated solution, according Theorem 3, which gives tighter and more detailed
16

---

Published as a conference paper at ICLR 2022
bound than in Theorem 1. Then, the total required computational complexity for sampling is
Tast
B2
log
(
A
(AT))
d..
B2
Teu
0
log
exp(37)
T=1
dL2
B2
exp
(3Tend)
=0
log
CLS
2
AAEC
12
dL2
B2 Bg Co)
=
log
npp
2
A2 + m97
where Ô hides log log terms.
Algorithm with Option B: According Theorem i, is sufficient to run outer-loop up Tend 2 times to acquire ap-appooximated solution. You can the detailed definition of C2- = O(max A 3/222 -//2)-1/3)) in Eqs. (52a)-(52c) of Theorem 3 and Proof of Theorem 1. Therefore, the total required computational complexity for sampling is
log ( nyy
Ea
cp
Σ
0
dL2
log
B2
exp exp (ST)
(RT)
Azc
(nI)
T=1
G.S
EC
dL2
B2
Tmd
50
log
exp (31')
A2co
T=1
dI.
B₂
exp (ST (nedd)
log
duco (TT))
1-
5/2
=ō max
1/3/2
1/
( (nnda+1)
1/2)
1/222
(1+ ayda
log
B2(B1 + Bg + Cu)
n3/2ep
A2
172
where Ô hides log log terms as well.
B.3 SAMPLING COMPLEXITY OF MALA
Here, adopt the convergence rate analysis of MALA obtained by Ma et al. (2019). They showed the convergence under the setting where Ú is justrrongly convex outside a compact domain and is L-smooth. Although our situation does not necessarily satisfy this condition, the proof of Ma et al. (2019) essentially utilizes it to show the LSI of the target density, Fortunately, we have already seen that the target density p" satisfies the LSI in Lemma 3. Thus, the analysis of Ma et al. (2019) can be applied to our setting.
From now on, we denote (HT) by p" Let R(s) be a positive real such that PM [|[[ 2 R(s)) < for 0 < < 1 where ps is the probability measure corresponding to p= and Ć > 0 be the smoothness parameter of log(p"): |/V log(p"(0) V log(p" (O)))(()')) & i. 0'|] Then, Ma et al. (2019) showed the following proposition.
Proposition 4 (Lemma 9 of Ma al. (2019) with modification). Let the initial density po satisfy
(e sup
< 8 PER" pr (0)
17

---

Published as a conference paper at ICLR 2022
for B > O. Then, for the step size satisfying
Ld' LA/3 LA/3R((/2390) 2812//33/3 }) [3/2R(c/28)
1 ==
min
and the number of iterations sattisfyng
1
2B
k = Ω
log
Ls
2 2
1
log
23
max R(e/23) Ld, L Riee (16222))/00//0))
(9)
LS
holds that
1p" - PklITy Sc.
We apply this proposition to our setting. From now on, we assume Assumptions and hold. Then, it can be easily seen that
L LB2 2\1
A2
employ Po( exp( X/ 11 (a Gaussian measure), then we may set 3 S exp(2BB2 /2) because
(0)
/2
p/2
eeep
Aum
2i=1 h₂(0") (TT)
X 111111)) de
(0)
2x
1
Ei= h0)gg (AT)
exp
dnn
D/2
/ exp
exp
(2B2/X2)
( d0'
2m
221/A2
P/2
exp (2 A2)
exp
( |00)0 do
2x
exp (2B2/A2).
(10)
Moreover, the proof of Lemma 16 of Ma et al. (2019) showed that we may set
R(s)
2 log
+ 1B-111011).
cLS
The second term in the right hand side can be evaluated as
/ 110122" (0)a0
/11011200 (0) (0) (0) po(()
Epp [1101]]]
do
S exp(2B22/2) / (10)13p((0)c0
A2 d,
exp((22//2)) A2)
2
where the inequality in the second line can be shown in the same manner Eq. (10). Then, we obtain
d
d
2B
A d.
R(s)
2 log
exp
CTS
2
Therefore, substituting these evaluations into Eq. (9) yields that
1
2B2
)
1
k 20
+ log
x
A2
1/3
2d
2B2
d
2d
2B2
Je/3
max
13/2
log
+
Ld.
4/3
log
A
CLS
A
1/3
18

---

Published as a conference paper at ICLR 2022
Algorithm 3 Definition of g(t)
1: p(ar) (1g
2: (AT+1) t argmax -i (gtns) )+ hiar (0))
(a2)
(RI) p
-iar
1
2nd2
War
3: for 1 nf 1,.... 7(T + 1) do
4: 7(0) plg(1)
((+1) - argmax 11, (gt.) + (0)0000 - -bh.))
((1)
5: 9i,
91.
-in
2nd2
Die
6: end for
sufficient to ensure l|p" - ps lv Σ c, where CLS can see that the right hand side can be bounded by
2 X2
exp( (4B2/A2) and İ LB,+2X, We
1 5/2
B2 + L) ()(1+)))+ + 1og(1/c)
3/2
0
d'
A2
CLS
which gives the assertion of Proposition
C PROOF OF THEOREM
In this section, we give the convergence rate analysis of our algorithm. We first show the convergence rate Algorithm I (the continuous limit version), then we prove Theorem 1 by bounding the discretiza- tion eπror. Before we begin the proofs, we need some auxiliary variables that interpolate between the continuous limit solution and the discretized one.
The key ingredient of our proof is to show the discrepancy between the solution with particle discretization and the one with the continuous limit sufficiently small throughout each inner loop from = 7I the continuous limit update but is initialized from the discretized one (f at the every starting point of the inner loop (see its definition in Algorithm 3). To fill in the gap between the continuous limit (()) and the discretized version g()), we also define gi()) (1 = i.....,,TT T( + 1)) which lies in between (see Algorithm 4). Basically, g(t) and g(() are updated by the same rule, but are initialized in different ways. Indeed, the primal variables corresponding to these two sequences
ÜT 1). To rigorously prove this intuition, we introduce g(t) which follows
updated as 1 he, (0)ōg(() and p((+1) x ())) xxxp (- 1 hi, (0)5g()).
jg (0) x (0) xxxp
nda
respectively. However, different initial primal variables p(xT) (TT) and p(nT) (nT)
employed for the two schemes (remember that p(22') is an approximation of (1g by the sampling scheme). We denote difference of the updated coordinates before and after updates as
(t) 9h. Di. (g(++)) def g(++1) g()), (g) g(t+1) (0)) (t (t=nT, iT........ T(n 1) 1), 5g def (t+1)
respectively.
At the same time, define some variants of r (t) in order to fill the gap more finely between g(9)
and (t), We define vfo direct discretization of p(t) using the sampling point 02 (see Algorithm
5). Let p(!)) be a numerator of (t) p (see Algorithm 6). We expect that -
(())
plays role of
interpolation between (t) and
In the convergence rate analysis, the smoothness of the dual objective plays an important role. The following lemma shows the smoothness of the dual of the regularization term.
19

---

Published as a conference paper at ICLR 2022
Algorithm 4 Definition of g(t)
1: p(aT) ((I)
2: Pig ((T++)) t argmax
-1 (Tiar) )+
hiar (0)pliz) (0)d0
(AT)
1ta
9
(RT)
2nd2
Bar
3: for t- = iI (0) exp
iT + +1
do
st-1 h4 (0, (0,))g{.))
4: (10(0))
nd2
J (0)) exp
122 E hi (0 (0)."'''") (T)
d0
5:
9i
(++1)
argmax 94.
-1, -1ggaa-11,,00 (gt.)
hi, (0)d0)) -
g1.
9.,
2nd2
6: end for
Algorithm 5 Definition of f(()
7 (RT) 1
(j .M)
2: if then
3: for = v.,', nT + 1,, 1..1,,,,((7 1) 1) do
-2 n2g Σi=nr hi (0, 16g, (-(
M exp
4: f p(++1)
[plnt) exp
nAg Σ1 - hi (0 (0))o (T)
d0'
5: end for
6: end if
Algorithm Definition of 2((()
(TT)
(j -......,,M
M
2: if then
3: for = 7T, üT + I,....
TT + - 2 do
4: ((++1
Σ h4 (0,)
exp
M
nd2 ==
5: end for
6: end if
Lemma 4. The first and second order derivatives of the second term of the dual objective D(g) can be evaluated as
/ cll(()al
- n
1
V Ve log
h((0)pi)(0))e)
(11)
/ alo(())do
V2X2 log
1
1=(0)),(000bl))))a h.(0))ll())00 / [1100)11000)0))
(12)
n 22
20

---

Published as a conference paper at ICLR 2022
In particular, the dual of the regularization term can be bounded by quadratic function follows: (a) Lı-norm bound,
22 log If [lg + 5gl < A2 log
[ [(o1()aa))
- Σ /
hu(0)pls((0)d00g + 2n X2
lóglli:
1=1
(13)
and (b) Lz-norm bound,
22 log f dlg 55)00)d
/ dlo())
/ (=())))lg())dd8g + 2nd2
I 22 log
l
(14)
The L1 -oomm bound is useful to derive the dual coordinate ascent method (Algorithms I and 2). On the other hand, the L2-no -o-morm bound also useful in the following convergence rate analysis.
Proof. First of all, it is easy to see that d2 log [f (O)) is two times differentiable by y,. Since direct calculation yields
0g. els --
a /
1
/ hz(0) exp -
- -1:1011 dB.
(15)
ndz we immediately obtain Eq. (13):
22
0 Ao log
1 (1(
h:(0)pls(0)a).
Ogi
J (0)
By differentiating the right hand side by gy again and using Eq. (15), we get the desired equality (12)
for the Hessian.
Note that the right hand side of the equation (12) for the Hessian evaluation is given by the covariance between hy and h, with respect to the probability distribution plg](0). Hence, each element is
bounded by मーズ due to the boundedness assumption of h; (Assumption (A2)), which yields Eq. (13), As for the L2-norm bound (14), it is easily obtained by the relation ||ōg]] S vnllóglla. O
C.I CONVERGENCE OF THE IDEAL UPDATE g(t) (ALGORITIIM 1)
First, we prove the convergence of the ideal update g(t) Convergence analysis in the continuous limit the basis for that of discretized updates; more concretely, we can use results and proofs for continuous limit in this section Theorem (convergent proof for discretized updates).
reference to derive Lemma 8, which is one of the key lemmas for
Lemma 5. Suppose that & is 1/yy-mmooth for all E (1....m} (Assumption (Al)). Consider the update of g(t) given in Algorithum 1 (continuous limit update). then, for any iteration t and any s E [0, 1], it holds that
(P((())) D(2(1))) - (-) 2 (())
E2, ((g(tt)) - D(G(1))
(16)
n
2X²
where Tf(t) = pl
(t)) def (1 7(1 22
Es, (ulf) - ((0))]]
(17)
and u. del E. (O)),()))) Here, the expectation was taken with respect to the choice of the coordinate it. Accordingly, we also have the same bound for the ideal update in Algorithm 2 (T)) and g(t) (1 nT +1,..., (T + 1))), considering the definition of g(t)
Proof. Since we will focus on the update of one step, i is omitted as in the following. At each update, the improvement the dual objective can be written as
n (1(+)))))) ())
21

---

Published as a conference paper at ICLR 2022
-G + ( - log f [10)))111(0000 + A2 log / (0
( + h:(0))00 (0)d0 (o( ((1) 2md2
I9c 1/ + E9 (
-0. s(u - -) + / (10) (u )(")) 2n12
s²
(1) g/003 + )o."),
(18)
where we used Eq. (13) in Lemma 4 for the first inequality, and the second inequality follows from the definition of g(++1)
Since e is y-strongly convex (7 = 0 is also allowed), it holds that
(o + s(u (2) )) s st (u) + (1 ()) ?
2 20(1
(19)
From the property of the Legendre transform, - = (J (O)):((0))) implies that
& (f000))000))) + (1())
/ h. (0)0))"00))))))
(20)
Similarly, we also have the following relation from the duality of the regularization term:
A1 / 1 (1)111 + A2 / p00) (o) + do log
h, (0))))0))))0)))))
(21)
Combining Eqs. (18), (19) and (20), we have that n m[D(g((+1)) ((
((()))) - f (n(0)) (())
(2(1-5)-
2(1 - 5) s)
n22 ) 1u, (2) g(00
8
)
- -//0) - [1((0)=100))0))
4 (f (B)):;(())) + +0(00
8
y(1 s) ndp
u (1)
9. 1p.
(22)
By taking expectation of both sides of (22) with respect to i, it follows that
E
n [D((g D(
e,
(f
+-
Σ
(1
1 "
Σ
/
ha ho(())p00
n
n1
8
7(1 -s)
11)
2n
nda
1=11
1
e (f
(0))
1 +
Σ
0 (0 (1)
1m1
1=1
+ A1 /
(fab)1(00))
1000 (0)]0011)3 - X2
1) log ( (0)do + + log
8
(11- ((--)) 5) nde
s Σ
+ 2n
1u
9.
1=1
22

---

Published as a conference paper at ICLR 2022
=s(((B))) D() s² G
(23)
2X2
where we used Eq. (21) for the first equality. Dividing the both sides of Eq. (23) by n, we obtain the desired inequality (16).
Lemma 6. 9, in Algorithm is uniformly bounded by B2, that is,
0
B2 (Vt= -...... nTand, Vi (I....,m)).
In the same way, 94 9:
(t)
(t)
and
9
(t)
in Algorithm 2 satisfy
19.
19. - B2 (Vt = 1,... n'Tnnd, Vi nL).
Proof. Let p be an arbitrary probability distribution. Let us consider the update
2 arg argmax ER
+2' (2'
(
(24)
for an arbitrary E {I,..., n.. Since hy is bounded by (Assumption (A2)), the sub-differential of E at z satisfies that
1] #
By the assumption (A3), i.c.,
Bg sup {14 (x)) |2| S 1] sup flyl ([-1.11 * 0) co.
1.
8.
we have that z in Eq. (24) is bounded by B2. Next, we consider its proximal-version update:
argmax { (-)7(0))+ h. h,(0)p(O)d0(2' - v) E'R
1
1/ y
(25)
2nA₂
In this proximal-version update, we added the term 2nAg (z y) whose maximizer with respect given by y. Since the maximum point of the sum of two concave functions is located between the maximizers of them, 2 in Eq. (255) satisfies that
a I max {B2,
From this and the initialization (ā!0) < B2, we have that Ig! ( in Algorithm is uniformly bounded by Bd by induction.
In the same way, di (1) Di and gi in Algorithm minimize the sum of the average h2,))) and quadratic term well. Note that the absolute value of the the average of (4,) is bounded with 1, regardless methods averaging, that is, weighted particle for ( and probability distribution for (t) and g( Thus, 9; (1) 0.0) and g!" uniformly bounded with B2, as well as (1) in Algorithm 1.
Remark 4. In the Section C.2, we sometimes face the situation where we should upper-bound increment of the coordinates (t) 5g and 8g
Lemma directly implies that (,) (g(t) and 8g(t) are all bounded by 2B, for all t 0, 0,1,...,it onn- Thus, for future convenience, we define
C1 def exp
2x
A2
with which exp
(2)) X2
exp
2nd 2xe(6g(1)
exp
2x
S C1 hold, where is is defined as (+) X2
Note that, however, we could have much tighter evaluations of (g) 5g and 6g((), which are dependent of residual D(g") - (t) or duality gap P(p(g (t)) ((g(t)), since the dual problem
strongly convex due the smoothness of primal problem.
23

---

Published as a conference paper at ICLR 2022
Lemma 7. Under Assumptions (A2) and (A3), it holds that
1) D(g I B1 + B2.
for Algorithm and
1( )() s Bı + B2.
Algorithm
Proof. By Assumptions (A2) and (A3),
that | (ʃ p (O)),;00)))))) inf, (x) is bounded by B1- Since inf, (x(x)) - inf (Assumption (A3)), we also have (0) sup, 10
(x(x)} = inf (x)) < co. In Algorithm 1, every initial coordinate g. is bounded by B3. Also, (0) in Algorithm 2 is similarly bounded by B2. Therefore, in the following, will be used without distinguishing between and ý. Moreover, in similar manner to Eq. (21) in Lemma 5, we have that
9i
A1 f 10010111 A (0) (0))0) + X2 log / 565 (00)0
ha h=(0)pig"010))a)gg""
from the duality.
Combining these observations, we have that
P(rl ())
Σ n i=1
1
( 1 100 (8))
n Σ
(0:00) X1 ʃ (0)
i=1
+ A2 / (0) log 10) (0)d0 X2 log
/ 1/g 10))a
(f (00) (0)h,(00d)
/
+
(o:
(0)
Σ
(0)) 9.
i=1
Σ 4 (ʃ p (0)),(()))
inf (x(x)
+ n
1
Σ
9 (
x
1=1
n Σ Bj + Bg B + B2.
where we used Lemma 6 in the last incquality.
Theorem 2. Lel ep(> 0) be a desired accuracy. To obtain an expected duality gap
E [P(p(faaa) - ((g(ta))))] I Ep in Algorithm 1, it suffices to have the total number Tend of iterations as
n )++2))) log (n+
1 1
B₁ + Bg
fend
7dve
7222
EP
Furthermore, if we randomly choose ((1))g(())))) from t = tend
+ I,..., bend in Algorithm 1, then
(1+ ny 1
Bi B2
Lend
2
n
log
+ n
7d22
sufficient to achieve ep duality gap expectation.
Proof. In Lemma 4, choosing
n7da I+nyda
[0. 1). then we have that
(1 - s)nTd2
-1-1=0 O.
24

---

Published as a conference paper at ICLR 2022
Hence, we have G 0 for all Thus, Lemma 5 implies that
E
> _ -
(26)
Therefore, we have that
- E D(g") - D(G
S - - =E =E <((D())) - -E D(g") D(g
< - (( -
This implies that
D(g") - D( - (1- -(g))))))))]
(27)
By applying this inequality repeatedly, we obtain that
E -D D(G(t)) -1- L - S SB1 + B2) exp
D(2(0)))
st (-m)
(28)
where the last inequality is obtained by [D(g") D( s E [P - and Lemma
7.
Moreover, making use of Eq. (26), the duality gap can be bounded follows:
E (()( D(g()) - " [( D(g(o) - " [D(0*) D(((()) (29)
Hence, combining Eqs. (28) and (29) yields
(B1 B2)
exp (- 72t
E 1(p()))
(- 1 noyd2
s
Therefore, in order to have an ep-approximated solution in expectation, it suffices to let
((n [
1
B + Bg
tend 2
n
log
7x2
7d2
Ep
Finally, show the second part. By summing up Eq. (26) over t tend 1,... bends we have that
(P((0))D(G((1))
< E ((g(tamat))) - (mag
E
t= Lend m1+1
1 D(g")) D(g
B1 + B2 exp
Fde(tend - -n)
1 myd2
1 1+ n7d2
Note that although we do not actually calculate g(fena++1), virtually define g(faaa+!) if the algorithm continued.
Therefore, in order to have ap-approximated solution, it suffices to let
1 > log ((1+ ny/d2 6P
1
Bi + Bg
fend
n +
+ n,
when we randomly choose
from lend 1,..., Lend-
C.2 CONVERGENCE OF THE PARTICLE-VERSION UPDATE (( (ALGORITIM 2)
Next, we consider the convergence in Algorithm 2.
In order to prove Theorem and give a convergence proof for Algorithm 2, we adopt the following strategy. First we prove Lemma 8, by rewriting Theorem 2 to allow the discretization error in each
25

---

Published as a conference paper at ICLR 2022
series of inner loops. Then, the assumptions of Theorem 2 is verified, by Theorem 3. There are two major streams leading up to Theorem 3. The first is about the error originated from the sampling accuracy (Langevin algorithms), and the second is to evaluate the error accumulated during the re-weighting stage. The former is considered in Section C.2.1 and the latter is in Section C.2.2, both of which combined in Section C.2.3 and finally Theorem 3 is proved.
Apart from the sources of the error, we bound the discrepancy of the duality gap between continuous limit and discretized updates (P) - )) - (P(g())) - -(g(()) with the following two parts,
A
def
E
- -
E (() - P()
B
each of which corresponds to the discrepancy in the primal problem and dual problem, respectively.
As mentioned above, these two terms attributed the two sources, that is, sampling and re-
(8) (0()) represent how much the two sequences g(t) and g(t) diverge
(6)
weighting. These variables C fB
during the updates of the inner-loop. Intuitively, in each inner-loop, the absolute values of c and cB are increasing sequences between = 7π + and ū(T + 1). Then, they are "reset" at = (T +
(t)
according to the initialization rule of g (((T+))+1 and again will increase from = n( + 1) + to i 2).
+
The following lemma states that under the assumption about CA, Algorithm achieves linear nndg
convergence with regard to the number of coordinate update. From now we fix 14 Also, we define ă = K
ands where = A
T+wryg"'
Lemma Suppose that Assumptions I and 2 hold. If ( is smaller than Coe-ST for ev-
(RT) ery T = 1,..., Tond with constant Co you can arbitrarily take, an expected duality gap
[ (nTTm))) S Ep can be achieved by total number Tend of outer itera- tions given by
( 1
1 log (n+
(Bi + B2 + I-exp(--)
Co
Tend 2
2n n
(30)
ny
7d2
Furthermore, if A and 6B are uniformly bounded by er for all n and - 12 + I...., TTen in addition to the assumptions above, then it also suffices to sel
(t)
(t)
n2 (1+ Tond 2 -
( log ((++
(B1 + B2 + 1-oxp( -3)
Co
+
(31)
n7/22
6P
to achieve ep-duality gap when we randomly choose output ())) from L. ü(Tend 1...., nTennd
( +
Remark Comparing the speed of convergence with regard coordinate update it exp(- st/ (i) for Theorem 2 and exp(-3T") & exp(-šš/í) for Lemma 8, which can be interpreted as follows. In the descretized version, Lemma allows discretization error to appear, the cost of a slight decrease in convergence speed. Indeed, you can choose arbitrary & from the open interval (0, Ks). Our choice Ka is just for simplicity. If you choose another much closer to Es, Tend can be smaller.
of 2
Remark 6. If you change the length of inner loop in the middle, the convergence rate after the change is obtained by simply replacing Bi + B2 with the duality gap at the time of the change.
Proof. Since the update of z(0) is the same as that of (0 in Algorithm during the inner loop. holds that
E D(g") D(g(RT+1))] s s01 (1-
8 ;) E D -D(g
n
for T = 0, 1..,.... Tend 1. We also have
D(g")= - D(g(+t1))] <
8
D(g") -D(G
n
26

---

Cp Published as a conference paper at ICLR 2022
fort + 1...., " ( + 1) - 1, in the same way as the derivation of Eq. (27).
By applying these inequalities above iteratively, we have that
E "
$ (1- -1- -1==)= -D((s")DD(OATT)) D(g(AT)))
S < (D(0") D(
(32)
Bounding ((g( by D(g(ūCr+1))) (1(T+1)) we have that E ( Dig (T)+ Se E p(gg)) ((g A (((T+))
EA
Therefore, by the Gronwall's lemma ("vanilla version" of Lemma 18), we obtain that
T
[D(0') - (()()T)) Se 25T - - + 237
Σ
25T (nI')
e
A
T=
Applying the assumption (n7)
6 S Coe-s7 -T this implies that
e 57 - c 25T
E D(g") - -[D('")-((G)
D( S Se- - D()
+ Co
1 1
Co
Se FT
E - |
Co
Bi B2+ 1-6-8
(33)
where we used Lemma 7 for the last inequality.
Moreover, the same reasoning as Eq. (29) yields that
E (TT)) D ((g(nm)
s "E D(g") [D('' (ar(a))))
(34)
Therefore, combining Eqs. (33) and (34), we can bound the duality gap
EE (T) - ((g((T)))
< e --- T2
B + B2 +
Co
This inequality provides the sufficient number of iterations attain the desired duality gap. In order to have an ap-appooxdmteed solution, it is sufficient to let
Tond 2 2/h 11
1+ nyd22
log
)
1
Bi + B2 + I I-expl
7d2
Next, show the convergence of the averaging version. Similarly Theorem have that for A
1,..., [.
A(Tond k+1)
E
Σ
(Plvlo")) D(G)))
< E D(g") - D(g D(gti(Tmma +)+11]]]]]]
t= kj+1
Co
<- 3(Taud -[1/xx-11/x Bi 1 B2+ +
1 1-e-8
Note that, if Eq. (31) holds, we have that - (1/x]] B + B3 + I-e
Cn
Hence, if
EA and (t) B are uniformly bounded by cn for all L. = ñ (Tend - [1/xl] + 1 1,..., TT it holds that
T.Tnnd
E
1
Σ
-
1/
t=n t= (Tand- [1/K]411
27

---

Published as a conference paper at ICLR 2022
R (T and
k+11
E
[1/k]
k=1
1 R (Tend k)-1
if
iFond
(t)
[1/k]n = [1/x])+
[1/k]n
BB
i [1/+])++
6P EP Ep.
5
This concludes the proof.
In the following, we derive the sufficient conditions for the assumptions Lemma 8 by considering the error induced by the sampling stage (Section C.2.1) 1) and the error induced at the re-weighting stage (Section C.2.2), respectively. For that purpose, we first prepare lemmas that will be used in the both proofs.
The following Lemma 9 evaluates how much the difference of two solutions in the previous step incurs difference of them in the next update in Algorithm 2. This will be used to bound the accumulated error through each inner loop.
Lemma 9, For arbitnary E [n], consider the following two updates:
21 = argmax = Elk
(x((x) + H2 - v1)
- |
811
2n12
22 = argmax
(f((2)) + H2(222 - V22
12/2 y21?
EG
2n12
Then, we have that
1
1
122 21 S
[112 H1l +
1/2 81
+ nda
n12
Proof. From the definition of ø1, it holds that
(x1) ) + H1
(21 - -y11 = 0,
(35)
nd
and similarly,
(22) H2
(22 422
n22
where (2) is an element in the lubdiiffernntal
0e, (2). From (35), we have that
1
1
(f'( + 112
(21 1/2) 11 H1 +
(2/2 31).
nd
nd2
Since E is --stonggly convex, - (z) H2(= U2) 2mAag 1= 12l is (r + mxa)
strongly concave.
Thus, we obtain that
22 21
1
1
V
(E] + H2(2z 42)
1z 1/2 2
7 + nda
2nd2
2222
v ( (f](z) + H2(2 -12)
1
1z yal2
2nd2
H2 Hi +
1
(122 2i)
7 nAg
nd2
1
<
H2 H
182 - vi
ndg
nua
where v/(=)]==z// is element of the sub-differential of z'_ This concludes the proof. n
28

---

Published as a conference paper at ICLR 2022
Now are ready to get into the main proofs which evaluate the discretization error.
C.2.1 CONVERGENCE OF AUXILIARY VARIABLES g( AND g(t)
We now fix some T, and consider the sequances {
/tAT 1EST71 10 Nt=TT+1
Lemma 10. Suppose that S (0) p(g assumptions, it holds that
ab - cc
(BT)
holds in Algorithm 2. Under these
9i 9i S Cz C2
(8T)
(36)
n(T + for = nT + 1....,
1). where C2 is defined as follows:
C2 def 20
n222 + +1)
exp
+ nda
nyd2 +1
Proof. By Lemma 9, we have that
10 (1+1)
([+11)
Dic
( he, (1.(")))))))) - / hi, (()p1"")(())d))
+ 1
10!00
-g!00)
7 + nAa
nd
(/1 plo))
nd2 gl - 0
(37)
7 nAg
so we evaluate each terms of Eq. (37).
First, we bound the first term of the right hand side of Eq. (37). Using Lemma 6, we have the following bounds:
C S exp
Σ hi, (0))g(())
SCI. Ci
sexp
Σ
h4, (0)6g(.)
SCI. (38)
nd ==RT
nd2 -=
Using (38), we have that
/
p (RT)) (O)
exp
(-
1
Σ
hi (0)6g/!)
-plo (x) exp
-plg
(-11
(-ndee
Σ
hi (0)66( do
nd2
nd
T=NT
-=RT
/
<G
plo
(0) (
do
+
(AT
exp
Σ
his (0)66(")
exp
Σ
hi (0)66//
d0
nd
nd22
T=ūT
<Ci
(aT)
(++1)
g1. s1.tt1))
(39)
nd
T=AT
Thus, combining Eqs. (38) and (39) yields we have the following bound for the first term of Eq. (37):
(() - (() do
(
mig E uz hir (0)6g(")
( exp
R his (0)6g!.)
d0
E1-1 B.r (8') )68! (r)
ply
Etelr vr Ho ( (805))
de'
E11r hip (0) (0))5(")
(TT)
E1- hip (0)6g(
d0
((
D1=1 nr h. (8) (o9(")))))))))))))}))
de'
D Rip (0')50
de'
(TT) exp
Σ(-1 u Nor (0)6g! (r)
plg ( exp
Ei r hop (0)
d0
plg(ar))(00)) exp mag E. h. (0)62) do'
2i Σイ (0'))){
no ( expp
nh, Σ! h4p (0 )0')8'''" 40
I ((g) exp(
Sxz Ei=1 Mu, (0')6))7)
do"
29

---

Published as a conference paper at ICLR 2022
p (0)
ndg Dieer Ri, (0)6s)))) -ply (0g)) expl - nig Erear he (0)6g,,) do
+
I [1lg(ar::100)) exp
ndz Etzar Ne, (0))6g))) ))) do
t-1
S2G
p
(1T) (0) exp
Σ Nr, (0)6g!
plg exp (-nde
Σ
hi, (0)8g
nd T=NT
nd2
T=NT
1-1
$207
(6T) cC
94,
(++1)
- Dir (
(40)
nda
T=AT
where we used the boundedness of each term following Eq. (38) for the second from last inequality and (39) for the last.
Finally, we bound D0 Di.
(1+1)
(++1)
by applying (40) to (37):
((+1)
1
91.
5
1
207
cC
(AT)
gr - gir
(++1)
-gir
+
1 (r-1)) - gir
lsk
+ ndp
nd2 ==7l
nd2
1
207c0
(BT)
207 + 1
t-1
Fi.
(1+1) - -,t- s
(41)
7
nd2
nd2
R=
Moreover, we have that
(TT+1)
(AT+1)]
/
NMir Wiirr (0)) (0))00 - - /hear (000
Par
S
I
7 ndg
(T)
(0) - p0g (0)) v
(42)
7 712
Combining Eqs. (41) and (42), we can apply Gronwall's lemma ("summation version" of Lemma 18): we have that, for t = nl, nT + 1, (( + 1) 1.
10 (++1)
((-1)
The
2C7
209 +1
(RT) (1+ (1+ nyd2
207 +1
R-2
cc
+ md2
(x+ + ndo
+
nd2
2C7
207 +1
â-1
(AT)
+ nuv
1
17d2 +1
CC
207
n(2C? +1)
(XT)
7 + nd2
I exp
nyd2 +1
C
where we used C1 2 1 for the second inequality. Remember that we defined C2 as C2 20?
exp (a(c(++) +1) By the definition of g(t) and g((), ()) = g((t) holds for = nT my
r nd2
1..... nT + 1) unless the coordinate i is chosen during the inner loop. Otherwise, we have (nT) for all t 2 r + 1 where 7 is the first time at which the coordinate is chosen as we have seen above. Thus, Eq. (36) holds for all t KT + 1,..., 7(T + 1) and I,...,n. n.
0.th g:
< C2c
C.2.2 PARTICLE APPROXIMATION, AND ERROR ANALYSIS OF (t))
From now on, bound the error induced by particle sampling in the re-weighting stage. The following three Lemmas relate, in three steps, the update of the continuous limit to the approximation by particle weighting.
30

---

Published as a conference paper at ICLR 2022
Lemma 11. With probability 1 8, it holds that
M Σ
(0,)r - hi, (0)p") 2 7C2
1 M
2n
ha
log
8
1=1
uniformly over all t nT, ..... n(T 1) 1 and the choice of the coordinates
(1+)==/T
((T++)-11
Proof. For t = =ī + " + 1) - 1, let
(T'
[-1
H = f2,a() -hi,())xx) exp
Σ
hs, (:)a CI,
nd
T=AT
14 [m], I =
E [m]
larl S 2B2 (r E ffT,,...,t ., 1})}
where Cz,a (1))))(0))) exp
(-nto E hi, (0')ar) do' for I i-11 and
(aaT..... 11-1). We also define
H4 == fI,a()) = -h;(-)
Σ
hi, ()a, CT
nd2 T=AT
I== = (laT, 11-1) E [m]r-äT, |arl S 2B2 (r E (iT....., 11)}
Then, we can see that H = Use[n] Hi. Then, Σ/ hi, (0)) can be rewritten as
Σ h.. (0,))
M Σ (0,)
In similar way, also have that
f (n,(())'") Eo (o).
To check this, first note that (, (0,),,00)))) (1))) j are i.idd random variables for each t = 7T + 1...., A(T + 1) 1. Indeed, pf4) is determined by g( which is updated by using continuous limit and thus is independent of the choice of particles (0, 1/) Their expectations can be evaluated as
(() exp
1 Σ h;, (0, (0,)8g
EØ, (B, = - (0)
ndg
d6
(() exp
nāo E1=1 h1 (0))g( (T) d0'
f h4 (0))(
Moreover, they are uniformly bounded
nd2 Ei hi, (0)))
exp
(M,;;00)))) & (0,)) ʃ (((T)) (0') exp
$ SC..
(43)
(-nto Σt=Ns h4 00' (09. (7) d0'
ndo
The same results are also true for t = iT.
Note that it suffices bound Z suppr. EH MM tion. For that purpose, we define the Rademacher complexity function class H' as
M E/= Fz. Lz, (0)] obtain the asser-
R(H') = E
sup SEH' M
Σ e(
1=1
where (s,)Mi is an i.i.d. sequence of the Rademacher random variables (P(s, = 1) = P(e, = -1) 1/2) and the expectation is taken with respect to both (ez) and (0, (0,) Then, the
31

---

Published as a conference paper at ICLR 2022
Rademacher concentration inequality (Mohri et al., 2012, Theorem 3.1) tells that the uniform bound on a function class H1' can be obtained as
M
2
P
sup S'EW
(f(0,) - Eelf(())) 2 2R(H)) + sup
|Ifl
log
S ö, (44)
M 3
JEH
2M
for 0 < 8 < 1. Here, let Z; sup SEH.
IM Ey= f(0,) - Eolf(()]] and then we also see that
= [a|] Zz Hence, we have that
n - it
PIZ
P(Z 2 8),
for s > O. If we substitute 8 +- maX4e()) 2R(H,) + supJeH, If v 2M log (2r), then by Eq. (44)
that the right hand side be bounded by d. Note that have already that || so I Cf for any f E H (see Eq. (43)). Therefore, we just need to bound R(H;). By the contraction property of the Rademacher complexity (Bouuchroon al., 2013, Theorem 1 11.6) or (Ledoux & Talagrand, 1991, Theorem 4.1 and its proof) and remembering the fact that the Rademacher complexity of function class is same as that its convex hull (Mohri al., 2012. Theorem 6.2), we obtain
R(H:.)
M
1-1
< I.a
E
sup
Ej exp
Σ
h4 (0,)a, - log Cz,a
(: Ihi (-1 s
M
ni22
j=1
T=AT
C1
M
1-1
sup
8j
Σ T=MT
h4, (0,)a, A2 log Ozaa
A2
M
n
j=1
2C1B2 'E
A2 2B² {
Σ Eg sh,((,)) j=1
max (fl1),e((=1):<//l55looog
A2
M
40,
2
V
M log(2n),
A2
where we used the contraction properties in the first and second inequalities, we used the convex hull argument in the third inequality, and the last inequality is by Massart's Lemma (Mohri et al., 2012, Theorem 3.3). Summarizing these evaluations, we obtain that
8C,B2
2n
))
P Z
V M
log(2n) CP
log 8
S 6.
A2
2M
The same bound also holds for = öT.
Then, by taking the uniform bound over L iT..,... ...., n(T 1) and simpifying the bound using à & n, we yield that
M
ha
(0,)5 - hv. (o) S 7C²
V
log M
2n
1=1
uniformly over = πT. ū(T + 1) and the choice of the sequence probability 1 - 5, which yields the assertion.
([T+1)-1
with
Lemma 12. With probability 1 - 8, it holds that
M
(e)
M
2n
Σ
(O,),))
Σ
S S70 (1
log
Σ =1
M
1-1
1=1
uniformly over all = öT. iT....,,fi A(T-1)-1
(T + 1) 1, = 1 1...., n and the choice of coordinates
(1+)
32

---

Published as a conference paper at ICLR 2022
Proof. Let Cp = J (())) exp exp (-nta Eisnr h4 (0)6g(r) d0. Then, noticing that r/)
ndz
(70) /Cp and |/4llo I I, we have that
((0)
(g)
((,))r))))
M
(0)
j'=
p10)
200
hg()))
Σ (1
Σj'=1 ff
M
(2)
Cp
p10)
Σ M
(1)
hz(0,)
1
j=1
2jr=1 1
Cp
M
(t)
<
1
We apply the same argument as Lemma 11 to evaluate the far right hand side 1 - -ド x((1)
by
replacing hi, (0) with in the definition of H. Then, we obtain the assertion.
From now we define eD
def
4n
ED
7
log
M
Then, we have the following bound.
Lemma 13. We have that, with probability 1 8 with respect realization of particles (0,)
(t) C2CReD
(45)
for all = hT +1..... +lT+1), 1, ..... and the choice of coordinates (i+)
(((+1)--1
uniformly.
Proof.
From Lemma 1l and 12, with probability 1 - 6, we have that
Σ hs, (0,)7 - -1
- (Op < CleD.
and that
M
M
his (0₂)
Σ hr (0,)) < <CEDD
f'==1
for all
nT 1. Therefore, by Lemma have that
19 (1+1)
9 (1+11)
M
(
Σ
(1)),)
(1(0))""00)dd0
+
1
19:
(9)
9
(1)
7 + ndg
Σj1 =1==1 3''
(3)
nd
1=1
M
(t)
(())
20
Σ
(0,)
19
(t)
9/01
(46)
7
M
Σ
(1)
nd
mAz
The second term of the right hand side of Eq. (46) can be further bounded as
M
(0)
+(0)
Σ (e,)
Ex= 3
1=1
Σx=ır>
33

---

Published as a conference paper at ICLR 2022
M
11)
(0)
M
3
+
(9)
Σ 1 3
2(0)
1=1
51 3'
M
M
Σ M (t)
Σ M
1
Σ
1 - 1
(0)
+ Σ
j'=1
]=1
E=1
f =1=1" 3
(1)
Σø'=1 f1=1 3
(2)
<C Σ
(0) 100 +
Σ
(0)
(t)
1
S'=1
SC1 Σ
(0)
10 (1)
C Σ
(1)
3
j=1
5'
M
<2C Σ
M
22C1
Σ
1
exp
1
Σ
h4,
(0,)6g(
(r)
1
exp
Σ
h4 (0, (0))
M
nd2 T=ET
M
nda T=AT
207
1-1
Σ
lg4,
(T-1)
Di,
(r+1)
(47)
n22 ==AT
fort = ūT.
(T + 1) - 1. This and Eq. (46) yields
(1+1)
Di, (1+1))
1
2CfED +
2C7 +1
Σ
t-1
19,,
(T+1)
Di, ((1)
+ nda
nd2
T=āT
fort = 7l' + I,..., 7(T' + 1) - 1. Also, it holds that
C7
(RT+1)
(8T+1)
99
9i
ED:
7 + nda
from Lemmall.
Therefore, applying Gronwall's lemma ("summation version" of Lemma 18) yields that for t ni.....,ā( + 1) 1.
2C3
2C7 +1
2C + R-2
(1+1) 1
Ti.
1
+
1+
CEED
+ mda
(7+
1
nd2
ny/2 +1
207
207 +1
1+
CRED
7 + nda
1md2+11
2C?
(m(2C9
+11)
+1)
ny/2+ +1 CieD
exp
7 + ndo
<C2CleD.
which concludes the proof of Eq. (45).
C.2.3 INTEGRATION OF TWO SOURCES OF DISCRETIZATION ERROR
From now on, integrate the convergence analysis from Section C.2.1 and Section C.2.2. Then, we extend our analysis from that of each inner loop to that of the outer loop over T = ,,n..,Tendt So far, we have been arguing only within one inner loop, while we may choose a different number
34

---

2B, Therefore, under the assumption Published as a conference paper at ICLR 2022
M of particles in each outer-loop T. Accordingly, we denote the number of particles at the T-th outer-iteration by M(T)) and add a subscript to én as
(RT) def P
1 M
An (A(ne
log
Moreover, we assume that the bound of Lemma 13 holds for every T' ................................ -I We denote by & this even. We know that P(E) 2 1-5. In the following, we analyze the convergence under the event where the bounds in Lemma 13 hold, so that |g!" 9; 1 (C256 (RT') + CC (2)) (T' 1
holds owing to Lemma 10 and Lemma 13. Therefore, the expectation with respect to the choice of coordinates taken under the condition of this event. We note that the bound in Lemma 13 holds uniformly over the choice of the coordinates and thus the distribution of the coordinates is not affected by being conditioned by the event.
In order to evaluate the effect of the discrepancy of coordinate updates on the difference of 6,, we present the following Lemma 14 on the local smoothness of dual space.
Lemma 14. Suppose that 2 2Ba and Assumption 1 holds. Then, D(g) is nba (127' + 1)-smooth with respect to g E Rn on the set A defined by
A = g E R" min min 93 (1)
min De S S max { max gi mfxx!"") (i= Ii=
where takes 1 = 0, 1.... iTonnd.
Proof.
Suppose that the coordinate i is chosen the 1-th iteration. By the definition of (0) we have
that
(8)
(4
(0) ) (h,(0,)
94
9i 11
(48)
M -1
nd
where é,' is an element of the sub-differential. The first term of the right hand side can be bounded by I as follows:
M
Σ
(1,(0,))
S max (4,(0,))
1
Σy
M
(8)
Also, Lemma 6 implies that the second term can be bounded by ndz 28, we have that
n Ag
10: (1 = 2.
Applying the same argument, we also have that
(0 < 2.
Therefore, for E [ri] that has already chosen in the algorithm until the 1-th step, we see that E 7'ssmohh on the interval between g: and ge For the coordinates that have not yet chosen,
we have that 9; (t) g 9; (0) and (e) (o) < 1. Therefore E is 7'-smooth the box A. By
I (A2n' + 1)-smooth. O summarizing this argument and applying Lemma 4, we see that D(g) is mdz
Combining the above results, we obtain the evalution of CA and sufficient conditions for the first statement of Lemma 8.
Lemma 15. Suppose that n 2 2B2 and fix an arbitrary T. if E(D((())) ((T() s
Co exp (-3T") is satisfied for all T" < T, then we have that
()) -D((") s 2
Co
7 ) (Cxx (nT)) + +C2C7 (n:))) exp ( + 1)) + (2 - 4d
(49)
35

---

the Cauchy-Schwarr inequality yields Published as a conference paper at ICLR 2022
for all t RT + 1,.... ".( + 1), where μ
2
B1+
1-0
In particular, if
(RT) I C3 exp
(AT) D
C3
3T
C
exp
(-2) 2
2
are satisfied for all T ..... Tend where
V Col2 - 11p
C3 S
(50)
C2 C2 (7 + 1/222
then it holds that for all 0,1,.... ........, .Tenn..
E[D(g")) D(g s exe (-8T),
where T [ Ail
Proof. Under the assumptions that n 2 2B, Lemma 14 holds. Fix some T and we will show the first assertion for nT + 1,...,
da
, (T + 1).
Let A, { (1 (1)(1) 8 E [0, 1]}. then Lemma 15 implies that D(g) is ndz (A27' + 1)- smooth on the set. Hence, for y0= 26 (BL+B2- co ----)
2
D(g(0)) - ((g(t))
(VD(2) -
sup gEA,
n1
2
A2n' +
+
[1 - -0
1 122" +
sup
+
2
4 nd2
gEAt
1
sup [D(g") D + (v+ 1.) (Cse (RT)
+ C2C7e(iT)) 27
$ 2
gEA.
If D(g(0)) (1) 2 0, then we have that
()) ())
D(o $ M (D(O')D(((0)))
1
+ 1 (v+ 1 (Csel (AT) + 120f ) a
Rearranging the terms we have that
D(G (0) $ / (D(o") (100))D(G(0)) (2 u)p
( xe 1 (Cze (iT) + )
(ar)
This holds as well when (() (D) < O.
Here, Lemma 5 implies that [1) is monotonically non-decreasing within the same inner loop, and thus we have that E [D - (() < E [D(g") (gg Moreover, the right hand side E [gg)) ((g(( can be further bounded by (B1 + Bg I-e exp(-3T) Co 2-1 exp(-3(T + 1)) by using the assumption and applying Eq. (33) Lemma Therefore, 2
(() (D)
Co exp (-3 3T 1)) (2 - p2 (x' 1)) (C2e (AT) (BT) + C)C7(
1
This concludes the first assertion.
36

---

Published as a conference paper at ICLR 2022
We will show the second assertion by induction. It is obvious that this is true at initialization. Suppose that Eq. (50) holds for t = 0, 1.... A.. Then the assumption of the first statement of Lemma 8 holds for T" ...... .T, which leads to Eq. (49) for t nf + 1...., n(T 1). Then we have that Eq. (50) also holds for = nT
A + 1), since we chose C3 as
VC 1)
C3
C2 22 + 1//2)
so that the right hand side of (49) is bounded by Co exp (-8(7 + 1)). Thus, the proof is completed by induction.
Although The proof for the first statement of Lemma is almost done, have some proof left for the second claim. The next two Lemmas measure the effect of discrepancy of coordinates in the dual space on the value of the main problem. Lemma 16. P(p(g) is
I mx (7 + 3)2 + 9B2) with respect to g on the set A.
Proof. We show the assertion by a direct calculation. Let Eplsl [bx] def [1:(0)) (O)dO. First, the second derivative of the loss function can be calculated
82
4, (Epis] [hvl)
0gg9k
0
nA 0gk e,' (Episl [hel) (Epia) [h,] Epial [hy] Eptol [heh)])
1
x² [e." ("risl (b+]) (Epial [bd] Epiil [by] - Perlol (1alk))) (Epial [bx] Eeiiol [ - Peill (hght))
+ 4,' (Bpio) [h:]) ( (lplll [he]] EEpis [ha] Bpis) [b,/tx]) Epist [hy]
P1plil [h.] (Epio) [hy] Epioi [hw] Eplsl (hr (EEpill [hehy] Epiia [ha] lEpiol (hattyhel)))
1
- nº (7 + 3B2),
where we used smoothness of & for the last inequality.
The regularization term can be split into the following two terms:
A1 /
plo(00)1011ι01 + X2 / plg)(0) log (O)
((())lol(0)d0g + Ao log / (lo(())o
We already confirmed nda smoothness of the second term in Lemma 4. As for the first term, have that
02
1
Eplai [helg: 9s
0gigk
n
1
0²
X2TEppo [hy]
(Eplo) [hs] Eplo) [hy] - Eplol (hrh)]) gi
ndą 0gk
1 a Eepll (nglhu)) n i=1
2\a (Eplsl [by) Eplal [Ma]]
Σ ( (Bepie) [hs] Epiol [ha] Eppol (lache)) Epiol [hg]
x²
-Eplgl [he] (Epis) [hy] Epia) [hy] Eptgl [hyhel) (Epto) [hahy] Epiol [ha] Epta) (h,hyhul) gil
? (2)2 + 6B2),
where we used Lemma 6 for the last inequality to bound gi-
Combining the two inequalities and the result from Lemma obtain the assertion.
37

---

Published as a conference paper at ICLR 2022
Lemma 17. Suppose Eq. (50) of Lemma 15 and its assumption hold for t 0, 1,... nT Then, we have that
E -
27 + 6Λ2 -18B2
Co VChe* (2 1il
Bi + B2 + Co +
- V2(7222 1) exp(-87)
\ V
where
T
[c]
Proof.
By using the smoothness of P(pig)) by Lemma 16, we have that
P - ())
| 1 g
sup SEA.
nd 7 + 3)2 + 9B2
+ 3A2 + +9B2
sup /VP(ll))]22
(g(0) - 1
nd2
sup P(p(g)) P()lo1))
7 + 3d2 + +9B2
(Czel (RT) + C2CRe (RT)
X2
where T = [ - 1. P(rlo(0)) - P(pla(01) 2 0, then we have that
1) - 1))
$ P(plg(()) - P(p(g*))
+ 312 + -92
(C2E6 (RT) + +22(76
(aT)
X3
7 3A2 + +9B2 |19(0
g'lla (C2e (RT)
C2C
AByn
Since D is -stronggl convex, we have - - S V (D(o)) D(g Therefore, it
holds that
<² 7 + 3A2 9B2
V(D()) D(g(o)) (C2e (RT) CC2 (m))
X3
which holds as well when ((p(g(01) P(plo(01) s 0. By taking expectations, yields that
E ()) - P(plo"0D]
7 + 312 + -9B2 EE
(10(())-D((()))
(cre (RT) + +2C966m) (AT)
7: + 3\z + +9B2
E
(D(gg)) D(G
(C2e (RT) + C2C7-600 (K2)
7: + 312 + 9B2
/E
[D(g)) D(G()) + D(G()) D(g(t)] (C2E (AT)
+ +2277000
(T))
27 6 18B,
Co 1-e-a + CoC2Cs exp (-5T),
B1 + B2 +
where we used Lemma and Lemma 15 for the last inequality.
Coll2 Col2-Pu0 p)a
Furthermore, substituting C2C3
we have the assertion.
V 2e (y
38

---

Published as a conference paper at ICLR 2022
Finally, will prove the Theorem 3, which gives sufficient conditions about sampling accuracy and number particles for Theorem
Theorem 3. Suppose that Assumptions 1 and 2 hold. When we choose an Option (A), suppose that set and that
Tond 2
1
log (n+ 722
(Bi + B2
expl 8)
(51)
n2222
Cp
in Eq. (30) in Lemma 8. Furthermore, when choose an Option (B), where we choose solution randomly from t = ñ(Tend - -1= I)..... i.Tend. suppose that we set Tend so that
1 (1+
1 ny222
Co
Tend 2 - h
log
+
CP
K
where
(2+ n7d2
2
Co
C max
Bi + B2 +
(52a)
1 exp (-8)
4Co,
(52b)
(87 (87+2200+72BB)) Bi +o+ Co v Coe (2-N)H
(52c)
N (2ylo/2-11)
(t) Then, the condition imposed E and E in Lemma 8 can be satisfied for both options under the event & if the following inequalities are satisfied:
(6)
2B2
11 2
A2
C
(nT)
< C3 exp
ST
(53)
2
(nT)
<
C3
exp
(-2 2
5T
M M(T)
(TT)
49C4
log
[AnTend
exp (3T),
C7
C
for T' , Tend 1. That is, under the conditions above, E [P((g] - D())) s ep is satisfied for the solution 9 of either Option A or B.
Proof. As for the Option (A), according to Lemma 15, Eqs. (51) and (53) are the sufficient conditions for the first statement of Lemma the event E.
Next, we will show the convergence of Option (B). In the same way, Eqs. (52a) and (53) assure (RT) is smaller than Che-37 for every T 1,..., Tend, according Lemma 15. Furthermore, that CA
Eqs. (52b) and (52c) states that both e and ég
(t) (t) uniformly bounded by 3 for all t ä(Tend
(t) [1]) + ,,..., nTeend Note that the evaluation of E follows from Lemma 15, and that of ég is derived in Lemma 17.
Therefore, all assertions are proved.
Theorem essentially shows Theorem 1. More detailed explanations can be found in the following.
Proof of Theorem 1. Theorem 3 gives the sufficient conditions for Lemma 8. Thus, we have that, for Option (A), with Egs. (51) and (53), an expected duality gap E (TT - Dgg cp can be achieved, conditioned by the event & whose probability no less than same way, for Option (B), with Eqs. (52a)-(52c) and (53), an conditional expectation of duality gap ((naa)) ( < 6p can be achieved. The sufficient condition of fend is then
(MTema)) { <
In the
obtained by noticing lend 7.. and setting č C3/7, O₂ = Of and Ó₂ 1+ me4g
2C,
39

---

Published as a conference paper at ICLR 2022
Next, evaluate the constants follows. First, note that C2 Ci by their definitions. Moreover, [((207+])
C
(2t1///
C2 = 0x,,0'1/0// [ exp my
= (22 exp
(((02+1)) my
/C)(2-4)
20 exp [ +1] and Of Cg. Finally, the evaluation of O3 can be
where we used C2
obtained by bounding Cy using 1/(1 exp( -)) ( and s + nnd2 I 1 + 24
0(1).
C.2.4 DISCRETE GRONWALL'S LEMMA
In the proof' of Lemmas 8, 10 and 13, used discrete version of the Gronwall's lemma. We give its formal statement in the following lemma.
Lemma 18 (Discrete Gronwall's 's lemma).
(I) (Vanilla version) Suppose that a sequence of real numbers (un) satisfies
Un-1 S athn + bu41 (vn 2 20),
for a positive real a > 0 and a sequence of real numbers (bn). Then, holds that
Un S a un
by.
(2) (Summation version) Suppose that sequence of real numbers (un) satisfies
Uy auk B (Vn 2 1),
for a positive real a > 0 and a real number B. Then, holds that
Un S (1 a1+a)"~1 + aun).
Proof. The assertion is just an application of the standard discrete version of the Gronwall's lemma. However, give a proof for completeness.
The first assertion is direct consequence of Lemma 5.1 of Mischler (2019). The second assertion is also proven by Lemma 5.2 of Mischler (2019). Indeed, Lemma 5.2 of Mischler (2019) yields that
n-1
Un S +
a(1 + a)m-1-kp B a(1 + a) -
A=1
The right hand side is equivalent to
B
1
(1+ a)*
+ +a (11++a)=-1u) a)n-1u
k=0
a) n- -
1
a(l + -)
+ a
-(1 + 1
(B + auo);
which yields the second assertion.
D ADDITIONAL EXPERIMENTS AND DETAILS OF EXPERIMENTAL SETTINGS
Here, give additional experiments and more detailed explanations of our experimental settings. All experiments were conducted under Google Colaboratory. All codes of the experiments in the main text and the following additional experiments are provided in the supplementary material.
40

---

Published as a conference paper at ICLR 2022
Additional experiments for single neuron teacher network teacher network consists of a single neuron:
We consider setting where the
y: [ (etTTT) x;)) Ee,
where the imput dimensionality is d = 10, wi E R is generated from the uniform distribution the unit sphere Sd-1 in the d-dimensional Euclidean space, Xi E R4 is distributed from a standard normal distribution on Rd, 6, E I oheys the Gaussian distribution with mean and standard deviation 0.5 and the activation function a is ReL.U. We generated n. = 100 trainig data points from this model and trained a student network by using P-SDCA, PDA, and a naive SGD. Here, as the number of particles, we employed M = 200 and the student network is the tanh neural network (h)(2) tanh(w' x + b) with 6 (w, b)). We employed the squared loss {(f. (y u) for the loss function. We conducted experiments with three patterns of regularization parameters: A = 2 = {0.01, 0.001, 0.0001). The naive SGD was trained to minimize the primal objective for the student network with M 200 without the entropy regularization, i.e..
1
M Σ
A M
M
ho (3s), Vi
+
10ml1
M
m11
We employed the ULA for sampling the particles for PDA and P-SDCA where we ran the algorithm with 200 steps in each re-sampling stage. The step-sizes v used for sampling (ULA) and SGD were set as 10
We plotted the excess primal objective value against the number of gradient evaluations for each method in Figure 2. The excess primal objective value is computed by subtracting the minimum objective value throughout the optimization and all optimization methods from the current primal objective. Since we are approximating the primal variable by finite particles, we cannot determine the entropy term exactly. To handle this issue, we employed the k-NN entropy estimator (Kozachenko & Leonenko, 1987; Brodersen, 2020) to estimate the entropy term from the finite particles. This is also applied SGD to compute the primal objective with the regularization term for the solution obtained by SGD. From Figure 2, we can see that P-SDCA and PDA behave almost identically for large regularization parameter X. In that regime, SGD once decreases the primal objective rapidly but increases it afterward. This is because SGD is minimizing an objective without the entropy regularization. For = 10 10-2 P-SDCA and PDA does not decrease the excess primal objective below 10 This would be due to the instability of the entropy estimator. For small X, P-SDCA outperforms other methods after some iterations because it yields linear convergence. This observation supports our theoretical analyses.
that
Additional experiments for a teacher network with multiple neurons We consider a setting where the teacher network consists of multiple neurons:
Mf Σ awwil x.) + Es,
Vi
where the width of the teacher network is M, 10, the input dimensionality is d = 5, the activation function the teacher network is (--) = tanh(-), the noise drawn from Normal distribution and
P-DDCA
P-SOCA PDA
P-SOCA
Docesp traliing ESE
training
10
10
65
1.0
1.5
2.0
25
0.6
0.5
1.0
3.5
2.0
2.5
20
0.5
05
1.6
1.5
2.0
2,5
Gradient evalueationn
Gracient evaluestions
Gradient evalueations
(a) A 0.01
(b) À = 0.001
(c) = 0.0001
Figure 2: The number of gradient evaluations versus the excess primal objective for P-SDCA, PDA and SGD for teacher network with single ReLU neuron.
41

---

Published as a conference paper at ICLR 2022
its Signal to Noise Ratio (SNR) is 7.0. For each method, selected the optimal step size: 10-4 for ULA in P-SDCA and PDA, and 10-5 for SGD. In the following, "step size of P-SDCA or PDA" refers to the step size for ULA in the sampling scheme. In addition to that, we plotted SGD with its step size 10-1 to see how the step size affects the convergence of SGD. The purpose to include SGI) as well as P-SDCA and PDA is to how the entropy term will behave when we do not explicitly include the entropic regularization term into the objective function.
Roguartudd tosd lua
trching
6 100
G 10
0x10
Phsh PAGDCA
ASDCA
10
regudatted
4x10-1
3 10
3x10
x -0
xx-0
2x1001
1.0
00
Gradienn
105
Gradient tralcaators
3xe
Grad ann woe
(a) A2
0.01
(b) A2 0.001
(c) A2
0.0001
Figure 3: The number of gradient evaluations versus the excess primal objective for P-SDCA, PDA and SGD for teacher network with multiple tanh neurons.
We compared our method with SGD and PDA with the regularization parameters A = 0.01 (fixed) and A2 0.01, 0.001, 0.0001. We repeated the experiments 10 times and took the average of primal objective values, which includes regularization terms. The primal objective value is plotted against the number of gradient evaluations in Figure 3 where the error-bar represents the standard error over the 10 repetitions.
For A2 = 0.01, we can see that particle methods (P-SDCA and PDA) showed faster convergence than that of SGD. Next, as for X2 0.001, we can see that P-SDCA converged faster than SGD and PDA. Although PDA and SGD with step size 1) = 10-5 also converged eventually, the behavior of PDA in the intermediate stage was quite unstable and SGD with 1) 10-5 showed quite slow convergence. These observations support the superiority of our method. Finally, as for X2 = 0.0001, SGD with step size 17- 10 5 was the fastest due to the small regularization parameter 22 for the particle methods which makes the convergence of them slower. However, we still see that P-SDCA outperformed PDA in terms of both convergence speed and stability. The convergence of SGD is strongly affected by its step size. In fact, SGD with larger step size (1- = 10 10-4) could decrease the objective rapidly at the early stage, but it eventually diverged. This would be partly for the same reason as in the single neuron setting, that is, due to the lack of the entropic regularization in the objective of SGD. In summary, we can see that P-SDCA properly minimizes the corresponding objective function in wide range of X2 and shows faster and more stable convergence than PDA as indicated by our theory.
Experiments on MNIST dataset We conduct numerical experiments to illustrate the "feature learning" aspect of mean field neural network. For this purpose, we run the algorithm in a binary classification task that separates "2" and "4" in the MNIST dataset. We subsampled n 1000 training examples and trained the two layer tanh network with width M = 2500. The step size for ULA in P-SDCA was set as 1) 10-5 and the regularization parameters were set as X 10-2 X2 10-4.
To visualize the adaptivity of the feature map, we first plot the evolution of singular values of the first layer's weight matrix, W [wi, W2,..., WMM RdxM until convergence, when the solution achieves 99.8% training accuracy. We can see that the singular values of the network trained by P-SDCA drastically changed during optimization. This indicates that the mean-field neural network does not "freeze" and can adaptively learn features during optimization by P-SDCA. We also plotted the same quantity for SGD updates in the NTK regime. In contrast, we observe only small change of the singular values. This is because the first layer's parameters do not change much during optimization in the NTK regime; in other words, SGD the NTK regime does not perform adaptive feature extraction. In that sense, we expect that training in the mean field regime is effective in setting where adaptive feature extraction required.
In addition to the singular value evolution, we evaluate how well the extracted features are "aligned" to the label. For that purpose, define kernel function kw (I, x²) >M o ( for weight matrix W [wi, W2:--.. WM] E R Then, we define the kernel alignment (Cristianini
42

---

Published as a conference paper at ICLR 2022
100.0
6.0
80.0
60.0
50.0
Singuat valuee
30.0
20.0
233
10.0
3,0
SV3
8.0
3Y4
2.5
5V4
5V5
SVS
5.0
others
2.0
others
steps
(a) P-SDCA (mean field regime)
(b) SGD (NTK regime)
Figure 4: Evolution of singular values of the first layer's weight matrix during optimization. The figure shows top 5 largest singular values in addition to the & x 10-th largest singular values for k 2,3, 4, ...; 9 and k x 100-th largest singular values for k = 1, 1,2,3.
et al., 2002) on the training data (24, 9:) as
(Kw: My )r
A(kw) = VKW Kw Flyy yy 1EF
where B) : Bij. Kw = (kw (Ii, +1) and y (y),..., Un) We also define the kernel alignment on the test data the same manner. We can see that the kernel alignment represents how strongly the kernel function defined by the features after the first layer is aligned to the target signal.
The kernel alignment of the solutions obtained by each method (P-SDCA, SGD NTK regime, SGD in mean field regime) is depicted in Figure 5. We can see that P-SDCA (and SGD in the mean field regime) properly improves the kernel alignment, which means that P-SDCA can extract informative features adaptively depending increase the kernel alignment due to the parameters almost "frozen" at initialization. This experiment highlights how well the mean field neural networks can execute feature learning.
the data. On the other hand, SGD in the NTK regime does not
1.0
0.8
0.8
Allgmment
0.6
0.2
SED (MTK)
02
SGD (NTK)
SGD (MF)
0.0
0
50
100
150
200
250
300
0.0
50
100
150
200
250
300
Outer loop steps Steps
Outer loop steps Stops
(a) Kernel alignment on training data
(b) Kernel alignment on test data
Figure 5: Evolution of kernel alignment of the extracted features during optimization.
43


Reviews:
[{"review_id": "nNPGk-txkK", "reviewer": "Reviewer_pSnQ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Thanks for the clarification. I'd like to keep my score.\n\n======================\n\nThe paper is clear in its exposition, despite involving many parameters. In the context of MF two-layer neural nets, the idea is new. The proofs seem correct. I also appreciate that the authors have made some honest remark concerning $\\lambda_2$, which is clearly a deficiency in this approach.\n\nI do not have much to complain, though a few things could be considered:\n- An explanation of the role of the re-weighting stage is needed. It is said to reduce the number of resampling times $\\tilde{n}$ down from $n$ and hence brings down the complexity. It would be good to be more quantitative about this. Would the result be $O((n/\\lambda_2\\gamma)\\log(1/\\epsilon_P)$ as mentioned in page 7?\n- It is said that Mei et al 2018 has the error growing exponentially with time. I think here the same also applies. One can take a different view that is to fix a target error to argue against this exponential dependency, but the same view applies to Mei et al 2018 as well.", "summary_of_the_paper": "The paper considers optimization of MF two-layer neural nets (the infinite-width version with entropic regularization). Specifically the paper proposes to do dual coordinate ascent, which allows for exponential convergence rate, together with a particle approximation scheme. The paper shows that in the discrete-time setting, an improved convergence rate is obtained (w.r.t. desired error $\\epsilon_P$, the number of data points $n$, the regularization $\\lambda_2$).", "summary_of_the_review": "The paper proposes an interesting variant of dual coordinate ascent that achieves improvements in the convergence rate. There is not much to criticize since the paper has optimized the complexity to a good extent, except for the inherent difficulty with the approach that depends crucially on $\\lambda_2$.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "ZdyggwKGG7H", "reviewer": "Reviewer_CmPk", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper concerns optimizing mean-field neural networks. The problem is inherently challenging because the infinite width limit leads to a functional formulation of the representation and, consequently, the optimization must be carried out over the space of probability measures by solving a PDE. Particle-based approaches alleviate some of the challenges, but, as the authors note, the original mean-field papers deal with discretization simply by taking the limit $\\Delta t\\to 0$. Here, the authors propose to instead work directly with a space-time discretized dynamics.\n\nThe authors examine a particular variant of the mean-field objective function in which an entropic regularization term is added.  However, unlike these previous works, the authors take advantage of Fenchel duality to convert the problem into a maximization problem in $\\mathbb{R}^n$ as opposed to an infimum over $p\\in \\mathcal{P}$. \n\nThe proposed algorithm relies on first sampling parameters from a proxy probability measure $\\rho^{(t)}(\\theta)d\\theta$ that is close in total variation distance from the target measure $\\rho[g^{(t)}]$. How is it ensured that the proxy is actually close? This strikes me as a challenging thing to verify.\n\nBecause the algorithm requires sampling a high-dimensional distribution, this step is performed infrequently and the particles are reweighted iteratively. The sampling with ULA or MALA is the basis of the convergence results, using a LS inequality. Of course, the exponential decay rate will be unfavorable if the target distribution is metastable. There's no real discussion of the nature of the parameter distribution. Is there some reason we should expect that it's easy to sample?", "summary_of_the_paper": "This paper presents an optimization alogrithm for mean-field shallow neural networks. The algorithm optimizes the parameter measure using a particle-based implementation of the stochastic dual coordinate ascent algorithm. The authors establish convergence results, namely an exponential convergence rate owing to the convexity of the dual problem.", "summary_of_the_review": "The paper outlines an algorithm and carefully analyzes it. The algorithm appears to be effective in a simple synthetic example. However, the assumption of efficient sampling of the parameter distribution requires substantial additional justification. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "YI45TZHIjk", "reviewer": "Reviewer_ny66", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Strengths \nThe paper is well-written, technically sound and I have no doubts about the correctness of the results (though I did not check the proofs in the appendix carefully). \n\nThe presented dual particle optimization scheme appears to be novel and might be useful for other problems on probability density than mean-field neural networks. \n\n### Weaknesses\nA) Motivation\n* Since the main contribution is an algorithm (and its analysis) for Eq. (1), it would be great to motivate this problem a bit more for the uninitiated reader. For example, is it expected to be preferable over standard neural network training (eg. due to the regularization or convexity) or is the goal rather to gain a better understanding on standard training through it? In the experiments, the entropic regularization seems to help generalization, so perhaps this could be mentioned as a motivation (with additional explanations). \n \nB) Clarity\n* It would be helpful if the paper could provide an explicit definition of the set P of probability density functions on R^{\\tilde d} and a proof of Proposition 1. I am confused about this point, since in order for Theorem 15.23 of Bauschke et al. (2011) to be applicable in Proposition 1, it seems we have to be in a Hilbert space setting. But I am unsure whether the set of probability density functions considered here is a subset of a Hilbert space.  As a remark, in case the Hilbert space setting does not suffice, note that the original work by Rockafellar \"Duality and stability in extremum problems involving convex functions\", Pacific J. Math, 1967 already proves the duality theorem in the more general setting of paired topological vector spaces.\n\n* It was a bit tedious to follow the proof of Lemma 1 (Lemma 3 in the appendix on page 19). I can see that the elements of the Hessian are bounded by one, but it would be nice if the paper could add a clarifying sentence how this yields Eq. (13).  Also, the sentence \"Note that the bound (12) of the Hessian in the form of the covariance matrix regarding h_i\" did not make sense to me, and Eq. (12) appears not to be a bound but rather an equality. \n\nC) Experimental evaluation \n* Would it be possible to run the method on a small binary classification example, e.g. on MNIST (2 vs 4)? The current experiment in Figure 1 suggests that the particle method generalizes much better than SGD (mean field / NTK), and it would be interesting to confirm this on a small real world dataset rather than in the teacher/student setup. Perhaps one could even show the learned 28x28 image-filters to illustrate the \"feature learning\" aspect of the mean-field algorithm (opposed to the NTK one). \n", "summary_of_the_paper": "The problem of training an infinitely wide neural network with one hidden layer can be written as an optimization problem in the space of probability distributions, where each atom in the distribution corresponds to a hidden neuron. The paper proposes to solve instead the dual problem (of an entropy regularized variant), which in the case of finite data is a finite-dimensional optimization problem.  The main contribution is the conception and analysis of an efficient dual coordinate ascent scheme to solve the dual, which in theory and practice requires fewer iterations to find an epsilon-accurate solution than existing methods. \n", "summary_of_the_review": "The paper proposes and studies an interesting optimization method for problems in the space of probability density functions. The main application considered are mean-field neural networks, but  I can see how this algorithm will be very useful also for other problems over probability densities which are ubiquitous in ML.  \n\nThere are several smaller issues mostly regarding the technical clarity and motivation in the paper (see main review). If these are clarified/addressed in a revision, I am also willing to increase my score. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "2ymXKLc90mE", "reviewer": "Reviewer_M8tQ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Comments:\nThis paper introduces techniques to directly solve the infinite-dimensional mean-field model for two-layer NN. It allows finite step size and improves the convergence rate to linear by solving the dual problem and assuming primal smoothness. My main comments are as follows:\n\n* Proposition 1 needs more careful treatment. The authors claim the Fenchel–Rockafellar strong duality holds without formal justification. Note that, the strong duality holds when certain domain inclusion condition [Eq. 15.33, Bauschke 2011] is verified, which is not trivial for Eq. 2. One should not think the domain of l_i^* would be the whole R as the loss could be a logistic loss, which is not coercive (see [Theorem 11.8(d), R1].\n\n* Intuitively, the update 4+5 in Algorithm 1 is very similar to running one-step SGD for the i-th data point. After computing step 4, the \"weights\" are updated by step 5. Is there any connection with a primal algorithm given this similarity?\n\n* On stepsize for P-SDCA: it seems the step size length in the Algorithm is n*lambda_2, which is from the lower estimation in Lemma 1. However, on Page 40, additional experimental results, they said \"we selected the optimal step size ... for P-SDCA and PDA\". If n=100 as on Page 39, then the step sizes for Fig 3(a-c) would be actually different. This should be made clear.\n\n---\n[R1] Rockafellar, R.T. and Wets, R.J.B., 2009. Variational analysis (Vol. 317). Springer Science & Business Media.", "summary_of_the_paper": "This work solves the entropy regularized mean-field model for a two-layer neural network. As the Fenchel–Rockafellar dual of that problem is a finite dimension problem, they consider an SDCA-type scheme to solve the dual problem, which still contains an integral term and they use the Langevin iterate procedure to solve the subproblem approximately. In contrast to the existing analysis for SGD in the mean-field regime, the new algorithm has a fixed length dual step size and has a fast linear convergence rate.", "summary_of_the_review": "This work introduces a dual algorithm to solve the infinite-dimensional mean-field training problem. The update step is solved approximation by estimating an integral term in every iteration. The new technique is interesting and the new linear convergence rate improves the prior result.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This is a solid paper and considers the problem of training a wide neural network with a single hidden layer. This can be framed as an optimization problem in the space of probability distributions with a suitable entropy regularization, where each atom in the distribution corresponds to a hidden neuron. The dual of this problem (for finite data) is a finite-dimensional optimization problem and the paper proposes a particle based coordinate ascent scheme.
The paper provides some convergence rate results. After the rebuttal, the authors have also included more experimental/numerical results.

The authors have answered the concerns raised by the reviewers and overall, the paper can be accepted:
The presented approach appears to be sufficiently novel and might be useful in other settings.
The presentation is clear and easy to follow for such a technical paper; the paper is well organized.
The limitations of the approach are clearly stated (dependence on the regularization parameter for entropy term that may be hard to select)
Decision: Accept (Poster)