Paper: Improved statistical and computational complexity of the mean-field Langevin dynamics under structured data (ICLR-2024)

Reviewer comment (Theory):
In my opinion, the claim that "anisotropy helps" seems too strong, a bit specific to the problem in the paper, and not necessarily a general statement that can be made for neural networks based on the assumptions and results of the paper. It would be great if the authors could better motivate a connection between this problem and more general neural networks with real-world datasets.

Author response:
**Connection to more general neural networks with real-world datasets**

First, we note that our primary goal is to develop theory for representation learning – to understand when it is beneficial in terms of statistical and computational complexity, and how it interacts with different aspects of the learning problem / algorithm, such as structure of the data and target function, model architecture, and optimization algorithm. More specifically, the focus of this work is to rigorously examine the interplay between feature learning and the anisotropic structure of input data, which, as remarked in the introduction, is relatively underexplored especially in the context of classification problems. While our theoretical setting is idealized, it is motivated from practical observations that real-world data exhibits low-dimensional structure (see below), and we believe that our analysis serves as an important step towards a concrete understanding of feature learning in more complicated and realistic settings. 

On the other hand, our problem setting is motivated by empirical observations on realistic data. Real-world data, including image and text, can be high-dimensional, yet neural networks often efficiently learn from data and avoid the “curse of dimensionality”, This success can be attributed to the fact that real-world data exhibits some underlying structure, and intrinsic low dimensionality is considered as one of the prominent factors. Specifically, low intrinsic dimensionality has the following two aspects.  
(1) *Low-dimensionality of ground truth (target function)*. This means that not all directions of the input features are important for predicting $y$.   
(2) *Anisotropy of input data*. This means that the features already contain low-dimensional structures despite the large ambient dimension.   
The $k$-sparse parity classification problem is a classical example of low-dimensional function acting on high-dimensional data, and most prior theoretical studies for this problem setting considered isotropic input data, which accounts for (1) but not (2). By introducing this “generalized” version of the sparse parity problem on anisotropic input, we aim to theoretically study the interplay between structured (anisotropic) data and the efficiency of feature learning via gradient descent. We show that neural networks trained by gradient-based algorithms indeed exploit such low-dimensional structure, as evident in the improved statistical and computational complexity.  
As for application to real-world data, our Theorem 1 suggests that neural networks can leverage low-dimensional structure to improve generalization performance, whereas Theorem 2 suggests that the structure of target function can be revealed by the gradient covariance matrix. We believe that these insights can be transferred to practical settings at a high level.

Reviewer comment (Novelty):
The authors state on page 2, under "Feature learning under structured data", that 

> in certain regression settings with low-dimensional target, structured data with a spiked covariance structure can improve the performance of both kernel methods and optimized NNs (Ghorbani et al., 2020; Ba et al., 2023; Mousavi-Hosseini et al., 2023; Suzuki et al., 2023b). However, these regression analyses do not directly translate to the binary classification setting which the k-parity problem belongs to.

This implies that the intuition "anisotropy helps" is perhaps already observed in the literature, which may weaken the impact of the paper. If the authors could expand on why the $k$-parity problem is worth studying separately, that would really help motivate things.

Author response:
**Difference from prior results on anisotropic data** 

We have revised the Introduction to highlight the difference from prior works and new technical challenges. We make the following remarks. 
1. Classification and regression problems have fundamentally different structures: it is possible that a predictor achieves vanishing classification error but has large regression loss. Concretely, the following aspects of our analysis do not translate from the regression setting: 
* For the neural network results, we exploited properties of the logistic loss, and more importantly, margin conditions on the optimal classifier. Such analysis does not follow from reducing the classification problem to regression (e.g., with the squared loss). 
* For the kernel lower bound, establishing a lower bound on the classification error is much more challenging than on the regression loss, because it is clear that a kernel model can perfectly predict the sign of the labels $y$ but do not achieve small squared error. 
2. Regarding the optimization dynamics, our analysis differs significantly from prior results (Ghorbani et al., 2020; Ba et al., 2023; Mousavi-Hosseini et al., 2023). (Ghorbani et al., 2020) did not provide any optimization guarantees for neural networks. (Ba et al., 2023; Mousavi-Hosseini et al., 2023) considered a narrow neural network in the so-called “rotation” regime, where the statistical complexity is governed by the *information exponent* of the ground truth. Roughly speaking, to learn a degree-$k$ parity function in this regime, the sample complexity may scale with $d^{\Theta(k)}$. In contrast, we focused on neural networks in the mean-field regime, which allows us to “decouple” the degree from the exponent of the dimension dependence, albeit at a cost of higher computational complexity (Theorem 1). Finally, we realize that (Suzuki et al., 2023b) does not handle the anisotropic setting – we have removed this citation in the paragraph on structured data.

Reviewer comment (Experiments):
How applicable is the coordinate transformation used in the paper, to conventional neural network training for real-world datasets?

Author response:
**"How applicable is the coordinate transformation used in the paper, to conventional neural network training for real-world datasets?"** 

The coordinate transform can be implemented in practical neural networks, since the gradient covariance estimation can be done by standard auto-differentiation analogous to the computation of second-order statistics for preconditioned gradient update. We intuitively expect such transformation to be beneficial when the batch size is sufficiently large. It is also possible that the covariance of SGD noise is also implicitly amplifying certain informative directions, thus achieving a similar effect as explicit coordinate transformation; we leave such investigation to future work.

Reviewer comment (Novelty):
I am not well-positioned to comment on the significance of the results within this research area. However, it is not clear to me why we should be interested in this formal setting beyond the fact that some previous works have considered it. The assumptions seem highly contrived and no attempt is made to link them to any practical task in a meaningful way. Could the authors explain why progress in this research direction is worth pursuing? I would argue that if this work is being submitted to a broad conference such as ICLR, more effort should be made to broaden its appeal by explaining why its setting is relevant.

Author response:
We thank the reviewer for the helpful feedback. 

The reviewer’s main concern is regarding the relevance of our theoretical findings to the ICLR community. We motivate our problem setting from the following three perspectives.  

**1. Why do we care about the theory of representation (feature) learning?** 

It is clear that representation learning is a central topic of ICLR, as indicated by the title of the conference. We argue that it is imperative to develop theory for representation learning – to understand when it is beneficial in terms of statistical and computational complexity, and how it interacts with different aspects of the learning problem / algorithm, such as structure of the data and target function, model architecture, and optimization algorithm. The focus of this work is to rigorously examine the interplay between feature learning and the anisotropic structure of input data, which, as remarked in the introduction, is relatively underexplored especially in the context of classification problems. While our theoretical setting is idealized, it is motivated from practical observations that real-world data exhibits low-dimensional structure (see point below), and we believe that our analysis serves as an important step towards a concrete understanding of feature learning in more complicated and realistic settings, and is certainly of interest to the ICLR community. In fact, there are at least 5 papers accepted to ICLR 2023 that theoretically analyzed the feature learning dynamics under similar idealized settings: [Mousavi et al.](https://openreview.net/pdf?id=6taykzqcPD), [Bordelon and Pehlevan](https://openreview.net/pdf?id=nZ2NtpolC5-), [Akiyama and Suzuki](https://openreview.net/pdf?id=6doXHqwMayf), [Telgarsky](https://openreview.net/pdf?id=swEskiem99), [Tian](https://openreview.net/pdf?id=s130rTE3U_X).

**2. Why is the sparse parity problem worth studying?** 

The $k$-sparse parity classification problem is a classical example of low-dimensional function acting on high-dimensional data. This function often serves as a testbed for theory of representation learning, and has been extensively studied in the isotropic setting  – see references in the Introduction section. The reason is twofold: (i) The target function only depends on a few coordinates of the input features, and hence we expect that the trained neural network representation can “zoom-in” to the relevant k-dimensional subspace. (ii) For methods that do not exploit such low-dimensional structure, such as neural networks in the lazy (NTK) regime, we expect the sample complexity to be inferior to the feature learning alternative – this provides a simple setting to illustrate the benefit of representation learning. 

**3. What new insights do we gain from introducing anisotropy?** 

Real-world data, including image and text, can be high-dimensional, yet neural networks often efficiently learn from data and avoid the “curse of dimensionality”, This success can be attributed to the fact that real-world data exhibits some underlying structure, and intrinsic low dimensionality is considered as one of the prominent factors. Specifically, low intrinsic dimensionality has the following two aspects (see [Pope et al.](https://arxiv.org/pdf/2104.08894.pdf) for example).  
(1) *Low-dimensionality of ground truth (target function)*. This means that not all directions of the input features are important for predicting $y$.   
(2) *Anisotropy of input data*. This means that the features already contain low-dimensional structures despite the large ambient dimension.    

As noted in the previous point, prior theoretical studies on the isotropic parity problem accounted for (1) but not (2). By introducing this “generalized” version of the sparse parity problem on anisotropic input, we aim to theoretically study the interplay between structured (anisotropic) data and the efficiency of feature learning via gradient descent. We show that neural networks trained by gradient-based algorithms indeed exploit such low-dimensional structure, as evident in the improved statistical and computational complexity. Similar problem setups and motivations have appeared in [Ghorbani et al. (2020)](https://proceedings.neurips.cc/paper/2020/hash/a9df2255ad642b923d95503b9a7958d8-Abstract.html) and [Ba et al. (2023)](https://openreview.net/forum?id=HlIAoCHDWW) but for the regression setting. 

We have revised the Introduction to highlight the motivation of our problem setting and the relevance of our theoretical results.

Reviewer comment (Writing):
Could the authors provide a definition of anisotropic input data? I don't think this was clearly defined. Does it exactly refer to the setting defined mathematically at the beginning of Sec 1.1?

Author response:
**"Could the authors provide a definition of anisotropic input data?"**

We follow the standard definition of isotropy/anisotropy, that is, a random vector is anisotropic if its covariance is not identity. 

We would be happy to clarify any further concerns/questions in the discussion period.

Reviewer comment (Theory):
Why can the existing results for the regression setting not be directly applied in the binary classification setting by converting it to a regression problem?

Author response:
**"Why can the existing results for the regression setting not be directly applied in the binary classification setting by converting it to a regression problem?"**

Note that classification and regression problems have fundamentally different structures: it is possible that a predictor achieves vanishing classification error but has large regression loss. Concretely, the following aspects of our analysis cannot be obtained by converting the classification problem into regression: 
* For the neural network results, we exploited properties of the logistic loss, and more importantly, margin conditions on the optimal classifier. Such analysis does not follow from reducing the classification problem to regression (e.g., with the squared loss). 
* For the kernel lower bound, establishing a lower bound on the classification error is much more challenging than on the regression loss, because it is clear that a kernel model can perfectly predict the sign of the labels $y$ but do not achieve small squared error.

Reviewer comment (Theory):
It appears that the whole work assumes that the matrix $A$ is known a priori or prespecified. This assumption might be too strong and make the problem of more theoretical interest than very practically relevant, although the motivation of studying anisotropic data seems to arise from the case of realistic datasets as mentioned in the abstract.

Author response:
**"The whole work assumes that the matrix $A$ is known a priori or prespecified"** 

We make the following remarks.  
* In the data generating process, we assume a fixed transformation $A$ (which is not known to the learning algorithm). This problem setting is motivated by the empirical observation that real-world features exhibit low-dimensional structure, which is not captured by many prior works on feature learning assuming isotropic input. By introducing this “generalized” version of the sparse parity problem on anisotropic input, we aim to theoretically study the interplay between structured data and the efficiency of feature learning via gradient descent. To further illustrate the benefit of feature learning in the anisotropic setting, in the revised manuscript, we introduced a stronger kernel lower bound that applies to a wide range of (spiked) anisotropic data. 
* We emphasize that the transformation $A$ is not given to the learning algorithm; instead, the algorithm only has access to i.i.d. observations from the data generating process, and our goal is to characterize how anisotropy (controlled by the matrix $A$) enhances the performance of MFLD. Similar problem setups and motivations have appeared in [Ghorbani et al. (2020)](https://proceedings.neurips.cc/paper/2020/hash/a9df2255ad642b923d95503b9a7958d8-Abstract.html) and [Ba et al. (2023)](https://openreview.net/forum?id=HlIAoCHDWW) for the regression setting. 
* As for implications for real-world data, since the learning algorithm does not know $A$ a priori, intuitively speaking, what Theorem 1 suggests is that neural networks can leverage low-dimensional structure to improve generalization performance, whereas Theorem 2 suggests that the structure of target function can be revealed by the gradient covariance matrix. We believe that these insights can be transferred to real-world settings.

Reviewer comment (Theory):
If we have real data while $A$ is generally unavailable or can only be estimated from data, what can we say about this case with the results of this paper?

Author response:
**"The whole work assumes that the matrix $A$ is known a priori or prespecified"** 

We make the following remarks.  
* In the data generating process, we assume a fixed transformation $A$ (which is not known to the learning algorithm). This problem setting is motivated by the empirical observation that real-world features exhibit low-dimensional structure, which is not captured by many prior works on feature learning assuming isotropic input. By introducing this “generalized” version of the sparse parity problem on anisotropic input, we aim to theoretically study the interplay between structured data and the efficiency of feature learning via gradient descent. To further illustrate the benefit of feature learning in the anisotropic setting, in the revised manuscript, we introduced a stronger kernel lower bound that applies to a wide range of (spiked) anisotropic data. 
* We emphasize that the transformation $A$ is not given to the learning algorithm; instead, the algorithm only has access to i.i.d. observations from the data generating process, and our goal is to characterize how anisotropy (controlled by the matrix $A$) enhances the performance of MFLD. Similar problem setups and motivations have appeared in [Ghorbani et al. (2020)](https://proceedings.neurips.cc/paper/2020/hash/a9df2255ad642b923d95503b9a7958d8-Abstract.html) and [Ba et al. (2023)](https://openreview.net/forum?id=HlIAoCHDWW) for the regression setting. 
* As for implications for real-world data, since the learning algorithm does not know $A$ a priori, intuitively speaking, what Theorem 1 suggests is that neural networks can leverage low-dimensional structure to improve generalization performance, whereas Theorem 2 suggests that the structure of target function can be revealed by the gradient covariance matrix. We believe that these insights can be transferred to real-world settings.

Reviewer comment (Writing):
Typo on page 6, after Proposition 2: “otherwise, we still”.

Author response:
We thank the reviewer for the positive evaluation and insightful comments. We have revised the manuscript and corrected the typos you pointed out.

Reviewer comment (Writing):
Typo on page 6, after Proposition 3: “Under this condition”.

Author response:
We thank the reviewer for the positive evaluation and insightful comments. We have revised the manuscript and corrected the typos you pointed out.