A Mean-Field Analysis of Neural Stochastic Gradient Descent-Ascent for Functional Minimax Optimization

Minimax Optimization. Our work is closely related to the literature on first-order methods for solving minimax optimization problems. These works establish the convergence rate or iteration complexity of first-order methods under various assumptions on the objective function. In particular, most of the existing works focus on finite-dimensional parameter spaces and one of the following objective functions: (i) convex-concave (Lin et al., 2020b; Ibrahim et al., 2019; Ouyang and Xu, 2021; Alkousa et al., 2019; Luo et al., 2021; Xie et al., 2020a; Han et al., 2024; Li et al., 2023; Jin et al., 2022), (ii) nonconvex-concave (Jin et al., 2019; Lin et al., 2020a; Lu et al., 2020a; Ostrovskii et al., 2021b; Zhao, 2023; Huang et al., 2022; Luo et al., 2020; Zhang et al., 2021a; Nouiehed et al., 2019; Thekumparampil et al., 2019), and (iii) nonconvex-nonconcave (Li et al., 2022; Diakonikolas et al., 2021; Ostrovskii et al., 2021a; Yang et al., 2022; Grimmer et al., 2022; Hajizadeh et al., 2024; Grimmer et al., 2023; Yang et al., 2020).

Our work can be viewed as an extension of convex-concave minimax optimization to the infinite-dimensional functional space. In particular, our objective is a regularized quadratic functional with respect to the input functions, which is then restricted to the class of overparameterized neural networks. Note that the objective of interest is in fact nonconvex-nonconcave in the neural network parameter space. Compared with the work on general nonconvex-nonconcave minimax optimization problems, our setting has a better underlying structure in the functional space in terms of convexity. This structure enables us to lift the network parameter updates to the Wasserstein space and analyze the gradient flow in the space of distributions. Our approach leverages the hidden convexity-concavity behind the seemingly nonconvex-nonconcave objective function and thus achieves better results in terms of algorithm convergence and complexity.

Mean-field Analysis in Deep Learning. Our work is closely related to the recent study of neural network training via gradient-based methods. One line of research establishes the convergence of gradient-based algorithms for training overparameterized neural networks under the “lazy training” regime, where the neural networks behave similarly to random kernel functions. Such a regime is also known as the as the neural tangent kernel regime (Jacot et al., 2018; Allen-Zhu et al., 2019a, b; Chen et al., 2020a; Frei and Gu, 2021; Zou and Gu, 2019; Du et al., 2018, 2019; Arora et al., 2019a, b; Huang and Yau, 2020). Our work is more related to another line of research based on the perspective of mean-field approximation (Mei et al., 2018, 2019; Chizat and Bach, 2018; Sirignano and Spiliopoulos, 2020b, a, 2022; Chen et al., 2020b; Fang et al., 2021b; Chen et al., 2019). Under the mean-field view, the neural network parameters are identified as a distribution over the parameter space. As a result, the evolution of parameters by gradient-based updates is captured by a differential equation that governs the evolution of the corresponding distribution. By elevating the training dynamics to an infinite Wasserstein space, the optimization objective often enjoys a benign landscape, which yields admits a more tractable analysis and global convergence. See, e.g, Zhang et al. (2020, 2021b); Fang et al. (2021b); Lu et al. (2020b); Fang et al. (2019); Chizat (2022); Hu et al. (2021); Nitanda et al. (2022) and the references therein. Also, see Fang et al. (2021a) for a recent survey.

Our work is especially related to the mean-field analysis of the Neural Temporal Difference (TD) (Zhang et al., 2020) and the Neural Actor-Critic (AC) (Zhang et al., 2021b) in reinforcement learning. These previous works have provided an analysis of the global convergence of the TD and AC algorithm with two-layered overparameterized neural networks. The optimization problem in these two tasks is the minimization of an objective where only one neural network is involved. Rather different from these works, we focus on minimax optimization, which requires neural network parameterization of both the primal function and the dual function. This brings new challenges to the analysis as the gradient dynamics of the primal and dual neural networks give birth to a coupled system of PDEs. To the best of our knowledge, our paper is the first to apply the mean-field limit to study the convergence of algorithms in solving the general form of functional conditional moment equations using neural networks.

Adversarial Estimation. Our work is also related to the literature on adversarial estimation, a method that solves a functional conditional moment equation by introducing a dual function and reformulating the original problem into a minimax optimization. Our work studies this type of minimax optimization with overparameterized neural networks. Thus, our work is more related to the study of adversarial estimation within neural network function classes (Dikkala et al., 2020; Chernozhukov et al., 2020; Bennett et al., 2019; Xu et al., 2021). Compared with our work, these studies focus on statistical errors pertinent to neural networks, assuming the optimization problem is solved perfectly. We instead study the optimization algorithm and establish the convergence of stochastic gradient-descent-ascent with neural networks.

Several previous works have also explored the convergence of optimization dynamics in adversarial estimation with neural networks. In particular, Neural GTD (Wai et al., 2020) and Neural SEM (Liao et al., 2020) analyze respectively the convergence for off-policy evaluation and structural equation models estimation with overparameterized two-layered neural network. However, their analyses are based on the idea of neural tangent kernel (NTK), where the employed neural network has a fixed representation during training, and the representation is completely determined by the initialization. In contrast, our work adopts the mean-field approach, which enables learning a data-dependent representation.

For any function $f\in\mathcal{F}$ and any $z\in\mathcal{Z}$, we define a functional $\bar{\delta}\colon\mathcal{Z}\times\mathcal{F}\rightarrow\mathbb{R}$ as

$\displaystyle\bar{\delta}(z;f):=\mathbb{E}_{X|Z}\bigl{[}\Phi(X,Z;f){\,\big{|}\,}Z=z\bigr{]},\qquad\forall f\in\mathcal{F},z\in\mathcal{Z}.$

(2.2)

In other words, the conditional moment equation problem in (2.1) boils down to finding a function $f_{0}\in\mathcal{F}$ such that $\bar{\delta}(\cdot;f_{0})$ is a zero function on $\mathcal{Z}$. Therefore an equivalent way to solve $f_{0}\in\mathcal{F}$ in (2.1) is by solving $\inf_{f\in\mathcal{F}}\mathbb{E}[(\bar{\delta}(Z;f))^{2}]$ (Ai and Chen, 2003; Chen and Pouzo, 2012). To control the complexity of the function class $\mathcal{F}$, Ai and Chen (2003) propose to use flexible sieve spaces $\mathcal{F}_{k(n)}$ that becomes dense in $\mathcal{F}$ as the sieve dimension $k(n)$ grows to infinity with data sample size $n$, and proposed the so-called sieve minimum distance criterion $\min_{f\in\mathcal{F}_{k}(n)}\mathbb{E}\bigl{[}\bar{\delta}(Z;f)^{2}\bigr{]}/2.$ In particular, Ai and Chen (2003) allow for two-layer NNs, splines, wavelets, Fourier series, and all kinds of polynomial sieves $\mathcal{F}_{k(n)}$ to approximate functions in $\mathcal{F}\subseteq L^{2}(\mathcal{W})$. Alternatively Chen and Pouzo (2012) propose the following penalized (or regularized) minimum distance criterion:

$\displaystyle\min_{f\in\mathcal{F}}\mathbb{E}\bigl{[}\bar{\delta}(Z;f)^{2}\bigr{]}/2+\lambda\cdot\mathcal{R}(f),$

(2.3)

where $\lambda\geq 0$ is a regularization parameter, $\mathcal{R}(f)$ is a regularizer on function $f\in\mathcal{F}$. They allow that $\mathcal{R}(f)$ to be any convex or lower-semicompact regularizer. In the minimum distance approach, for any fixed $f$, the authors first estimate $\bar{\delta}(z;f)$ by the following least squares criterion:

$\displaystyle\mathop{\mathrm{argmin}}_{\delta\in L^{2}(\mathcal{Z})}\mathbb{E}\Big{[}1/2\cdot\big{(}\Phi(X,Z;f)-\delta(Z)\big{)}^{2}\Big{]}=\mathop{\mathrm{argmax}}_{\delta\in L^{2}(\mathcal{Z})}\mathbb{E}\Big{[}\Phi(X,Z;f)\delta(Z)-1/2\cdot\delta(Z)^{2}\Big{]}$

Furthermore, we assume that the functional $\Phi$ is affine in $f$, which captures several important applications in machine learning and causal inference listed in Section 2.2. Specifically, we define $\widetilde{\Phi}(x,z,f)=\Phi(x,z,f)-\Phi(x,z,0)$, where $0$ stands for the zero function on $\mathcal{W}$. Then for any two functions $f_{1},f_{2}\in\mathcal{F}$ and any $a,b\in\mathbb{R}$, we have

$\displaystyle\widetilde{\Phi}(x,z;af_{1}+bf_{2})=a\cdot\widetilde{\Phi}(x,z;f_{1})+b\cdot\widetilde{\Phi}(x,z;f_{2}),\qquad\forall(x,z)\in\mathcal{X}\times\mathcal{Z}.$

(2.4)

Solving (2.1) with Overparameterized Neural Networks. In the sequel, we aim to solve the problem in (2.1) based on i.i.d. data points sampled from $\mathcal{D}$, with $\mathcal{F}$ being a class of overparameterized neural networks. In this case, it is possible that (2.1) does not have a solution within $\mathcal{F}$. Furthermore, for the choice of regularizer, we consider the following specific form of $\mathcal{R}(f)$:

$\displaystyle\mathcal{R}(f)=\mathbb{E}_{\mathcal{D}}[\Psi(X,Z;f)]$

(2.5)

where for any given $(x,z)\in\mathcal{X}\times\mathcal{Z}$, $\Psi(x,z;f):\mathcal{F}\rightarrow{\mathbb{R}}_{+}$ is a convex functional of $f$ that maps each function $f$ to a scalar. Moreover, $\Psi$ satisfies

	$\displaystyle\Psi(x,z;0)$	$\displaystyle=0,\qquad\Psi(x,z;f)\geq 0,\qquad\forall f\in\mathcal{F},$		(2.6)
	$\displaystyle\frac{\delta\Psi(x,z;af_{1}+bf_{2})}{\delta f}$	$\displaystyle=a\cdot\frac{\delta\Psi(x,z;f_{1})}{\delta f}+b\cdot\frac{\delta\Psi(x,z;f_{2})}{\delta f},\qquad\forall f_{1},f_{2}\in\mathcal{F},\;a,b\in\mathbb{R}.$		(2.7)

Equation (2.6) requires that $\Psi(X,Z;f)$ is a non-negative functional of $f$ that is equal to $0$ if and only $f=0$. Equation (2.7) requires that the functional derivative of $\Psi(X,Z;f)$ with respect to $f$, is linear in $f$. One example of $\Psi$ is the $L_{2}$-regularizer of the following type, $\Psi(x,z;f)=f(w)^{2}$. Here $w\in\mathcal{W}$ is a subset of variables that contain values from both the endogenous variables $x$ and exogenous variables $z$.

Minimax Estimation. To solve the optimization problem in (2.3), we first transform it into a unconditional moment formulation by introducing a dual function. By Fenchel duality, we can rewrite the objective function $J(f)$ as follows,

	$\displaystyle J(f)$	$\displaystyle=\mathbb{E}_{\mathcal{D}}\Bigl{[}1/2\cdot\bar{\delta}(z;f)^{2}+\lambda\Psi(X,Z;f)\Bigr{]}$		(2.8)
		$\displaystyle=\mathbb{E}_{\mathcal{D}}\Bigl{[}\max_{g:\mathcal{Z}\rightarrow\mathbb{R}}\left(g(z)\cdot\mathbb{E}\bigl{[}\Phi(X,z;f){\,\big{|}\,}z\bigr{]}-1/2\cdot g(z)^{2}\right)+\lambda\Psi(X,Z;f)\Bigr{]}$
		$\displaystyle=\max_{g:\mathcal{Z}\rightarrow\mathbb{R}}\mathbb{E}_{\mathcal{D}}\Bigl{[}g(Z)\cdot\Phi(X,Z;f)-1/2\cdot g(Z)^{2}+\lambda\Psi(X,Z;f)\Bigr{]}.$

The formulation in (2.8) leads to the following minimax optimization problem:

$\displaystyle\min_{f}\max_{g}\mathcal{L}(f,g)=\mathbb{E}_{\mathcal{D}}\Bigl{[}g(Z)\cdot\Phi(X,Z;f)-1/2\cdot g(Z)^{2}+\lambda\Psi(X,Z;f)\Bigr{]}.$

(2.9)

We note that $\mathcal{L}$ is a convex-concave functional with respect to function $f$ and $g$. We denote by $(f^{*},g^{*})$ the unique saddle point of (2.9). Here the uniqueness of $f^{*}$ comes from the convexity of regularization $\Phi(X,Z;f)$, and $g^{*}(z)=\mathbb{E}[\Phi(X,Z;f^{*})|Z=z]$ implies the uniqueness of $g^{*}$. Without the regularization, i.e., $\lambda=0$, the saddle point of (2.9) is $f^{*}=f_{0}$ and $g^{*}=0$.

Policy Evaluation. We consider a Markov decision process given by $({\mathcal{S}},\mathcal{A},\mathcal{P},r,\gamma)$, where ${\mathcal{S}}\subseteq\mathbb{R}^{d}$ is the state space, $\mathcal{A}$ is the action space, $\mathcal{P}:{\mathcal{S}}\times\mathcal{A}\rightarrow\mathscr{P}({\mathcal{S}})$ is the transition kernel, $r:{\mathcal{S}}\times\mathcal{A}\rightarrow[0,1]$ is the reward function, $\gamma\in(0,1)$ is the discount factor. Given a policy $\pi:{\mathcal{S}}\rightarrow\mathscr{P}(\mathcal{A})$, an agent interacts with the environment in the following manner. At a state $s_{t}$, the agent takes an action $a_{t}\sim\pi(\cdot{\,|\,}s_{t})$ and receives a reward $r_{t}=r(s_{t},a_{t})$. Then, the agent transits to the next state $s_{t+1}\sim\mathcal{P}(\cdot{\,|\,}s_{t},a_{t})$. We denote the transition kernel induced by policy $\pi$ by $\mathcal{P}^{\pi}(s^{\prime}{\,|\,}s)=\int_{\mathcal{A}}\mathcal{P}(s^{\prime}{\,|\,}s,a)\pi(a{\,|\,}s)\mathrm{d}a$ for any $s,s^{\prime}\in{\mathcal{S}}$. In policy evaluation, we aim to estimate the value function $V^{\pi}:{\mathcal{S}}\rightarrow\mathbb{R}$ defined as follows,

$\displaystyle V^{\pi}(s)=\mathbb{E}_{\pi}\Bigl{[}\sum_{i=0}^{\infty}\gamma^{i}r(s_{i},a_{i}){\,\Big{|}\,}s_{0}=s\Bigr{]},$

where the expectation $\mathbb{E}_{\pi}$ is taken with respect to $a_{t}\sim\pi(\cdot{\,|\,}s_{t})$ and $s_{t+1}\sim\mathcal{P}(\cdot{\,|\,}s_{t},a_{t})$ for $t\geq 0$. By the Bellman equation (Sutton and Barto, 2018), it holds for any $s\in{\mathcal{S}}$ that

$\displaystyle V^{\pi}(s)-{\mathcal{T}}^{\pi}V^{\pi}(s)=0,\quad{\mathcal{T}}^{\pi}f(s)=\mathbb{E}_{a\sim\pi(\cdot{\,|\,}s)}\bigl{[}r(s,a)\bigr{]}+\gamma\mathbb{E}_{s^{\prime}\sim\mathcal{P}^{\pi}(\cdot{\,|\,}s)}\bigl{[}f(s^{\prime})\bigr{]}.$

(2.10)

Corresponding to the Bellman equation in (2.10), let $\mathcal{D}$ denotes the joint distribution of the state-action tuple $(s,a,s^{\prime})$ under policy $\pi$, the value function $V^{\pi}$ satisfies the following functional conditional moment equation,

$\displaystyle\mathbb{E}_{s^{\prime}|s}\Bigl{[}r(s,a)-V^{\pi}(s)+\gamma\cdot V^{\pi}(s^{\prime}){\,\Big{|}\,}s\Bigr{]}=0.$

(2.11)

We notice that (2.11) is a special case of the functional conditional moment equation in (2.1) by setting the exogenous variable $Z$ to be the current state $s$, the endogenous variable $X$ to be the next state $s^{\prime}$ and the function to be estimated $f:{\mathcal{S}}\rightarrow\mathbb{R}$ to be defined on the state space ${\mathcal{S}}$. In this case, the functional is $\Phi(X,Z;f)=r+\gamma\cdot f(X)-f(Z)$, where $r$ is the reward function. We remark that the reason function $f$ can be evaluated simultaneously on $X$ and $Z$ is that both $X$ and $Z$ are variables defined on ${\mathcal{S}}$. Following the same derivation of (2.8), policy evaluation can be formulated as the following minimax optimization problem,

$\displaystyle\min_{f}\max_{g}\mathcal{L}(f,g)=\mathbb{E}_{\mathcal{D}}\Bigl{[}g(Z)\cdot\bigl{(}r+\gamma\cdot f(X)-f(Z)\bigr{)}-1/2\cdot g(Z)^{2}+\lambda\Psi(X,Z;f)\Bigr{]}.$

Nonparametric Instrumental Variables Regression. The nonparametric instrumental variables model is common and useful in statistics and economics. The model can be described simply by a line of equation

$\displaystyle Y=f_{0}(X)+\varepsilon,\quad\mathbb{E}\bigl{[}\varepsilon{\,\big{|}\,}Z\bigr{]}=0.$

where $Y$ in an observed outcome, $X$ is the endogenous variable, $Z$ is the exogenous variable, $f_{0}$ is the true model that characterize the relationship between $Y$ and $X$ and is also the function we want to estimate. In this model, $\varepsilon$ is a noise possibly correlated with the endogenous $X$ but uncorrelated with the exogenous $Z$. It’s straightforward to see that NPIV model fits into the framework of the functional conditional moment equation by plugging the model equation into the equation about $\varepsilon$,

$\displaystyle\mathbb{E}_{\mathcal{D}}\Bigl{[}Y-f_{0}(X){\,\big{|}\,}Z\Bigr{]}=0.$

(2.12)

We notice that (2.12) is a special case of functional conditional moment equation in (2.4) by identifying $X$, $Z$ with the endogenous and exogenous variable respectively and setting the functional as $\Phi(X,Z;f)=Y-f(X)$. Following the same derivation of (2.8), the problem of NPIV is equivalent to the following minimax optimization problem,

$\displaystyle\min_{f}\max_{g}\mathcal{L}(f,g)=\mathbb{E}_{\mathcal{D}}\Bigl{[}g(Z)\cdot\bigl{(}Y-f(X)\bigr{)}-1/2\cdot g(Z)^{2}+\lambda\Psi(X,Z;f)\Bigr{]}.$

Asset Pricing. Asset pricing refers to the process of determining the fair value of financial assets. This field is fundamental in finance and underpins much of the work in investment, portfolio management, and risk assessment. Semiparametric Consumption Captial Asset Pricing Model (CCAPM) is a foundational model in asset pricing that describes the relationship between systematic risk and expected asset returns, which also incorporates the influence of the consumption preference of investors over time. Moreover, CCAPM can be characterized through a functional conditional moment equation (Chen et al., 2014; Chen and Ludvigson, 2009). To describe the model, let $C_{t}$ denote the consumption level at time $t$, $c_{t}\equiv C_{t}/C_{t-1}$ the consumption growth. The marginal utility of consumption at time $t$ is given by $\text{MU}_{t}=C_{t}^{-\gamma_{0}}f_{0}(c_{t})$, where $\gamma_{0}>0$ is the discount factor, $f_{0}:\mathcal{C}\to\mathbb{R}$ is the nonparametric structural demand function, which is an unknown positive function of our interest and is defined on $\mathcal{C}$, the space of consumption growth. The unknown function $f_{0}$ can be understood as a taste shifter that describes how the marginal utility of consumption changes with the state of the economy in terms of consumption growth.

Now, consider the growth-return tuple $(c_{t},\widetilde{r}_{t+1},c_{t+1})$ for $t\in\mathbb{N}^{+}$ with joint distribution $\mathcal{D}$, where $c_{t}$ is the consumption growth at the current time $t$, and $c_{t+1}$ is the consumption growth at the next time $t+1$. $\widetilde{r}_{t+1}$ is a modified return observed in this period, which is a known function of the actual return $r_{t+1}$ and the consumption growth $c_{t+1}$ at time $t+1$. We consider the scenario where the time series of consumption growth $\{c_{t}\}_{t\geq 0}$ follows a time-homogenous Markov chain with a smooth transition kernel. That being said, both conditional transition probabilities $c_{t+1}|c_{t}$ and $c_{t}|c_{t+1}$ admit a smooth density function. The CCAPM model captures the behavior of $f_{0}$ through the following equation:

Adversarial Riesz representer Estimation. Many problems in statistics, causal inference, and finance involve the task of learning a continuous linear functional in the following form,

$\displaystyle\mathcal{V}(g)=\mathbb{E}\bigr{[}m(V;g)\bigl{]}.$

(2.14)

where function $g\in\mathcal{G}:\mathcal{X}\rightarrow\mathbb{R}$, $\mathcal{F}$ is defined on a function space $\mathcal{G}$, and $V$ is a random vector of which we have access to observations and represents the source of randomness in the functional. Moreover, suppose such continuous linear functional $\mathcal{F}(\cdot)$ is also mean-square continuous with respect to $L^{2}$ norm. In that case, it can be written in a more benign and useful manner, which is also often the case. Formally speaking, for such linear functional $\mathcal{V}$, there exists function $f_{0}$ such that for any $g\in\mathcal{G}$,

$\displaystyle\mathcal{V}(g)=\mathbb{E}\bigl{[}f_{0}(X)g(X)\bigr{]}.$

(2.15)

The function $f_{0}$ here is called the Riesz representer of the linear functional $\mathcal{V}$, and the equation (2.15) is known as the Riesz representation theorem. Information about the Riesz representation of such linear functional is crucial to numerous applications and learning tasks. Therefore, we aim to estimate $f_{0}$ by exploiting the relationship characterized by the equation. We have the following trivial observation that the true Riesz representer $f_{0}$ can be recovered by solving the following equation,

$\displaystyle\mathbb{E}\Bigl{[}f_{0}(X)-f(X){\,\big{|}\,}X\Bigr{]}=0.$

(2.16)

Of course, $f=f_{0}$ will solve the equation above, and therefore the true Riesz representer is achieved. We remark that this is indeed a special case since the expectation taken in (2.16) is unconditioned. In the equation, we only involve the endogenous variable $X$, which also indicates that the exogenous variable $Z$ coincides with $X$. While special, the problem still fits in the framework discussed here. By setting $\Phi(X,Z;f)=f(X)-f_{0}(X)$, we recovered the intractable formulation of Riesz representer estimation.

However, unlike the previous examples where have access to observations of each term in the equation, here we have no direct access to values of $f_{0}$, making the problem seemingly intractable. Fortunately, the alternative formulation of the original problem as a minimax optimization problem solves this difficulty. When written in the minimax formulation, we will again see the linear functional $\mathcal{V}$ show up in the equation in the form of (2.14), which can be approximated using empirical values calculated from accessible observations of the random vector $V$. Following the same derivation of (2.8) and the definition of Riesz representer in (2.15), the problem of adversarial Riesz representer estimation is equivalent to the following minimax optimization problem,

$\displaystyle\min_{f}\max_{g}\mathcal{L}(f,g)=\mathbb{E}_{\mathcal{D}}\Bigl{[}m(V;g)-f(X)\cdot g(X)-1/2\cdot g(X)^{2}+\lambda\Psi(X,X;f)\Bigr{]}.$

(2.17)

Again, we stress that in (2.17), the absence of $Z$ is due to the fact both the endogenous and exogenous variables are described by $X$ and the objective is computationally tractable since we have access to both observations of $X$ and $V$.