Learning Neural Networks with Distribution Shift: Efficiently Certifiable Guarantees

Thank you for your time and feedback!

Regarding the runtime in unbounded case, note that there is evidence that the dependence on the parameter $\epsilon$ must be exponential, even without distribution shift and assuming that the training marginal is Gaussian (see lines 198-204 and Diakonikolas et al., 2020b; Goel et al., 2020b). In the bounded case, it is an interesting open direction to achieve better runtime, but this is beyond the scope of this work.

Overall, obtaining provable guarantees for learning often requires algorithms with high time complexity. The value of such results, especially in the TDS setting where we make no assumptions on the test distribution, is that one can be confident in the performance of the algorithms. This is crucial in critical settings where errors might lead to catastrophic results, and heuristic approaches are not appropriate.

Regarding the limitations of our work, we thank the reviewer for their suggestions and we will consider including them in future versions of our paper, to make it more accessible to a general audience. However, we stress that our results are a part of a long line of works on computational learning theory. In this context, obtaining a polynomial-time algorithm (even if the exponent is large) is the gold standard for problems where only exponential time results were previously known. Reducing the exponent in these cases is usually reserved for follow up work.

Questions:

• We believe that an important open question is whether one can remove the hypercontractivity assumption from the results on bounded training distributions or show that this assumption is necessary. In the classical setting (where there is no distribution shift), similar results can be obtained without assuming hypercontractivity. It is not clear whether the need for hypercontractivity can be alleviated with a different analysis, or whether it is an inherent limitation of the Kernel method for TDS learning.
• Our tests will actually accept whenever the training and test distributions have matching low-degree moments. In order to obtain testers that will accept when the two distributions are close in statistical distance, one might combine our results with prior work on TDS learning (Goel et al., 2024). See point 5 in our response to reviewer ovXB for more details.
• You are correct that we only care about the test error. However, in order to obtain a bound on the test error and since we do not have access to labels from the test distribution, we use the training error as a proxy. The error of this proxy depends on the joint error $\lambda$ (see line 058) plus some excess error depending on the amount of distribution shift (usually measured in terms of some notion of distance between distributions, like the discrepancy distance, see, e.g., Mansour et al., 2009). Our tests essentially certify that the final excess error term is small.

We wish to thank the reviewer for their feedback and for appreciating our work. We would like to stress that this is a theory paper and the main contribution is to a long line of theoretical research on provably efficient learning algorithms. We believe that experimental evaluation of our algorithms is an interesting and promising direction but more suitable for a separate project.

We thank the reviewer for their feedback.

We would first like to clarify that we do not make any structural assumptions on the noise on the labels (the reviewer seems to suggest that the label noise is Gaussian).

Assumptions on the training marginal. We stress that our results not only capture bounded distributions, but also arbitrary subgaussian distributions (Section 4), which is a significantly milder assumption. For bounded and hypercontractive distributions, we essentially match the best known guarantees for the standard setting (without distribution shift). For unbounded distributions, our results give new bounds even for the standard setting, where all of the known results require much stronger distributional assumptions (like Gaussianity). Note also that there is strong evidence that there are no dimension-efficient learning algorithms with provable guarantees under no distributional assumptions even for a single ReLU or sigmoid neuron (see, e.g., [DKMR’22]).

In this context, our work provides a positive message: the algorithmic approaches proposed in standard learning theory (like Kernel methods and polynomial regression) can be enhanced with appropriate testers (like Algorithm 1) to provide strong guarantees even in the presence of distribution shift.

[DKMR'22] Diakonikolas, I., Kane, D., Manurangsi, P., & Ren, L. (2022, May). Hardness of learning a single neuron with adversarial label noise. In International Conference on Artificial Intelligence and Statistics (pp. 8199-8213). PMLR.

Clarity questions. Thank you for your suggestions to improve the presentation of our work. However, we kindly disagree that our paper lacks clarity in general (see also the Strengths section of Reviewer ovXB).

• In section 1.2 we provide a high-level explanation of the main technical challenges and we provide more details in section 3 (particularly, subsection 3.1) and section 4. Regarding Algorithm 1 (for bounded distributions), as we stress in lines 110-117, the straightforward idea of matching pairwise correlations of the expanded features provides suboptimal runtime. We instead follow a data-dependent testing approach (see, e.g., lines 118-126 and section 3.1), which crucially uses the kernel trick in a non-standard way. In particular, we use test examples to form a data-dependent feature map $\phi$ (see lines 298-318 for the definition of $\phi$) and test the pairwise correlations of the corresponding features.
• For an explanation on the approximation theory ideas for our applications on neural networks with unbounded training marginals, see lines 474-485.
• We explain how we use the polynomial kernels in lines 127-130. However, motivated by your comment, we will consider putting this explanation in the beginning of section 1.2 for future versions of our work. The definition of the concise feature map $\phi$ is given in lines 298-318. The way it is used in Algorithm 1 is through matrices $\hat\Phi, \hat \Phi’$ (see lines 361-371 for the relationship between $\phi$ and $\hat\Phi, \hat\Phi’$).

We wish to thank the reviewer for appreciating our work and for their thorough feedback!

Question 1. One motivating example for the setting we consider is when one is, say, given a model that is trained on data from some hospital A in some part of the world and then is asked whether the predictions of the model can be trusted in order to decide the correct treatments for the patients in some other hospital B elsewhere. In this case, the learner does have unlabeled examples from hospital B: each patient has a file with the results of their exams, as well as any experienced symptoms. The crucial question is whether we can trust our model to label patients from hospital B, and knowing even when to reject is very useful. In general, TDS learning is particularly important for critical tasks, where obtaining even a small number of labels is very expensive.

Question 2. Thank you for catching this, we have included the full statement in the revision.

Question 3. Note that our algorithms only require uniform convergence with respect to the training distribution, which is structured. We do not require uniform convergence with respect to the test distribution, which can be arbitrary. This enables, for example, a sample complexity upper bound which is independent from the input dimension $d$ in the case of bounded training marginals (see Algorithm 1). We will highlight this in future revisions of the paper.

Relaxing this requirement (uniform convergence with respect to the training distribution) would be an interesting result, as it would imply an information-theoretic separation between TDS learning and a closely related learning primitive (testable agnostic learning, see Rubinfeld & Vasilyan, 2023). In particular, Rademacher complexity is known to characterize the sample complexity of testable agnostic learning for any function class (see Gollakota et al., 2023). However, proving a lower bound for the sample complexity of TDS learning in terms of the Rademacher complexity or showing that no such bound exists remains open.

Question 4. We use the overline notation to refer to a set of labeled examples and distinguish it from its unlabeled counterpart ($\{(x,y)\}$ versus $\{(x)\}$). It seems that the notation is consistent in lines 330, 333 (Algorithm 1).

Question 5. The reviewer is raising a very interesting question. Our tests are formally guaranteed to accept only when the test distribution is equal to the training distribution. However, there are recent techniques in TDS learning which yield tests that accept much wider classes of test distributions. For example, Goel et al. (2024) showed that, in the case of binary classification, it is possible to accept all distributions that are close to the training distribution in total variation distance. We expect this result to extend to our setting with an appropriate modification of our algorithms, but we did not work it out. Moreover, Chandrasekaran et al. (2024) showed that, for some function classes like halfspaces, there are TDS learners that accept whenever the test marginal is sufficiently structured (i.e., has some mild concentration and anti-concentration properties). Such distributions can even be far from the training marginal in statistical distance. It is interesting to explore whether such results can also be obtained in the regression setting.

Finally, note in our algorithms, the candidate output hypothesis does not depend on the examples from the test distribution. This means that if we wanted to test a number of different test distributions, we would not need to train the model from scratch, but rather only use test examples to decide whether to accept or reject.