On the Complexity of Learning Sparse Functions with Statistical and Gradient Queries

Summary

This paper studies the problem of learning juntas from the perspective of restricted statistical query algorithms. Concretely, for some d-dimensional product distribution \mu_x^d and some conditional distribution \mu_{y|z}, the learner is given access (via samples or, in this paper, certain queries) to the distribution of (y,x) where x ~ \mu_x^d and y|x has distribution \mu_{y|z} where z is some tuple of P unknown coordinates of x. That set of coordinates that influence y is called the support of the distribution, and the goal is to learn the support. A Q-restricted statistical query algorithm is one that only makes queries to E[phi(y,x)] for functions phi in Q. Examples include SQ and CSQ; this paper introduces, for any fixed loss function \ell, another class DLQ_\ell, where the allowable queries are gradients of \ell(f(x,w),y) (w.r.t. w) for arbitrary functions f. The main result of this paper is a characterization for any {SQ, CSQ, DLQ} of the query complexity of the support learning problem, when P, \mu_x, \mu_{y|z} are fixed but d is growing: up to lower-order terms it has the form d^k for a "leap exponent" k that can be determined from \mu_x,\mu_{y|z}, and the query class Q. As a consequence, the paper gives separations and equivalences among these classes: notably, CSQ is equivalent to DLQ for \ell_2 loss (Proposition 6.2a), but SQ is equivalent to DLQ for \ell_1 loss (Proposition 6.1a). Additionally, the paper has some results showing that if one performs SGD on a two-layer neural network in the mean-field regime and with loss function \ell, then DLQ_\ell (rather than SQ or CSQ) is what generically characterizes success of the dynamics (more precisely, Leap(DLQ_\ell) = 1 is what characterizes success with O(d) samples).

Strengths

- This paper extends prior work by Abbe, Boix-Adsera, and Misiakiewicz to a more general setting (specifically, the prior work considered the special case of CSQ and the uniform distribution on the hypercube). While CSQ has long been motivated by the fact that it captures gradient descent, a key takeaway of this paper is that this is only for \ell_2 loss, and for other losses a different characterization is needed. This seems like a useful observation for the community. - The paper also contrasts the power of adaptive queries versus non-adaptive queries, and gives a characterization of query complexity for both. - The paper is written carefully and rigorously.

Weaknesses

- I spent a while trying to parse the notation for the proof of Theorem 5.1(a). Ultimately I think it is not super complicated, but some intuitive, high-level explanation of the construction would have been very helpful in addition to the formal definitions. Also, the proof outline explains the boilerplate elements that appear in any SQ lower bound argument but gives no intuition/proof sketch for the crucial and presumably novel inequality (25).

Questions

- Does the characterization extend beyond just {SQ, CSQ, DLQ} to any query class Q, or if not what breaks down? - The equivalence between SQ and DLQ for ell_1 loss rests on a universal approximation result. I am curious if there is intuition for why the analogous approximation result for ell_2 does not hold? Is it just that there is no non-linearity in the gradient? If this is the intuition, it seems to me that in some sense "most" loss functions besides ell_2 ought to induce an equivalence between DLQ_ell and SQ? Proposition 6.2(c) is interesting because it gives a loss function that does not degenerate to SQ nor CSQ: is there no universal approximation here because it's a polynomial? - I assume that CSQ and SQ provably *do not* capture SGD on neural networks, but I didn't see a formal statement to this effect; does this follow from the separations in Section 6? - This is a minor technical confusion: in lines 505-506, I see why the null responses are compatible with nu^\sigma, but since nu_{S^\star} is not in the hypothesis class (since it's a partially "decoupled" distribution), I don't see why this immediately gives failure of the learning algorithm? I would have thought one needs two hypothesis in the hypothesis class to be indistinguishable. Of course I think this should also follow from (26) but I just wanted to check if I am misunderstanding.

Rating

6

Confidence

3

Soundness

4

Presentation

3

Contribution

3

Limitations

Yes

Summary

This submission studies the computational complexity of learning low-dimensional functions (juntas) and examines a new type of statistical query termed DLQ, which captures the difficulty of learning using gradient-based queries. The authors identify quantities that dictate the computational lower- and upper-bounds based on the properties of the link function and show that the DLQ complexity may match either SQ or CSQ depending on the loss function. This generalizes the earlier leap exponent in (Abbe et al. 2023), which is specific to CSQ models. Finally, in a simplified hypercube setting, the authors show that a mean-field neural network can learn functions with DLQ-leap exponent 1 using $n=O(d)$ samples, which matches the DLQ complexity.

Strengths

This submission makes a number of contributions to the learning of low-dimensional functions on product distributions: * The proposed DLQ "interpolates" between SQ and CSQ depending on the structure of the loss function. This provides insights into the class of low-dimensional functions that can be learned by gradient descent algorithms. * The authors also distinguish between adaptive and non-adaptive queries, providing a separation between algorithms with one gradient step versus consecutive steps. * The mean-field analysis extends (Abbe et al. 2022) to different loss functions and confirms that gradient descent training of a two-layer neural network learns juntas with DLQ leap exponent 1.

Weaknesses

I have the following concerns/questions: 1. The class of functions with the prescribed DLQ leap exponent is rather opaque, even in the simple hypercube setting. For Gaussian single-index models, the SQ-leap exponent is completely characterized in (Damian et al. 2024), and it is known that all polynomial links have SQ exponent 1 or 2. But in the current submission, it is not clear what functions have DLQ leap exponent 1. The authors provided one example in Proposition 6.2(c) for one specific loss and link function, but there is no general principle that determines the DLQ hardness. Intuitively, if the output is discrete, then the detectable set for SQ versus CSQ may not differ too much, so it is not easy to see that the SQ/DLQ-learnable function class is significantly larger on hypercube input. Can the authors clarify? 2. The current submission does not take the rotational symmetry into account, and the query dependence in the lower bound is polynomial as opposed to logarithmic in the Gaussian case (Damian et al. 2022). For algorithms that are not axis-aware, such as SGD from rotationally invariant initialization, do we expect a lower bound that is stronger in terms of query dependence? 3. Assumption 2.1 needs more explanation beyond "*it always holds when $\mathcal{X}$ is discrete*". Can the authors comment on how to translate this assumption to a condition on the link function when the input is continuous, such as Gaussian? 4. How do the authors relate the number of queries in the SQ lower bound to the sample size in the mean-field analysis? --- The authors’ response addressed most of my concerns. I have updated my evaluation accordingly.

Questions

See Weaknesses.

Rating

7

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

See Weaknesses.

Summary

This paper studies the complexity of learning sparse functions (i.e., juntas) using statistical and gradient queries. For that, it introduces Differentiable Learning Queries (DLQ) that can be seen as a generalization of correlation statistical queries beyond square loss. With the notion, the paper gives a characterization of the query complexity for learning the support of a sparse function over generic product distributions. The characterization is in terms of leap exponent introduced from Abbe et al earlier. Finally, the paper also introduces cover exponent that characterizes non-adaptive query complexity.

Strengths

The paper studies a well-motivated question of query complexity of learning juntas, a long established question in ML theory. The paper builds upon the prior work of Abbe et al and provides a generalization beyond CSQ and square loss function. I find the results strong at a technical level. To my knowledge, the main theorems are novel. The work proves general relations between SQ, CSQ and DLQ. The paper also provides experimental validation of the theoretical results, though it's somewhat limited. The paper is well written.

Weaknesses

Figure 1 is an experimental study. I believe it'd be nice to expand this into a small section that extends beyond learning y_2 (2) and using the 2 loss functions here. I also can't find more details: what is the architecture of this two-layer net? Abbe et al [2023] has some nicer numerical results. I'd recommend check it out. Compared with Abbe et al 2023, the paper only tackles the case when the target function is junta. It doesn't extend much beyond sparse functions. Minor comments follow: * Line 18: remove “that is” * Line 24: it should be clarified that Abbe et al. [2022, 2023] studies learning juntas using neural networks via SGD

Questions

Any upper bound results or conjectures on how to match the lower bound in section 5?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The paper lists several limitations.