Two-layer neural network on infinite-dimensional data: global optimization guarantee in the mean-field regime
Anonymous Author(s)
AffiliationAddress email
Abstract
Analysis of neural network optimization in the mean-field regime is important asthe setting allows for feature learning .
Existing theory has been developed mainlyfor neural networks in finite dimensions, i.e., each neuron has a finite-dimensionalparameter.
However, the setting of infinite-dimensional input naturally arises in machine learning problems such as nonparametric functional data analysis and graphclassification.
In this paper, we develop a new mean-field analysis of two-layerneural network in an infinite-dimensional parameter space.
We first give a generalization error bound, which shows that the regularized empirical risk minimizerproperly generalizes when the data size is sufficiently large, despite the neuronsbeing infinite-dimensional.
Next, we present two gradient-based optimizationalgorithms for infinite-dimensional mean-field networks, by extending the recentlydeveloped particle optimization framework to the infinite-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 finite sample setting, respectively.
Toour knowledge this is the first quantitative global optimization guarantee of neuralnetwork on infinite-dimensional input and in the presence of feature learning.
Introduction
A variety of machine learning problems need to handle input from infinite-dimensional spaces.
Forinstance, Ling and Vieu (2018); Ferraty et al.
(2007) studied non-parametric function regressionproblems where the input is a function in an infinite-dimensional functional space and Kriege et al.(2020) studied graph classification problems using kernel method, which can be cast as a learningproblem with inputs from an infinite-dimensional reproducing kernel Hilbert space.
Among various models that deal with those infinite-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 thefixed basis functions in Rossi and Conan-Guez (2005) with adaptive bases, in order to leveragethe power of representation (feature) learning and the flexibility of deep learning.
While neuralnetwork based models with infinite-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 finite-dimensional data.
It is experimentally observed that the gradient-based optimization methods can yield sufficiently 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 considersufficiently wide neural network compared with the data size, and the model parameterization can bedivided into the mean-field 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-field 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-field analysis, the dynamics of gradient descent is described by Wasserstein gradientflow in the space of probability distributions on the parameters.
Mei et al.
(2018) showed that themean-field Langevin dynamics converges to the global optimal solution, and Hu et al.
(2019) provedlinear convergence with respect to the objective function with sufficiently strong KL divergenceregularization.
However, those studies mainly analyzed the algorithm in continuous time, and do notestablish quantitative convergence rate for discrete time and finite 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 finite width setting.
Furthermore, in the case of finitesum 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 infinite-dimensional input, then the optimizationguarantee becomes meaningless due to the dimension dependence.
Therefore, we need to construct anew methodology and theory for the infinite-dimensional problem.
Our contributions.
In this work, we extend the mean-field analysis for neural network training to infinite-dimensional parameter space, which covers two-layer neural network with infinitedimensional inputs.
First, we establish a generalization error bound which entails that our modelgeneralizes properly when the number of training data is sufficiently large.
Next, we propose twogradient-based algorithms that globally optimize the KL-regularized objective, corresponding to theinfinite-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 infinite-dimensional Hilbert space.
We therefore adapt an infinite-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-field regime whose input is infinite-dimensional,and establish a generalization error bound of the regularized empirical risk minimizer.•
We propose two optimization algorithms for our infinite-dimensional neural network by incorporating the infinite-dimensional Langevin dynamics in the inner loop of PDA and P-SDCA.• We prove that our infinite-dimensional extension of PDA achieves polynomial order convergence,and that of P-SDCA achieves exponential (outer loop) convergence even in an infinite-dimensionalsetting (while there is exponential dependency on the regularization parameter as in usual Langevinmethods).
To the best of our knowledge, this is the first quantitative global optimization guarantee of neural network in the infinite-dimensional mean-field 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 infinite-dimensional input
In this section we introduce the mean-field two-layer neural network with infinite-dimensional input.
We first 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 infinitedimensional) separable Hilbert space H .
Since H is separable, there is a complete orthonomal system( e j ) ∞ j =0 for which it holds that H = (cid:110)(cid:80) ∞
j =0 a j e j (cid:12)(cid:12)(cid:12) (cid:80) ∞
j =0 a 2 j < ∞ (cid:111) , and H is equipped with theinner product given by ⟨ f, g ⟩ H :
[-P 1 j =
-]= (cid:80) ∞
j =0 a j b j , where f = (cid:80) ∞ j =
a j e j , g = (cid:80) ∞
j =
b j e j ∈ H .
Let
H 0 be a subset of H defined by H 0 :
= { 0 · e 0 + (cid:80) ∞
j =1 a j e j | (cid:80) ∞
j =1 a 2 j < ∞} .
We consider aneural network model that takes an infinite-dimensional input x ∈ X ⊆ H 0 and gives an outputy ∈ Y ⊆ R , where X and Y are domains of input and output, respectively.
Each neuron in thenetwork has parameter θ := ae 0 + w ∈ H for a ∈ R and w ∈ H 0 , and is represented byh θ ( x ) :
= σ 1 ( a ) σ 2 ( ⟨ w, x ⟩ H ) , 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-field regime, the two layer neural network model(with infinite-dimensional input) can be expressed as the average over the M neurons: [Equation] [Equation] :=
[Equation]
M
[Equation] [Equation] This can be seen as a finite 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 define a Gaussian process taking its value in H .
For that purpose, we define a subspace H K γ ⊂ H for γ ≥ 0 as follows.
First, we define a linearoperator T K : H → H as T K f = (cid:80) ∞
j =0 µ j a j e j for f = (cid:80) ∞
j =
a j e j , where ( µ j ) ∞ j =0 correspondsto the spectrum of T K satisfying µ j > 0 ( ∀ j ) and is sorted in decreasing order ( µ 1 ≥ µ 2 ≥ · · · ).
Note that the orthonormal system ( e j ) ∞ j =0 is taken so that it coincides with the eigen-functions of
T K .
Then, we define H K γ :
= T γ 2 K H =
{ T γ 2 K h | h ∈ H} and we can naturally equip H K γ with anorm ∥ f
∥ H Kγ := (cid:80) ∞ j =0 µ − γ j a 2 j for f = (cid:80) ∞ j =
a j e j ∈ H K γ .
For concise notation we simply write
H K :
= H K 1 .
Next, define an operator A = λ 1 T − 1 K for a positive regularization parameter λ 1 > 0 ,i.e.,
Af = λ 1 (cid:80) ∞
j =0 µ − 1 j a j e j for f = (cid:80) ∞ j =
a j e j ∈ H K , and let ν be a Gaussian measure in Hwhose mean is 0 and its covariance is A − 1 (i.e., ⟨ z, x ⟩ H ( z ∼ ν ) obeys the Gaussian distribution withmean 0 and covariance ⟨ x, A − 1 x ⟩ H for any x ∈ H ) 2 .
Finally, we define 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 || ν
[Equation]
[Equation]
Objective.
Now we define 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-field 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 θ ∼ π [ ∥ θ ∥ 2 H ] <
∞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 first term of the objectivefunction is written as 1 n (cid:80) 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 finite-dimensionalsetting where A = λ 1 I , we can easily see that the KL divergence is decomposed into the sumof ℓ 2 -regularization λ 1 E θ ∼ π [ ∥ θ ∥ 2 ] and negative entropy E θ ∼ π [log d π ( θ ) / d θ ] .
Here the negative entropy naturally arises from the mean-field 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 ∈ P 2 by the mean-field 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 satisfiesq ∗ ∝ exp( − 1 λ 2 δL ( q ∗ ) δq ) , where δL ( q ) δq : H → R is the derivative with respect to the distribution q suchthat (cid:82) δL ( q ) δq ( θ )d( p − q )( θ ) =
lim ϵ → 0 1 ϵ [ L ( ϵp + (1 − ϵ ) q ) − L ( q )] for any distribution p ∈ P 2 .
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 definition 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 ) :
= (cid:82) K ( t, s ) f ( s )d µ ( s ) for a positive definite kernel
K (Caponnetto and De Vito, 2007; Steinwart and Christmann, 2008).
The goal is to estimate somenon-linear functional from the input x ∈ H K to output y ∈ 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 confirms that the model properly generalizes whenthe data size is sufficiently large.
Specially, we show the generalization bound for two problemsettings: binary classification and regression.
The following theorem stats the generalization boundfor binary classification 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-1loss by ℓ 01 ( z, y ) , i.e., ℓ 01 ( z, y ) :
= 1 [ 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 π ◦ ∈ P 2 such that h π ◦ ( x ) y ≥ 1 / 2 for all ( x, y ) ∈ supp D , and that | h θ ( x ) | ≤ 1holds for all θ ∈ H , x ∈ X .
Then, the following holds with probability 1 − δ with respect to thechoice of S ⊂ X :
[Equation]
The KL -term is bounded when π ◦ is a variation of the Gaussian measure ν such that d π ◦ d ν log( d π ◦ d ν ) isintegrable with respect to ν .
Note that if we set λ 2 = 1 / √ n , the generalization error is O (1 / √ n ) .
This order matches that in the finite-dimensional input setting shown in Nitanda et al.
Mean-field optimization with infinite-dimensional Langevin algorithm
As previously remarked, the core idea of the mean-field analysis is to directly optimize the distributionof the parameters (2).
For that purpose, the ideal situation would be to maintain infinitely manyparticles (neurons) to represent the distribution; however, in a tractable algorithm we need to usean approximation with finite number of particles.
This finite-particle approximation may lead toinstability due to the interaction between particles, which is known to be difficult to control.
To overcome this difficulty, 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 first-order approximation of the loss function
L ( π ) : = E ( X,Y ) ∼D [ ℓ ( E θ ∼ π [ h θ ( X )] , Y )] as ( π − π t ) δL ( π t ) δπ :=
(cid:82) δ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 (finitedimensional 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 infinite-dimensionalsampling problem, which cannot be directly extended from the finite-dimensional counterpart.
Let π be a distribution on H that satisfies d π d ν ( x ) ∝ exp ( − G ( x ))
(e.g., G = λ − 1 2 δL ( π t ) /δπ ).
Thenwe know π is the stationary distribution of the infinite-dimensional stochastic differential equation,
[Equation]
where ( W t ) t ≥ 0 is a cylindrical Brownian motion on H (Da Prato et al., 1996) and ∇ 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., X k +1 = S η (cid:0) X k −
[Equation]
∇ G ( X k ) + √ 2 ηζ k (cid:1) , 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 / ∈ H in the infinite-dimensional setting, and thus the naive
Euler-Maruyama discretization does not ensure X k +1 ∈ H while the semi-implict scheme does.
In practice, we need to approximate the infinite-dimensional vector X k by a finite-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 ∇ G N ( x ) : =
P N ( ∇ G ( P N x )) .
Then the finite-dimensionalapproximation of the semi-implicit Euler scheme is given as
[Equation]
A convergence rate analysis of the infinite-dimensional gradient Langevin dynamics (3) with timediscretization and finite-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).
Infinite-dimensional Particle Dual Averaging
In this section, we present our first algorithm to solve the risk minimization problem (2) which is aninfinite-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 finite-dimensional setting.
Extending the algorithm and convergence analysis to infinite-dimensional input is non-trivial, and todo so we make use of the weak convergence property of the infinite-dimensional gradient Langevindynamics (GLD) shown in Muzellec et al.
Algorithm description of infinite-dimensional PDA.
Here we provide the algorithm descriptionof PDA with infinite-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ˆ π ( t ) =
1 M (cid:80) Mm =1 δ θ ( t ) m
(i.e., h Θ ( t ) =
E θ ∼ ˆ π ( t ) [ h θ ( · )] ).
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 Θ ( t ) ( x ( t ) ) , y t ) h
( θ, x t ) . The idea of the Nesterov’s dual averaging method is to take its weighted sum of the gradient over itshistory: ¯ g ( t ) ( θ ) :=
[Equation]
sg
[Equation]
s )
[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 π ( t +1) ∗ d
ν ( θ ) ∝ exp (cid:0) − ¯ g ( t ) ( θ ) (cid:1) .
Hence in the inner loop, we sample particles from π ( t +1) ∗ usinginfinite-dimensional gradient Langevin dynamics (3) 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 ∇ ¯ g ( t ) N ( · ) :
= P N ( ∇ ¯ 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 [-+
-]+1 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 , . . .
, 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 infinitedimensional PDA.
Let P X be the distribution of input x ∈ X and G :=
supp P X .
Our analyses relyon 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, forall y, z ∈ Y , | ∂ z ℓ ( z, y ) | ≤ 2 holds.(A3)
Both σ 1 and σ 2 are thrice continuously differentiable.
Moreover,
it holds thatmax {∥ σ 1 ∥ ∞ , ∥ σ ′ 1 ∥ ∞ , ∥ σ ′′ 1 ∥ ∞ , ∥ σ 2 ∥ ∞ , ∥ σ ′ 2 ∥ ∞ , ∥ σ ′′ 2 ∥ ∞ } ≤ ˆ b .(A4)
There exists 12 < γ < 2 , B X > 0 such that ∥ x ∥ H K 1+ γ < B X for all x ∈ G .
Remark.
(A1) is a sufficient 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 satisfied by common loss functions such as square lossand logistic loss.
(A3) is satisfied 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 infinite-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]
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 κ ∈ (0 , 1 / 2) ,there exists ˆ C
[Equation]
[Equation]
such that the following holds
[Equation]
Theorem 3 establishes the convergence rate of sampling via infinite-dimensional GLD in the PDAalgorithm.
According to this theorem, to achieve sufficiently small sampling error ϵ ( t ) A < ϵ , itsuffices to take µ N =
O ( ϵ 2 / (1 − 2 κ ) ) , η t =
O ( ϵ 2 / (1 − 2 κ ) ) , T t = Ω( ϵ − 2 1 − 2 κ log(1 /ϵ ) / Λ) with someκ ∈ (0 , 1 / 2) .
This result is comparable to the finite-dimensional counterpart in Nitanda et al.
(2021),where η t =
O ( ϵ 2 ) and T t = Θ( ϵ − 2 log(1 /ϵ ) / Λ) is shown to be sufficient, which is close to our boundup to a factor of κ despite the infinite dimensionality.
It is worth noting that the finite-dimensionalanalysis in Nitanda et al.
(2021) is based on the convergence result of GLD under KL-divergence(Vempala and Wibisono, 2019).
However, in our infinite-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 thefinite-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 and let ¯ ϵ ( T ) A := 2 T ( T +3) (cid:80) T +1 t =2 ϵ ( t ) A .
Under (A1) , (A2) , (A3) , for alldistributions π ∗ ∈ P 2 , the following inequality 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 first-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 infinite-dimensional setting.
We overcame this difficulty by reformulating theupdate using the Radon–Nikodym density with respect to the Gaussian measure ν .
Interestingly, thisnew technique yields a refined 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 ( π [Equation] ˆ [Equation]
∗
[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 ∞ -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 infinite-dimensional PDA to obtain ϵ -accurate solution is
O (1 / (Λ ϵ 3 − 2 κ 1 − 2 κ ) log(1 /ϵ )) for some κ ∈ (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-field method).
We emphasize that our analysisguarantees the convergence to the global optimal solution, not the stationary point or local minimizer.
Infinite-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 finite-sum property.
For that purpose, we propose an infinitedimensional 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 finite-dimensional setting.
Especially, the lack of density function in the infinite-dimensional parameter space necessitates acareful modification of the update rule.
Indeed, we reconstruct the algorithm based on the densitywith respect to the Gaussian measure ν .
Algorithm description of infinite-dimensional P-SDCA.
Here we give an algorithmic descriptionof our infinite-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 isbounded for i = 1 , . . . , n .
We denote the conjugate function of ℓ ( · , y i ) by ℓ ∗ ( · , y i ) 4 .
Define the dualobjective D : R n → R as
[Equation]
Then, if inf π ∈P 2 L ( π ) > −∞ holds, it holds that inf π ∈P 2 L ( π ) =
sup g ∈ R n D ( g ) .
We construct the update rule based on this duality theorem.
For a dual variable g ∈ 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 ) ⌘-][Equation]
(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 ∈ [ 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 ∈ [ n ] is chosen uniformly at random:
[Equation]
[-g 
-]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
=
[Equation]
for j ̸ = i t .
We can see that the update formula requires the expectation E ν (cid:2)
h ( θ, x i t ) · p [ g ( t ) ] ( θ ) (cid:3) .
We approximate this expectation by a weighted average of M particles as E ν (cid:2) h
( θ, x i t ) · p [ g ( t ) ]( θ ) (cid:3) ≈
[Equation]
[Equation]
t ) m 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 )
Mm =1 .
In particular: (i) at the resampling steps (once in every ˜ n iterations), we runthe infinite-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 ∈
[ M ] ; (ii) in non-refreshing steps,we update the weights ( r m )
Mm =1 as r ( t +1) m
[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 infinite-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 κ ∈ (0 , 1 / 2) , it holds that [Equation] (cid:17)
[Equation]
, where ˇ C 1 , ˇ C 2 , Λ are constants depending on ˆ b, λ 1 , λ 2 and ϕ : H → R is any test function satisfying∥ ϕ
[Equation]
≤ exp
[Equation] exp (cid:16) 4˜ [Equation] 2 (cid:17)(cid:17) with sufficient 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 ′ ℓ ( x ′ , y i ) | ≤ B 1 for all x ∈ [ − 1 , 1] andi ∈ [ n ] .
Let ˜ s :
= ˜ nλ 2 1+ nλ 2 and take any T end ≥ 2 n ˜ n (cid:16) 1+ 1 nλ 2 (cid:17) log (cid:18)(cid:16) n + 1 λ 2 (cid:17) B 1 +2+ 1 1 − exp( − ˜ s ) ϵ P (cid:19) forϵ P > 0 .
Suppose that ϵ (˜ nT ) C and M satisfy ϵ (˜ nT ) C ≤ ˆ C 1 exp (cid:0) − ˜ sT 2 (cid:1) ( ∀ T ∈ [ T end ]) and M ≥ˆ C 2( ϵ (˜ nT ) C ) 2 log (cid:0) 4 nT end δ (cid:1) where ˆ C − 1= O (cid:18)
[Equation] (cid:16) [Equation]
n (
[Equation]
{+2 + [Equation] nλ 2 +
 (cid:17)(cid:19) and+} ˆ C 2 = exp (cid:16) 4 B 1 ˜ n λ 2 n (cid:17) .
Then, after ˜ nT end iterations of infinite-dimensional P-SDCA, the duality gap can be bounded as 332 E (cid:2)
[Equation]
E (cid:3) ≤ ϵ 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.
Here, the expectation istaken over the choice of the coordinates in each step and the event E corresponds to measure concentration induced by the M -particle sampling.
Consequently, P-SDCA can achieve a duality gap smallerthan ϵ P with O (log(1 /ϵ P )) outer loop iterations.
In combination with the sampling complexity, wesee that the total iteration cost to obtain a solution with ϵ 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 complexitymatches the finite-dimensional case in Oko et al. except for the log factor, which again derivedthe complexity of sampling from GLD based on the 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 defined on [0 , 1] , and H K be the Sobolev space H 1 ([0 , 1]) , which is a set of absolutely continuous functions defined on [0 , 1] satisfyingu (0)
= [-
-]0 and
u ′ ∈ L 2 ([0 , 1]) .
Then, H is a Hilbert space equipped with inner product defined by⟨ x 1 , x 2 ⟩ H =
(cid:82) 1 0 x 1 ( t ) x 2 ( t )d t , and H K is a Hilbert space equipped with inner product defined by⟨ x 1 , x 2 ⟩ H K
= (cid:82) 1 0 x ′ 1 ( t ) x ′ 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 ( ⟨ w ◦ , x 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 ◦ = (cid:80) ¯
Nj = 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 ∥ x i ∥ H K 2 < B X with some B X .
In our implementation, we only used ¯ Nbases of H to define w ◦ (note that w ◦ ∈ 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 benefitof feature learning.
For that purpose, we compare the performance of the two-layer mean-fieldneural 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 = ⟨ (cid:80) Nj =
α j e j , x ⟩ H , where ( α 1 , . . . , α N ) ∈ R N are chosen tominimize 1 n (cid:80) ni
=1 (cid:16)(cid:80)
Nj =
1 x ( j ) i α j − y i {+(cid:17) 
+} + λ (cid:80) Nj
=
α 2 j with regularization parameter λ > 0 .
As forthe Nadaraya-Watson estimator, the estimator ˆ y for an input x is given by ˆ y = (cid:80) ni =
y i k ( ∥ x i − x ∥ H /h ) (cid:80) ni
=1 k ( ∥ x i − x ∥ H /h ) ,where h > 0 is a bandwidth and k ( u ) = max { 1 − u 2 , 0 } for u ∈ R (Epanechnikov kernel). [Figure]
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 functionsto 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 − Regularization parameter for ridgeregression: λ =
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-field two-layer neural network has bettergeneralization ability than non-adaptive methods.
This is mainly due to the feature learning ability ofthe mean-field networks, whereas the non-adaptive methods are not able to find informative features.
Conclusion
In this work we studied two-layer neural networks in the mean-field regime as a model for machinelearning problems with infinite-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 infinite-dimensional parameter space.
Leveraging the convergence guarantee ofinfinite-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 finite-sum setting.
Numerical experiments on synthetic data confirmed that mean-fieldneural network outperforms linear estimators that do not learn features. Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[No](d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] All assumptionsare stated in the each theorem statements and Assumption 2.
(b) Did you include complete proofs of all theoretical results?
[Yes] All of the proofs of
theoretical results are shown in the appendices.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[No](b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] The details of the experiments are shown in Section 6.
(c) Did you report error bars (e.g., with respect to the random seed after running experi-
ments multiple times)?
[Yes] See Figure
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[N/A]
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] The codes we used in
the experiments are based on supplemental materials of Oko et al.
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]