Robust Feature Learning for Multi-Index Models in High Dimensions

TL;DR

This paper demonstrates that in high-dimensional multi-index models, robust feature learning can be achieved efficiently by first learning features normally and then tuning a linear layer for adversarial robustness, with sample complexity independent of dimension.

Contribution

It proves that hidden directions in multi-index models provide optimal low-dimensional projections for adversarial robustness, enabling robust learning with minimal additional samples.

Findings

Robust feature learning is as straightforward as standard learning.

Sample complexity for robustness does not depend on ambient dimension.

Hidden directions serve as optimal projections for adversarial robustness.

Abstract

Recently, there have been numerous studies on feature learning with neural networks, specifically on learning single- and multi-index models where the target is a function of a low-dimensional projection of the input. Prior works have shown that in high dimensions, the majority of the compute and data resources are spent on recovering the low-dimensional projection; once this subspace is recovered, the remainder of the target can be learned independently of the ambient dimension. However, implications of feature learning in adversarial settings remain unexplored. In this work, we take the first steps towards understanding adversarially robust feature learning with neural networks. Specifically, we prove that the hidden directions of a multi-index model offer a Bayes optimal low-dimensional projection for robustness against $\ell_2$-bounded adversarial perturbations under the squared loss, assuming that the multi-index coordinates are statistically independent from the rest of the coordinates. Therefore, robust learning can be achieved by first performing standard feature learning, then robustly tuning a linear readout layer on top of the standard representations. In particular, we show that adversarially robust learning is just as easy as standard learning. Specifically, the additional number of samples needed to robustly learn multi-index models when compared to standard learning does not depend on dimensionality.