Mean-Field Analysis for Learning Subspace-Sparse Polynomials with Gaussian Input

Thanks a lot for your helpful and insightful comments and questions. Please find our responses below:

• Connection/difference between the reflective property (equations (3.1)-(3.2)) and $\textup{IsoLeap}(h^*)>1$ (Abbe et al. 2023):
  • Our conditions (3.1) and (3.2) are more general than $\textup{IsoLeap}(h^*)>1$ in the sense that (3.1) and (3.2) can be extended straightforwardly to general function $h^*$, while $\textup{IsoLeap}(h^*)$ can only be defined for polynomial $h^*$.
  • If we only consider a polynomial $h^*$, then it can be prove that $\textup{IsoLeap}(h^*)>1$ implies (3.2). However, we currently do not know whether (3.2) implies $\textup{IsoLeap}(h^*)>1$ or not.
  • (3.1) is different from (3.2) and $\textup{IsoLeap}(h^*)>1$ when $\sigma$ is less expressive. In our original submission, we give a counter-example with $\sigma(\zeta) = \zeta$, but this is not the only example. There exist other counter-examples with more expressive $\sigma$. In fact, for any $n\geq 2$, consider $h^*(z) = \textup{He}_1(z_1)+\textup{He}_n(z_1) \textup{He}_1(z_2)$ and $\sigma(\zeta) = \zeta^n$. Then it can be verified that (3.1) is satisfied while $\textup{IsoLeap}(h^*)=1$.
• Discretized dynamics and sample complexity: There have been standard results (see e.g. Mei et al. (2019)) that bound the distance between the SGD and the mean-field dynamics. To apply the results of Mei et al. (2019), one needs the boundedness of $f^*$ and $\sigma$. However, in our work, $f^*$ and the second activation function $\hat{\sigma}$ in Algorithm 1 are bounded, which is the reason that we did not give a sample complexity result in our original submission. Recently, we realized that a sample complexity result can be derived as follows and we will include such analysis in our revision.
  • Step 1: We slightly modify the mean-field dynamics by replacing $f^*$ and $\hat{\sigma}$ by $\text{sign}(f^*)\min(|f^*|,C_f)$ and $\text{sign}(\hat{\sigma})\min(|\hat{\sigma}|,C_{\sigma})$ that are both bounded. $C_f$ and $C_{\sigma}$ are dimension-free constants guaranteeing that the distance between the original dynamics and the modified dynamics is smaller than $\epsilon$. Let us remark that $C_f$ and $C_{\sigma}$ can be chosen as dimension-free since $f^*$ is subspace-sparse and dynamics in Algorithm 1 does not learn any information about $V^\perp$.
  • Step 2: Apply the standard results in Mei et al. (2019) to the modified dynamics and then obtain the sample complexity for SGD.
• Some explanations for Algorithm 1: We agree with you that taking the average in Algorithm 1 is non-standard. The reason is that we want to lift the linear independence to the algebraic independence. For the change of activation function in Algorithm 1, we find it is somehow acceptable in the existing literature (see Abbe et al. (2022)).
• Consistency with Lee et al. (2024) and Arnaboldi et al. (2024):
  • Lee et al. (2024) require a sample complexity $\mathcal{O}(d\cdot \text{polylog} d)$. So the training time is not constant, due to the $\text{polylog}d$ term. Our work always considers a finite training time $T$ and we definitely agree that SGD may learn a larger family of functions if we allow the training time to scale with the dimension $d$.
  • Arnaboldi et al. (2024) analyze a variant of SGD, not the original SGD. Our necessary condition considers the mean-field dynamics corresponding to the original SGD.

Thanks a lot for your helpful and insightful comments and questions. Please find our responses below:

• Connection/difference between the reflective property (equations (3.1)-(3.2)) and $\textup{IsoLeap}(h^*)>1$ (Abbe et al. 2023):
  • Our conditions (3.1) and (3.2) are more general than $\textup{IsoLeap}(h^*)>1$ in the sense that (3.1) and (3.2) can be extended straightforwardly to general function $h^*$, while $\textup{IsoLeap}(h^*)$ can only be defined for polynomial $h^*$.
  • If we only consider a polynomial $h^*$, then it can be prove that $\textup{IsoLeap}(h^*)>1$ implies (3.2). However, we currently do not know whether (3.2) implies $\textup{IsoLeap}(h^*)>1$ or not.
  • (3.1) is different from (3.2) and $\textup{IsoLeap}(h^*)>1$ when $\sigma$ is less expressive. In our original submission, we give a counter-example with $\sigma(\zeta) = \zeta$, but this is not the only example. There exist other counter-examples with more expressive $\sigma$. In fact, for any $n\geq 2$, consider $h^*(z) = \textup{He}_1(z_1)+\textup{He}_n(z_1) \textup{He}_1(z_2)$ and $\sigma(\zeta) = \zeta^n$. Then it can be verified that (3.1) is satisfied while $\textup{IsoLeap}(h^*)=1$.
• Definition 3.2 and Theorem 3.5: We understand that it is a bit confusing that we use $u+v^\top z_S^\perp$ in Definition 3.2 and use $w^\top x_S^\perp$ without bias term in Theorem 3.5. The reason can be seen in line 479-480, where we write $w^\top x_S^\perp = w^\top x_V^\perp + w^\top (x_V-x_S)$. Then $w^\top x_V^\perp$ is the bias term $u$ and $x_V-x_S$ is the orthogonal complement of $x_S$ in $V$, corresponding to $z_S^\perp$ in Definition 3.2. We will make this point more clear in our revision.
• Delta measure on a subspace: For a subspace $S$, the delta measure $\delta_S$ is a probability measure on $S$ such that for any continuous and compactly supported function $f:S\to\mathbb{R}$, we have $\int_S f(x)\delta_S(dx) = f(0)$. We will add the definition in our revision.

Thanks a lot for your helpful and insightful comments and questions. Please find our responses below:

• Connection/difference between the reflective property (equations (3.1)-(3.2)) and $\textup{IsoLeap}(h^*)>1$ (Abbe et al. 2023):
  • Our conditions (3.1) and (3.2) are more general than $\textup{IsoLeap}(h^*)>1$ in the sense that (3.1) and (3.2) can be extended straightforwardly to general function $h^*$, while $\textup{IsoLeap}(h^*)$ can only be defined for polynomial $h^*$.
  • If we only consider a polynomial $h^*$, then it can be prove that $\textup{IsoLeap}(h^*)>1$ implies (3.2). However, we currently do not know whether (3.2) implies $\textup{IsoLeap}(h^*)>1$ or not.
  • (3.1) is different from (3.2) and $\textup{IsoLeap}(h^*)>1$ when $\sigma$ is less expressive. In our original submission, we give a counter-example with $\sigma(\zeta) = \zeta$, but this is not the only example. There exist other counter-examples with more expressive $\sigma$. In fact, for any $n\geq 2$, consider $h^*(z) = \textup{He}_1(z_1)+\textup{He}_n(z_1) \textup{He}_1(z_2)$ and $\sigma(\zeta) = \zeta^n$. Then it can be verified that (3.1) is satisfied while $\textup{IsoLeap}(h^*)=1$.
• Connection between Theorem 3.4 and Abbe et al. (2022): The goal of Theorem 3.4 is to generalize the result in Abbe et al. (2022) to the rotation-invariant setting. Therefore, the statement of Theorem 3.4 is similar to Abbe et al. (2022). However, the proofs are significantly different since Abbe et al. (2022) analyze the dimension-free dynamics and we do not use that. We directly show that the flow cannot learn any information about $S$ by connecting it with a flow in $\mathbb{R}\times S^\perp$.
• Some explanations for Algorithm 1:
  • In Algorithm 1, we need to repeat Step 2 for $p$ times, and for each time the flow is an interacting particle system, where the whole system is described by a PDE and the dynamics of a single particle in the interacting system is described by an ODE. Therefore, Step 3 does not train $p$ particles independently, but trains $p$ interacting particle systems. Compared to training only one interacting particle system as in Abbe et al. (2022), we need to repeat this step to lift the linear independence to the algebraic independence.
  • You are correct that we change the activation function in Step 4. We want to mention that this is somehow acceptable as in existing literature (see Abbe et al. (2022)).
• Exponential convergence in Theorem 4.3: You are correct that the exponential convergence is from linear dynamics, i.e., the second phase. This idea is standard as in many existing works including Abbe et al. (2022). The most difficult part is to prove that the eigenvalue of the kernel matrix has a positive lower bound.
• Verification of Assumption 4.1: It is very difficult to directly verify Assumption 4.1 in general. But there are several cases that are manageable and we will let you know if we figure out a more general case during the discussion period.
  • If $h^*(z) = c_1 z_1+c_2 z_1 z_2 + \cdots + c_p z_1 z_2\cdots z_p$, then Assumption 4.1 can be proved with nonzero coefficients $c_1,c_2,\dots,c_p$.
  • If $h^*(z) = \sum_{s\in\mathcal{S}} c(s) \prod_{i=1}^p He_{s_i}(z_i)$, where $\mathcal{S}$ is a finite subset of $\mathbb{Z}_{\geq 0}^p$ with $(1,0\dots,0),(1,1,0,\dots,0),\dots,(1,1,\dots,1)\in\mathcal{S}$, then Assumption 4.1 can be proved in the almost sure sense when $c(s)$ are randomly generated (e.g., from iid Gaussian) for all $s\in\mathcal{S}$.

I thank the authors for their detailed response and for clearing out some of my confusion. Here are come comments:

• Concerning isoleap: The isoleap is defined in terms of the Hermite coefficients in the orthonormal decomposition of the target function, which always exist for squared integrable functions (which is always the case when using squared loss), so it is not limited to polynomials. (The isoleap is always defined and finite for any squared integrable function.)

• Concerning Assumption (3.2): Let's consider $S$ to be one direction $z_1$. Then $h_* (z) = \sum_{k \geq 0} He_k (z_1) h_k (z_{-1}).$ Therefore the condition (3.2) shows that $h_1(z_{-1}) = 0$ for all $z_{-1}$. Hence, in this coordinate basis, $z_1$ only appears in $He_k$ with $k\geq2$, so its leap $\geq 2$ in that basis, and therefore the isoleap is $\geq 2$ (this generalizes straighforwardly to $S$ of any dimension). Hence Assumption (3.2) is equivalent to isoleap $\geq 2$ (if I am not missing anything!). Indeed, Assumption 3.1 takes into account the expressivity of the activation function. However, note that in your example, one cannot fit $h_*$ anyway (degree $n$ activation and degree $n+1$ activation). More generally, for any $\sigma$ degree $k$ polynomial, I would expect (didn't check) $\sigma (u + v^T z_{S^\perp})$ to span all polynomials of degree-$k$ and to be equivalent to Assumption (3.2) whenever $h_*$ is degree $\leq k$, which is anyway the only part of the target function that can be fitted by the neural network. Hence, for $\sigma$ degree-$k$ polynomial, we could imagine the following: decompose $h_* = h_{\leq k} + h_{>k}$, assume Assumption (3.2) on $h_{\leq k}$ and disregard $h_{>k}$, which does not contribute to the dynamics by orthogonality. (Note that $h_{\leq k}$ can only have higher leap than $h_*$.) All in all, I am still not convinced that Assumption (3.1) is significantly different from isoleap to claim ``we propose a basis-free generalization of the merged-staircase property in [1]''. I apologize for being nit-picking (and I do not consider it to be the most important part of your paper), but several other groups, including [2,7,9,10] cited in your paper, have worked on this basis-free leap complexity. Please let me know if I misunderstood anything.

• Concerning Algorithm 1 and Assumption 4.1: Thank you for the clarification, I mistakenly thought the initialization was a point mass. I acknowledge that the Gaussian case is much harder than the hypercube case, due to the need of an algebraic independence. Thank you for making the effort of verifying Assumption 4.1 in both these cases, I believe it already makes a strong case for Assumption 4.1!

Thank you for taking into account my comments. I increased my score to 5. The setting is interesting and can be of interest to the NeurIPS community. However, I believe that the paper still falls short of solving the main difficulty of going from hypercube to Gaussian, which is the algebraic independence, and requires an unatural construction (repeating step 2).

Thank you very much for your encouraging comments! We appreciate your time and consideration.