Robust Feature Learning for Multi-Index Models in High Dimensions

We thank the reviewer for their careful assessment and valuable feedback, and address their comments and questions below.

Tight characterization of robust test risk:

While we agree with the reviewer that a tight characterization of the robust test risk would be highly valuable, we would like to highlight that even in the standard setting there is no tight characterization of the risk under feature learning, unlike its random feature counterpart. This is due to the inherent complexity of analyzing feature learning, and it is unlikely that the current techniques will be able to provide a sharp characterization of test risks in feature learning, especially in the distribution-free setting of our Theorems 4-7.

Data assumption: having anisotropic covariance and non-Gaussian data:

As we show in the general response above, it is possible for anisotropic covariance matrices to satisfy Assumption 1. Further, Assumption 1 is necessary for Theorem 1, and without it the optimal representation may no longer be $U$. Based on our discussion in the general response, it is an interesting direction to consider other algorithms that achieve optimal performance beyond Assumption 1.

Please note that Theorem 4-7 only require Assumption 1 and do not depend on Gaussianity or isotropic covariance. In particular, since gradient-based feature learning algorithms are known for anisotropic covariance (e.g. [1,2]), and for non-Gaussian distributions (e.g. [3]) one can combine these with our results to obtain end to end robust learnability guarantees.

Limitation of $\ell_2$ perturbation:

As we discuss in the general response above, the $\ell_2$ attack norm is also necessary for Theorem 1, as the optimal representation may change with the $\ell_\infty$ attack. This is another interesting open direction, and our discussion in the general response above can provide a starting point for developing algorithms in that setting.

Assumption on the perturbation budget:

We would like to highlight that scaling $\varepsilon$ with $\sqrt{d}$ in this setting will be undesirable, since it will encourage the optimal predictor given by $x \mapsto h(\langle u, x\rangle)$ to have a fixed prediction over an input range which grows with high dimensions. This is because $\langle u, x\rangle$ is $\mathcal{O}(1)$ with high probability, yet $\langle u, \delta\rangle$ can be $\mathcal{O}(\sqrt{d})$ where $\delta$ is the perturbation. Therefore, the prediction of $h$ will collapse to a constant regardless of input in high dimensions, which is undesirable. We believe the natural scaling of $\varepsilon$ in our setting is constant with respect to dimension.

Only polynomial and ReLU activations are considered:

We remark that in principle our results should be applicable to other activation functions. However, please see the general response above for a discussion on the technical difficulties of considering other activations.

We thank the reviewer for their careful assessment and valuable feedback, and address their comments and questions below.

Technical challenges in considering activations beyond ReLU and polynomial:

We refer to the general response above for a discussion on these challenges.

Restructuring the proof of Theorem 1 and including proof outline for other results:

We thank the reviewer for this suggestion, and will restructure the proof of Theorem 1 and add a high-level overview of the proof technique of other results in the final version.

The dependence of the optimal adversarial risk on dimension:

Assuming the zero predictor is in $\mathcal{F}$, one can notice that $\mathrm{AR}^* \leq \mathbb{E}[y^2]$ which is independent of dimension. Moreover, even in an unnatural setting where $\mathrm{AR}^*$ grows with dimension, our final bounds will be dimension-independent, since in that case $\tilde{\epsilon} = \epsilon$ in all theorem statements.

Lack of definition for Landau notation:

We thank the reviewer for expressing their concern about the notation, in the final version, we will provide further clarification in the notation section of the manuscript. All our bounds are non-asymptotic, and the Landau notation only hides absolute constants (or polylogairthmic terms when used with tilde). We also assume certain quantities such as $k$ are absolute constants in this paper, and we mention such assumptions when introducing those quantities.

Guarantees from robust training:

Currently there are no theoretical guarantees for the representation learned by adversarial training on the first layer in our setting. We believe investigating the sample and computational complexity arising from this adversarial training is an interesting open direction.

Curse of dimensionality in $k$:

The reviewer’s interpretation of the exponential dependence on $k$ as a curse of dimensionality for approximating functions in a $k$-dimensional subspace is indeed correct. The reason this dependence does not show up when considering polynomials is that the space of polynomials of bounded degree is a much smaller space. For example, the space of $k$-index polynomials of degree $p$ is a finite-dimensional space with dimension $k^p$ which only scales polynomially with $k$. The space of all (pseudo-Lipschitz) $k$-index functions is infinite-dimensional, and approximating it with a finite-dimensional space using an $\epsilon$-covering requires $(1/\epsilon)^\mathcal{O(k)}$ many parameters, exponential in $k$.

The choice of activation from Lee et al., 2024:

We thank the reviewer for pointing out this potential confusion. The algorithm in Proposition 8 makes use of a special polynomial activation with a sufficiently large degree. In the revised version, we have provided further clarification in Corollary 9 and 11 in terms of the assumption on activation.

Norm of $x$ and scale of $\varepsilon$:

When considering weights on $\mathbb{S}^{d-1}$, the natural scaling of $\Vert x \Vert$ is $\sqrt{d}$ since it leads to $\langle w, x\rangle$ being of order 1. This is indeed the same reason why $\varepsilon$ is assumed to be of order 1. We will further clarify this in the final version of the manuscript.

Further recommendations for clarity:

We thank the reviewer for their recommendations to improve clarity. We will incorporate them in the final version. In particular, Definitions 2 and 3 indeed assume the rows to be unit-norm.

We would be happy to clarify any concerns or answer any questions that may come up during the discussion period.

We thank the reviewer for their careful assessment and valuable feedback, and address their comments and questions below.

The validity of Theorem 1 beyond Assumption 1:

As we show in the general response above, Assumption 1 is necessary for Theorem 1, and without it the statement of the theorem may no longer hold. However, as also shown in the general response, in certain settings one can still identify closed form transformations on $U$ that lead to optimal presentations provided that they are able to estimate the covariance matrix. Developing algorithms that are optimal beyond Assumption 1 based on this observation is an interesting direction for future work.

Using Theorems 4-7 with Anisotropic Gaussians:

We demonstrate in the general response above that Gaussians with certain anisotropic covariance matrices can satisfy Assumption 1, in which case we can use Theorems 4-7.

Further, we thank the reviewer for noticing the typo that we fixed in the revised version.

We would be happy to clarify any concerns or answer any questions that may come up during the discussion period.

We thank the reviewer for their careful assessment and valuable feedback, and address their comments and questions below.

Exponential dependency on $k$:

We would like to highlight that without additional assumptions, the exponential dependency on $k$ is predicted by minimax lower bounds [1, Section 2.5], therefore unfortunately not improvable in general.

Improving Notation:

We thank the reviewer for their suggestion, and will improve the notation for an easier understanding in the final version of the manuscript.

The meaning of $\beta$ in the SFL definition:

A lower bound on $d\mu/d\tau_k$ is asking for the stochastic weights sampled from $\mu$ to be sufficiently present across all directions on the low-dimensional sphere $\mathbb{S}^{k-1}$. Note that when $d\mu/d\tau_k \geq \beta$, then if the uniform measure assigns some probability $\tau_k(A)$ to some $A \subseteq \mathbb{S}^{k-1}$, then $\mu(A) \geq \beta \tau_k(A)$, thus we are sufficiently covering $A$ when $\beta > 0$.

Validity of Assumption 1:

We refer to the general response for examples of anisotropic/low-rank covariance structure where Assumption 1 is verified. Please note that in the example provided in the general response, both $\Sigma_1$ and $\Sigma_2$ can be low-rank, leading to a low-rank overall covariance matrix $\Sigma$. However, we agree with the reviewer in that the independence between $Ux$ and $\tilde{U}x$ is still necessary for Theorem 1, as we show in the general response that without this assumption, the theorem may no longer hold.

We would be happy to clarify any concerns or answer any questions that may come up during the discussion period.

References:

[1] A. Tsybakov. “Introduction to Nonparametric Estimation.” 2009.

We thank the reviewer for their careful assessment and valuable feedback, and address their comments and questions below.

Discussion on extension to $\ell_\infty$ norm:

We refer to the general response above about a discussion on extension to $\ell_\infty$ attack norm, and the necessity of the $\ell_2$ assumption for our theoretical result.

Going from squared to logistic loss:

We agree with the reviewer that our proof technique relies crucially on a decomposition of the squared loss, and considering other loss functions is an important open direction. We have added a simple experiment on the MNIST dataset in Appendix D which demonstrates that the intuitions from our theory carry over to the case of cross entropy loss.

Improving the structure of the paper:

We thank the reviewer for raising their concern about the structure of the paper, and we will improve the presentation in the final version of the manuscript.

Typos in the proof of Lemma 16:

We thank the reviewer for noticing the typos. We have fixed them in the revised version, and made the proof more easily readable.

Having W fixed at SD training vs. continuing with AD training:

The additional improvement from continuing with adversarial training in the first layer weights comes from the fact that given the finite number of samples in the standard training phase, it is only able to approximately recover the low-dimensional information in the target function. Further training allows exploiting further samples and more information about the low-dimensional projection.

Why Algorithms 2 and 3 are not used in the numerical experiments:

Our goal in the numerical experiments section was to show that one does not need to consider the contrived feature learning algorithms of prior works to take advantage of the observation in Theorem 1, and even simpler algorithms such as vanilla SGD that are used in practice can be sufficient. We will further clarify this in the final version.

We would be happy to clarify any concerns or answer any questions that may come up during the discussion period.

We thank the reviewer for their careful assessment and valuable feedback, and address their comments and questions below.

Lack of extensive experiments:

We thank the reviewer for this suggestion. We have added a simple experiment on MNIST in Appendix E to demonstrate that the intuitions from our theoretical results carry over to settings with neural networks with multiple layers, cross entropy loss, and $\ell_\infty$ norm for attack. We must highlight however that this experiment is not meant to perform an extensive study. In particular, we retrain all layers from standard training, as in this case it is not entirely clear which layers are performing representation learning. We believe that in practice, retraining all layers will be useful nevertheless, since the algorithm will not be able to perfectly learn the optimal representation during standard training due to finite samples and compute. It will be an interesting direction for future work to perform an extensive empirical study over modern architectures to consider the effect of fixing certain layers from standard training and performing adversarial training on the rest.

Confusion in the statement of Theorem 1:

Thank you for your feedback. We will accordingly rewrite the statement of Theorem 1 in the final version to make it more readable.

Discussing relevant works:

Thank you for pointing out these relevant works. Collins et al., 2024 consider learning multiple tasks of the multi-index model form using the feature learning framework. Wang et al., 2024 consider learning a generalization of the single-index model where instead of having a linear feature, we have a degree $k$ polynomial of the input as the feature. We believe that intuitions from our work can extend to these settings as well, i.e. one can combine standard representation learning with robust fine-tuning. We have added a discussion of these works in the revised version of our manuscript.

Is the first inequality in the proof of Theorem 1 an equality?:

The inequality can be strict because we are moving the maximum outside of the expectation. Equality would only hold when $f$ is a minimizer of $\mathrm{AR}$ in $\mathcal{F}$, and additionally $h(U\cdot) \in \mathcal{F}$.

How did you select the direction of $u$ in the experiments?:

In this section, we look at single-index models of the form $x \mapsto g(\langle u, x\rangle)$. We randomly sample $u$ from the unit sphere in $\mathbb{S}^{d-1}$. Since we use the same seed in all experiments, this means $u$ is the same across experiments. We will make this more explicit in the final version.

Actionable insights from this work:

Indeed, while our focus is on the theoretical part, one actionable insight from our work is that standard training yields a better initialization than the usual random initialization for adversarial robustness. Furthermore, in certain settings, fine-tuning the final layers of a network might be enough to achieve adversarial robustness, although this still requires a full backpropagation over the network to generate adversarial examples. We will highlight these insights in the final version of our manuscript.

We would be happy to clarify any concerns or answer any questions that may come up during the discussion period.

We thank the reviewer for their careful assessment and valuable feedback, and address their comments and questions below.

Technical challenges and novelties:

Unexpectedness of the result: One reason for the unexpectedness of the result of Theorem 1 is that we simply do not always have $U$ as an optimal low-dimensional projection of the input, unless we take into account additional assumptions. Some examples where $U$ is not the optimal direction are given in the general response above, which are due to using $\ell_\infty$ norm or having certain anisotropic covariances. This is in contrast with the standard setting, where $U$ always provides an optimal projection regardless of additional assumptions. This highlights the importance of the assumptions and the proof of Theorem 1, and that one can not always take the statement of Theorem 1 for granted without taking into account the assumptions. We will add this discussion to the final version of the manuscript.

Technical challenges: We agree with the reviewer that once the rather unexpected result of Theorem 1 is proved, it is expected that we can end up with dimension-independent bounds for sample and computational complexity. However, we would like to highlight some other aspects of our analysis here that differs from that of prior works:

• We introduce DFL and SFL as a formalism for feature learning that unifies the previous approaches to establishing alignment with target directions and performing dimension reduction. This might be of interest for future works since one only needs to prove that the DFL or SFL assumption holds in some setting, and then they can readily use our approximation and generalization guarantees to end up with learnability results.
• The Rademacher complexity analyses for two-layer neural networks are typically performed using Talagrand’s contraction lemma, see e.g. [1, Proposition 7]. This approach is not compatible with the adversarial loss, and we instead take the covering number approach with Maurey’s sparsification lemma as used in [2] for our Rademacher complexity analysis.
• Our approximation analysis goes beyond that of the recent literature such as Damian et al., 2022 in the sense that we cover pseudo-Lipschitz functions with a universal approximation argument inspired from [1]. However, the universal approximation argument of [1] requires trained biases, whereas we prove our results under random biases, which helps keep the training of the second layer a convex problem.
  We will further discuss these differences in the final version of our manuscript.

Assumptions on $\mathcal{F}$ in Theorem 1:

Theorem 1 does not require any particular assumptions on $\mathcal{F}$, beyond the ones that are required for the mathematical objects in the theorem to make sense such as measurability and finite expectation.

Is $h$ assumed to be in $\mathcal{F}$:

$h(U\cdot)$ does not need to be in $\mathcal{F}$. However, if $h(U\cdot)$ is additionally in $\mathcal{F}$, such as in the examples, then $\mathrm{AR}(h(U \cdot)) = \mathrm{AR}^*$.

Constraint on the norm of $w_i$ in DFL:

Yes, DFL and SFL both assume the weights $w_i$ to be on the unit sphere $\mathbb{S}^{d-1}$, which is also the setting in the rest of the paper, including Theorem 4. We thank the reviewer for pointing out that we had missed to mention this, and have fixed it in the revised version.

We would be happy to clarify any concerns or answer any questions that may come up during the discussion period.

References:

[1] F. Bach. “Breaking the Curse of Dimensionality with Convex Neural Networks.” JMLR 2017.

[2] T. Zhang. “Covering number bounds of certain regularized linear function classes.” JMLR 2002.