Robust Generalization of Quadratic Neural Networks via Function Identification

1 INTRODUCTION . Recent work has demonstrated that neural networks are not robust to shifts in the underlying data , including both distribution shifts ( i.e. , where the data comes from a new distribution independent of the neural network parameters ) ( Hendrycks & Dietterich , 2019 ; Taori et al. , 2020 ) as well as adversarial shifts ( i.e. , where the shift can depend on the neural network parameters ) ( Szegedy et al. , 2013 ) . Accordingly , there has been a great deal of interest in better understanding why neural networks fail to be robust ( Tsipras et al. , 2018 ; Ilyas et al. , 2019 ) and on improving robustness ( Goodfellow et al. , 2014 ; Raghunathan et al. , 2018 ; Cohen et al. , 2019 ) . From the perspective of learning theory , there is little reason to expect neural networks to be robust , since generalization bounds typically assume that the test examples are from the same distribution as the training examples . PAC-Bayesian generalization bounds allow for a limited amount of robustness , but only if the support of the target distribution q is contained in that of the source distribution p , since it requires that the KL divergence DKL ( q ‖ p ) is finite . Yet , distribution shifts ( Hendrycks & Dietterich , 2019 ) often shift probability mass to inputs completely outside the source distribution . Instead , the reason we might expect neural networks to be robust to these shifts is that humans are robust to them ; for instance , small pixel-level shifts considered in adversarial examples are typically unnoticeable to humans , yet these shifts can move the image completely off of the distribution of natural images . This fact indicates a gap in our theoretical understanding of neural networks . In particular , the key question is understanding settings under which we may expect neural networks to be robust to distribution shifts that are “ large ” ( e.g. , in terms of KL divergence ) . We study a strategy for closing this gap based on the statistical concept of identifiability ( Hsu et al. , 2012 ) . At a high level , this concept assumes that the true model belongs to the model family ; then , in the limit of infinite training data , the learning algorithm can exactly recover the parameters of the true model . For instance , in linear regression , the data is generated according to the model y = 〈θ∗ , x〉+ξ , where ξ is σ-subgaussian noise . Then , under mild assumptions on the training data Z = ( X , Y ) , the ordinary least squares ( OLS ) estimator θ̂ ( Z ) recovers the true parameter—i.e. , in the limit of infinite data , θ̂ ( Z ) = θ∗ . With finite samples , OLS satisfies high-probability convergence rates of the form ‖θ̂ ( Z ) − θ∗‖2 ≤ . ( 1 ) The connection to robustness is that if ( 1 ) holds , then for any input x such that ‖x‖2 ≤ xmax , we have |〈θ̂ ( Z ) , x〉 − 〈θ∗ , x〉| ≤ ‖θ̂ ( Z ) − θ∗‖2 · ‖x‖2 ≤ · xmax . ( 2 ) Thus , for any distribution q ( x ) with support onB2 ( 0 , xmax ) = { x ∈ X | ‖x‖2 ≤ xmax } , θ̂ ( Z ) obtains bounded error—i.e. , we have Eq ( x ) [ ( 〈θ̂ ( Z ) , x〉 − 〈θ∗ , x〉 ) 2 ] ≤ 2x2max with high probability . Thus , a natural question is whether we can obtain similar kinds of parameter identification bounds for neural networks . A key complication is that practical neural networks are often over-parameterized , possibly to facilitate optimization ( Du & Lee , 2018 ; Jacot et al. , 2018 ) . In this setting , identification is impossible , since multiple parameters can yield the same model . Nevertheless , it may be possible to obtain bounds of the form ( 2 ) —in particular , even if we do not recover the true parameters θ∗ , we can still recover the function fθ∗ ( x ) . We refer to this notion as function identification . Furthermore , we show that quadratic neural networks satisfy function identification bounds under mild conditions . To demonstrate its utility , we show how function identification can be leveraged to obtain regret guarantees for a bandit ( Rusmevichientong & Tsitsiklis , 2010 ) where each arm is a quadratic neural network . Linear bandits fundamentally involve covariate shift since their “ covariates ” are arms , which are adaptively chosen through the learning process as a function of past observations ; thus , existing approaches have all operated in the setting where there is a unique and identifiable global minimizer . Similarly , we build on recent work proving bounds on transfer learning in the setting of bounded label shift and unbounded covariate shift ( Bastani , 2020 ; Xu et al. , 2021 ) ; again , we show that we can leverage function identification to easily transfer learn quadratic neural networks . Our results suggest that one strategy for improving robustness of neural networks is to design models that can represent the true data generating process . However , doing so can be challenging due to the complexity of most real-world data generating processes mapping covariates to labels—e.g. , mapping natural language to semantic meaning or images to object detections and labels . As a consequence , we study how our results can connect to neural module networks ( Andreas et al. , 2016 ) , which are designed to break down complex tasks into smaller ones that can each be solved by a neural network . For instance , we may break down the task “ count the number of red balls in image x ” into ( i ) detecting balls , ( ii ) detecting red objects , and ( iii ) intersecting the two sets , and ( iv ) summing the results . Intuitively , neural modules can generalize more robustly since ( i ) it is more likely that an individual neural module designed to solve a simple task can be identified from training data , and ( ii ) even if module composition is not itself identifiable , shifts in the compositional structure tend to be smaller than shifts in the underlying data distribution . We study a simplified form of neural module networks , where modules are quadratic neural networks composed in sequence according to a given input ; for simplicity , we assume they can be trained in a supervised way , ensuring robust generalization . Then , we show that under certain conditions , compositions of these models are also robust , including the case where there are shifts in the distribution over compositions . Related work . Prior work has connected misspecification ( i.e. , the true model is in the model family ) and robustness to covariate shift ( Shimodaira , 2000 ; Wen et al. , 2014 ) ; however , having a correctly specified model is insufficient if the true parameters are not identifiable—e.g. , in linear regression , if the covariance matrix Σ = Ep ( x ) [ xx > ] is singular , then θ is not identifiable ; thus , the estimated model may not be robust . Quadratic neural networks can not be identified even if the model is correctly specified since the parameters have a continuous symmetry ( i.e. , orthogonal transformations ) . Recent work has studied learning under adversarial examples ( Goodfellow et al. , 2014 ; Raghunathan et al. , 2018 ; Cohen et al. , 2019 ) and corrupted training data ( Steinhardt et al. , 2017 ; Diakonikolas et al. , 2019 ) . In contrast , we are interested in robustness to covariate shift ; there has been recent work empirically showing that neural networks are sensitive to distribution shift ( Hendrycks & Dietterich , 2019 ; Taori et al. , 2020 ; Ruis et al. , 2020 ; Ribeiro et al. , 2020 ; Koh et al. , 2020 ) . Distributionally robust optimization enables training of models robust to small shifts ( Duchi & Namkoong , 2018 ) , but we are interested in potentially large shifts . Unsupervised domain adaptation ( Ben-David et al. , 2007 ; Blitzer et al. , 2008 ) learns a model on a covariate shifted target distribution ; however , they rely on unlabeled examples from the target domain , whereas we do not . There has been recent theory on robustness to adversarial perturbations—e.g. , showing there may be a tradeoff between robustness and on-distribution generalization ( Tsipras et al. , 2018 ) , and that non-robust algorithms tend to learn predictive but brittle representations compared to adversarially robust ones ( Ilyas et al. , 2019 ) . In contrast , we show that these tradeoffs are mitigated when the true model function can be identified despite over-parameterization . Furthermore , adversarial shifts are typically bounded ( e.g. , small ` ∞ norm ) , whereas the shifts we consider may be large . There has been a great deal of recent work on deep learning theory , including on quadratic neural networks ; however , it has largely focused on optimization ( Ge et al. , 2017b ; Jacot et al. , 2018 ; Du et al. , 2019 ; Gao et al. , 2019 ; Soltanolkotabi et al. , 2018 ; Li et al. , 2018 ) , and on-distribution generalization ( Neyshabur et al. , 2017 ; Du & Lee , 2018 ; Jacot et al. , 2018 ; Arora et al. , 2018 ; Long & Sedghi , 2019 ) . In contrast , we are interested in out-of-distribution generalization . We discuss additional related work on matrix factorization and multi-armed bandits in Appendix A , as well as a discussion of the novelty of our results . 2 PROBLEM FORMULATION . We consider a model fθ : X → Y , with covariates X ⊆ Rd , labels Y ⊆ R , and parameters θ ∈ Θ ⊆ Rm . A generalization bound from learning theory typically has the form Pp ( Z ) [ Lp ( θ̂ ( Z ) ) ≤ ] ≥ 1− δ where Lp ( θ ) = Ep ( x ) [ ( fθ ( x ) − fθ∗ ( x ) ) 2 ] , ( 3 ) where , δ ∈ R > 0 , Z = { ( x1 , y1 ) , ... , ( xn , yn ) } ⊆ X × Y with yi = fθ∗ ( xi ) + ξi is a training set of i.i.d . observations from a distribution p ( i.e. , p ( Z ) = p ( x1 , y1 ) · ... · p ( xn , yn ) ) , ξi is bounded random noise independent of xi with |ξi| ≤ ξmax , θ∗ ∈ Θ are the true parameters , and θ̂ ( Z ) = arg min θ∈Θ L̂ ( θ ; Z ) where L̂ ( θ ; Z ) = 1 n n∑ i=1 ( fθ ( xi ) − yi ) 2 is an estimator based on the training data Z.1 In particular , they assume that the training inputs xi ∼ p are i.i.d . samples from the same distribution as the test example x ∼ p. Definition 2.1 . The model fθ and distribution p satisfy function identification if for any , δ ∈ R > 0 , we have Pp ( Z ) [ ∀x ∈ X . ( fθ̂ ( Z ) ( x ) − fθ∗ ( x ) ) 2 ≤ ] ≥ 1− δ for n = |Z| sufficiently large . Function identification implies generalization bounds even when the test data comes from a different distribution q . In particular , we say fθ robustly generalizes if for any q with support on X , we have Pp ( Z ) [ Lq ( θ̂ ( Z ) ) ≤ ] ≥ 1− δ , ( 4 ) where the difference from ( 3 ) has been highlighted in red . It is easy to see that function identification implies ( 4 ) . Note that the true model fθ∗ does not change , so there is no label shift .

Summary: The paper considers uncertainty estimation in overparameterized shallow neural networks with quadratic activation function. In particular, the paper assumes a non-linear regression model where labels are generated by an aforementioned neural net f* and noise is bounded. Then, uncertainty estimation is in giving a confidence interval around f*(x) for any given x. This differs from the usual uncertainty estimation in the linear regression, where one is concerned with deviation of parameters. The bound is a combination of the uniform convergence type bound on the excess risk and an interesting observation about strong convexity of the risk in the parameters of the network. The strong convexity here combined with Lipschitzness of the predictor allows to "convert" excess risk bound into the deviation between predictors.  The uncertainty estimation developed here is later used to give an explore-then-commit type of bandit algorithm with $T^{\\frac{2}{3}}$ regret.