Matching the Statistical Query Lower Bound for $k$-Sparse Parity Problems with Sign Stochastic Gradient Descent

Summary

This paper studies the problem of learning the $k$-sparse parity problem over the $d$-dimensional boolean hypercube, using a two-layer neural network. The main result is that a specific modification of online SGD, called "sign SGD," can learn the k-parity problem $n = \tilde O(d^{k-1})$ samples and a network width of $2^{\Theta(k)}$; the number of scalar queries to the model is $m\cdot n \cdot d = \tilde O(d^k)$, and thus this algorithm matches the SQ lower bound of $d^k$.

Strengths

- The k-sparse parity problem has attracted much recent interest in the deep learning theory community, with many prior works focusing on both lower bounds (via the SQ/CSQ framework) and learning guarantees for neural networks. This paper makes good progress by showing that the SQ lower bound of $d^k$ queries can indeed be achieved by a gradient-based algorithm. - The paper is well written and easy to understand, and I was able to follow the proof sketch. - I have skimmed the appendices, and the proofs appear to all be correct.

Weaknesses

- My main issue with this paper is the choice of the Sign SGD algorithm. Unlike vanilla SGD, Sign SGD is not invariant to the choice of the basis, and is implicitly taking advantage of the structure of the problem being k-parity (where in the ground truth solution each neuron weight is either -1, 0, or 1). For example, the Sign SGD algorithm here would fail if the task was instead k-parity in an unknown basis. Many of the prior works discussed in Tables 1 and 2, however, are rotationally invariant. - In particular, the paper (Glasgow, 2023) had the more ambitious goal to show that vanilla SGD, with minimal modifications, can learn the 2-parity problem. The algorithm in that paper is rotationally invariant and will succeed if the target was XOR in an unknown basis. I thus find the comparisons to (Glasgow, 2023) in the paper, such as "our result matches the best sample complexity for solving the XOR (i.e., 2-parity) problem (Glasgow, 2023)" (line 84) to be a bit misleading. - On a more minor note, I also find the choice of activation $\sigma(z) = z^k$ to be unrealistic, as this again is essentially hardcoding the fact that the target problem is $k$-parity. As seen in Tables 1 and 2, most of the prior work uses the ReLU activation.

Questions

- Can you include more intuition in the main text for Lemma 5.3, specifically why a batch size of $\tilde O(d^{k-1})$ is sufficient for gradient concentration?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

N/A

Summary

This paper considers the problem of learning a $k$-parity function on the $d$-dimensional hypercube using SGD on a two-layer neural network. They consider a specific choice of activation ($\sigma (x) = x^k$) and training dynamics (correlation loss, online sign SGD, fixed second layer weights) and show the following: for $m = 2^{\Theta (k)}$ neurons, batch size $B = \Theta (d^{k-1} \text{polylog} (d))$, then $T = \Theta (log (d))$ steps of online sign SGD achieves small $0/1$-test error. In particular, the total computation time scales as $\Theta (mTB) = \Theta (d^{k} \text{polylog} (d))$. This is noteworthy as it matches the lower bound for learning parities from a large class of learning algorithms, namely statistical queries.

Strengths

- The paper is well written, clear, and a very pleasant read (thank you!). Remarks and discussions provide background and comparison with other works, and help outline the contributions. - The theoretical analysis and statements are clean, and the proof quite compact and elegant, thanks to some simplifying assumptions in the model.. - The fact that SGD on a (quite regular) neural network can match the runtime complexity of the SQ class of algorithms is a nice message. There is a growing literature trying to understand the power of learning with SGD on neural networks, compared to other classes of algorithms, and I believe this paper is a useful addition to the literature.

Weaknesses

I am not sure the paper offers striking novel insights compared to existing literature, or that the analysis is particularly interesting in terms of proof techniques. This is the reason for my rating, but I remain open to changing it during the discussion period. - I think it was already quite widely believed that SGD on NNs can match the SQ lower bound (this is indicated in many papers on learning sparse functions on the hypercube/multi-index functions on Gaussian data). The case of $k$-parities was not written/proven before, and I think this paper makes the necessary effort to settle it. However, I believe some existing analysis can be extended to this case (e.g., [Abbe2023a], [Margalit et al., 2023]), even though they will have drawbacks specific to their choice of simplifications. - The analysis and proofs are quite nice, but it feels like a game of “there exists a specific choice of hyperparameters so that an analysis goes through”. Given that the paper doesn’t put forth particularly novel insights/phenomena, the analysis of a more vanilla training setting (non-monomial activation, more general loss, and standard GD) would have felt more substantial. - There are several questions about this setting that remain completely not understood (e.g., what happens when reusing batches) which further contributes to the impression that this work is incremental.

Questions

About the conclusion on going beyond SQ: the paper “On the power of differentiable learning versus PAC and SQ learning” shows that it is indeed possible with enough machine precision: we can emulate Gaussian elimination. Of course, this is not a very interesting result, as it uses a highly non-standard network that encodes an algorithm in its architecture. There are other noise-tolerant algorithms that do slightly better than SQ, it would be interesting to show that regular architectures prevent their emulation (besides low machine precision).

Rating

6

Confidence

4

Soundness

4

Presentation

4

Contribution

2

Limitations

Yes

Summary

The paper addresses the k-sparse parity problem, a fundamental one in computational complexity and algorithmic theory, by using stochastic gradient descent (SGD) on two-layer fully-connected neural networks. The authors demonstrate that SGD can efficiently solve the problem on a d-dimensional hypercube with a sample complexity that matches the established lower bounds of Statistical Query (SQ) models.

Strengths

The paper provides a solid theoretical foundation by matching the SQ lower bound for learning k-sparse parity problems, which is a significant achievement in this area. Given the recent interest in SQ and CSA bounds and the amount of works on single and multi-index models folllowing the work of Abbe et al, I find these results pertinent.

Weaknesses

The results, both in their scope and sytle, are more aligned with theoretical computer science (CS) than machine learning (ML). While the mathematical contributions are significant, the immediate applicability to practical ML problems is very limited.

Questions

A variant of the parity problem that is not on the hypercube would be if one uses y = sign(z1 z2 z3) with z_mu = x*w_mu, with w_mu a hidden direction on the hypercube. I was wondering how different would the SG bounds and the results of the algorithm be in this case.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Summary

This work considers learning k-sparse parities with two layer neural networks trained with stochastic signed gradient descent, where the second layer is fixed at initialization and the first layer is trained. It is shown that with the activation $\sigma(z) = z^k$, a network with width $2^{\Theta(k)}$ can solve the k-sparse parity problem with $\tilde{O}(d^{k-1})$ samples in $O(\log d)$ iterations, almost matching the Statistical Query lower bound for this problem in terms of the number of queries.

Strengths

The theoretical analysis of the paper is conceptually simple and clean, and there is sufficient discussion on the intuition behind the proof techniques in the main text, easily conveying the ideas to the reader.

Weaknesses

While the goal of this work is to understand the training dynamics of modern neural networks, given the use of a polynomial activation and signed gradients, it is not clear to what extent the intuitions from this work extend to more mainstream algorithms, such as ReLU networks trained with SGD. Further, it is possible to more explicitly discuss some limitations of the work. For example, it appears that the algorithm is not rotationally invariant, therefore is using the knowledge of the coordinate system, in comparison with the work of Glasgow (2023) which does not require this knowledge (for the simpler XOR problem). It might also give readers a more accurate picture if “signed SGD” is used in the title and abstract instead of “SGD”.

Questions

1. Is there a way to trade off compute (iterations/neurons) against samples in upper bounds while maintaining the optimal number of queries? 2. The authors mention non-i.i.d. data as room for future research. There are perhaps even more accessible settings to be investigated. For example, can these results be extended to non-isotropic data, similar to the setting of [1]? 3. The summation notation in Line 183 might need clarification. I might have missed this, but I didn't see the definition of $s$ used here. [1] Nitanda et al. "Improved statistical and computational complexity of the mean-field Langevin dynamics under structured data." ICLR 2024.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Some limitations could be more explicitly discussed as described in the “Weaknesses” section.

Response

Thank you to the authors for your detailed response, and for agreeing to add a discussion of the limitations of SignGD to the paper and revise comparisons to prior work. While I still believe that the choice of algorithm and activation limit the impact of this paper (this sentiment seems to be shared by other reviewers), I agree that matching the SQ lower bound with a gradient-based algorithm on a standard neural network is a nice contribution, and I thus raise my score to "Weak Accept."