P ARTICLE S TOCHASTIC D UAL C OORDINATE A SCENT : E XPONENTIAL CONVERGENT ALGORITHM FOR MEAN FIELD NEURAL NETWORK OPTIMIZATION
A BSTRACT
We introduce Particle-SDCA, a gradient-based optimization algorithm for two-layer neural networks in the mean ﬁeld regime that achieves exponential convergence rate in regularized empirical risk minimization.
The proposed algorithm can be regarded as an inﬁnite 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 particle method that approximately solves the dual problem, and an importance re-weighting technique to reduce the computational cost.
The convergence rate of our method is veriﬁed by numerical experiments.
I NTRODUCTION
A major challenge in developing an optimization theory for neural network is to handle the nonconvexity of loss landscape.
Recent works showed that overparameterization is one key property that makes global optimization possible.
In particular, the effectiveness of overparameterization has been extensively investigated in the theory of Neural Tangent Kernel (NTK) and the mean ﬁeld analysis .
The NTK analysis considers a linear approximation of neural networks and casts the gradient descent dynamics to that in a corresponding reproducing kernel Hilbert space (Jacot et al., 2018).
Since the whole argument can be carried out essentially in 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 ﬁxed kernel fails to 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 et al., 2019).
In contrast, the mean ﬁeld 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 is 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-Eijnden, 2018; Sirignano & Spiliopoulos, 2020).
A beneﬁt of the PDE formulation is that the objective function becomes convex in the space of measures.
This being said, since the dynamics is essentially inﬁnite dimensional, most convergence results are shown in the continuous limit (inﬁnite-width limit).
Also, many existing guarantees only handle the continuous time dynamics, due to difﬁculty in analyzing the time discretization of the PDE.
Indeed, Nitanda & Suzuki (2017);
Chizat & Bach (2018); Rotskoff & Vanden-Eijnden (2018) proved convergence to a global optimal solution for essentially inﬁnite width networks under continuous time settings.
On the other hand, global optimality may be shown under less restrictive settings if we consider the noisy gradient descent algorithm; the corresponding stochastic dynamics can be formulated as McKean–Vlasov dynamics, and its convergence has been studied in Mei et al.
; Hu et al.
However, most of these works establish convergence result for the discrete-time ﬁnite-width method by bounding the difference from its continuous-time and inﬁnite-width limit counterpart, which usually incurs a large discretization error.
Hence, it is important to develop a practical algorithm with optimization guarantees for ﬁnite 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 speciﬁcally, the algorithm realized time discretization by extending the dual averaging technique for ﬁnite dimensional convex optimization to the inﬁnite dimensional setting, and realized the space discretization by sampling ﬁnite number of particles via the gradient Langevin dynamics.
Eventually, the method achieves O (1 /t ) convergence with respect to the outer-loop iteration t .
Although O (1 /t ) is optimal in an online optimization setting, there is still room for improvement for a ﬁnite sample size setting.
Motivated by this observation, we propose a 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 ﬁeld regime that achieves linear convergence (w.r.t.
the outer loop steps) for regularized empirical risk minimization 1 .
We adopt the stochastic dual coordinate ascent (SDCA) technique (Shalev-Shwartz & Zhang, 2013a; Taká ˇ c 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 a ﬁnite 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 ˜ n -iterations, which provides signiﬁcant computational speedup.
Furthermore, unlike many existing analyses concerning the time-space discretized version of McKean–Vlasov dynamics (Mei et al., 2018; Bou-Rabee & Schuh, 2020; Bou-Rabee & Eberle, 2021), 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 ﬁxed (non-vanishing) step size.
Our contribution can be summarized as follows (for comparison with existing works see Table 1):
• We propose a new optimization method for mean ﬁeld representation of neural networks that is efﬁcient in the ﬁnite 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 0.
• The integral required in each update is approximated by an average over ﬁnite 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-weighting scheme for particles, we only resample particles once in ˜ n updates, where ˜ n can be as large as the total sample size n ; this signiﬁcantly reduces the computation cost.
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
P RELIMINARIES
We ﬁrst 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 a twolayer neural networks, where the input and output spaces are denoted by X ⊂ R d and Y ⊂ R , respectively.
Let h θ : X → Y be one neuron with the parameters θ ∈ R ˜ d .
For example, we may set h θ ( x ) = tanh( r ) σ ( w (cid:62) x ) for θ =
( r, w ) ∈ R × R d and an activation function σ (the tanh operation is merely to ensure boundedness of the output).
In the mean ﬁeld regime, we consider neural networks represented as an average of neurons:
f Θ ( x ) = 1 M (cid:80) Mm =1 h θ m ( x ) where
Θ =
( θ m )
Mm =1 ⊂ R ˜
d is a set of parameters.
The continuous limit in the mean ﬁeld regime is obtained by taking M → ∞ , and by an analogy to the law of large numbers, the function f Θ converges to the following integral representation of a two-layer neural network:
[Equation]
where p : R ˜ d → R is a probability density function on R ˜ d
representing the weight of parameters.
Here, we denote by P the set of probability density functions on R ˜ d .
As in the typical mean ﬁeld analysis (Nitanda & Suzuki, 2017; Mei et al., 2018; Chizat & Bach, 2018), we aim to optimize the density function p ∈ P so that the neural network f p accurately predicts the output y ∈ Y from the input x ∈ X .
To deﬁne the empirical risk and predictive risk, we let (cid:96) ( z, y ) :
Y × Y → R be a convex loss function, such as the squared loss (cid:96) ( z, y ) =
( z − y ) 2 / 2 for regression, and the logistic loss (cid:96) ( z, y )
= log(1 + exp( − yz )) for classiﬁcation.
For each ( x i , y i ) , we use the notation h
i ( θ ) and (cid:96) i ( f ( x i )) to indicate h θ ( x i ) and (cid:96) ( f ( x i ) , y i ) respectively.
To estimate the density function p ∈ P , we minimize the regularized empirical risk deﬁned by
[Equation]
where λ > 0 is a regularization parameter and KL( p, N (0 , σ 2 I ))
represents the KL divergence from the normal distribution N (0 , σ 2 I ) with mean 0 and covariance σ 2 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:
[Equation]
where λ 1 , λ 2 > 0 are the regularization parameters.
Importantly, the objective is convex with respect to the density p ∈ P although it is non-convex with respect to each parameter θ .
We make full use of this convexity in the following analysis.
The main difﬁculty to optimize (2) stems from the following two factors: (i) there is an integral with respect to p that should be approximated by a computationally tractable scheme, and (ii) there is a summation over n data points, which could lead to costly gradient evaluations.
To overcome the former difﬁculty, we sample a ﬁnite number of particles from the density p .
As for the later, we adopt the stochastic dual coordinate ascent (SDCA) technique, which we explain in the next section.
P ROPOSED METHOD : P ARTICLE -SDCA
We now explain our proposed Particle-Stochastic Dual Coordinate Ascent (P-SDCA) method.
3. D UAL PROBLEM AND ALGORITHM DESCRIPTION
The SDCA method relies on 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.
Proposition 1 (Fenchel’s duality theorem) .
Suppose that (cid:96) i : R → R is a proper convex function, h i : R d → R is bounded, and the primal problem satisﬁes inf p ∈P P ( p ) >
−∞ , then the following duality holds:
[Equation]
where q [ g ]( θ ) := exp (cid:110) − 1 λ 2 (cid:0) 1 n (cid:80) ni
=1 h i ( θ ) g i + λ 1 (cid:107) θ (cid:107) 2 (cid:1)(cid:111) and (cid:96) ∗ i is the convex conjugate 2 of (cid:96) i .
Moreover, p ∗ ∈ P and g ∗ ∈ R n are both optimal solutions of primal and dual problems if and only if f p ∗ ( x i ) ∈
∂(cid:96) ∗
i ( g ∗ i ) and p [ g ∗ ]
= p ∗ , where p [ g ]( θ ) :
=
q [ g ]( θ ) (cid:82) q [ g ]( θ (cid:48) )
d θ (cid:48) and ∂(cid:96) ∗ i denotes the sub-differential of (cid:96) ∗ i .
This proposition can be derived from Theorems 3 and 8 of Rockafellar (1967) (its ﬁnite dimensional version can be found in 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 ∗ =
p [ g ∗ ] .
In the dual problem, updating one coordinate g i can be carried out by picking up one data points ( x i , y i ) .
The main strategy of SDCA is to randomly pick the sample index i ∈ [ n ] 3 and update the corresponding coordinate g i by maximizing the dual (i.e., dual coordinate ascent).
Before we state the algorithm, we assume the following condition.
Assumption 1.
(A1) (cid:96) i : R → R is (1 /γ ) -smooth 4 for all i ∈ [ n ] .
(A2) sup θ ∈ R ˜
d | h i ( θ ) | ≤ 1 for all i ∈ [ n ] .
The ﬁrst assumption (A1) is satisﬁed, for example, by the squared loss (cid:96) i ( u )
=
( y i − u ) 2 / 2 with γ = 1 and the logistic loss (cid:96) i ( u ) =
log(1 + exp( − yu )) with γ = 1 .
As for the second assumption (A2), our analysis is valid for bounded neurons, and the upper-bound 1 is chosen just for simplicity and can be replaced by another value.
In fact, if we assume sup θ ∈ R ˜
d | h i ( θ ) | ≤ R instead, the analysis can be reduced to the setting of R = 1 by rescaling h i and the loss function, i.e., R can be absorbed into the smoothness parameter, which then becomes R/γ .
As stated above, the strategy to optimize the dual problem is to (i) randomly pick i ∈ [ n ] and (ii) update g i by approximately maximizing D ( g ) with respect to g i with other coordinates ﬁxed.
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 ∈ R n and δg ∈ R n , the second term of D ( g ) can be bounded as
[Equation]
Using this inequality, we can derive the update rule as a maximizer of an lower bound of the objective as follows.
When we update only the i -th component (i.e., δg j = 0 ( ∀ j (cid:54) = i ) ) for randomly chosen index i , the dual objective can be bounded as follows:
[Equation]
Hence, by substituting ˇ g ← ˇ g ( t ) and maximizing the right hand side with respect to δg i =
g
i − ˇ g ( t ) i , we obtain the following ideal version of the update rule of ˇ g ( t ) :
[Equation]
while the other components are not updated: ˇ g ( t +1) j
= ˇ g
[Equation]
t )
j ( j (cid:54) = i ).
Note that this update requires only one data point while the full gradient descent requires the whole n data points.
Along with the updated dual variable ˇ g ( t ) , we update the primal variable as ˇ p ( t ) =
p [ˇ g ( t ) ] .
The whole optimization procedure based on this update rule is summarized in Algorithm 1.
Require: training data { ( x i , y i ) } ni =1 and number of iterations t end 1: Choose ˇ g (0) i
s.t.
| (cid:96) ∗ i (cid:48)
(ˇ g (0) i ) | ≤ 1 ( i = 1 , . . . , n ) and (cid:96) ∗ i (ˇ g (0) i ) ≤ (cid:96) ∗ i (0) 2: for t = 0 , 1 , . .
, t end do 3:
Randomly choose i t from { 1 , 2 , . .
. , n } 4: ˇ g ( t +1)
i t ← argmax g it ∈ R (cid:110) − (cid:96) ∗ i t
( g i t )
+ (cid:82) h
i t ( θ )
ˇ p ( t ) ( θ )d θ ( g i t − ˇ g ( t ) i t ) − 1 2 nλ 2 |
g i t − ˇ g ( t ) i t | 2 (cid:111) .
Update the primal solution as ˇ p ( t +1)
=
p [ˇ g ( t +1) ] .
6: end for 7: return g ( t end )
3. O UR PROPOSAL : P ARTICLE UPDATE METHOD
The update (3) requires an integration with respect to p [ˇ g ( t ) ] .
We approximate the integral by a weighted sum over M particles:
[Equation]
where ( θ m )
Mm =1 are particles that approximately cover the support of p [ˇ g ( t ) ] , and ( r m )
Mm =1 are the weights of each particle.
Therefore, the key point of designing an efﬁcient algorithm is to construct a “good” set of particles ( θ m )
Mm =1 and their weights ( r m )
Mm =1 .
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 ˜ n ∈ [ n ] iterations, and perform the re-weighting stage in the other iterations.
The detailed description of each stage is given below.
(1) Sampling stage.
In the sampling stage, we sample M i.i.d.
particles ( θ m )
Mm =1 from the density function p [ g ( t ) ] by using a Monte Carlo sampler, and place an even weight r ( t ) m
= 1 /M ( m = 1 , . . .
, M ) on the sampled particles ( θ m )
Mm =1 .
As for the sampler, we may employ the unadjusted gradient Langevin algorithm (ULA), or the Metropolis-adjusted Langevin algorithm (MALA), both described in 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 p ( t ) that is in (cid:15) ( t ) C distance from the target density p [ g ( t ) ] measured by the total variation distance 5 .
Our convergence rate analysis will take this error into account.
(2) Re-weighting stage.
Since the sampling procedure requires relatively heavy computation, it is desirable to perform this operation once in a while.
We consider performing the sampling step once in ˜ n 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 p [ g ( t ) ]
( θ ) ∝ q [ g ( t ) ]( θ ) =
exp (cid:110) − 1 λ 2 (cid:16) 1 n (cid:80) n i =1 h
i ( θ ) g ( t ) i + λ 1 (cid:107) θ (cid:107) 2 (cid:17)(cid:111) which yields q [ g ( t +1) ] =
q [ g ( t ) ] exp[ − h i t
( θ )( g ( t +1)
i t − g ( t ) i t ) / ( nλ 2 )] for the index i t
chosen at the t -th iteration.
We adapt this relation to update the weight r ( t ) m of particles as
[Equation]
during the re-weighting stage because r ( t ) m represents “importance” of θ m which is proportional to p [ g ( t ) ]( θ m ) /p
[ g (˜ nT ) ]
( θ m ) so that the particle approximation (4) is (nearly) unbiased.
We allow the number of particles to depend on the outer-loop iteration T , and write M (˜ nT ) as the number of particles in the T -th outer-loop.
The entire procedure is summarized in Algorithm 2.
Finally, we output the solution of the last iterate.
We call this output as Option (A) and write its solution as g (A)out .
Alternatively, there is another choice Option (B) , which randomly selects the stopping time t (cid:48) end from the ﬁnal n -iterations, i.e., that is randomly chosen from { ˜ nT end − n + 1 , . . .
, ˜ nT end } , and returns the solution g (B)out = g ( t (cid:48) end ) .
This technique improves the accuracy up to O ( n ) factor.
Therefore, we can reduce the number of particles by O ( n ) times.
Require
[Table]
y i ) } ni
[Table]
1 and numbers of inner
[Table]
loop iterations ˜ n and outer
[Table]
loop iterations
[Table]
, M (˜ nT ) ) .
9: end for 10: end for 11: return Option (A): g (A)out = g (˜ nT end ) ; Option (B): g (B)out = g ( t (cid:48) end ) for t (cid:48) end that is randomly chosen from ˜ n ( T end − (cid:100) n ˜ n (cid:101) )
+ 1 , . .
. , ˜ nT end } .
Sampling algorithms.
Here, we present two examples that can be applied to sampling the particles ( θ m )
Mm =1 from p [ g ]( θ ) ∝ exp( − U ( θ )) where U ( θ ) is the potential function deﬁned as U ( θ ) :
=λ 2 (cid:0) 1 n (cid:80) ni
[Equation]
(1) ULA (unadjusted gradient Langevin algorithm):
The algorithm generates samples from the following noisy gradient descent update:
[Equation]
where η > 0 is the step-size and ξ k ∼ N (0 , I ) is a random noise generated from the ˜ d -dimensional standard normal distribution.
The sampling efﬁciency 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 ˜ θ k +1 by ULA, i.e., ˜ θ k +
=
θ k − η ∇ U ( θ k ) + √ 2 ηξ k , which is accepted with probability α given by
[Equation]
where q ( θ (cid:48) | θ ) ∝ exp( −(cid:107) θ (cid:48) − θ − η ∇ U ( θ ) (cid:107) 2 / 4 η ) .
If ˜ θ k +1 is not accepted, we set ˜ θ k +
= θ k .
The convergence rate of MALA has been studied by Bou-Rabee & Hairer (2013); Ma et al.
; Chen et al.
(2020a), to name a few.
Due to the Metropolis-Hastings step, MALA allows for larger step size than ULA, which leads to an improved iteration complexity with respect to the desired accuracy (cid:15) C .
C ONVERGENCE ANALYSIS
We present the convergence analysis of our method under the following assumptions.
Assumption 2.
(A3) inf x (cid:96) i ( x ) >
−∞ holds for all i ∈ [ n ] .
There exists a constant γ (cid:48) such that ∀ i ∈
[ n ] , (cid:96) ∗ i is γ (cid:48) -smooth
in the interval { y ∈ dom( (cid:96) ∗ i ) | ∂(cid:96) ∗ i ( y ) ∩ [ − 2 , 2] (cid:54) = ∅} .
B 1 :=
sup {| (cid:96) i ( x ) − inf x (cid:48) (cid:96) i ( x (cid:48) ) | | | x | ≤ 1 , i ∈ [ n ] } and B 2 :=
sup {| (cid:96) (cid:48) i ( x ) | | | x | ≤ 1 , i ∈ [ n ] } are both ﬁnite.
(A4) h
i is L -smooth for all i = 1 , . . . , n .
It is straightforward to verify the ﬁniteness of B 1 , B 2 and inf x (cid:96) i ( x ) in (A3) for several practical loss functions such as the squared loss and logistic loss.
The smoothness assumption ( γ (cid:48) ) on (cid:96) ∗ i is required to ﬁll in the gap between g ( t ) and ˇ g ( t ) , and is also satisﬁed by the loss functions listed above:
the squared loss satisﬁes
it with γ (cid:48)
= 1 and the logistic loss satisﬁes it with γ (cid:48)
= ( e 2 + 1) 2 /e 2 .
The 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 satisﬁed, for example, by a (truncated) two-layer neuron h θ ( x ) = tanh( r ) σ ( w (cid:62) x ) for θ =
( r, w ) ∈ R × R d (which also satisﬁes the boundedness assumption (A2)).
We are now ready to present the convergence rate of our proposed method Algorithm 2.
Here, we deﬁne ˜ s def = ˜ nγλ 2 2(1+ nγλ 2 ) .
Theorem 1.
Suppose that Assumptions 1 and 2 hold.
Let t end =
˜ nT end .
There exist constants ˆ C 1 , ˆ
C 2 and ˆ C 3 which can depend on ˜ s and λ 2 such that, if
[Equation]
for 0 < δ < 1 and λ 2 ≥ 2 B/n , then there exists an event E 6 satisfying P ( E ) ≥ 1 − δ in which the conditional expectation of the duality gap can be bounded by E (cid:104) P ( p [ g (A)out ]) − D ( g
[Equation]
A
[Equation]
out )
(cid:12)(cid:12)
[Equation]
(cid:105) ≤ (cid:15) P as long as the number of iterations t end satisﬁes the following conditions:
[Equation]
Moreover, the constants have the following dependency on n, ˜ n, λ 2 , ˜ s : ˆ C 2 = exp( 4 B
[Equation]
n λ 2 n ) ,
[Equation]
=
O (cid:16) λ − 1 / 2ˆ C
[Equation]
exp (cid:16)
[Equation]
n (2 ˆ C 2 +1) nγλ
2 +1 (cid:17)(cid:17) and ˆ C 3 =
O (max { ˜ s − 1 , λ −
[Equation] [Equation]
[Equation]
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 ﬁnite dimensional version of SDCA (Shalev-Shwartz & Zhang, 2013a).
In contrast, a (non-stochastic) full-batch gradient descent requires
O (( n/λ 2 γ ) log(1 /(cid:15) P ))
total gradient evaluations to optimize a convex objective with a condition number 1 / ( λ 2 γ ) .
Importantly, SDCA (and our method) “decouples” the sample size n and the condition number 1 / ( λ 2 γ ) yielding only O (( n + 1 /λ 2 γ ) log(1 /(cid:15) P )) iterations, which is beneﬁcial in a setting of big data and a small regularization parameter.
The difference between options (A) and (B) is only an O (log( n )) factor in terms of the number of iterations.
However, this improvement affects the number of particles given by Eq.
Indeed, the log( n ) factor improvement of t end allows us to take M (˜ nT end )
≥ Ω( ˆ C
ˆ C 3 / ( ˆ C 21 (cid:15) P ))
in the last inner loop, whereas option (A) requires M (˜
nT end )
≥ Ω( n ˆ C
ˆ C 3 / ( ˆ C 21 (cid:15) P )) .
Either way, the number of particles can be of linear order in terms of the inverse of the target precision (cid:15) P .
Remark 1.
Depending on the regularization strength λ 2 , the re-sampling interval ˜ n should be chosen differently.
For constant order (non-vanishing) λ 2 , we may take the re-sampling interval as ˜ n = n so that all constants are bounded.
On the other hand, in the small λ 2 regime, say λ
2 =
O (1 /n ) , we have ˆ C 2 =
exp
( O (˜ n )) , and there also appears an exponential dependency on ˆ C 1 w.r.t.
˜ n .
In this case, ˜ n = 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 a trade-off.
In practice, re-sampling can be executed in parallel and thus efﬁcient training (using GPU) is possible.
Importantly, to obtain the usual O (1 / √ n ) -generalization error, λ 2 = 1 / √ n is sufﬁcient (Nitanda et al., 2021) and thus we may take ˜ n = Ω( √ n ) in that situation.
We however conjecture that this trade-off is unavoidable without additional assumptions.
Sampling complexity.
In Theorem 1, it is assumed that the density p (˜ nt ) from which the particles are generated is within (cid:15) (˜ nT ) C from the target density p
[ g (˜ nT ) ] measured by the total variation distance.
The next proposition provides an upper bound of the number of iterations required to obtain the (cid:15) (˜ nT )
C 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 a constant c LS (see Deﬁnition 2 for its precise deﬁnition).
particles are sampled from the sampler to approximate the integral.
Proposition 2.
Under Assumptions 1 and 2, the target density p [ g (˜ nT ) ] at each outer iteration satisﬁes the log-Sobolev inequality with constant c LS = 2 λ 1 λ 2 exp( − 4 B 2 /λ 2 ) .
Accordingly, the number of iterations K ( T ) ∗ to obtain the (cid:15) (˜ nT ) C -accurate distribution can be evaluated as follows:
[Equation]
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 simpliﬁed as K ( T ) ∗ ≥ Ω (cid:0) ( c LS (cid:15) (˜ nT ) C ) − 2 log( n ) (cid:1) .
On the other hand, MALA requires only poly-log order complexity with respect to (cid:15) (˜ nT ) C as
[Equation]
2 (cid:1) , which is therefore more preferable to obtain higher accuracy.
Combining the required number of particles M (˜ nT ) 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-Sobolev constant c LS also depends exponentially on 1 /λ 2 .
While the exponential factor can be avoided in the outer loop (Theorem 1) by appropriately setting ˜ n , such dependence is known to be unavoidable for c LS in the most general setting (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 complexity of our algorithm (Algorithm 2) with the MALA sampler can be roughly bounded as
[Equation]
where we assumed ˆ C − 1 1 , ˆ C
= 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.
A DDITIONAL RELATED WORKS
Hu et al.
; Sirignano & Spiliopoulos (2020) showed a linear convergence rate of the McKeanVlasov dynamics.
However, the theory only handles the inﬁnite width limit and continuous time dynamics, and it is from trivial to covert their results to a discrete time and ﬁnite width setting.
Mei et al.
(2018) also analyzed space-time discretization of a similar dynamics.
However, a quantitative convergence rate is not provided, and the time discretization error grows exponentially with the time horizon.
Compared to these results, we deal with general (1 /γ ) -smooth
loss functions and made the dependency on the parameters λ 2 and γ 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 ﬁeld models.
When specialized to mean ﬁeld neural networks, our formulation covers a 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 (log(1 /(cid:15) P ) /(cid:15) P ) , whereas our bound requires only poly-log order cost due to the ﬂexible choice of sampling algorithm.
Jabir et al.
(2021) also considered spatial-time discretization of mean ﬁeld dynamics of deep neural network.
However, it contains the time discretization lag h in their convergence rate, while our method is purely discrete time and is not affected by such a time discretization error.
Moreover, their analysis requires a strong regularization (which corresponds to λ 1 in our setting) so that the objective becomes convex-like with respect to each particle to show the geometric ergodicity; this is not required in our analysis.
N UMERICAL 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 i =
σ ( w ∗(cid:62) x i + b ∗ )
+ (cid:15) i , where x i and (cid:15)
i are independently drawn from Gaussian distributions.
We optimize a tanh neural network ( h θ ( x ) =
tanh( w (cid:62) x + b ) with θ = ( w, b ) ) of width M s as a student using Algorithm 2.
[Figure]
P
[Figure]
SDCA
[Figure]
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 ReLU activation for σ and (note that even though this is a simple setting, the teacher model is “not” included in the student model).
The number of examples and dimensionality of the input space are n = 500 and d = 50 , respectively.
The other hyperparameters are set to λ 1 = 10 − 4 , λ 2 = 10 − 4 , and M s = 200 .
The left side of
Figure 1 depicts the convergence of the primal objective to the optimal value.
The x -axis is the number of outer-iteration and the y -axis is the excess primal objective P ( · ) − inf p P ( p ) including the regularization terms on a training set and a 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 ( σ ( · ) = sign( · ) ), and the student is an overparameterized tanh network of width
M s = 200 .
The number of examples and input dimensionality are n = 1000 and d = 50 , and we set λ 1 = 10 − 3 , λ 2 = 10 − 4 .
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 ( f Θ ( x ) = 1 √ M (cid:80)
Mm =1 h θ m ( x )) for SGD optimization in addition to the mean ﬁeld scaling.
The test error with respect to the number of gradient evaluations is displayed in the right side of Figure 1.
We make the following observations: (i) P-SDCA achieves a 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 ﬁeld 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.
C ONCLUSION
In this paper, we proposed a new stochastic optimization technique for two-layer neural networks with the mean ﬁeld representation.
Our algorithm can be seen as an inﬁnite dimensional extension of SDCA that maintains its advantage, namely, that it converges linearly with respect to the outeriteration, and is much more efﬁcient 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 ﬁnite 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 signiﬁcantly different from existing numerical approximation approaches of McKean-Vlasov 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 in this work is not rate optimal – the optimal algorithm can be obtained by combining Nesterov’s acceleration with stochastic methods (Allen-Zhu, 2017; Murata & Suzuki, 2017).
Extending our algorithm to such an optimal accelerated method is an interesting future direction.
In addition, while the exponential dependence on 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.
A CKNOWLEDGMENT
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 (JPMJPR1928)
T ABLE OF C ONTENTS 1 Introduction
2 Preliminaries 3 Proposed method: Particle-SDCA 3.
Dual problem and algorithm description . .
Our proposal: Particle update method . . . . . . . . . . . . . . . . . . . . . . . .
4 Convergence analysis 5 Additional related works 6 Numerical experiments 7 Conclusion
A Fenchel’s duality theorem (Proof of Proposition 1)
B Computational complexity of sampling algorithms (Proof of Proposition 2)
B.1 Log-Sobolev inequality of the target function . . . . . . . . . . . . . . . . . . . .
B.2 Sampling complexity of ULA . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.3 Sampling complexity of MALA . . . . . . . . . . . . . . . . . . . . . . . . . . .
C Proof of Theorem
C.1 Convergence of the ideal update ˇ g ( t ) (Algorithm 1) . . . . . . . . . . . . . . . . .
C.2 Convergence of the particle-version update g ( t ) (Algorithm 2) . . . . . . . . . . .
C.2.
Convergence of auxiliary variables ˜ g ( t ) and ¯ g ( t ) . . . . . . . . . . . . . . .
C.2.
Particle approximation, and error analysis of g ( t ) . .
C.2.
Integration of two sources of discretization error . . . . . . . . . . . . . .
C.2.
Discrete Gronwall’s lemma . . . . . . . . . . . . . . . . . . . . . . . . . .
D Additional
——Appendix——
A F ENCHEL ’ S DUALITY THEOREM (P ROOF OF P ROPOSITION 1)
To prove Proposition 1, we utilize the inﬁnite dimensional Fenchel’s duality theorem by Rockafellar (1967).
Two real vector spaces E and E ∗ with locally convex Hausdorff topology are called topologically paired vector spaces if all the elements of each space can be identiﬁed with all the continuous linear functionals on the other.
Therefore, for a continuous linear functional y : X → R , there exists
y ∈
E ∗ such that we can write y ( x ) =
(cid:104)
y, x (cid:105) E ∗ ,E as a bilinear form (cid:104)· , ·(cid:105) E ∗ ,E .
Deﬁnition 1 (Fenchel-Legendre convex conjugate) .
Let X, X ∗ be topologically paired vector spaces.
For f : X → R ∪ {∞} , its Fenchel-Legendre convex conjugate f ∗ : X ∗ → R ∪ {∞} is deﬁned by
[Equation]
Remark 3.
When X = X ∗ =
R d equipped with the usual inner product, the deﬁnition of the
Fenchel-Legendre convex conjugate coincides with the deﬁnition introduced in the main text.
Lemma 2 (Theorems 3 and 8 of Rockafellar (1967)) .
Let ( X, X ∗ ) and ( Y, Y ∗ ) be two pairs of topologically paired vector spaces.
Let f : X → R ∪ {∞} and g :
Y → R ∪ {∞} are proper lower semi-continuous convex functions, and let A : X → Y be a bounded linear operator.
Assume that ∃ x ∈ X such that f ( x ) is ﬁnite and g is ﬁnite and continuous at Ax .
Then,
[Equation]
where A ∗ is
the adjoint operator of A such that (cid:104) y, Ax (cid:105) Y ∗ ,Y
= (cid:104)
A ∗
y, x (cid:105) X ∗ ,X .
Moreover, ˆ x and ˆ y are the optimal solutions of the primal and dual problems respectively if and only if A ˆ x ∈ ∂g ∗ (ˆ y ) and A ∗ ˆ y ∈ ∂f (ˆ x ) .
Now, we are ready to prove Proposition 1.
Proof of Proposition 1.
We utilize Lemma 2 to show the assertion.
We let Y =
R n .
Then, we know that Y is self-dual, that is, its topological dual Y ∗ is Y itself where the linear functional is given by the standard inner product (cid:104) u, v (cid:105) Y ∗ ,
Y =
(cid:80) ni
=1 u i v i for u
= ( u i ) ni
=1 ∈ Y and v = ( v i ) ni
=1 ∈ Y ∗ .
We let X be the set of L 1 -integral functions with respect to the Lebesgue measure, i.e., X =
L 1 ( R ˜ d ) =
L 1 ( R ˜ d ; µ ) , where µ is the Lebesgue measure on ( R ˜ d , B ( R ˜ d )) .
It is known that its topological dual is the set of essentially bounded functions equipped with L ∞ -norm, i.e.,
X ∗ =
L ∞ ( R ˜ d )(=
L ∞ ( R ˜ d ; µ ))
=
{ f : R ˜ d
→ R | f is B ( R ˜ d ) -measurable , (cid:107) f (cid:107) ∞ < ∞} , where (cid:107) f (cid:107) ∞ =
inf { M > 0 | µ ( { x | | f ( x ) | ≥ M } )
= 0 } for the Lebesgue measure µ .
The bilinear functional (cid:104)· , ·(cid:105) X ∗ ,X is given by (cid:104) g, f (cid:105) X ∗ ,X = (cid:82) fg d µ .
Then, X, X ∗ become topologically paired vector spaces by equipping X ∗ with its weak ∗ topology with respect to the inner product (cid:104)· , ·(cid:105) X ∗ ,X .
We let A : L 1 ( R ˜ d ) → R n be Af =
( (cid:82) h i ( θ ) f
( θ )d θ ) n i =1 ∈ R n for f ∈ L 1 ( R ˜ d ) .
Then, we can see that A is a bounded linear functional by the boundedness of h i ( · ) .
It is also known that the KL divergence is jointly lower semi-continuous with respect to the weak convergence topology (Posner, 1975, Theorem 1).
Since the L 1 -norm 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, KL divergence is lower semi-continuous with respect to the L 1 -norm of the corresponding density functions.
It is also well known that KL divergence is jointly convex and it is easy to check that p ∈ P (cid:55)→ KL( p, N (0 , σ 2 I )) is proper, i.e., ∃ p ∈ P such that KL( p, N (0 , σ 2 I )) < ∞ .
Under these settings, if we set g ( u ) = 1 n (cid:80) ni
=1 (cid:96) i ( u i ) for u
= ( u i ) ni
=1 ∈ Y = R ˜ d and
[Equation]
∞
[Equation]
for p ∈
L 1 ( R ˜ d ) , then we can see that P ( p ) =
g ( Ap ) + f ( p ) for p ∈ P and inf p ∈P P
( p )
=
inf h ∈ X
{ g ( Ah ) + f ( h ) } .
We can also verify the condition in Lemma 2.
Indeed, if we let p 0 ∈ P be the density function of N (0 , λ 1 /λ 2 I ) , then f ( p 0 ) = 0 < ∞ and g is continuous at Ap 0 because a convex function is continuous when its effective domain is the whole space R n (Rockafellar, 1976,
Theorem 10.1), and the domain of g is the whole space R n (remember that the range of (cid:96) i is R ).
Therefore, by applying Lemma 2 and going through a simple calculation to obtain f ∗ , we obtain the desired assertion.
B C OMPUTATIONAL COMPLEXITY OF SAMPLING ALGORITHMS (P ROOF OF P ROPOSITION 2) B. L OG -S OBOLEV INEQUALITY OF THE TARGET FUNCTION
In each outer iteration, we sample M particles from the density p [ g (˜ nT ) ] so that the particles are sampled from an approximated density p (˜ nT ) with TV( p (˜ nT )
|| p [ g (˜ nT ) ]) ≤ (cid:15) ( T ) S .
To estimate the convergence of Monte-Carlo methods such as ULA and MALA, the log-Sobolev inequality associated with the target density plays the central role.
Deﬁnition 2 (Log-Sobolev inequality) .
Let p ( θ ) be a smooth probability density function on R ˜ d .
Then, we say p ( θ ) (or its corresponding probability measure on ( R , B ( R ˜ d )) ) satisﬁes the log-Sobolev inequality (LSI) with a constant c LS >
if and only if, for any smooth function φ : R p
→ R with E p [ φ 2 ] < ∞ , it holds that
[Equation]
The following lemma shows that the target density p [ g (˜ nT ) ] satisﬁes the LSI.
Lemma 3.
Suppose max i | g (˜ nT ) i | ≤ B
(this is true by Lemma 6 below with B = B 2 ) and assume (A2) of Assumption 1, then p [ g (˜ nT ) ] satisﬁes the LSI with the constant
[Equation]
Proof.
Remember that p [ g (˜ nT ) ] is given as
[Equation]
Therefore, p [ g (˜ nT ) ]( θ ) can be written as
[Equation]
where f ( θ )
= − 1 λ 2 1 n (cid:80)
ni
=1 h i ( θ ) g (˜ nT ) i and q ( θ ) is the probability density function of a Gaussian measure N (0 , λ 2 2 λ 1 I ) .
Note that (cid:107) f (cid:107) ∞
≤ B/λ 2 by the assumption (A2) and max i | g (˜ nT ) i | ≤ B .
Then, it is well known that q satisﬁes the LSI with a LSI constant c LS ( q )
= 2 λ 1 λ 2 and Holley & Stroock (1987) showed that the bounded perturbation as in Eq.
(8) preserves the LSI condition where the
LSI constant of p [ g (˜ nT ) ] is bounded by c LS
= c LS ( q ) / exp(4 (cid:107) f (cid:107) ∞ ) ≥ 2 λ 1 λ 2 exp( − 4 B/λ 2 ) .
This concludes the proof.
By Lemma 3, we show that existing convergence rate analyses of ULA and MALA such as Vempala & Wibisono (2019);
Ma et al.
(2019) can be applied to our target density p [ g (˜ nT ) ] .
From now on, we denote our target density p [ g (˜ nT ) ] by p ∗ , and the negative log density of p [ g (˜ nT ) ] by U , that is, p ∗ ( θ ) = exp( − U ( θ )) .
In addition, let p k be the marginal distribution of θ k (sample in the k -th iteration of ULA or MALA).
B. S AMPLING COMPLEXITY OF ULA
First, we apply the convergence rate analysis of ULA developed by Vempala & Wibisono (2019) to our setting.
We utilize the following proposition given by Vempala & Wibisono (2019).
Proposition 3 (Theorem 1 of Vempala & Wibisono (2019)) .
Assume that p ∗ ( θ )d θ satisﬁes the log-Sobolev inequality with a constant c LS and U is L -smooth.
Under these assumptions, for the step size η satisfying η ≤ c LS 4 L 2 , we have that
[Equation]
Proposition 3 implies that if the step size satisﬁes η
[Equation]
min
[Equation]
and the number of iterations satisﬁes
[Equation]
then we obtain an (cid:15) -accurate distribution p k in terms of the KL-divergence.
If we employ p
( θ ) ∝ exp( − λ 1 λ 2 (cid:107) θ (cid:107) 2 ) (a Gaussian measure), then 2KL ( p 0 (cid:107) p ∗ ) is bounded by 4 B 2 λ 2 , from Lemma 6:
[Equation]
[Table]
[Equation]
Moreover, using Pinsker’s inequality bounds total variation norm by KL-divergence:
[Equation]
Therefore, in order to have an (cid:15) C -accurate distribution in total variation
[Table]
it sufﬁces to let the step size satisfy η ≤ c LS 4 L 2 min (cid:110) 1 , (cid:15)
2 C 8˜ d (cid:111) and the number of iterations satisfy
[Equation]
Hence, the total computational complexity can be bounded by O (cid:16) ˜ dL 2 c 2LS (cid:15) 2 C log B 2 λ 2 (cid:15) C (cid:17) .
sampling by ULA According to Theorem 1, sampling accuracy in total variation (cid:15) (˜ nT )
C should be smaller than O (exp (cid:0) − ˜ sT 2 (cid:1) ) in each sampling step T .
Therefore, the total computational complexity for sampling can be calculated as follows:
Algorithm 2 with Option A: Suppose we run outer-loop up to
T end ≥
[Equation]
times.
Indeed it is sufﬁcient to have an (cid:15) Papproximated solution, according to Theorem 3, which gives a tighter and more detailed bound than in Theorem 1.
Then, the total required computational complexity for sampling is
[Equation]
where ˜ O hides log log terms.
Algorithm 2 with Option B: According to Theorem 1, it is sufﬁcient to run outer-loop up to T end ≥ 1˜
[Equation]
log
[Equation]
+
[Equation]
ˆ C 3 (cid:15) P (cid:17) times to acquire an (cid:15) P -approximated solution.
You can see the detailed deﬁnition of ˆ C
= O (max { ˜ s − 1 , λ − 3 / 2˜ s − 1 / 2 } ) in Eqs.
(52a)-(52c) of Theorem 3 and in Proof of Theorem 1 .
Therefore, the total required computational complexity for sampling is
[Equation]
where ˜ O hides log log terms as well.
B. S AMPLING COMPLEXITY OF MALA
Here, we adopt the convergence rate analysis of MALA obtained by Ma et al.
They showed the convergence under the setting where U is µ -strongly convex outside of 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 ∗ satisﬁes the LSI in Lemma 3.
Thus, the analysis of Ma et al.
(2019) can be applied to our setting.
From now on, we denote p [ g (˜ nT ) ] by p ∗ .
Let R ( s ) be a positive real such that P ∗ ( (cid:107) θ (cid:107) ≥ R ( s ))
≤ s for 0 < s < 1 where P ∗ is the probability measure corresponding to p ∗ and ˜ L > 0 be the smoothness parameter of log( p ∗ ) :
(cid:107)∇ log( p ∗ ( θ ))
− ∇ log( p ∗ ( θ (cid:48) ))
(cid:107) ≤ ˜ L (cid:107) θ − θ (cid:48)
(cid:107) .
Then, Ma et al.
(2019) showed the following proposition.
Proposition 4 (Lemma 9 of Ma et al.
(2019) with modiﬁcation) .
Let the initial density p 0 satisfy
[Equation]
for β > 0 .
Then, for the step size η > 0 satisfying
[Equation]
and the number k of iterations satisfying
[Equation]
it holds that
[Equation]
We apply this proposition to our setting.
From now on, we assume Assumptions 1 and 2 hold.
Then, it can be easily seen that
[Equation]
If we employ p
( θ ) ∝ exp( − λ 1 λ 2 (cid:107) θ (cid:107) 2 ) (a Gaussian measure), then we may set β ≤ exp(2 B 2 /λ 2 ) because
[Equation]
Moreover
[Equation]
the proof of Lemma 16 of Ma et al.
(2019) showed that we may set
[Equation]
The second term in the right hand side can be evaluated as
[Equation]
where the inequality in the second line can be shown in the same manner as Eq.
Then
[Equation]
we obtain
[Equation]
Therefore
[Equation]
substituting these evaluations into Eq.
(9) yields that
[Equation]
1: ˜ p (˜ nT
[Equation]
← p [ g (
[Equation]
nT ) ]
[Equation]
3: for t = ˜
nT + 1 , . . .
, ˜ n ( T + 1) − 1 do 4: ˜
p ( t ) ← p
[˜ g ( t ) ] 5:
˜ g ( t +1)
i t ← argmax
[Equation]
g it −
[Equation]
(cid:90) h
[Equation]
( θ
[Equation]
p
[Equation]
( θ
[Equation]
d θ (˜ g
i
[Equation]
− ˜ g ( t
[Equation]
i
[Equation]
− 2 nλ 2 | ˜ g
i t − ˜ g ( t ) i t |6: end for is sufﬁcient to ensure (cid:107) p ∗ − p k (cid:107) TV ≤ (cid:15) , where c LS = 2 λ 1 λ 2 exp( − 4 B 2 /λ 2 ) and ˜ L =
LB 2 +2 λ 1 λ.
We can see that the right hand side can be bounded by
[Equation]
which gives the assertion of Proposition 2.
C P ROOF OF T HEOREM
In this section, we give the convergence rate analysis of our algorithm.
We ﬁrst show the convergence rate Algorithm 1 (the continuous limit version), then we prove Theorem 1 by bounding the discretization error.
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 is sufﬁciently small throughout each inner loop from t = ˜
nT to t = ˜ n ( T + 1) .
To rigorously prove this intuition, we introduce ˜ g ( t ) which follows the continuous limit update but is initialized from the discretized one g (˜ nT ) at the every starting point of the inner loop (see its deﬁnition in Algorithm 3).
To ﬁll in the gap between the continuous limit ˜ g ( t ) and the discretized version g ( t ) , we also deﬁne ¯ g ( t ) ( t = ˜
nT, . .
, T (˜ n + 1)) which lies in between (see Algorithm 4).
Basically, ¯ g ( t ) and ˜ g ( t ) are updated by the same rule, but are initialized in different ways.
Indeed, the primal variables corresponding to these two sequences are updated as p [˜ g ( t +1) ]
( θ ) ∝ p [˜ g ( t ) ]( θ ) × exp (cid:16) − 1 nλ 2 h
i t ( θ ) δ ˜ g ( t ) (cid:17) and ¯ p ( t +1)
∝ ¯ p ( t ) ( θ ) × exp (cid:16) − 1 nλ 2 h
i t ( θ ) δ ¯ g ( t ) (cid:17) , respectively.
However, different initial primal variables ˜ p ( nT ) =
p [ g ( nT ) ] and ¯ p ( nT ) = p ( nT ) are employed for the two schemes (remember that p ( nT ) is an approximation of p [ g ( nT ) ] by the sampling scheme).
We denote difference of the updated coordinates before and after updates as
[Equation]
1) , respectively.
At the same time, we deﬁne some variants of r ( t ) j , in order to ﬁll the gap more ﬁnely between ¯ g ( t ) and g ( t ) .
We deﬁne ¯ r ( t ) j as a direct discretization of ¯ p ( t ) using the sampling point θ j (see Algorithm 5).
Let ˆ r ( t )
j be a numerator of ¯ r ( t ) j (see Algorithm 6).
We expect that
[Equation]
plays a role of interpolation between ¯ p ( t ) and
[Equation]
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.
[Table] [Table]
[Equation]
5: end for 6: end if
Algorithm 6 Deﬁnition of ˆ r ( t ) j 1: ˆ r
˜ nT )
j ← 1 M ( j = 1 , . . . , M ) 2: if ˜ n > 1 then 3: for t = ˜
nT, ˜ nT + 1 , . . .
, ˜ n ( T + 1) − 2 do 4: ˆ r (
[Equation]
←
[Equation]
(cid:33) 5: end for 6: end if
Lemma 4.
The ﬁrst and second order derivatives of the second term of the dual objective D ( g ) can be evaluated as
[Equation]
In particular, the dual of the regularization term can be bounded by a quadratic function as follows: (a) L 1 -norm bound,
[Equation]
and (b) L 2 -norm bound,
[Equation]
2 log
[Equation]
The L 1 -norm bound is useful to derive the dual coordinate ascent method (Algorithms 1 and 2).
On the other hand, the L 2 -norm bound is also useful in the following convergence rate analysis.
Proof.
First of all, it is easy to see that λ 2 log (cid:2)(cid:82) q
[ g ]( θ )d θ (cid:3) is two times differentiable by g i .
Since a direct calculation yields
[Equation]
we immediately obtain Eq.
[Equation]
By differentiating the right hand side by g j 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 h i and h j with respect to the probability distribution p [ g ]( θ ) .
Hence, each element is bounded by 1 n 2 λ 2 due to the boundedness assumption of h i (Assumption (A2)), which yields Eq.
As for the L 2 -norm bound (14), it is easily obtained by the relation (cid:107) δg (cid:107) 1 ≤ √ n (cid:107) δg (cid:107) 2 .
C. C ONVERGENCE OF THE IDEAL UPDATE ˇ g ( t ) (A LGORITHM 1)
First, we prove the convergence of the ideal update ˇ g ( t ) .
Convergence analysis in the continuous limit is the basis for that of discretized updates; more concretely, we can use results and proofs for continuous limit in this section as a reference to derive Lemma 8, which is one of the key lemmas for
Theorem 1 (convergent proof for discretized updates).
Lemma 5.
Suppose that (cid:96) i is 1 /γ -smooth for all i ∈ { 1 . . . , n } (Assumption (A1)).
Consider the update of ˇ g ( t ) given in Algorithm 1 (continuous limit update), then, for any iteration t and any s ∈ [0 , 1] , it holds that
[Equation]
where ˇ p ( t ) = p [ˇ g ( t ) ] ,
[Equation]
and u ( t ) i def = (cid:96) (cid:48) i (cid:0)(cid:82) ˇ p ( t ) ( θ ) h
i ( θ )d θ (cid:1) .
Here, the expectation was taken with respect to the choice of the coordinate i t .
Accordingly, we also have the same bound for the ideal update in Algorithm 2 ( g (˜ nT ) and ˜ g ( t ) ( t = ˜
nT
+ 1 , . .
. , ˜ n ( T + 1)) ), considering the deﬁnition of ˜ g
( t ) .
Proof.
Since we will focus on the update of one step, i t is omitted as i in the following.
At each update, the improvement in the dual objective can be written as
[Equation] [Equation]
where we used Eq.
(13) in Lemma 4 for the ﬁrst inequality
[Equation]
and the second inequality follows from the deﬁnition of ˇ g ( t +1) .
Since (cid:96) ∗ i is γ -strongly convex ( γ = 0 is also allowed)
[Equation]
it holds that
[Equation]
From the property of the Legendre transform, u ( t ) i = (cid:96)
[Equation]
i
[Equation]
p
[Equation]
(cid:1) implies that
[Equation]
Similarly
[Equation]
we also have the following relation from the duality of the regularization term:
[Equation]
Combining Eqs.
(18), (19) and (20)
[Equation]
we have that
[Equation]
By taking expectation of both sides of (22) with respect to i
[Equation]
it follows that
[Equation] [Equation]
where we used Eq.
(21) for the ﬁrst equality.
Dividing the both sides of Eq.
(23) by n , we obtain the desired inequality (16).
Lemma 6.
ˇ g ( t ) i in Algorithm 1 is uniformly bounded by B 2 , that is,
[Equation]
In the same way, g ( t ) i , ˜ g ( t ) i and ¯ g ( t ) i in Algorithm 2 satisfy
[Equation]
Proof.
Let p be an arbitrary probability distribution.
Let us consider the update
[Equation]
for an arbitrary i ∈ { 1 , . . . , n } .
Since h i is bounded by 1 (Assumption (A2)), the sub-differential of (cid:96) ∗ i at z satisﬁes that
[Equation]
By the assumption (A3), i
[Equation]
e.,
[Equation]
∞ , we have that z in Eq.
(24) is bounded by B 2 .
Next, we consider its proximal-version update:
[Equation]
In this proximal-version update, we added the term − 1 2 nλ 2 ( z − y ) 2 whose maximizer with respect to z is given by z = y .
Since the maximum point of the sum of two concave functions is located between the maximizers of them, z in Eq.
(25) satisﬁes that
[Equation]
From this and the initialization | ˇ g (0) i | ≤ B 2 , we have that | ˇ g ( t ) i | in Algorithm 1 is uniformly bounded by B 2 by induction.
In the same way, g ( t ) i , ˜ g ( t ) i and ¯ g ( t ) i in Algorithm 2 minimize the sum of the average of h i ( θ j ) and quadratic term as well.
Note that the absolute value of the the average of h i ( θ j ) is bounded with 1 , regardless of methods of averaging, that is, weighted particle for g ( t ) i and probability distribution for ˜ g
( t ) i and ¯ g ( t ) i .
Thus, g ( t ) i , ˜ g ( t ) i and ¯ g ( t ) i are uniformly bounded with B 2 , as well as ˇ g ( t ) i in Algorithm 1.
Remark 4.
In the Section C.2, we sometimes face the situation where we should upper-bound increment of the coordinates δg ( t ) ,
δ ˜ g ( t ) and δ ¯ g ( t ) .
Lemma 6 directly implies that δg ( t ) ,
δ ˜ g ( t ) and δ ¯ g ( t ) are all bounded by 2 B 2 for all t = 0 , 1 , . . . ,
˜
nT end .
Thus, for future convenience, we deﬁne
C 1 def = exp (cid:18) 2 κBλ(cid:19) with which exp (cid:16) 2 κ | δg ( t ) | λ(cid:17) , exp (cid:16) 2 κ | δ ˜ g ( t ) | λ(cid:17) , exp (cid:16) 2 κ |
δ ¯ g ( t ) | λ(cid:17) ≤ C 1 hold, where κ is deﬁned as κ = ˜ nn .
Note that, however, we could have much tighter evaluations of δg ( t ) , δ ˜ g ( t ) and δ ¯ g ( t ) , which are dependent of residual D ( g ∗ ) − D ( g ( t ) ) or duality gap P [ p ( g ( t ) )] − D ( g ( t ) ) , since the dual problem is strongly convex due to the smoothness of primal problem.
Lemma 7.
Under Assumptions (A2) and (A3), it holds that
[Equation]
for Algorithm 1 and
[Equation]
in Algorithm 2.
Proof.
By Assumptions (A2) and (A3), we can see that (cid:12)(cid:12) (cid:96) i (cid:0)(cid:82) p (0) ( θ ) h
i ( θ )d θ (cid:1) − inf x (cid:96) i ( x ) (cid:12)(cid:12) is bounded by B 1 .
Since inf x (cid:96) i ( x ) > − inf (Assumption (A3)), we also have (cid:96) ∗ i (0)
=
sup x { 0 · x − (cid:96) i ( x ) } = − inf x (cid:96) i ( x ) < ∞ .
In Algorithm 1, every initial coordinate ˇ g (0) i is bounded by B 2 .
Also, g (0) i in Algorithm 2 is similarly bounded by B 2 .
Therefore, in the following, g will be used without distinguishing between g and ˇ g .
Moreover, in a similar manner to Eq.
(21) in Lemma 5, we have that
[Equation]
from the duality.
Combining these observations, we have that
[Equation]
n
[Equation]
where we used Lemma 6 in the last inequality.
Theorem 2.
Let (cid:15) P ( > 0) be a desired accuracy.
To obtain an expected duality gap
E (cid:2) P (ˇ p ( t end ) )
− D (ˇ g ( t end ) ) (cid:3) ≤ (cid:15) P in Algorithm 1, it sufﬁces to have the total number T end of iterations as
[Equation]
Furthermore, if we randomly choose (ˇ p ( t ) , ˇ g ( t ) ) from t = t end − n + 1 , . .
, t end in Algorithm 1, then t end
[Equation]
is sufﬁcient to achieve (cid:15) P duality gap in expectation.
Proof.
In Lemma 4, choosing s = nγλ 2 1+ nγλ 2 ∈ [0 , 1] , then we have that
[Equation]
Hence, we have G ( t ) ≤ 0 for all t .
Thus, Lemma 5 implies that
[Equation]
Therefore, we have that
[Equation]
This implies that
[Equation]
By applying this inequality repeatedly, we obtain that
[Equation] [Equation]
where the last inequality is obtained by E (cid:2)
D
[Equation]
g ∗ ) − D
[Equation]
ˇ g
[Equation] [Equation]
) (cid:3) ≤ E
[Equation]
P (ˇ p
[Equation]
− D (ˇ g
[Equation]
(cid:3) and Lemma 7.
Moreover, making use of Eq.
(26), the duality gap can be bounded as follows:
[Equation]
Hence, combining Eqs.
(28) and (29) yields
[Equation]
Therefore, in order to have an (cid:15) P -approximated solution in expectation
[Equation]
it sufﬁces to let
[Equation]
Finally
[Equation]
we show the second part.
By summing up Eq.
(26) over t =
t end − n + 1 , . .
, t end , we have that
[Equation]
Note that although we do not actually calculate ˇ g ( t end +1) , we can virtually deﬁne ˇ g ( t end +1) as if the algorithm continued.
Therefore, in order to have an (cid:15) P -approximated solution, it sufﬁces to let
[Equation]
when we randomly choose (ˇ p ( t ) , ˇ g ( t ) ) from t = t end − n + 1 , . . .
, t end .
C. C ONVERGENCE OF THE PARTICLE - VERSION UPDATE g ( t ) (A LGORITHM 2)
Next, we consider the convergence in Algorithm 2.
In order to prove Theorem 1 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 series of inner loops.
Then, the assumptions of Theorem 2 is veriﬁed, by Theorem 3.
There are two major streams leading up to Theorem 3.
The ﬁrst 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 are combined in Section C.2.3 and ﬁnally 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 ( g ( t ) )
− D ( g ( t ) )) − ( P (˜ g ( t ) )
− D (˜ g ( t ) )) with the following two parts,
[Equation]
each of which corresponds to the discrepancy in the primal problem and dual problem, respectively.
As mentioned above, these two terms are attributed to the two sources, that is, sampling and reweighting.
These variables (cid:15) ( t ) A , (cid:15) ( t ) B represent how much the two sequences g ( t ) and ˜ g ( t ) diverge during the updates of the inner-loop.
Intuitively, in each inner-loop, the absolute values of (cid:15) ( t ) A and (cid:15) ( t ) B are increasing sequences between t = ˜
nT
+ 1 and ˜ n ( T + 1) .
Then, they are “reset” at t = ˜ n ( T + 1) according to the initialization rule of ˜ g
(˜ n ( T +1)+1)
and again will increase from t =
˜ n ( T + 1) + 1 to ˜ n ( T + 2) .
The following lemma states that under the assumption about (cid:15) ( t ) A , (cid:15) ( t ) B , Algorithm 2 achieves linear convergence with regard to the number of coordinate update.
From now on, we ﬁx s = nγλ 2 1+ nγλ 2 .
Also, we deﬁne ˜ s
= κ · s 2 = ˜ nγλ 2 1+ nγλ 2 , where κ = ˜ nn .
Lemma 8.
Suppose that Assumptions 1 and 2 hold.
If (cid:15) (˜ nT )
A is smaller than C 0 e − ˜
sT
for every T = 1 , . . .
, T end with a constant C
you can arbitrarily take, an expected duality gap E (cid:2)
P ( p [ g ( nT end ) ]) − D ( g ( nT end ) )
[Equation]
≤ (cid:15) P can be achieved by a total number T end of outer iterations given by
[Equation]
Furthermore, if (cid:15) ( t ) A and (cid:15) ( t ) B are uniformly bounded by (cid:15) P 4 for all t = ˜ n ( T end −
[Equation]
1 κ
[Equation]
+
[Equation]
, nT end in addition to the assumptions above, then it also sufﬁces to set
[Equation]
to achieve (cid:15)
P -duality gap when we randomly choose output ( p ( t ) , g ( t ) ) from t =
˜ n ( T end − (cid:100) 1 κ (cid:101) ) + 1 , . .
˜
nT end .
Remark 5.
Comparing the speed of convergence with regard to coordinate update t , it is exp( − st/ ˜ n ) for Theorem 2 and exp( − ˜ sT )
∝ exp( − s 2 t/ ˜ n ) for Lemma 8, which can be interpreted as follows.
In the descretized version, Lemma 8 allows discretization error to appear, at the cost of a slight decrease in convergence speed.
Indeed, you can choose arbitrary ˜ s from the open interval (0 , κs ) .
Our choice of ˜ s = κs 2 is just for simplicity.
If you choose another ˜ s much closer to κs , T end can be smaller.
Remark 6.
If you change the length of inner loop in the middle, the convergence rate after the change is obtained by simply replacing B 1 + B 2 with the duality gap at the time of the change.
Proof.
Since the update of ˜ g ( t ) is the same as that of ˇ g ( t ) in Algorithm 1 during the inner loop, it holds that
[Equation]
for T = 0 , 1 , . . .
, T end − 1 .
We also have
[Equation]
for t = ˜
nT + 1 , . . .
, ˜ n ( T + 1) − 1 , in the same way as the derivation of Eq.
By applying these inequalities above iteratively, we have that
[Equation]
Bounding D (˜ g (˜ n ( T +1)) ) by D ( g (˜ n ( T +1)) ) − (cid:15) (˜ n ( T +1))
A , we have that
[Equation]
Therefore, by the Gronwall’s lemma (“vanilla version” of Lemma 18), we obtain that
[Equation]
Applying the assumption (cid:15) ( nT )
A ≤ C 0 e − ˜ sT , this implies that
[Equation] [Equation]
where we used Lemma 7 for the last inequality.
Moreover, the same reasoning as Eq.
(29) yields that
[Equation]
Therefore, combining Eqs.
(33) and (34)
[Equation]
we can bound the duality gap as
[Equation]
This inequality provides the sufﬁcient number of iterations to attain the desired duality gap.
In order to have an (cid:15) P -approximated solution, it is sufﬁcient to let
[Equation]
Next, we show the convergence of the averaging version.
Similarly to Theorem 2, we have that for k = 1 , . . .
, (cid:100) 1 κ (cid:101) ,
[Equation]
Note that, if Eq.
(31) holds, we have that 1 s e − ˜ s ( T end −(cid:100) 1 /κ
(cid:101) )
(cid:16) B 1 + B 2 +
C 0 1 − e − ˜ s (cid:17) ≤ (cid:15) P 2 .
Hence, if (cid:15) ( t ) A and (cid:15) ( t ) B are uniformly bounded by (cid:15) P 4 for all t = ˜ n ( T end − (cid:100) 1 /κ (cid:101) ) + 1 , . . .
, ˜ nT end , it holds that
[Equation] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
This concludes the proof.
In the following, we derive the sufﬁcient conditions for the assumptions in Lemma 8 by considering the error induced by the sampling stage (Section C.2.1) and the error induced at the re-weighting stage (Section C.2.2), respectively.
For that purpose, we ﬁrst 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 arbitrary i ∈ [ n ] , consider the following two updates:
[Equation]
Then, we have that
[Equation]
Proof.
From the deﬁnition of z 1 , it holds that
[Equation] [Equation]
and similarly,
[Equation]
where (cid:96) ∗ i (cid:48)
( z ) is an element in the sub-differential ∂(cid:96) ∗
i ( z ) .
From (35), we have that
[Equation]
Since (cid:96) ∗ i is γ -strongly convex, − (cid:96) ∗ i ( z )
+ H 2 ( z − y 2 ) − 1 2 nλ 2 |
z − y 2 | 2 is (cid:16) γ + 1 nλ 2 (cid:17) -strongly concave.
Thus, we obtain that
[Equation]
where ∇ f ( z ) | z = z (cid:48) is an element of the sub-differential of f at z
= z (cid:48) .
This concludes the proof.
Now we are ready to get into the main proofs which evaluate the discretization error.
C.2. C ONVERGENCE OF AUXILIARY VARIABLES ˜ g ( t ) AND ¯ g ( t )
We now ﬁx some T , and consider the sequances { g ( t ) } ˜ n ( T +1) t =˜
nT , { ¯ g ( t ) } ˜ n ( T +1) t =˜
nT +1 , { ˜ g ( t ) }
˜ n ( T +1) t =˜
nT +1 .
Lemma 10.
Suppose that (cid:82) (cid:12)(cid:12) p (˜ nT ) ( θ ) − p [ g (˜ nT ) ]
( θ ) (cid:12)(cid:12) d θ ≤ (cid:15) (˜ nT )
C holds in Algorithm 2.
Under these assumptions, it holds that
[Equation]
) for t = ˜
nT + 1 , . . .
, ˜ n ( T + 1) , where C 2 is deﬁned as follows:
[Equation]
Proof.
By Lemma 9, we have that
[Equation]
so we evaluate each terms of Eq.
[Equation]
First
[Equation]
we bound the ﬁrst term of the right hand side of Eq.
Using Lemma
[Equation]
we have the following bounds:
[Equation]
Using (38)
[Equation]
we have that
[Equation]
Thus
[Equation]
combining Eqs.
(38) and (39) yields we have the following bound for the ﬁrst term of Eq.
[Equation]
[Table]
[Equation]
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 | ¯ g ( t +1)
i t − ˜ g ( t +1)
i t | by applying (40) to (37):
[Equation]
g
[Equation]
Moreover
[Equation]
we have that
[Equation]
Combining Eqs.
(41) and (42)
[Equation]
we can apply Gronwall’s lemma (“summation version” of Lemma 18): we have that, for t = ˜
nT, ˜ nT + 1 , . . .
, ˜ n ( T + 1) − 1 ,
[Equation]
, where we used C 1 ≥ 1 for the second inequality.
Remember that we deﬁned C 2 as C 2 = 2 Cγ + 1 nλ 2 exp (cid:16) ˜ n (2 C 21 +1) nγλ
2 +1 (cid:17) .
By the deﬁnition of ˜ g ( t ) and ¯ g ( t ) , ¯ g ( t ) i = ¯ g ( t ) i holds for t = ˜
nT + 1 , . . .
, ˜ n ( T + 1) unless the coordinate i is chosen during the inner loop.
Otherwise, we have (cid:12)(cid:12)(cid:12) ˜
g ( t ) i − ¯ g ( t ) i (cid:12)(cid:12)(cid:12) ≤ C 2 (cid:15) ( nT ) C for all t ≥ τ + 1 where τ is the ﬁrst time at which the coordinate i is chosen as we have seen above.
Thus, Eq.
(36) holds for all t = ˜
nT
+ 1 , . .
˜ n ( T + 1) and i = 1 , . . . , n .
C.2.
P ARTICLE APPROXIMATION , AND ERROR ANALYSIS OF g ( t )
From now on, we 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.
Lemma 11.
With probability 1 − δ , it holds that
[Equation]
uniformly over all t = ˜
nT, . .
˜ n ( T + 1) − 1 and the choice of the coordinates ( i τ ) ˜ n ( T +1) − 1 τ =˜
nT . .
Proof.
For t = ˜
nT + 1 , . . .
, ˜ n ( T + 1) − 1 , let
H :
= (cid:40)
f
[Equation]
) (cid:9) , where C I ,a := (cid:82) p (˜ nT )
( θ (cid:48) )
exp (cid:16) − 1 nλ 2 (cid:80) t − 1 τ =˜
nT h
i τ ( θ (cid:48) )
a τ (cid:17) d θ (cid:48) for I =
( i ˜ nT , . .
, i t − 1 ) and a = ( a ˜ nT , . . .
, a t − 1 ) .
We also deﬁne
[Equation]
Then, we can see that H = ∪ i ∈
[ n ] H i .
Then, (cid:80) Mj
=1 h i t ( θ j )¯ r ( t )
j can be rewritten as
[Equation]
In a similar way, we also have that
[Equation]
To check this, ﬁrst note that ( Mh i t ( θ j )¯ r ( t ) j ) Mj
=1 are i.i.d.
random variables for each t = ˜
nT
+ 1 , . .
˜ n ( T + 1) − 1 .
Indeed, ¯ r ( t ) j is determined by ¯ g ( t ) which is updated by using a continuous limit and thus is independent of the choice of particles ( θ j ) Mj =1 .
Their expectations can be evaluated as
[Equation]
Moreover, they are uniformly bounded as
[Equation]
The same results are also true for t = ˜
nT .
Note that it sufﬁces to bound Z = sup f I ,a ∈H | 1 M (cid:80) Mj
=
f I ,a ( θ j ) − E θ
[ f I ,a ( θ )] | to obtain the assertion.
For that purpose, we deﬁne the Rademacher complexity of a function class H (cid:48) as
[Equation]
where ( ε j ) Mj =1 is an i.i.d.
sequence of the Rademacher random variables ( P ( ε j = 1) =
P ( ε j = − 1) = 1 / 2 ) and the expectation is taken with respect to both ( ε j ) M j =1 and ( θ j ) M j =1 .
Then, the
Rademacher concentration inequality (Mohri et al., 2012, Theorem 3.1) tells that the uniform bound on a function class H (cid:48) can be obtained as
[Equation]
for 0 < δ < 1 .
Here, let Z i =
sup f ∈H i | 1 M (cid:80) Mj =
1 f ( θ j ) − E θ [ f ( θ )]
| and then we also see that
Z = max i ∈ [ n ] Z i .
Hence, we have that
[Equation]
for s > 0 .
If we substitute s ← max i ∈ [ n ] 2 R
[Equation]
(cid:1) , then by Eq.
(44) we can see that the right hand side can be bounded by δ .
Note that we have already seen that (cid:107) f (cid:107) ∞ ≤ C 21 for any f ∈ H (see Eq.
Therefore, we just need to bound R ( H i t ) .
By the contraction property of the Rademacher complexity ((Boucheron et al., 2013, Theorem 11.6) or (Ledoux & Talagrand, 1991, Theorem 4.12) and its proof) and remembering the fact that
the
Rademacher complexity of a function class is same as that of its convex hull (Mohri et al., 2012,
Theorem 6.2), we obtain
[Equation]
where we used the contraction properties in the ﬁrst and second inequalities
[Equation]
we used the convex hull argument in the third inequality
[Equation]
and the last inequality is by Massart’s Lemma (Mohri et al., 2012,
Theorem 3.3).
Summarizing these evaluations
[Equation]
we obtain that
[Equation]
The same bound also holds for t = ˜
nT .
Then, by taking the uniform bound over t = ˜ nT, . . .
, ˜ n ( T + 1) − 1 and simpifying the bound using ˜ n ≤ n , we yield that
[Equation]
uniformly over t = ˜
nT, . .
˜ n ( T + 1) − 1 and the choice of the sequence ( i τ ) ˜ n ( T +1) − 1 τ =˜
nT with probability 1 − δ , which yields the assertion.
Lemma 12.
With probability 1 − δ , it holds that
[Equation]
uniformly over all t = ˜
nT, . .
˜ n ( T + 1) − 1 , i = 1 , . . . , n and the choice of coordinates ( i τ ) ˜
n ( T +1) − 1 τ =˜
nT .
Proof.
Let C p
= (cid:82) p (˜ nT )
( θ ) exp (cid:16) − 1 nλ 2 (cid:80) t − 1 τ =˜
nT h
i τ ( θ ) δ ¯ g ( τ ) (cid:17) d θ .
Then, noticing that ˆ r ( t ) j
= ˆ r ( t ) j /C p and (cid:107) h
i (cid:107) ∞ ≤ 1 , we have that (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
M (cid:88)
j = h i ( θ j ) ˆ r ( t ) j (cid:80) Mj (cid:48)
=
1 ˆ r ( t ) j (cid:48) −
M (cid:88)
j = h i ( θ j )¯ r ( t ) j (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) =
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
M (cid:88)
j = h i ( θ j ) ˆ r ( t ) j (cid:80) Mj (cid:48)
=
1 ˆ r ( t ) j (cid:48) −
M (cid:88)
j = h i ( θ j ) ˆ r ( t )
j
C p (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) = (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
M (cid:88)
j = h i ( θ j ) ˆ r ( t ) j (cid:80) Mj (cid:48)
=
1 ˆ r ( t ) j (cid:48) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
We apply the same argument as Lemma 11 to evaluate the far right hand side (cid:12) (cid:12)(cid:12) 1 − (cid:80) Mj
=1 ¯ r ( t ) j (cid:12)(cid:12)(cid:12) by replacing h i t ( θ ) with 1 in the deﬁnition of H .
Then, we obtain the assertion.
From now on, we deﬁne (cid:15) D as
[Equation]
Then, we have the following bound.
Lemma 13.
We have that, with probability 1 − δ with respect to realization of particles ( θ j ) Mj =1 ,
[Equation]
for all t = ˜
nT +1 , . . .
, ˜ n ( T +1) , i = 1 , . . . , n and the choice of coordinates ( i τ ) ˜ n ( T +1) − 1 τ =˜
nT uniformly.
Proof.
From Lemma 11 and 12, with probability 1 − δ , we have that
[Equation]
¯ p ( t ) ( θ )d θ (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) ≤ C 21 (cid:15) D , and that (cid:12) (cid:12) (cid:12)(cid:12)(cid:12)(cid:12)
M (cid:88)
j = h
i t ( θ j )
ˆ r ( t ) j (cid:80) Mj (cid:48)
=
1 ˆ r ( t ) j (cid:48) −
M (cid:88)
j = h
i t ( θ j )¯ r ( t ) j
(cid:12) (cid:12) (cid:12)(cid:12)(cid:12)(cid:12) ≤ C 21 (cid:15) D , for all t = ˜
nT, . .
˜ n ( T + 1) − 1 .
Therefore, by Lemma 8, we have that | g ( t +1)
i t − ¯ g
( t +1)
i t | ≤ γ +
[Equation]
nλ  (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
M (cid:88)
j = h i ( θ j ) r ( t ) j (cid:80) Mj (cid:48)
=1 r ( t ) j (cid:48) − (cid:90) h
[Equation]
( θ )¯ p ( t ) ( θ )d θ (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) + 1 nλ 2 | g ( t )
i − ¯ g ( t ) i | 
 ≤ γ + 1 nλ
[Equation]
 
[Equation]
C
[Equation]
(cid:15) D + (cid:12)(cid:12)(cid:12) (cid:12)(cid:12)(cid:12)
M (cid:88)
j = h i ( θ j ) (cid:32)
[Equation]
The second term of the right hand side of Eq.
(46) can be further bounded as (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
M (cid:88)
j = h i ( θ j )
[Equation]
r
[Equation]
−
[Equation]
(cid:33)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Table]
) i τ − ¯ g
( τ +1)
i τ | (47) for t = ˜
nT, . .
˜ n ( T + 1) − 1 .
This and Eq.
(46) yields
[Equation] [Equation]
τ =˜
nT | g ( τ +
[Equation]
− ¯ g
( τ +1)
[Equation]
τ | (cid:33) for t = ˜
nT + 1 , . . .
, ˜ n ( T + 1) − 1 .
Also, it holds that | g (˜ nT +1)
i −
[Equation]
g (˜ nT
[Equation]
i
[Equation]
from Lemma11.
Therefore, applying Gronwall’s lemma (“summation version” of Lemma 18) yields that for t = ˜
nT, . .
˜ n ( T + 1) − 1 ,
[Equation]
i
[Equation]
which concludes the proof of Eq.
C.2. I NTEGRATION OF TWO SOURCES OF DISCRETIZATION ERROR
From now on, we 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 = 0 , . . .
, T end .
So far, we have been arguing only within one inner loop, while we may choose a different number
M of particles in each outer-loop T .
Accordingly, we denote the number of particles at the T -th outer-iteration by M (˜ nT ) and add a subscript to (cid:15) D as
[Equation]
Moreover, we assume that the bound of Lemma 13 holds for every T = 0 , . .
, T end − 1 .
We denote by
E this even.
We know that P ( E ) ≥ 1 − δ .
In the following, we analyze the convergence under the event
E where the bounds in Lemma 13 hold, so that | ¯ g ( t )
i
[Equation]
i
[Equation]
holds owing to Lemma 10 and Lemma 13.
Therefore, the expectation with respect to the choice of coordinates is taken under the condition of this event.
We note that the bound in Lemmaholds 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 (cid:15) A , we present the following Lemma 14 on the local smoothness of dual space.
Lemma 14.
Suppose that n ≥ 2 B 2 λ 2 and Assumption 1 holds.
Then, D ( g ) is 1 nλ 2 ( λ 2 γ (cid:48) + 1) -smooth with respect to g ∈ R n on the set
A deﬁned by
[Equation]
. . , n ) (cid:111) , where t takes t = 0 , 1 , . . . ,
˜
nT end .
Proof.
Suppose that the coordinate i is chosen at the t -th iteration.
By the deﬁnition of g ( t ) , we have that
[Equation]
where (cid:96) ∗ i (cid:48) is an element of the sub-differential.
The ﬁrst term of the right hand side can be bounded by 1 as follows:
[Equation]
Also, Lemma 6 implies that the second term can be bounded by 2 B 2 nλ 2 .
Therefore, under the assumption n ≥ 2 B 2 λ 2 , we have that
[Equation]
Applying the same argument, we also have that
[Equation]
Therefore, for i ∈ [ n ] that has already chosen in the algorithm until the t -th step, we see that (cid:96) ∗ i is γ (cid:48) -smooth on the interval between g ( t ) i and ˜ g ( t ) i .
For the coordinates i that have not yet chosen, we have that g ( t ) i =
˜
g ( t ) i
= g (0) i and | (cid:96) ∗ i (cid:48)
( g (0) i ) | ≤ 1 .
Therefore (cid:96) ∗ i is γ (cid:48) -smooth on the box A .
By summarizing this argument and applying Lemma 4, we see that D ( g ) is 1 nλ 2 ( λ 2 γ (cid:48)
+ 1) -smooth.
Combining the above results, we obtain the evalution of (cid:15) A and sufﬁcient conditions for the ﬁrst statement of Lemma 8.
Lemma 15.
Suppose that n ≥ 2 B 2 λ 2 and ﬁx an arbitrary T .
if E [ D (˜ g (˜ nT (cid:48) ) ) − D ( g (˜ nT (cid:48) ) )] ≤
C 0 exp ( − ˜ sT (cid:48) ) is satisﬁed for all T (cid:48) ≤ T , then we have that
[Equation]
for all t = ˜
nT
+ 1 , . .
˜ n ( T + 1) , where µ =
[Equation]
In particular, if
[Equation]
are satisﬁed for all T = 0 , . .
, T end , where
[Equation]
then it holds that for all t = 0 , 1 , . . .
, ˜ nT end ,
[Equation]
where T = (cid:6) t ˜ n (cid:7) .
Proof.
Under the assumptions that n ≥ 2 B 2 λ 2 , Lemma 14 holds.
Fix some T and we will show the ﬁrst assertion for t = ˜
nT + 1 , . . .
, ˜ n ( T + 1) .
Let A t =
(cid:8) s ˜ g
( t )
+ (1 − s ) g
( t ) (cid:12)(cid:12) s ∈ [0 , 1] (cid:9) , then Lemma 15 implies that D ( g ) is 1 nλ 2 ( λ 2 γ (cid:48)
+ 1) - smooth on the set.
Hence, for µ =
[Equation]
, the Cauchy-Schwarz inequality yields
[Equation]
If D (˜ g ( t ) ) − D ( g ( t ) ) ≥
[Equation]
then we have that
[Equation]
Rearranging the terms we have that
[Equation]
D
[Equation]
This holds as well when D (˜ g ( t ) ) − D ( g ( t ) ) ≤ 0 .
Here, Lemma 5 implies that E [ D (˜ g ( t ) )] is monotonically non-decreasing within the same inner loop, and thus we have that E (cid:2) D ( g ∗ ) − D (˜ g ( t ) ) (cid:3) ≤ E (cid:2) D ( g ∗ ) − D ( g (˜ nT ) ) (cid:3) .
Moreover, the right hand side E (cid:2) D ( g ∗ ) − D ( g (˜ nT ) ) (cid:3) can be further bounded by (cid:16) B 1 + B 2 +
[Equation]
s
[Equation]
2 − µµ exp( − ˜ s ( T + 1)) by using the assumption and applying Eq.
(33) of Lemma 8.
Therefore,
[Equation]
This concludes the ﬁrst assertion.
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 , . . .
˜
nT .
Then the assumption of the ﬁrst statement of Lemma 8 holds for T (cid:48)
= 0 , . .
, T, which leads to Eq.
(49) for t = ˜
nT + 1 , . . .
, ˜ n ( T + 1) .
Then we have that Eq.
(50) also holds for t = ˜
nT + 1 , . . .
, ˜ n ( T + 1) , since we chose C 3 as
[Equation]
so that the right hand side of (49) is bounded by C 0 exp ( − ˜ s ( T + 1)) .
Thus, the proof is completed by induction.
Although The proof for the ﬁrst statement of Lemma 8 is almost done, we 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
nλ 22 ( γ + 3 λ 2 + 9 B 2 ) -smooth with respect to g on the set A .
Proof.
We show the assertion by a direct calculation.
Let E p
[ g ] [ h i ] def = (cid:82) h
i ( θ ) p [ g ]( θ )d θ .
First, the second derivative of the loss function can be calculated as
[Equation]
where we used smoothness of (cid:96) i for the last inequality.
The regularization term can be split into the following two terms:
[Equation]
We already conﬁrmed 1 nλ 2 − smoothness of the second term in Lemma 4.
As for the ﬁrst term, we have that
[Equation]
where we used Lemma 6 for the last inequality to bound g i .
Combining the two inequalities and the result from Lemma 4, we obtain the assertion.
Lemma 17.
Suppose Eq.
(50) of Lemma 15 and its assumption hold for t = 0 , 1 , . . .
˜
nT end .
Then, we have that
[Equation]
where T = (cid:6) tn (cid:7) .
Proof.
By using the smoothness of P ( p [ g ]) by Lemma
[Equation]
we have that
[Equation]
where T = (cid:6) tn (cid:7) − 1 .
if P ( p [ g ( t ) ]) − P ( p [˜ g ( t ) ]) ≥ 0 , then we have that
[Equation]
Since D is γn -strongly convex
[Equation]
we have that (cid:107) g
( t ) − g ∗ (cid:107) 2 ≤ (cid:113) nγ (cid:
) D ( g ∗ ) − D ( g ( t ) ) (cid:1) .
Therefore, it holds that
[Equation]
which holds as well when P ( p [ g ( t ) ]) − P ( p [˜ g ( t ) ]) ≤ 0 .
By taking expectations
[Equation]
it yields that
[Equation]
where we used Lemma 8 and Lemma 15 for the last inequality.
Furthermore, substituting C 2 C 3 = √ C 0 (2 − µ ) µ √ 2 e ˜ s ( γ (cid:48) +1 /λ 2 ) , we have the assertion.
Finally, we will prove the Theorem 3, which gives sufﬁcient conditions about sampling accuracy and number of particles for Theorem 1.
Theorem 3.
Suppose that Assumptions 1 and 2 hold.
When we choose an Option (A), suppose that we set T end so that
[Equation]
as in Eq.
(30) in Lemma 8.
Furthermore, when we choose an Option (B), where we choose a solution randomly from t = ˜ n ( T end − (cid:100) 1 κ (cid:101) ) , . .
. , ˜ nT end , suppose that we set T end
so that
[Equation]
where
C
[Equation]
Then, the condition imposed on (cid:15) ( t ) A and (cid:15) ( t ) B in Lemma 8 can be satisﬁed for both options under the event E if the following inequalities are satisﬁed:
[Equation]
for T = 0 , 1 , . . .
, T end − 1 .
That is, under the conditions above, E [ P ( p [ˆ g ]) − D (ˆ g )) |E ] ≤ (cid:15) P is satisﬁed for the solution ˆ g of either Option A or B.
Proof.
As for the Option (A), according to Lemma 15, Eqs.
(51) and (53) are the sufﬁcient conditions for the ﬁrst statement of Lemma 8 in the event E .
Next, we will show the convergence of Option (B).
In the same way, Eqs.
(52a) and (53) assure that (cid:15) (˜ nT )
A is smaller than C 0 e − ˜
sT
for every T = 1 , . . .
, T end , according to Lemma 15.
Furthermore,
Eqs.
(52b) and (52c) states that both (cid:15) ( t ) A and (cid:15) ( t ) B are uniformly bounded by (cid:15) P 4 for all t =
˜ n ( T end − (cid:100) 1 κ (cid:101) ) + 1 , . . .
nT end .
Note that the evaluation of (cid:15) ( t ) A follows from Lemma 15, and that of (cid:15) ( t ) B is derived in Lemma 17.
Therefore, all assertions are proved.
Theorem 3 essentially shows Theorem 1.
More detailed explanations can be found in the following.
Proof of Theorem 1.
Theorem 3 gives the sufﬁcient conditions for Lemma 8.
Thus, we have that, for
Option (A), with Eqs.
(51) and (53), an expected duality gap E (cid:2) P ( p [ g (˜ nT end ) ]) − D ( g (˜ nT end ) )
|E (cid:3) ≤ (cid:15) P can be achieved, conditioned by the event E whose probability is no less than 1 − δ .
In the same way, for Option (B), with Eqs.
(52a)-(52c) and (53), an conditional expectation of duality gap
E (cid:104) P ( p [ g ( t (cid:48) end ) ]) − D ( g ( t (cid:48) end ) )
|E
(cid:105) ≤ (cid:15) P can be achieved.
The sufﬁcient condition of t end is then obtained by noticing t end =
˜ nT end and setting ˆ C 1 = C 3 / 7 , ˆ C 2 =
C 21 and ˆ C 3 = 2 C q 1+ 1 nγλ 2 .
Next, we evaluate the constants as follows.
First, we note that ˆ C 2 = C 21 by their deﬁnitions.
Moreover,
C
[Equation]
where we used C 2 =
[Equation]
exp (cid:16) ˜ n (2 C 21 +1) nγλ 2 +1 (cid:17) and C 21 =
ˆ C 2 .
Finally, the evaluation of ˆ C 3 can be obtained by bounding C q using 1 / (1 − exp( − ˜ s ))
=
O (˜ s − 1 ) and 1 ≤ 1 + 1 nγλ 2 ≤ 1 + 1 2 B 2 γ =
O (1) .
C.
[Equation]
D ISCRETE G RONWALL ’ S LEMMA
In the proofs of Lemmas 8, 10 and 13, we used a discrete version of the Gronwall’s lemma.
We give its formal statement in the following lemma.
Lemma 18 (Discrete Gronwall’s lemma) .
(1) (Vanilla version) Suppose that a sequence of real numbers ( u n ) satisﬁes
[Equation]
, for a positive real a > 0 and a sequence of real numbers ( b n ) .
Then, it holds that
[Equation]
(Summation version) Suppose that a sequence of real numbers ( u n ) satisﬁes
[Equation]
B
[Equation]
n
[Equation]
, for a positive real a > 0 and a real number B .
Then, it holds that
[Equation]
Proof.
The assertion is just an application of the standard discrete version of the Gronwall’s lemma.
However, we give a proof for completeness.
The ﬁrst assertion is a 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
[Equation]
The right hand side is equivalent to
[Equation]
which yields the second assertion.
D A DDITIONAL EXPERIMENTS AND DETAILS OF EXPERIMENTAL SETTINGS
Here, we 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.
Additional experiments for a single neuron teacher network We consider a setting where the teacher network consists of a single neuron:
[Equation]
where the input dimensionality is d = 10 , w ∗ 1 ∈ R d is generated from the uniform distribution on the unit sphere S d − 1 in the d -dimensional Euclidean space, x i ∈ R
d is distributed from a standard normal distribution on R d , (cid:15) i ∈
R obeys the Gaussian distribution with mean 0 and standard deviation 0 .
5 and the activation function σ is ReLU.
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 θ ( x ) =
tanh( w (cid:62) x + b ) with θ = ( w, b ) ).
We employed the squared loss (cid:96) ( f, y ) =
( y − u ) 2 /for
the loss function.
We conducted experiments with three patterns of regularization parameters: λ
1 = λ 2 = λ ∈ { 0 . 01 , 0 .
The naive SGD was trained to minimize the primal objective for the student network with M = 200 without the entropy regularization, i.e.,
[Equation]
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 η used for sampling (ULA) and SGD were set as 10 − 3 .
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 a ﬁnite 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 ﬁnite particles.
This is also applied to 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 λ .
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 − 2 , P-SDCA and PDA does not decrease the excess primal objective below 10 − 2 .
This would be due to the instability of the entropy estimator.
For small λ , we see that
P-SDCA outperforms other methods after some iterations because it yields linear convergence.
This observation supports our theoretical analyses.
Additional experiments for a teacher network with multiple neurons We consider a setting where the teacher network consists of multiple neurons:
[Equation]
where the width of the teacher network is M t = 10 , the input dimensionality is d = 5 , the activation function of the teacher network is σ ( · ) = tanh( · ) , the noise is drawn from Normal distribution and
[Figure]
[Equation]
[Figure]
its Signal to Noise Ratio (SNR) is 7 . 0 .
For each method, we selected the optimal step size: 10 −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 − 4 to see how the step size affects the convergence of SGD.
The purpose to include SGD as well as P-SDCA and PDA is to see how the entropy term will behave when we do not explicitly include the entropic regularization term into the objective function.
[Figure]
We compared our method with SGD and PDA with the regularization parameters λ 1 = 0 .
01 (ﬁxed) and λ 2 = 0 .
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 λ 2 = 0 .
01 , we can see that particle methods (P-SDCA and PDA) showed faster convergence than that of SGD.
Next, as for λ 2 = 0 .
001 , we can see that P-SDCA converged faster than SGD and PDA.
Although PDA and SGD with step size η = 10 − 5 also converged eventually, the behavior of PDA in the intermediate stage was quite unstable and SGD with η = 10 − 5 showed quite slow convergence.
These observations support the superiority of our method.
Finally, as for λ 2 = 0 . 0001 , SGD with step size η = 10 − 5 was the fastest due to the small regularization parameter λ 2 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 ( η = 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 λ 2 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 ﬁeld neural network.
For this purpose, we run the algorithm in a binary classiﬁcation 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 η = 10 − 5 and the regularization parameters were set as λ 1 = 10 − 2 , λ 2 = 10 − 4 .
To visualize the adaptivity of the feature map, we ﬁrst plot the evolution of singular values of the ﬁrst layer’s weight matrix, W = [ w 1 , w 2 , . .
. , w M ] ∈ R d × M , 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-ﬁeld 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 ﬁrst layer’s parameters do not change much during optimization in the NTK regime; in other words, SGD in the NTK regime does not perform adaptive feature extraction.
In that sense, we expect that training in the mean ﬁeld regime is effective in a setting where adaptive feature extraction is required.
In addition to the singular value evolution, we evaluate how well the extracted features are “aligned” to the label.
For that purpose, we deﬁne a kernel function k W ( x, x (cid:48) )
= (cid:80) Mm =1 σ ( w (cid:62)
m x ) σ ( w (cid:62)
m
x (cid:48) ) for a weight matrix W =
[ w 1 , w 2 , . .
. , w M ] ∈ R d × M .
Then, we deﬁne the kernel alignment (Cristianini
[Figure]
[Equation]
where (cid:104) A, B (cid:105) F := (cid:80) ni
=1 (cid:80) nj =
A ij B ij , K W =
( k W ( x i , x j ) n,ni =1 ,j =1 ) and y =
( y 1 , . . .
, y n ) (cid:62) .
We also deﬁne the kernel alignment on the test data in the same manner.
We can see that the kernel alignment represents how strongly the kernel function deﬁned by the features after the ﬁrst layer is aligned to the target signal.
The kernel alignment of the solutions obtained by each method (P-SDCA, SGD in NTK regime, SGD in mean ﬁeld regime) is depicted in Figure 5.
We can see that P-SDCA (and SGD in the mean ﬁeld regime) properly improves the kernel alignment, which means that P-SDCA can extract informative features adaptively depending on the data.
On the other hand, SGD in the NTK regime does not increase the kernel alignment due to the parameters almost “frozen” at initialization.
This experiment highlights how well the mean ﬁeld neural networks can execute feature learning.
[Figure]