Two-layer neural network on infinite-dimensional data: global optimization guarantee in the mean-field regime
Abstract
Analysis of neural network optimization in the mean-field regime is important as the setting allows for feature learning .
Existing theory has been developed mainly for neural networks in finite dimensions, i.e., each neuron has a finite-dimensional parameter.
However, the setting of infinite-dimensional input naturally arises in machine learning problems such as nonparametric functional data analysis and graph classification.
In this paper, we develop a new mean-field analysis of two-layer neural network in an infinite-dimensional parameter space.
We first give a generalization error bound, which shows that the regularized empirical risk minimizer properly generalizes when the data size is sufficiently large, despite the neurons being infinite-dimensional.
Next, we present two gradient-based optimization algorithms for infinite-dimensional mean-field networks, by extending the recently developed particle optimization framework to the infinite-dimensional setting.
We show that the proposed algorithms converge to the (regularized) global optimal solution, and moreover, their rates of convergence are of polynomial order in the online setting and exponential order in the finite sample setting, respectively.
To our knowledge this is the first quantitative global optimization guarantee of neural network 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.
For instance, Ling and Vieu (2018); Ferraty et al.
(2007) studied non-parametric function regression problems 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 learning problem with inputs from an infinite-dimensional reproducing kernel Hilbert space.
Among various models that deal with those infinite-dimensional input problems, neural networks are particularly 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 of functional 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 fixed basis functions in Rossi and Conan-Guez (2005) with adaptive bases, in order to leverage the power of representation (feature) learning and the flexibility of deep learning.
While neural network based models with infinite-dimensional input have been applied in many applications, their optimization using gradient-based methods is not well-understood.
In particular, no optimization guarantee has been shown for any algorithm in the presence of feature learning — such theoretical result 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 consider sufficiently wide neural network compared with the data size, and the model parameterization can be divided 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 focus on the mean-field regime, as it captures the presence of feature learning, which is one of the main advantages 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 gradient flow in the space of probability distributions on the parameters.
Mei et al.
(2018) showed that the mean-field Langevin dynamics converges to the global optimal solution, and Hu et al.
(2019) proved linear convergence with respect to the objective function with sufficiently strong KL divergence regularization.
However, those studies mainly analyzed the algorithm in continuous time, and do not establish 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 computational complexity, 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 respect to its outer loop iteration.
Both PDA and P-SDCA employ a double-loop structure, and make use of 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 Vempala and 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 optimization guarantee becomes meaningless due to the dimension dependence.
Therefore, we need to construct a new 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 model generalizes properly when the number of training data is sufficiently large.
Next, we propose two gradient-based algorithms that globally optimize the KL-regularized objective, corresponding to the infinite-dimensional extension of PDA and P-SDCA, respectively.
As mentioned above, PDA and P-SDCA require Monte Carlo sampling, and in our setting sampling is performed in an infinite-dimensional Hilbert space.
We therefore adapt an infinite-dimensional gradient Langevin dynamics (Debussche, 2011; Bréhier, 2014), the weak convergence of which has been established in Muzellec et al. (2022).
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-dimensional setting 1 .
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 the dimension 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 a function.
As previously discussed, although statistical properties (e.g., consistency) of these models have 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 theoretical analyses, 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 the inner product given by ⟨ f, g ⟩ H :
= (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 a neural network model that takes an infinite-dimensional input x ∈ X ⊆ H 0 and gives an output y ∈ Y ⊆ R , where X and Y are domains of input and output, respectively.
Each neuron in the network has parameter θ := ae 0 + w ∈ H for a ∈ R and w ∈ H 0 , and is represented by
[Equation] , where σ 1 , σ 2 : R → R are activation functions, which we assume to be bounded and smooth (such as tanh ) 2 .
We also denote h θ ( x ) by h ( θ, x ) .
Let the number of neurons be M and denote the set of parameters 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]
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 converge to the integral form as M increases.
This viewpoint motivates us to design optimization algorithms for 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 linear operator 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 corresponds to 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 a norm ∥
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 H whose mean is 0 and its covariance is A − 1 (i.e., ⟨ z, x ⟩ H ( z ∼ ν ) obeys the Gaussian distribution with mean 0 and covariance ⟨ x, A − 1 x ⟩ H for any x ∈ H ) 3 .
Finally, we define the Kullback-Leibler (KL) divergence from a probability measure ν 2 to a probability measure ν 1 that is absolutely continuous with respect to ν 2 as KL( ν 1 || ν 2 ) :
= E ν 1 (cid:104) log d ν 1 d ν 2 (cid:105) = E ν 2 (cid:104) d ν 1 d ν 2 log d ν 1 d ν 2 (cid:105) .
Objective.
Now we define the objective of our optimization method.
Suppose that the pairs of input and output are independently identically distributed from D , and let ℓ ( · , · ) :
Y × Y → R be a loss function.
In the mean-field regime, we optimize the following regularized risk minimization problem 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 the empirical measure of n training data points ( x i , y i ) ni
=1 as D , for which the first term of the objective function is written as 1 n (cid:80) ni =1 ℓ ( E θ ∼ π
[ h θ ( x i )] , y i ) .
In this setting, the problem can be regarded as a regularized empirical risk minimization problem.
As for the second term, in a finite-dimensional setting where A = λ 1 I , we can easily see that the KL divergence is decomposed into the sum of ℓ 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., 2022; Chizat, 2022).
Indeed, it is known that, if we optimize a loss function L ( q ) with respect to a probability 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 satisfies q ∗ ∝ exp( − 1 λ 2 δL ( q ∗ ) δq ) , where δL ( q ) δq : H → R is the derivative with respect to the distribution q such that (cid:82) [Equation] q [Equation]
[Equation] [ L [Equation] ϵp +
[Equation] L [Equation] q )] for any distribution p ∈ P 2 .
One important example that we can apply our model to is the non-parametric functional regression problem (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 .
One typical 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 some non-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 network in the empirical risk minimization setting, which confirms that the model properly generalizes when the data size is sufficiently large.
Specially, we show the generalization bound for two problem settings: binary classification and regression.
The following theorem stats the generalization bound for 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 0loss by ℓ 01 ( z, y ) , i.e., ℓ 01 ( z, y ) :
= 1 [ zy < 0] .
Let π ∗ be the optimal solution of the empirical risk minimization problem (2) for a given set S of n training data points (i.i.d.
from D ).
Suppose there exists a distribution π ◦ ∈ P 2 such that h π ◦ ( x ) y ≥ 1 / 2 for all ( x, y ) ∈ supp D , and that | h θ ( x ) | ≤holds for all θ ∈ H , x ∈ X .
Then, the following holds with probability 1 − δ with respect to the choice of S ⊂ X :
[Equation]
The KL -term is bounded when π ◦ is a variation of the Gaussian measure ν such that d π ◦ d ν log( d π ◦ d ν ) is integrable 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 distribution of the parameters (2).
For that purpose, the ideal situation would be to maintain infinitely many particles (neurons) to represent the distribution; however, in a tractable algorithm we need to use an approximation with finite number of particles.
This finite-particle approximation may lead to instability 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 PDA and P-SDCA algorithms (Nitanda et al., 2021; Oko et al., 2022).
We provide a short summary of this 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 + is 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 linearization and 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-dimensional sampling 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 ) /δπ ).
Then we 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 Riesz representer of the Fréchet derivative of G on H .
To discretize the continuous-time SDE, we employ the semi-implicit Euler scheme:
[Equation]
[Equation]
G
[Equation]
, where ζ k is a realization of standard Gaussian process on 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-dimensional approximation of the semi-implicit Euler scheme is given as
[Equation]
A convergence rate analysis of the infinite-dimensional gradient Langevin dynamics (3) with time discretization and finite-dimensional approximation was given by Muzellec et al.
(2022), which is an important 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 an infinite-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 is guaranteed 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 to do so we make use of the weak convergence property of the infinite-dimensional gradient Langevin dynamics (GLD) shown in Muzellec et al.
Algorithm description of infinite-dimensional PDA.
Here we provide the algorithm description of PDA with infinite-dimensional input.
Since we are optimizing the expected loss directly, we employ an online gradient descent approach: we draw one input-output pair ( x t , y t ) from D at the t -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 loop step: 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 its history: ¯ g ( t ) ( θ [Equation] :=
[Equation]
t
[Equation]
[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 by d π ( t +1) ∗ d ν ( θ ) ∝ exp (cid:0) − ¯ g ( t ) ( θ ) (cid:1) .
Hence in the inner loop, we sample particles from π ( t +1) ∗ using infinite-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 empirical distribution ˆ π t +1 but its “true” distribution).
We will show that π ( t +1)
( θ ) well approximates the ideal update π ( 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 ) =
[Equation]
+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 rely 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 ∈ 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 ∥ 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 present the result only for p = 2 for simplicity.
(A2) is satisfied by common loss functions such as square loss and 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 the difference 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]
Ω (cid:0)
[Equation]
[Equation]
[Equation]
λ
[Equation]
exp
[Equation]
λ
[Equation]
[Equation]
such that the following holds
[Equation]
Theorem 3 establishes the convergence rate of sampling via infinite-dimensional GLD in the PDA algorithm.
According to this theorem, to achieve sufficiently small sampling error
ϵ ( t ) A < ϵ , it suffices 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 bound up to a factor of κ despite the infinite dimensionality.
It is worth noting that the finite-dimensional analysis 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 do not expect convergence under the KL metric since the distribution is hardly absolutely continuous with respect to another.
Instead, our convergence guarantee is solely based on a weak convergence analysis, which allows us to establish computational complexity bound that is close to that of the finite-dimensional counterpart in Nitanda et al.
As for (i) , the following theorem evaluates the difference between π ∗ and π (ˆ t ) ∗ with respect to the objective function value.
The proof of this theorem can be found in Appendix C.2.
Theorem 4.
Suppose δ > 0 and let ¯ ϵ ( T )
A [Equation]
ϵ [Equation] ) A .
Under (A1) , (A2) , (A3) , for all distributions π ∗ ∈ 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 excess objective 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 strongly convex (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, this is not guaranteed in the infinite-dimensional setting.
We overcame this difficulty by reformulating the update using the Radon–Nikodym density with respect to the Gaussian measure ν .
Interestingly, this new 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 solution and 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 |
[Equation]
Notice that since the particle approximation does not have a density, the bound above does not control the 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 number of outer loops iterations is O (1 /ϵ ) , and the number of particles M is O
((1 /ϵ ) 2 log(1 /ϵ )) , then our proposed algorithm achieves ϵ -error in terms of the output of the model (or the loss function) from a solution 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 the desired accuracy at polynomial order of iteration complexity (but with exponential dependency on 1 /λ 2 , which would be unavoidable for any mean-field method).
We emphasize that our analysis guarantees 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, in the case of empirical risk minimization, we can improve the rate of convergence from polynomial to 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 Oko et 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 a careful modification of the update rule.
Indeed, we reconstruct the algorithm based on the density with respect to the Gaussian measure ν . We
Algorithm description of infinite-dimensional P-SDCA.
Here we give an algorithmic description of our infinite-dimensional extension of P-SDCA.
The main ingredient of SDCA is to execute optimization 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 bounded for i = 1 , . . . , n .
We denote the conjugate function of ℓ ( · , y i ) by ℓ ∗ ( · , y i ) 5 .
Define the dual objective 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 , the corresponding primal solution can be retrieved by p [ g ]( θ ) :=
[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, and update the corresponding coordinate so that it maximizes the dual objective.
Suppose that we have the dual solution g ( t ) at the t -th iteration, then we update its i t -th coordinate by the following formula where i t ∈ [ n ] is chosen uniformly at random:
[Equation]
for j ̸ =
[Equation]
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
[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 weightings of the particles ( r m )
Mm =1 .
In particular: (i) at the resampling steps (once in every ˜ n iterations), we run the infinite-dimensional GLD to generate M particles ( θ m )
Mm =1 from the distribution π ( t ) ∗ that is given by 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]
it ) nλ
[Equation]
(cid:21) .
This re-weighting scheme 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]
, where ˇ C 1 , ˇ C 2 , Λ are constants depending on ˆ b, λ 1 , λ 2 and ϕ : H → R is any test function satisfying ∥ ϕ ∥ ∞ ≤ 1 with sufficient smoothness (see Appendix D.3 for more details).
Using the 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 ∈ [ − ˆ b, ˆ b ] and i ∈ [ 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+ [Equation]
(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( ϵ (˜ nT ) C ) 2 log (cid:0) 4 nT end δ (cid:1) where ˆ C − 1= O (cid:18) λ − 12ˆ C 2 1+ 1 nλ 2 exp (cid:16) {+˜ n (2ˆ b ˆ C 2 +1) nλ 2 +
+}
[-2 + [Equation] nλ 2 +
-] (cid:17)(cid:19) and ˆ C 2 = exp (cid:16) 8max { ˆ b 2 , ˆ b 4 } ˜ n λ 2 n (cid:17) .
Then, after ˜ nT end iterations of infinite-dimensional P-SDCA, the duality gap can be bounded as E (cid:2)
L ( p [ g (˜ nT end ) ]) − D ( g (˜ nT end ) )
(cid:12)(cid:12) 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 is taken 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 smaller than ϵ P with O (log(1 /ϵ P )) outer loop iterations.
In combination with the sampling complexity, we see that the total iteration cost to obtain a solution with ϵ P duality gap is O ((1 /ϵ P ) 2 1 − 2 κ log(1 /ϵ P )) , which is much better than that of PDA shown in Section 4.
Numerical Experiments
We numerical evaluate our proposed methods on a non-linear functional regression problem.
Let H be 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] satisfying u (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) 
+}
0 x ′ 1 ( t ) x ′ 2 ( t )d t .
Consequently, H K is a reproducing kernel Hilbert space whose kernel 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 ◦ ) ∈ H are generated by a ◦ ∼ N (0 , 5 2 ) and w ◦ = (cid:80) ¯ Nj =1 w ( j ) e j , where w ( j ) ∼ N (0 , 5 2 ) .
We randomly generate x i from a Gaussian process whose covariance operator is a kernel function of H K 3 , and accept it when it holds that ∥ x i ∥ H K 2 < B X with some B X .
In our implementation, we only used ¯ N bases of H to define w ◦ (note that w ◦ ∈ H 0 still holds for this implementation).
We set ¯ N so that it is 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 benefit of feature learning.
For that purpose, we compare the performance of the two-layer mean-field neural 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 they do 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 to minimize 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 for the Nadaraya-Watson estimator, the estimator ˆ y for an input x is given by ˆ y = (cid:80) ni = [Equation] y [Equation]
[-=
-]∥ [Equation] x ∥ H [Equation]
h [Equation] ni = [Equation]
k [Equation] ∥
[Equation] 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 student model.
Our experiments were performed using the following hyperparameters: Upper bound of input K B X = of basis generate data
=
Hyperparameters we observed that λ
1 achieved the best performance over all choices of
As shown in Figure 1, PDA and P-SDCA converge and achieve small test loss.
Moreover, for both algorithms, 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 better generalization ability than non-adaptive methods.
This is mainly due to the feature learning ability of the 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 machine learning problems with infinite-dimensional input.
We proposed two optimization algorithms to learn the neural network, by extending the Particle Dual Averaging and Particle Stochastic Dual Coordinate
Ascent method to infinite-dimensional parameter space.
Leveraging the convergence guarantee of infinite-dimensional gradient Langevin dynamics, we showed that our proposed methods can globally optimize the training objective at a rate of polynomial order for the online setting and exponential order for the finite-sum setting.
Numerical experiments on synthetic data confirmed that mean-field neural network outperforms linear estimators that do not learn features.
Acknowledgments
NN was partially supported by CREST (JPMJCR2015).
TS was partially supported by Japan Digital
Design and CREST (JPMJCR2115).
AN was partially supported by JSPS Kakenhi (22H03650) and
JST-PRESTO (JPMJPR1928).
DW was partially supported by the Borealis AI Fellowship.