Two-layer neural network on inﬁnite-dimensional data: global optimization guarantee in the mean-ﬁeld regime
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Analysis of neural network optimization in the mean-ﬁeld regime is important asthe setting allows for feature learning .
Existing theory has been developed mainlyfor neural networks in ﬁnite dimensions, i.e., each neuron has a ﬁnite-dimensionalparameter.
However, the setting of inﬁnite-dimensional input naturally arises in machine learning problems such as nonparametric functional data analysis and graphclassiﬁcation.
In this paper, we develop a new mean-ﬁeld analysis of two-layerneural network in an inﬁnite-dimensional parameter space.
We ﬁrst give a generalization error bound, which shows that the regularized empirical risk minimizerproperly generalizes when the data size is sufﬁciently large, despite the neuronsbeing inﬁnite-dimensional.
Next, we present two gradient-based optimizationalgorithms for inﬁnite-dimensional mean-ﬁeld networks, by extending the recentlydeveloped particle optimization framework to the inﬁnite-dimensional setting.
Weshow that the proposed algorithms converge to the (regularized) global optimalsolution, and moreover, their rates of convergence are of polynomial order in theonline setting and exponential order in the ﬁnite sample setting, respectively.
Toour knowledge this is the ﬁrst quantitative global optimization guarantee of neuralnetwork on inﬁnite-dimensional input and in the presence of feature learning.
Introduction
A variety of machine learning problems need to handle input from inﬁnite-dimensional spaces.
Forinstance, Ling and Vieu (2018); Ferraty et al.
(2007) studied non-parametric function regressionproblems where the input is a function in an inﬁnite-dimensional functional space and Kriege et al.(2020) studied graph classiﬁcation problems using kernel method, which can be cast as a learningproblem with inputs from an inﬁnite-dimensional reproducing kernel Hilbert space.
Among various models that deal with those inﬁnite-dimensional input problems, neural networks areparticularly interesting due to their ability to learn features and model nonlinear data.
For example,
Rossi and Conan-Guez (2005) proposed the Functional Multi-Layer Perceptron, where the input offunctional data is converted to vector form by basis expansion and then fed into a neural network;the universal approximation and statistical consistency of the proposed model were also analyzed.
Recently, Yao et al.
(2021) proposed the Adaptive Functional Neural Network, which replaces theﬁxed basis functions in Rossi and Conan-Guez (2005) with adaptive bases, in order to leveragethe power of representation (feature) learning and the ﬂexibility of deep learning.
While neuralnetwork based models with inﬁnite-dimensional input have been applied in many applications, theiroptimization using gradient-based methods is not well-understood.
In particular, no optimizationguarantee has been shown for any algorithm in the presence of feature learning — such theoreticalresult can be challenging to establish even for ﬁnite-dimensional data.
It is experimentally observed that the gradient-based optimization methods can yield sufﬁciently small training error in training neural network models, despite the non-convexity of the landscape
Theoretical explanations of this observation often rely on overparameterization , that is, to considersufﬁciently wide neural network compared with the data size, and the model parameterization can bedivided into the mean-ﬁeld regime (Nitanda and Suzuki, 2017; Mei et al., 2018; Chizat and Bach,2018) and the neural tangent kernel regime (Jacot et al., 2018;
Du et al., 2018).
In this work we focuson the mean-ﬁeld regime, as it captures the presence of feature learning, which is one of the mainadvantages of neural network (Chizat et al., 2019; Yang and Hu, 2021).
In the mean-ﬁeld analysis, the dynamics of gradient descent is described by Wasserstein gradientﬂow in the space of probability distributions on the parameters.
Mei et al.
(2018) showed that themean-ﬁeld Langevin dynamics converges to the global optimal solution, and Hu et al.
(2019) provedlinear convergence with respect to the objective function with sufﬁciently strong KL divergenceregularization.
However, those studies mainly analyzed the algorithm in continuous time, and do notestablish quantitative convergence rate for discrete time and ﬁnite width settings.
To overcome this limitation, Nitanda et al.
(2021) proposed the Particle Dual Averaging (PDA)method that globally optimizes the KL-regularized objective with a polynomial order computationalcomplexity, in a completely discrete time and ﬁnite width setting.
Furthermore, in the case of ﬁnitesum objective, Oko et al.
(2022) proposed the Particle Stochastic Dual Coordinate Ascent (P-SDCA)method which further improves the computational complexity to an exponential order with respectto its outer loop iteration.
Both PDA and P-SDCA employ a double-loop structure, and make useof Monte Carlo sampling from an intermediate target distribution in the inner loop.
In particular,
Nitanda et al.
(2021) applied the convergence rate of the overdamped Langevin algorithm in Vempalaand Wibisono (2019) for their optimization analysis, whereas Oko et al.
(2022) also considered the
Metropolis-adjusted Langevin algorithm (MALA)
(Ma et al., 2019) for the inner loop sampling.
However, if we naively apply these algorithms to inﬁnite-dimensional input, then the optimizationguarantee becomes meaningless due to the dimension dependence.
Therefore, we need to construct anew methodology and theory for the inﬁnite-dimensional problem.
Our contributions.
In this work, we extend the mean-ﬁeld analysis for neural network training to inﬁnite-dimensional parameter space, which covers two-layer neural network with inﬁnitedimensional inputs.
First, we establish a generalization error bound which entails that our modelgeneralizes properly when the number of training data is sufﬁciently large.
Next, we propose twogradient-based algorithms that globally optimize the KL-regularized objective, corresponding to theinﬁnite-dimensional extension of PDA and P-SDCA, respectively.
As mentioned above, PDA and P-SDCA require Monte Carlo sampling, and in our setting samplingis performed in an inﬁnite-dimensional Hilbert space.
We therefore adapt an inﬁnite-dimensionalgradient Langevin dynamics (Debussche, 2011; Bréhier, 2014), the weak convergence of which hasbeen established in Muzellec et al.
Our contributions can be summarized as follows:
• We introduce two-layer neural network in the mean-ﬁeld regime whose input is inﬁnite-dimensional,and establish a generalization error bound of the regularized empirical risk minimizer.•
We propose two optimization algorithms for our inﬁnite-dimensional neural network by incorporating the inﬁnite-dimensional Langevin dynamics in the inner loop of PDA and P-SDCA.• We prove that our inﬁnite-dimensional extension of PDA achieves polynomial order convergence,and that of P-SDCA achieves exponential (outer loop) convergence, even in an inﬁnite-dimensional setting.
To the best of our knowledge, this is the ﬁrst quantitative global optimization guarantee ofneural network in the inﬁnite-dimensional mean-ﬁeld regime with the presence of feature learning.
Other related works.
Ferré and Villa (2006) proposed the SIR-NNr, which incorporates thedimension reduction method termed “sliced inverse regression” into a neural network.
In addition,
Rossi et al.
(2005) introduced a radial-basis function network as a nonlinear model whose input is afunction.
As previously discussed, although statistical properties (e.g., consistency) of these modelshave been studied, convergence guarantee of gradient-based optimization has not been established.
Two-layer neural network with inﬁnite-dimensional input
In this section we introduce the mean-ﬁeld two-layer neural network with inﬁnite-dimensional input.
We ﬁrst formulate mathematical notations and prepare assumptions required for our theoreticalanalyses, and then present a generalization error of the proposed model.
Neural network model.
We consider a setting where the input is included in a (possibly inﬁnitedimensional) separable Hilbert space H .
Since H is separable, there is a complete orthonomal system( e j ) 1 j =0 for which it holds that H =
[Equation]
o , and H is equipped with theinner product given by h f, g i H :=
{+P 1 j =
+}
a j b j , where f = P 1 j =
a j e j , g =
P 1 j =
b j e j 2 H .
Let
H 0 be a subset of H deﬁned by H 0 :
= { 0 · e 0 + P [Equation]
a j e j [Equation]
P 1 j =1 a 2 j < 1 } .
We consider aneural network model that takes an inﬁnite-dimensional input x 2 X ✓ H 0 and gives an outputy 2 Y ✓ R , where X and Y are domains of input and output, respectively.
Each neuron in thenetwork has parameter ✓ := ae 0 + w 2 H for a 2 R and w 2 H 0 , and is represented by
[Equation] where  1 ,  2 : R ! R are activation functions, which we assume to be bounded and smooth (such astanh ) 1 .
We also denote h ✓ ( x ) by h ( ✓ , x ) .
Let the number of neurons be M and denote the set ofparameters by ⇥ =
( ✓ r )
Mr =1 .
Then, in the mean-ﬁeld regime, the two layer neural network model(with inﬁnite-dimensional input) can be expressed as the average over the M neurons:
[Equation]
This can be seen as a ﬁnite sum approximation of the integral form E ✓ ⇠ ⇡ [ h ✓ ( x )] for a distribution ⇡on the parameter space.
Indeed, if ( ✓ r )
Mr =1 are i.i.d.
realizations from ⇡ , then h ⇥ ( x ) would convergeto the integral form as M increases.
This viewpoint motivates us to design optimization algorithmsfor the distribution of parameters ⇡ and utilize the convexity of the objective in the space of measures.
Notations.
To introduce our method, we need to deﬁne a Gaussian process taking its value in H .
For that purpose, we deﬁne a subspace H K  ⇢ H for   0 as follows.
First, we deﬁne a linear 107 operator T K : H ! H as T K f =
P 1 j =0 µ j a j e j for f =
P 1 j =
a j e j , where ( µ j ) 1 j =0 correspondsto the spectrum of T K satisfying µ j > 0 ( 8 j ) and is sorted in decreasing order ( µ 1  µ 2  · · · ).
109 Note that the orthonormal system ( e j )j =0 is taken so that it coincides with the eigen-functions of
T K .
Then, we deﬁne H K  :=
T  2 K H =
{ T  2 K h | h 2 H} and we can naturally equip H K  with anorm k f k H K  :=
P 1 j =0 µ   j a 2 j for f = P 1 j =
a j e j 2 H K  .
For concise notation we simply write
H K :
= H K 1 .
Next, deﬁne an operator A =  1 T  1 K for a positive regularization parameter  1 > 0 ,i.e.,
Af =
P 1 j =0 µ  1 j a j e j for f = P 1 j =
a j e j 2 H K , and let ⌫ be a Gaussian measure in Hwhose mean is 0 and its covariance is A  1 (i.e., h z, x i H ( z ⇠ ⌫ ) obeys the Gaussian distribution withmean 0 and covariance h x, A  1 x i H for any x 2 H ) 2 .
Finally, we deﬁne the Kullback-Leibler (KL)divergence from a probability measure ⌫ 2 to a probability measure ⌫ 1 that is absolutely continuouswith respect to ⌫ 2 as KL( ⌫ 1 || ⌫ 2 ) :
= E ⌫h log d ⌫ 1 d ⌫i = E ⌫h d ⌫ 1 d ⌫ 2 log d ⌫ 1 d ⌫
[Equation]
i .
Objective.
Now we deﬁne the objective of our optimization method.
Suppose that the pairs ofinput and output are independently identically distributed from D , and let ` ( · , · ) :
Y ⇥ Y ! R bea loss function.
In the mean-ﬁeld regime, we optimize the following regularized risk minimizationproblem with respect to the distribution of parameters where the regularization term is given by the
KL-divergence from the Gaussian measure ⌫ :
[Equation]
where  2 > 0 and P 2 is the entire set of probability measures that satisfy
E ✓ ⇠ ⇡ [ k ✓ k 2 H ] < 1 124 and are absolutely continuous with respect to the Gaussian measure ⌫ .
We may consider theempirical measure of n training data points ( x i , y i ) ni
=1 as D , for which the ﬁrst term of the objectivefunction is written as 1 n P ni
=1 ` ( E ✓ ⇠ ⇡ [ h ✓ ( x i )] , y i ) .
In this setting, the problem can be regarded asa regularized empirical risk minimization problem.
As for the second term, in a ﬁnite-dimensionalsetting where A =  1 I , we can easily see that the KL divergence is decomposed into the sum of ` 2 -regularization  1 E ✓ ⇠ ⇡ [ k ✓ k 2 ] and negative entropy E ✓ ⇠ ⇡ [log d ⇡ ( ✓ ) / d ✓ ] .
Here the negativeentropy naturally arises from the mean-ﬁeld Langevin dynamics (Hu et al., 2019; Nitanda et al.
Chizat, 2022).
Indeed, it is known that, if we optimize a loss function L ( q ) with respect to aprobability distribution q 2 P 2 by the mean-ﬁeld Langevin dynamics with temperature parameter  2 ,then its stationary distribution is given by q ⇤ = arg min q :density L ( q ) +  2 E q [log( q )] which satisﬁesq ⇤ / exp(  1  2  L ( q ⇤ )  q ) , where  L ( q )  q : H ! R is the derivative with respect to the distribution q suchthat
[Equation]
One important example that we can apply our model to is the non-parametric functional regressionproblem (Ling and Vieu, 2018; Ferraty et al., 2007).
Suppose that the input x is a function on R(such as time evolution of temperature) included in some reproducing kernel Hilbert space H K .
Onetypical deﬁnition of H K is given by H K = T 1 / 2 K H where H =
L 2 ( µ ) for a probability measure
µand T K is an integral operator given by T K f ( t ) :=
R K ( t, s ) f ( s )d µ ( s ) for a positive deﬁnite kernel
K (Caponnetto and De Vito, 2007; Steinwart and Christmann, 2008).
The goal is to estimate somenon-linear functional from the input x 2 H K to output y 2 R (using neural network (1)).
Generalization bound.
We now give a generalization error bound for our two-layer neural networkin the empirical risk minimization setting, which conﬁrms that the model properly generalizes whenthe data size is sufﬁciently large.
Specially, we show the generalization bound for two problemsettings: binary classiﬁcation and regression.
The following theorem stats the generalization boundfor binary classiﬁcation problems.
We defer the regression case to Appendix A due to space limitation.
Theorem 1.
Let D be a distribution of X ⇥ Y , ` be the smoothed hinge loss, and denote the 0-1 149 loss by ` 01 ( z, y ) , i.e., ` 01 ( z, y ) :=
[ zy < 0] .
Let ⇡ ⇤ be the optimal solution of the empirical riskminimization problem (2) for a given set S of n training data points (i.i.d.
from D ).
Suppose thereexists a distribution ⇡  2 P 2 such that h ⇡  ( x ) y  1 / 2 for all ( x, y ) 2 supp D , and that | h ✓ ( x ) |  1 152 holds for all ✓ 2 H , x 2 X .
Then, the following holds with probability 1   with respect to the 153 choice of S ⇢ X :
[Equation]
Note that if we set  2 = 1 / p n , the generalization error is O (1 / p n ) .
This order matches that in theﬁnite-dimensional input setting shown in Nitanda et al.
Mean-ﬁeld optimization with inﬁnite-dimensional Langevin algorithm
As previously remarked, the core idea of the mean-ﬁeld analysis is to directly optimize the distributionof the parameters (2).
For that purpose, the ideal situation would be to maintain inﬁnitely manyparticles (neurons) to represent the distribution; however, in a tractable algorithm we need to usean approximation with ﬁnite number of particles.
This ﬁnite-particle approximation may lead toinstability due to the interaction between particles, which is known to be difﬁcult to control.
To overcome this difﬁculty, we employ the linearization technique that was used in the original PDAand P-SDCA algorithms (Nitanda et al., 2021; Oko et al., 2022).
We provide a short summary ofthis idea: given the t -th step solution ⇡ t , we apply a ﬁrst-order approximation of the loss function
L ( ⇡ ) := E ( X,Y ) ⇠ D [ ` ( E ✓ ⇠ ⇡ [ h ✓ ( X )] , Y )] as ( ⇡  ⇡ t )  L ( ⇡ t ) ⇡ :=
R  L ( ⇡ t ) ⇡ ( ✓ )d( ⇡  ⇡ t )( ✓ ) .
Then, ⇡ t +1is updated as the minimizer of the linearized loss objective: argmin ⇡ ( ⇡  ⇡ t )  L ( ⇡ t ) ⇡ +  2 KL( ⇡ || ⌫ ) ,which can be written as d ⇡ t +1 =
exp(  1  2  L ( ⇡ t ) ⇡ ) · d ⌫ .
Hence we just need to draw particles from⇡ t +1 in the inner loop.
Importantly, the interaction between particles disappear due to the linearizationand we may sample in an i.i.d. manner.
The original algorithms of PDA and P-SDCA for (ﬁnitedimensional input) employed the gradient Langevin dynamics (GLD) and its variant such as MALA,the convergence rate of which have been extensively studied (Raginsky et al., 2017; Vempala and
Wibisono, 2019; Ma et al., 2019).
However, in our setting we need to solve an inﬁnite-dimensionalsampling problem, which cannot be directly extended from the ﬁnite-dimensional counterpart.
Let ⇡ be a distribution on H that satisﬁes d ⇡ d ⌫ ( x ) / exp (  G ( x ))
(e.g., G =   1 2  L ( ⇡ t ) / ⇡ ).
Thenwe know ⇡ is the stationary distribution of the inﬁnite-dimensional stochastic differential equation,
[Equation]
[Equation]
where ( W t ) t  0 is a cylindrical Brownian motion on H (Da Prato et al., 1996) and r G is the Rieszrepresenter of the Fréchet derivative of G on H .
To discretize the continuous-time SDE, we employthe semi-implicit Euler scheme:
[Equation]
i.e.,
[Equation]
, where ⇣ k is a realization of standard Gaussian processon H , and S ⌘ :=
(Id + ⌘ A )  1 where Id is an identity mapping.
We remark that our choice of semiimplicit scheme is due to the fact that ⇣ k / 2 H in the inﬁnite-dimensional setting, and thus the naive
Euler-Maruyama discretization does not ensure X k +1 2 H while the semi-implict scheme does.
In practice, we need to approximate the inﬁnite-dimensional vector X k by a ﬁnite-dimensional one.
For that purpose, we apply the Galerkin approximation .
Let P N be the orthogonal projection from
H to H N := span { e 0 , . . .
, e N } , and r G N ( x ) :=
P N ( r G ( P N x )) .
Then the ﬁnite-dimensional 186 approximation of the semi-implicit Euler scheme is given as
[Equation]
A convergence rate analysis of the inﬁnite-dimensional gradient Langevin dynamics (4) with timediscretization and ﬁnite-dimensional approximation was given by Muzellec et al.
(2022), which is animportant ingredient of our convergence rate analysis (the details can be found in Appendix B).
Inﬁnite-dimensional Particle Dual Averaging
In this section, we present our ﬁrst algorithm to solve the risk minimization problem (2) which is aninﬁnite-dimensional extension of the Particle Dual Averaging (PDA) method Nitanda et al.
PDA is a combination of Nesterov’s dual averaging (Nesterov, 2009) and particle sampling; it isguaranteed to converge globally at a sublinear rate, but is tailored to the ﬁnite-dimensional setting.
Extending the algorithm and convergence analysis to inﬁnite-dimensional input is non-trivial, and todo so we make use of the weak convergence property of the inﬁnite-dimensional gradient Langevindynamics (GLD) shown in Muzellec et al.
Algorithm description of inﬁnite-dimensional PDA.
Here we provide the algorithm descriptionof PDA with inﬁnite-dimensional input.
Since we are optimizing the expected loss directly, weemploy an online gradient descent approach: we draw one input-output pair ( x t , y t ) from D at thet -th iteration, and update the parameter by replacing the expected loss with a single-sample loss` ( E ✓ ⇠ ⇡ [ h ✓ ( x t )] , y t ) .
In addition, we apply the following linearization to the loss at each outer loopstep: suppose that we have M particles ⇥ ( t ) =
( ✓ ( t ) m )
Mm =1 whose empirical measure is denoted by [Equation]
.e. [Equation]
E [Equation]
⇠
[Equation] [Equation] ✓ ( · )] ).
Then the loss can be linearized around ˆ ⇡ ( t ) as
[Equation]
We denote the gradient component in the right hand side as g ( t ) ( ✓ ) :
= @
z ` ( h ⇥ ( [Equation] ) ( x (
[Equation] ) ) , y
[Equation]
h
[Equation]
The idea of the Nesterov’s dual averaging method is to take its weighted sum of the gradient over itshistory: ¯ g ( t ) ( ✓ ) :=
[Equation]
✓ ) , and update the distribution as
[Equation]
As we have seen above, the Radon-Nykodym deriatives of ⇡ ( t +1) ⇤ with respect to ⌫ is given byd ⇡ [Equation] t +
[Equation]
⇤
[Equation] ⌫ [Equation] ✓ ) / exp   ¯ g ( t ) ( ✓ )  .
Hence in the inner loop, we sample particles from ⇡ ( t +1) ⇤ usinginﬁnite-dimensional gradient Langevin dynamics (4) by setting G = ¯ g ( t ) ( ✓ ) .
That is, we obtain the
M -particle approximation by iteratively calculating the following updates:
[Equation]
where S ⌘ t :=
(Id + ⌘ A )  1 and r ¯ g ( t ) N ( · ) :
= P N ( r ¯ g ( t ) ( P N · )) .
We repeat the GLD iterations for
T t times in the inner loop and obtain the set of particles ( ✓ ( t ) m )
Mm =
=
( ˜ ✓ ( T t ) N,m )
Mm =1 for the next outer loop update.
We denote by ⇡ ( t +1)
( ✓ ) the distribution of each particle ˜ ✓ ( T t ) N,m (this is not the empiricaldistribution ˆ ⇡ t {++
+} but its “true” distribution).
We will show that ⇡ ( t +1) ( ✓ ) well approximates the idealupdate ⇡ ( t +1) ⇤ ( ✓ ) in the proof of Theorem 5 below.
Finally, after T outer loop iterations, the algorithm outputs M parameters ✓ (ˆ t ) 1 , . . .
, ✓ (ˆ t ) M whereˆ t 2 { 2 , . . .
, T +1 } is randomly chosen according to a distribution P ( t ) =
2 t T ( T +3) ( t = 2 , . . .
, T +1) .
The full algorithm is summarized in Appendix C.
Convergence Analysis.
Now we establish the global convergence rate of our proposed inﬁnitedimensional PDA.
Let P X be the distribution of input x 2 X and G :=
supp P X .
Our analyses rely 222 on the following assumptions.
Assumption 2.
(A1) µ k ⇠ 1 k 2 , i.e. there exists a constant ↵ 1 , ↵ 2 such that ↵ 1  k 2 µ k  ↵ 2 .
(A2) Y ⇢ [  1 , 1] .
In addition, ` ( z, y ) is 1-smooth convex function with respect to z .
Moreover, for all y, z 2 Y , | @ z ` ( z, y ) |  2 holds. (A3)
Both  1 and  2 are thrice continuously differentiable.
Moreover, it holds that
[Equation] (A4)
There exists 12 <  < 2 , B X > 0 such that k x k H K 1+  < B X for all x 2 G .
Remark.
(A1) is a sufﬁcient condition to guarantee convergence of gradient Langevin dynamics(Muzellec et al., 2022), in which the exponent k  2 may be generalized to k  p ( p > 1) , but we presentthe result only for p = 2 for simplicity.
(A2) is satisﬁed by common loss functions such as square lossand logistic loss.
(A3) is satisﬁed by several practical activation functions such as sigmoid and tanh .
Finally, (A4) is a regularity condition that is used in the convergence guarantee of GLD.
We show the convergence of inﬁnite-dimensional PDA by the following two steps: (i) bounding thedifference between the optimal solution ⇡ opt of the objective (2) and the optimal update ⇡ (ˆ t ) ⇤ at theˆ t -th iteration (Theorem 4), (ii) bounding the difference between the optimum auxiliary solution ⇡ (ˆ t ) ⇤and its particle approximation ˆ ⇡ (ˆ t ) (Theorem 5).
To begin with, let ✏ ( t )
A be the weak convergence error of sampling in terms of the network output, i.e.,
[Equation]
[Equation] This error can be bounded by the following theorem, proof of which can be found in Appendix C.1.
Theorem 3.
Under assumptions (A1) – (A4) , for all t = 1 , . . .
, T and  2 (0 , 1 / 2) ,there exists ˆ C=
O ⇣ 1 +   1 2 exp  O    1⌘ , ˆ C= O  exp  O    1 , ⇤ =⌦  min {  1 ,  12
[Equation]
} exp   O    1 such that the following holds
[Equation]
Theorem 3 establishes the convergence rate of sampling via inﬁnite-dimensional GLD in the PDAalgorithm.
According to this theorem, to achieve sufﬁciently small sampling error ✏ ( t ) A < ✏ , itsufﬁces to take µ N =
O ( ✏ 2 / (1  2  ) ) , ⌘ t =
O ( ✏ 2 / (1  2  ) ) , T t = ⌦ ( ✏  2 1  2  log(1 / ✏ ) / ⇤ ) with some 2 (0 , 1 / 2) .
This result is comparable to the ﬁnite-dimensional counterpart in Nitanda et al.
(2021),where ⌘ t =
O ( ✏ 2 ) and T t = ⇥ ( ✏  2 log(1 / ✏ ) / ⇤ ) is shown to be sufﬁcient, which is close to our boundup to a factor of  despite the inﬁnite dimensionality.
It is worth noting that the ﬁnite-dimensionalanalysis in Nitanda et al.
(2021) is based on the convergence result of GLD under KL-divergence(Vempala and Wibisono, 2019).
However, in our inﬁnite-dimensional setting, we generally donot expect convergence under the KL metric since the distribution is hardly absolutely continuouswith respect to another.
Instead, our convergence guarantee is solely based on a weak convergenceanalysis, which allows us to establish computational complexity bound that is close to that of theﬁnite-dimensional counterpart in Nitanda et al.
As for (i) , the following theorem evaluates the difference between ⇡ ⇤ and ⇡ (ˆ t ) ⇤ with respect to theobjective function value.
The proof of this theorem can be found in Appendix C.2.
Theorem 4.
Suppose  > 0 .
Under (A1) , (A2) , (A3) , for all distributions ⇡ ⇤ 2 P 2 , the followinginequality holds with probability 1   :
[Equation]
This theorem shows that when the number of outer iteration T is larger than O (1 / ✏ ) , the excessobjective function value of the auxiliary solution ⇡ (ˆ t ) ⇤ becomes smaller than ✏ .
It is known that this
O (1 /T ) rate is the minimax optimal for online ﬁrst-order methods when the objective is stronglyconvex (Agarwal et al., 2009).
We note that the analysis in the original PDA algorithm (Nitanda et al.
,2021) heavily relies on the existence of probability density function of the parameters; however, thisis not guaranteed in the inﬁnite-dimensional setting.
We overcame this difﬁculty by reformulating theupdate using the Radon–Nikodym density with respect to the Gaussian measure ⌫ .
Interestingly, thisnew technique yields a reﬁned analysis that is indeed dimension independent.
Next we control the particle discretization error (ii) , i.e., the difference between the auxiliary solutionand its particle approximation.
The proof can be found in Appendix C.3.
Theorem 5.
Under (A1) , (A2) , (A3) , (A4) , the following inequality holds with probability 1   :
[Equation]
In particular, we have | L ( ⇡ (ˆ t ) ⇤ )
L (ˆ ⇡ (ˆ t ) )
[Equation]
Notice that since the particle approximation does not have a density, the bound above does not controlthe objective L including the KL-regularization, but instead the L 1 -norm error and the loss function
L of the estimated function.
By combining the previous theorems, we know that when the numberof outer loops iterations is O (1 / ✏ ) , and the number of particles M is O
((1 / ✏ ) 2 log(1 / ✏ )) , then ourproposed algorithm achieves ✏ -error in terms of the output of the model (or the loss function) from asolution with ✏ -excess objective value.
Overall, the iteration complexity of our inﬁnite-dimensional PDA to obtain ✏ -accurate solution is
O (1 / ( ⇤ ✏ 3  2  1  2  ) log(1 / ✏ )) for some  2 (0 , 1 / 2) , i.e., it is possible to obtain a solution with thedesired accuracy at polynomial order of iteration complexity (but with exponential dependency on1 /  2 , which would be unavoidable for any mean-ﬁeld method).
We emphasize that our analysisguarantees the convergence to the global optimal solution, not the stationary point or local minimizer.
Inﬁnite-dimensional Particle Stochastic Dual Coordinate Ascent
In the previous section, we considered the expected risk minimization problem.
On the other hand, inthe case of empirical risk minimization, we can improve the rate of convergence from polynomialto exponential by making use of the ﬁnite-sum property.
For that purpose, we propose an inﬁnitedimensional extension of Particle Stochastic Dual Coordinate Ascent (P-SDCA) proposed in Okoet al.
As in PDA, the original P-SDCA is also tailored to the ﬁnite-dimensional setting.
Especially, the lack of density function in the inﬁnite-dimensional parameter space necessitates acareful modiﬁcation of the update rule.
Indeed, we reconstruct the algorithm based on the densitywith respect to the Gaussian measure ⌫ .
Algorithm description of inﬁnite-dimensional P-SDCA.
Here we give an algorithmic descriptionof our inﬁnite-dimensional extension of P-SDCA.
The main ingredient of SDCA is to executeoptimization in the Fenchel dual of the objective function that is derived in the following lemma.
Lemma 6.
Suppose that ` ( · , y i ) : R !
R is a proper convex function and h ( · , x i ) :
H ! R is 297 bounded for i = 1 , . . . , n .
We denote the conjugate function of ` ( · , y i ) by ` ⇤ ( · , y i ) 4 .
Deﬁne the dualobjective D : R n ! R as
[Equation]
Then, if inf ⇡ 2 P 2 L ( ⇡ ) > 1 holds, it holds that inf ⇡ 2 P 2 L ( ⇡ ) =
sup g 2 R n D ( g ) .
We construct the update rule based on this duality theorem.
For a dual variable g 2 R n , thecorresponding primal solution can be retrieved by p [ g ]( ✓ ) := {+exp ⇣  1 n  P ni =
+}
{+g i h ( ✓ ,x i ) ⌘ E ⌫ h exp ⇣  1 n  P ni =
 g i h ( ✓ ,x i ) ⌘+}
i (indeed,for the dual optimal solution g ⇤ , the primal optimal solution can be recovered by p [ g ⇤ ]d ⌫ ).
In the
SDCA algorithm, we randomly pick up one coordinate i t 2 [ n ] from the uniform distribution, andupdate the corresponding coordinate so that it maximizes the dual objective.
Suppose that we havethe dual solution g ( t ) at the t -th iteration, then we update its i t -th coordinate by the following formulawhere i t 2 [ n ] is chosen uniformly at random:
[Equation]
{+g 
+}
i
[Equation]
⌘
[Equation]
⇣ g 0 it  g ( t ) it ⌘ 2 n o , g
[Equation]
=
[Equation]
for j 6 = i t .
We can see that the update formula requires the expectation E ⌫
[Equation] ⇤ .
We approximate this expectation by a weighted average of M particles as E ⌫
[Equation]
[Equation]
where r ( t ) m is the weight of the m -th particle.
Here, the particles ( ✓ m )
Mm =1 are“refreshed” once in every ˜ n iterations, and in the remaining iterations, we only update the weightingsof the particles ( r m )
M m =1 .
In particular: (i) at the resampling steps (once in every ˜ n iterations), we runthe inﬁnite-dimensional GLD to generate M particles ( ✓ m )
Mm =1 from the distribution ⇡ ( t ) ⇤ that is givenby d ⇡ ( t )
⇤ d ⌫ = p [ g ( t ) ] , and reset the weights as r m = 1 /M for all m 2 [ M ] ; (ii) in non-refreshing steps,we update the weights ( r m ) Mm =1 as r ( t +1) m [Equation]
h it [Equation] ✓ m )
[Equation]
g
[Equation]
t
[Equation]
it  g
[Equation]
t )
it
[Equation]
n  [Equation]  . This re-weightingscheme allows us to reduce the number of sampling steps, which can be computationally demanding.
Convergence analysis.
Now we establish the convergence rate of the proposed inﬁnite-dimensional
P-SDCA.
Suppose at the T -th sampling stage, we run the GLD with step size ⌘ T for J T iterations.
Then the sampling error can be bounded as follows: denote by p (˜ nT ) the distribution of the particles,under (A1) , (A2) , (A3) , (A4) , for all  2 (0 , 1 / 2) , it holds that
[Equation]
where ˇ C 1 , ˇ C 2 , ⇤ are constants depending on ˆ b,  1 ,  2 and  : H ! R is any test function satisfying
[Equation]
⌘⌘ with sufﬁcient smoothness (see Appendix D.3 for more details).
Usingthe sampling error ✏ (˜ nT ) C , we can now derive the iteration complexity of the algorithm.
Theorem 7.
Assume (A1) – (A4) and | ` ( x, y i )  inf x 0 ` ( x 0 , y i ) |  B 1 for all x 2 [  1 , 1] andi 2 [ n ] .
Let ˜ s :=
˜ n  2 1+ n  2 and take any T end  2 n ˜ n ⇣ 1+ 1 n ⌘ log ✓⇣ n + 1 ⌘ B 1 +2+ 1 1  exp(  ˜ s ) ✏ P ◆ for✏ P > 0 .
Suppose that ✏ (˜ nT ) C and M satisfy ✏ (˜ nT ) C  ˆ C 1 exp   ˜ sT ( 8 T 2 [ T end ]) and M ˆ C( ✏ (˜ nT ) C ) 2 log  4 nT end   where ˆ C  1=
O ✓   1 2ˆ C 2 1+ 1 n  2 exp ⇣ ˜ n (2 ˆ C 2 +1) n  2 + ⌘◆
and
[-2 + [Equation] nλ 2 +
 (cid:17)(cid:19) and-] ˆ C 2 = exp ⇣ 4 B 1 ˜ n  2 n ⌘ .
Then, after ˜ nT end iterations of inﬁnite-dimensional P-SDCA, the duality gap can be bounded as
E ⇥ L ( p [ g (˜ nT end ) ])  D ( g (˜ nT end ) )
E ⇤  ✏ P , for a high-probability event E such that P ( E )  1   .
The detailed statement and proof of Theorem 7 can be found in Appendix D.4.
Consequently,
P-SDCA can achieve a duality gap smaller than ✏ P with O (log(1 / ✏ P )) outer loop iterations.
Incombination with the sampling complexity, we see that the total iteration cost to obtain a solutionwith ✏ P duality gap is O ((1 / ✏ P )(log(1 / ✏ P )) 2 ) , which is much better than that of PDA shown in Section 4.
Also, we note that the total complexity matches the ﬁnite-dimensional case in Oko et al.(2022) except for the log factor, which again derived the complexity of sampling from GLD based onthe KL-divergence instead of the weak convergence.
Numerical Experiments
We numerical evaluate our proposed methods on a non-linear functional regression problem.
Let Hbe L 2 ([0 , 1]) , which is a set of square integrable functions deﬁned on [0 , 1] , and H K be the Sobolev space H 1 ([0 , 1]) , which is a set of absolutely continuous functions deﬁned on [0 , 1] satisfyingu (0)
= {+
+}
and u 0 2 L 2 ([0 , 1]) .
Then, H is a Hilbert space equipped with inner product deﬁned byh x 1 , x 2 i H =
R 1 0 x 1 ( t ) x 2 ( t )d t , and H K is a Hilbert space equipped with inner product deﬁned byh x 1 , x 2 i H K =
R 1 0 x 0 1 ( t ) x 0 2 ( t )d t .
Consequently, H K is a reproducing kernel Hilbert space whosekernel function is given as K ( s, t ) =
min { s, t } .
We consider a teacher-student setup.
The data is generated as y i
= ˆ  1 ( a  )ˆ  2 ( h w  , x i i H )
+ ✏ i( i = 1 , . . . , n ) where ✏ i ⇠ N (0 , 1) and ˆ  1 ( · )
= ˆ  2 ( · ) = sign( · ) .
The true parameters ( a  , w  ) {+
+} Hare generated by a  ⇠ N (0 , 5 2 ) and w  =
P ¯ Nj =1 w ( j ) e j , where w ( j ) ⇠ N (0 , 5 2 ) .
We randomlygenerate x i from a Gaussian process whose covariance operator is a kernel function of H K 3 , andaccept it when it holds that k x i k H K 2 < B X with some B X .
In our implementation, we only used ¯ Nbases of H to deﬁne w  (note that w  2 H 0 still holds for this implementation).
We set ¯ N so that itis larger than N that is the number of bases for the Galerkin approximation.
In addition to demonstrating the validity of our proposed method, we also illustrate the beneﬁtof feature learning.
For that purpose, we compare the performance of the two-layer mean-ﬁeldneural network against that of non-adaptive estimators: the ridge regression estimator and the
Nadaraya-Watson estimator.
Here, we use the terminology “non-adaptive” to indicate that theydo not perform feature learning.
For ridge regression, we implemented the estimator using an
N -basis approximation, i.e., y = h P Nj =
↵ j e j , x i H , where ( ↵ 1 , . . . , ↵ N ) 2 R N are chosen tominimize 1 n P ni =
[Equation]
Nj =
[Equation] j ) i ↵ j  y i [-(cid:17) 
-]⌘ 2 +
P
Nj =1 ↵ 2 j with regularization parameter  > 0 .
As forthe Nadaraya-Watson estimator, the estimator ˆ y for an input x is given by ˆ y = P ni =
y i k ( k x i  x k H /h )
P ni =1 k ( k x i  x k H /h ) , 359 where h > 0 is a bandwidth and k ( u ) = max { 1  u 2 , 0 } for u 2 R (Epanechnikov kernel).
We used the squared loss for the optimization of twolayer neural network.
We employed tanh for the activation functions  1 ,  2 of the student model.
Hence,the teacher model is a bit out side of the studentmodel.
Our experiments were performed using thefollowing hyperparameters: Upper bound of inputdata in H K 2 : B X = 100 ; Number of basis functions 367 to generate data: ¯ N = 300 ; Number of training datan = 100 ; Number of basis functions for Galerkinapproximation: N = 150 ; Regularization parameters for PDA and P-SDCA:  1 = 10  2 ,  2 = 10  5 ;
Number of particles: M = 200 ; Hyperparametersfor PDA: T t = 10 , ⌘ t = 10  5 , minibatch size 50;
Hyperparameters for P-SDCA: J T = 10 , ˜ n = 1000 , ⌘ t = 10  4 ; Regularization parameter for ridgeregression:  = 0 .
1 , 1 , 10 (we observed that  = 1 achieved the best performance over all choices of ); Bandwidth for Nadaraya-Watson estimator: h = 0 .
As shown in Figure 1, PDA and P-SDCA converge and achieve small test loss.
Moreover, for bothalgorithms, the eventual test loss is smaller than those of linear estimators (Ridge-regression and
Nadaraya-Watson estimator), which illustrates that the mean-ﬁeld two-layer neural network has bettergeneralization ability than non-adaptive methods.
This is mainly due to the feature learning ability ofthe mean-ﬁeld networks, whereas the non-adaptive methods are not able to ﬁnd informative features.
Conclusion
In this work we studied two-layer neural networks in the mean-ﬁeld regime as a model for machinelearning problems with inﬁnite-dimensional input.
We proposed two optimization algorithms to learnthe neural network, by extending the Particle Dual Averaging and Particle Stochastic Dual Coordinate
Ascent method to inﬁnite-dimensional parameter space.
Leveraging the convergence guarantee ofinﬁnite-dimensional gradient Langevin dynamics, we showed that our proposed methods can globallyoptimize the training objective at a rate of polynomial order for the online setting and exponentialorder for the ﬁnite-sum setting.
Numerical experiments on synthetic data conﬁrmed that mean-ﬁeldneural network outperforms linear estimators that do not learn features.