Convergence of Mean-Field Langevin Stochastic Descent-Ascent for Distributional Minimax Optimization

Thank you for your thoughtful comments! Below, we respond to the concerns and clarify the novelty and contributions of our work.

On our analysis: We would clarify that both our proof technique and the resulting conclusions differ from the aforementioned papers in several fundamental ways. First, these prior works (Kim et al. (2024), Chen et al. (2022), Suzuki et al. (2023), Nitanda et al. (2022)) tackle the discrete-time MFLD problem by first establishing convergence in continuous time (typically via Wasserstein gradient flow) and then controlling the discretization error as a separate step. In contrast, our proof directly analyzes the per-step improvement using Taylor expansion and remainder estimates. This not only provides a more elegant and adaptable proof framework, but also yields stronger results as detailed in the next point.

Moreover, our results also differ in several key aspects. Kim et al. (2024) do not study MFL-SDA. To our best knowledge, their analysis of other algorithms does not extend to our established convergence rate of MFL-SDA. While Yang et al. (2020) study a minimax optimization in Euclidean space, and our analysis shares some structural similarities with theirs, the time-discretization error anlaysis of the probability functional on distributional space requires substantially more sophistication -- we devote pages of proofs to controlling this error, whereas in the Eucliean case, it follows directly from smoothness assumptions. The other cited references--Chen et al. (2022), Suzuki et al. (2023), and Nitanda et al. (2022)--focus on minimization problems, whereas minimax problems present more complexity due to the interaction between two players.

On the Lyapunov function: As you said, $\lambda$ can be fixed to 1. Yet, our chosen form of the Lyapunov function follows what is commonly used in previous works on minimax problems in Euclidean space (Yang et al. (2020)) and in continuous-time MFLD (Lu (2023)).

Comparison with MFL-AG and MFL-ABR: First, our analysis of stochastic gradient descent-ascent with last-iterate convergence and the inexact gradient analysis closely aligns with practical implementations.
Second, MFL-AG achieves a sample complexity $O(\epsilon^{-O(1/\alpha)})$ to reach an $O(\epsilon)$-approximation of the saddle point, where the bound on LSI constant $\alpha$ can be small--potentially even of order $1/(2+d/2)$. In contrast, our convergence analysis for MFL-DA establishes a sample complexity $O(\frac{1}{\epsilon}\log\frac{1}{\epsilon})$. MFL-ABR is a double-loop algorithm, it is shown that the outer loop has an $O(\frac{1}{\epsilon}\log\frac{1}{\epsilon})$ sample complexity; however, it needs an inner loop and the total complexity is not specified.

On Experimental Results and Applications: The primary aim of our work is to establish convergence guarantees for commonly used stochastic gradient descent-ascent algorithms. We acknowledge the value of empirical validation, we have included a numerical experiment based on Example 2 in our paper. We apply our algorithm to the nonlinear instrumental variable (NPIV) regression problem $$ \min_f \max_g E[g(Z)(Y-f(X))-\frac{1}{2}g(Z)^2+\lambda R(f)]

Thank you very much for your critical feedback, especially your sharp observations about the order of the bias term. Below, we address your concerns. \newcommand{\bE}{\mathbb{E}} \newcommand{\o}{\omega}

Addressing Weaknesses

Bias term of second moment: Indeed, the norm control in the submitted version (adopted from [3, Lemma 1]) is sub-optimal. Below, please find a stronger control, adopted from [2, Lemma 1]. The corrected expression for $r_{g2}$ should be $r_{g2}= 2M^2_1+2\tau^2\bE_{\mu_0}[\\|\\theta_0\\|^2_2]+4M^2_1+ \tau(4\eta_2M^2_1+4\tau d)$ and similarly for $r_{h2}$. They are of order $O(1)$, thereby the bias term $R_1$ in Theorem 1 is $O(\eta_1)$.

where the first equality uses the fact that $\xi_k^2$ is an independent zero-mean normal; the first inequality follows the fact that $\\|h_k+\tau\o_k\\| = \\|\nabla \frac{\delta J}{\delta \nu}$ \mu_{k+1}, \nu_k $(\o_k)\\| \leq M_1$; the second inequality is because $2M_1\bE_{\nu_k}[\\|\\o_k\\|]\leq \frac{\tau}{2} (\bE_{\nu_k}[\\|\o_k\\|])^2 + 2M^2_1/\tau \leq \frac{\tau}{2}\bE_{\nu_k}[\\|\o_k\\|^2]+2M^2_1/\tau$ and last inequality holds for $\tau<1/(4\eta_2)$. Hence, we obtain

The remaining part of the proof stays the same. Hence, we can directly check that the remaining constants $r_{gi}$, $r_{hi}$ as well as $\Gamma_0$, $\Gamma_1$, $\Gamma_2$ are all of order $O(1)$. Since $\eta_1$ and $\eta_2$ are of the same order, we conclude that the bias term $R_1= \frac{\lambda(\Gamma_2\eta_2^2+(1-2\eta_2\tau\alpha)(\Gamma_1+\Gamma_0)\eta_1^2)+\Gamma_1\eta_1^2}{\eta_1\tau\alpha}$ is of order $O(\eta_1)$.

Regarding your question on the definition of constants in (18), they are defined in the proof of Theorem 2 (Line 1214, etc.). We apologize for not explicitly displaying them.

Assumptions of LSI in examples: Indeed, we have provided two applications in Section 4, with the LSI being verified in Proposition 3 and Proposition 4, respectively. This follows a line of work like [1, proposition 5.1] and [2, Appendix A]. The setups therein are satisfied in our context.

Existence and uniqueness of optimal $\mu^\*,\nu^\*$ : The primal objective $J(\mu, \nu)$ is convex in $\mu$ and concave in $\nu$. With an additional KL (or entropy) regularization, the regularized objective $E(\mu, \nu)$ becomes strongly convex in $\mu$ and strongly concave in $\nu$. This ensures the existence and uniqueness of the mixed Nash equilibrium $(\mu^*, \nu^*)$; please refer to the detailed discussion before equation (3) in the manuscript as well as [2, Proposition 2.1].

Addressing Questions

A1. To get (24), the first line follows directly from equation (23); the second line follows from the definition of $T$; the third line follows from the Taylor's expansion of $\log {\rm det}$, where $\bar{\o}_k$ is the point that achieves the mean-value; and the fourth line follows from the property of trace operator and that ${\rm Tr}(h_k^2) =\\|\nabla h_k\\|_F^2 = (\\|\nabla^2 \frac{\delta J}{\delta \nu}\\|_F - \tau)^2 \leq (M_2 + \tau)^2$.

A2. Yes, you're absolutely right.

Thank you so much again for your insightful and constructive feedback. We hope our responses alleviate your concerns.

References:

[1]Chizat, L. Mean-field Langevin dynamics: Exponential convergence and annealing. Transactions on Machine Learning Research, 2022.

[2]Suzuki, T., Wu, D., and Nitanda, A. Mean-field Langevin dynamics: Time-space discretization, stochastic gradient, and variance reduction. Advances in Neural Information Processing Systems, 36, 2024.

[3] Nitanda, Atsushi, Denny Wu, and Taiji Suzuki. Convex analysis of the mean field langevin dynamics. International Conference on Artificial Intelligence and Statistics, 2022.

Thank you for your thoughtful and detailed feedback! Below, we address the comments and questions point-by-point.

Pessimistic LSI bound: We will explicitly comment on the weakness of this type of analysis, and we are trying to overcome this in our ongoing research.

Explicit dependence on other parameters: In the revised version, we will include a dedicated paragraph discussing the roles of parameters and here we provide a sketch.

On dimension $d$ and regularization parameters: We agree that a more explicit discussion of the role of parameters in the convergence rate would be beneficial. Viewing $\tau,\eta_1,\eta_2$ as small numbers, in Appendix A.3 on page 11, the remainder $r_{g2}, r_{h2}$ due to the second moment are $\max\\{O(1),\tau d\\}$, similarly $r_{g3}, r_{h3}$ are $\max \\{ O(1),\tau^{3/2}d^{3/2} \\}$. Substituting them into $\Gamma_0, \Gamma_1, \Gamma_2$, we can get $\Gamma_0=\max\\{O(1), \tau d,\tau^{2}d^{2},\frac{d}{\alpha^{1/2}\tau} \\}$, and $\Gamma_{1(2)}=\max\\{O(1),\tau d,\tau^2d^2\\}$. Substutiuting them into $R_1$ in (10), we can get $R_1=O(\frac{d^2\eta_1}{\tau^3\alpha^3})$. Let $R_1=\epsilon$ and choose $\eta_1= O(\frac{\epsilon \tau^3 \alpha^3}{d^2})$, we obtain a sample complexity $K=O(\frac{d^2}{\epsilon \tau^4 \alpha^4}\log \frac{1}{\epsilon})$. Comparing with sample complexity $K=O(\frac{d^2}{\epsilon \tau^2 \alpha^2}\log \frac{1}{\epsilon}$) in [1] for the MFLD of minimization problems, the higher orders of $\tau, \alpha$ is because the two-timescale algorithm MFL-SDA: the overall complexity of the algorithm depends on the slower descent step.

On other parameters: The variance $\zeta,\sigma^2$ only affect the $O(1)$ term in $\Gamma_i$. The parameters $\eta_1,\eta_2$ correspond to the fast descent and fast ascent regime similar to [3], and thus their choice depends on the instance and the user's emphasis on the descent/ascent part.

Another standard of suboptimality: Our problem is to find a saddle point rather than a maxima(mimima), hence $E(\mu_K,\nu_K)-E(\mu^*,\nu^*)$ may not apply to our analysis.

Addressing References Not Discussed:

Thank you for pointing this out! We will cite this paper in the revision and investigate if their Theorem 5.2 can inspire a better LSI bound for our problem.

Addressing Questions:

A1. Thank you for catching this: the Talagrand's inequality should be $W_2^2(\mu_k,\nu^*)\leq \frac{2}{\alpha}\mathrm{KL}(\mu_k|\mu^*)$.

A2. $\lambda$ corresponds to the Lyapunov function similar to [2] [3], and it is not an explicit hyperparameter in our algorithm.

A3. Yes, it holds point-wise for all $\theta$ and $\omega$.

A4. See our response on the dimension above.

A5, A6. Thank you for noticing this. We will correct the notation.

We thank you very much again for the constructive feedback and helpful suggestions!

Reference:

[1] Nitanda, A., Wu, D., & Suzuki, T. (2022). Convex analysis of the mean field Langevin dynamics. In International Conference on Artificial Intelligence and Statistics (pp. 9741-9757). PMLR.

[2] Lu, Y. (2023). Two-scale gradient descent ascent dynamics finds mixed Nash equilibria of continuous games: A mean-field perspective. In International Conference on Machine Learning (pp. 22790-22811). PMLR.

[3] Yang, J., Kiyavash, N., & He, N. (2020). Global convergence and variance-reduced optimization for a class of nonconvex-nonconcave minimax problems. arXiv preprint arXiv:2002.09621.

Thank you very much for your thorough and constructive feedback. Below, we address your main points raised. $\newcommand{\bs}{\boldsymbol}$

Addressing Weakness:

Particle Discretization: For this problem, we have verified that our method is feasible in the particle setting. As an illustration, we briefly summarize the adjustments needed in our proof technique under the particle approximation setting:

We set the joint distribution $\bs\mu,\bs\nu$ as $N$ finite particles $\bs\theta=[\theta^i]^N_{i=1}$ and $\bs \omega = [\omega^i]^N_{i=1}$. In this case, the LSI has the form adapted from [1, Lemma 7]

where $\bs\mu_k$ and $\bs\nu_k$ denote the distribution of $\bs\theta$ and $\bs\omega$ in the $k$-th iteration, respectively; $h_i$ denotes the partial derivative w.r.t the $i$-th particle; $\mathcal{L}_2^N$ denotes the Lynapunov function for the counterpart on the product space; $\bs\rho$ denotes the Gibbs distribution; $O(\frac{1}{N})$ is from the propagation of chaos. Other log-Sobolev inequalities used in our setting have similar forms. Thereby, most derivations in the manuscript involving Taylor expansions can be extended to analyze the joint distribution of particles, but with additional approximation error $O(1/N)$ due to weak particle interactions. The per-step improvement would be of the form

where the $O(\frac{\eta_1}{N})$ term comes from the propagation of chaos. Building on our current line of thought, we believe this approach can be extended to the stochastic gradient setting while maintaining the same sample complexity order as well. However, due to the extensive computations involved, we will present a complete proof of our results in an upcoming extended journal version.

Parameters: Due to the space limit, please refer to our response to Reviewer 5VGF for a detailed analysis of how the sample complexity in Theorems 1 and 2 depends on various parameters, including the dimension $d$. Notably, the sample complexity of MFL-AG in [2] is $O(\epsilon^{-O(1/\alpha)})$, which is exponential in the LSI constant $\alpha$, hence in $d$ under their conservative bound $O(1/(2+d/2))$ of $\alpha$, whereas our bound does not suffer from this exponential dependence. The Hessian bound $C_1$ occurs in bias $\Gamma_i=O(C_1)$. As for the condition number, adopting the notion of effective condition number in [3, Theorem 2.1], our bound has a similar dependence in the case of zero-sum game.

Addressing Questions:

Line 143 and Equation (9): Thank you for pointing this out. We will correct the issues in the updated version of the manuscript.

Boundedness of third-moment: The existence of this issue arises from the presence of $\mathbb{E}_{\mu_k}[\\|g_k\\|_2^3]$ in our remainder term. Therefore, in the stochastic gradient part, to ensure that this remainder term remains uniformly bounded, we need to impose a bound on the third-order moment. In fact, [1] assumed that the stochastic gradient is uniformly bounded with respect to any $\omega$ and $\theta$, which is stronger than ours.

Thank you again for your valuable feedback! We will correct the typos and formatting issues as you suggested.

References

[1] Suzuki, T., Wu, D., & Nitanda, A. (2023). Convergence of mean-field Langevin dynamics: time-space discretization, stochastic gradient, and variance reduction. Advances in Neural Information Processing Systems, 36, 15545-15577.

[2] Kim, J., Yamamoto, K., Oko, K., Yang, Z., & Suzuki, T. (2023). Symmetric mean-field langevin dynamics for distributional minimax problems. arXiv preprint arXiv:2312.01127.

[3] Lu, Yulong. Two-scale gradient descent ascent dynamics finds mixed Nash equilibria of continuous games: A mean-field perspective. International Conference on Machine Learning. PMLR, 2023.