Improved statistical and computational complexity of the mean-field Langevin dynamics under structured data

We thank the reviewer for the positive evaluation and helpful feedback. We address the technical comments below.

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.

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.

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., Bordelon and Pehlevan, Akiyama and Suzuki, Telgarsky, Tian.

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. 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) and Ba et al. (2023) 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.

We thank the reviewer for the positive evaluation and insightful comments. We have revised the manuscript and corrected the typos you pointed out. We address the technical points below.

"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) and Ba et al. (2023) 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.

"I also wonder if the problem can be tackled with a preconditioned version of MFLD"

Thank you for the interesting suggestion. At a high level, there is indeed a parallel between the coordinate transformation and preconditioning. However, we note that naively preconditioning the entire update yields the stationary distribution unchanged (though the optimization speed might improve), whereas our coordinate transform on the input feature yields a different stationary distribution on the parameter space, which enables us to establish a better LSI constant and generalization. It is an interesting question of whether we should apply the preconditioner to the gradient of the loss, the regularization, or the Gaussian noise separately, as here different design choices would result in different stationary distributions of MFLD. We intend to investigate this perspective in future work.

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