Robust Generalization of Quadratic Neural Networks via Function Identification
Authors: Kan Xu
Hamsa Bastani
Osbert Bastani
Venue: ICLR 2022 Submitted (2022)
Keywords: neural network
function identification
robust generalization

Abstract: A key challenge facing deep learning is that neural networks are often not robust to shifts in the underlying data distribution. We study this problem from the perspective of the statistical concept of parameter identification. Generalization bounds from learning theory often assume that the test distribution is close to the training distribution. In contrast, if we can identify the ``true'' parameters, then the model generalizes to arbitrary distribution shifts. However, neural networks are typically overparameterized, making parameter identification impossible. We show that for quadratic neural networks, we can identify the function represented by the model even though we cannot identify its parameters. Thus, we can obtain robust generalization bounds even in the overparameterized setting. We leverage this result to obtain new bounds for contextual bandits and transfer learning with quadratic neural networks. Overall, our results suggest that we can improve robustness of neural networks by designing models that can represent the true data generating process. In practice, the true data generating process is often very complex; thus, we study how our framework might connect to neural module networks, which are designed to break down complex tasks into compositions of simpler ones. We prove robust generalization bounds when individual neural modules are identifiable.

Under review as a conference paper at ICLR 2022
ROBUST GENERALIZATION OF QUADRATIC NEURAL NETWORKS VIA FUNCTION IDENTIFICATION
Anonymous authors Paper under double-blind review
ABSTRACT
A key challenge facing deep learning is that neural networks often not robust to shifts in the underlying data distribution. We study this problem from the perspective of the statistical concept of parameter identification. Generalization bounds from learning theory often assume that the test distribution is close to the training distribution. In contrast, if we can identify the "true" parameters, then the model generalizes to arbitrary distribution shifts. However, neural networks are typically over peee-paramtteized, making parameter identification impossible. We show that for quadratic neural networks, we can identify the function represented by the model even though we cannot identify its parameters. Thus, we can obtain robust generalization bounds even in the pver-parameterized setting. We leverage this result to obtain new bounds for contextual bandits and transfer learning with quadratic neural networks. Overall, results suggest that we improve robustness of neural networks by designing models that can represent the true data generating process. In practice, the data generating process is often very complex; thus, we study how our framework might connect to neural module networks, which are designed to break down complex tasks into compositions of simpler ones. We prove robust generalization bounds when individual neural modules are identifiable.
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 well adversarial shifts (i.c., where the shift can depend on the neural network parameters) (Szegeyy et al. 2013). Accordingly, there has been great deal of interest in better understanding why neural networks fail to be robust (Tsipras et al. 2018 [lyas et al., 2019) and on improving robustness (Goodfellow et al.
2014 Raghunathan al. 2018 Cohen 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 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 Dx g || 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 them: for instance, small lixel-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 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
exactly recover the parameters of the true
model. For instance, in linear regression, the data is generated according to the model y = (8" 2)) + E. where is o-subgaussian noise. Then, under mild assumptions on the training data Z (X, Y), the

---

Under review as a conference paper at ICLR 2022
ordinary least squares (OLS) estimator (()) recovers the true parameter- -.e., in the limit of infinite data, ((Z) = 0- With finite samples, OLS satisfies high-probability convergence rates of the form
(0(Z) 0 1|2 S
The connection to robustness is that if holds, then for any input x such that ||>|||2 S Umaxx we have
(((()),) 2) (0" 2>1 - [l((Z) 80 ||11||2 S c. Fmuxx
Thus, for any distribution ((x) with support on B2(0,, Tmax] = fI E X
|a]|l S xmax), ). O(2) obtains bounded error-i.e., we have By(x) ((())), - [0)] e max 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 peerrtarameeerized possibly to facilitate optimization (Du & Lee 2018 Jacot 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 6", we can still recover the function Je- (z)). 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 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 wnique 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 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- c.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 al. 2016). which are designed to break down complex tasks into smaller ones that can each be solved by 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 more likely that an individual neural module designed to solve 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 simplified form of neural module networks, where modules are quadratic neural networks composed in sequence according to a given input; for simplicity, assume they can be trained in 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 (Shimodairal 2000; Wen al. 2014);; however, having correctly in linear regression, if the covariance matrix Σ Exp()) [x=>" is singular, then is not identifiable; thus, the estimated model may not be robust. Quadratic neural networks cannot be identified even if the model is correctly specified since the parameters have a continuous symmetry (i.e., orthogonal transfornations).
specified model is insufficient if the true parameters are not identifiable-e.
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- et al. 2007]] Blitzer et al.) (0088 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
2

---

Under review as a conference paper at ICLR 2022
robustness to adversarial perturbations e.g., showing there may be tradeoff between robustness and on-distributton generalization (Tsipras et al. 2018). and that non-robust algorithms tend to learn predictive but brittle representations compared to adversarially robust ones (llyas et al., 2019). In contrast, we show that these tradeoffs are mitigated when the true model function can be identified despite over- preerparameterization. Furthermore, adversarial shifts typically bounded (e.g., small foc 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 al., 2019 Gao al., 2019 Soltanolkotabi
al., 2018; Li et al. 2018). and on-distribution
generalization (Neyshabur et al. 2017, Du & Lee, 20188 Jacot et al.) 2018) Arora et al. 20188 Long
Sedghi 2019). In contrast, we are interested out-of out-of-distributin generalization.
We discuss additional related work on matrix factorization and multi-armed bandits in Appendix A as well discussion of the novelty of our results.
PROBLEM FORMULATION
We consider a model fo - y, with covariates X S R2, labels y S R. and parameters 0 E C A generalization bound from learning theory typically has the form
P P212)(L((((2)) s sel2--
where
(0) Ep(=)((f((x) - for (x:))
where c, 8 R>0, 7 = (1,,,)),,... v1),... (#:: yn un) X y with y: fo- (I:) + E, training set of i..d. observations from distribution p (i.e., p(2) p(11, 11) ... p( Un & is bounded random
noise independent of x; with K&| S Emax, 0" E θ are the true parameters, and
n
0(Z) arg min i(; 2) 8Ee
where
- i(;:Z) - n.
- (fo(z)) -1))
an estimator based on the training data α In particular, they assume that the training inputs I; p are i.i.d. samples from the same distribution as the test example p.
Definition 2.1. The model fo and distribution satisfy function identification if for any e E Rso. we have Ppt P(Z))(V (Va X. (fo(2)()) - fe. (x))² S 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 So robustly generulizes if for any q with support on X, we have
P IL (0(Z)) < 1-8,
(4)
where the difference from (3 has been highlighted in red. is easy to see that function identification implies 4 Note that the true model fo- does not change, there is label shift.
FUNCTION IDENTIFICATION OF QUADRATIC NEURAL NETWORKS
Traditional statistical bounds parameter identification can provide guarantees for arbitrary covariate shift. In particular, suppose we have bound of the form
P Pp( [10(Z) 8" 1l2 s sc] 2 21
(5)
and assume that the model family fo is K K-Lipschitz continuous in 0; then, we have
Lg (0(2)) < K² [10(Z) <
(6)
with prohability at least 1 - 5 according to p(Z). In particular, this bound holds for any covariate distribution q. Our goal to extend these techniques to quadratic neural networks, which are over-parameterized so we cannot identify the true parameters ii..., (5) does not hold.
Note that in 3). the loss T.p omits the lahel errors including it would result in additive constant to T.p. This choice ensures that the optimal parameters have zero loss-i.e. (!") for any g-
3

---

Under review as a conference paper at ICLR 2022
3.1 QUADRATIC NEURAL NETWORKS
We consider a neural network Sa. where 0 E Rdx*, with single hidden layer with & neurons-i.e. fo(x)) = Σj Gj ())) We consider the over-parameterriztion case where & can be much larger than d. Following prior work (Du & Lee 2018), we assume that fo has quadratic activations and output weights equal to one ie., ((2) and a 1 for each j [kc] so we have
f0(x) = Σ 1=1
10,, x)
We assume the true (training) loss is the mean-squared error (p)(0) [ppc))((o((t) fo- (x))) (c)))) and we consider model trained using an empirical estimate of this loss on the training dataset:
((Z) arg min i(0;; Z) Ace
where
i(00 7) (vo(s,)) -.))..
Now, our goal is to obtain bound of the form (6); to this end, we assume the following:
Assumption 3.1. ||0|| S Tmax and ||0|| S 0max.
Assumption 3.2. There exists a R2 such that IE (E) ((x AA) 2 al|A||? for any symmetric E Rddd
Our second assumption standard: in particular, it closely related to the assumption in linear
regression that the minimum eigenvalue of the covariance matrix is lower bounded -i.e., E =
Eptu) [xx] > As an example, when x is i.i.d. uniform each component, e.g., p(x)
Ih Uniform(z;; (-122, 1/2)), then take 1/180; give proof in Appendix B.1
3.2
ROBUST GENERALIZATION
Our approach leverages the fact that fa(x) = x xx (00 )æ: thus, fe resembles a matrix factorization model. Recent work has leveraged this connection translate matrix factorization theory quadratic neural networks (Du & Lee 2018). We let g(0) 00 and fx( x oz. where & E $ S Rdd in
which case fx(x) fg(o) (z); in addition, we define (0) = Ep(z) - for (x)) where o*
g(0"), and [(: Z) n- E (f)( y so (p(0) (o()))) and i(: Z)
(g(()) Z).
Finally, we assume that || < maxi in general, we have @max S 0 by Assumption 3.1
We begin by stating several lemmas establishing the properties needed for function identification. Our first lemma says that the loss is strongly convex in Φ.
Lemma 3.3. Under Assumption
3.2
the loss (,()) is 2a-strongly convex in Φ.
We give a proof in Appendix
B.2
Our next lemma says that our model family Lipschitz in d.
Lemma 3.4. Under Assumptions
3.1 & 3.2 fo and L are K K-Lipschitz in é. where K = 4
B.3 Our final lemma says that our estimate of the loss function is a uniformly good approximation of the true loss.
We give a proof in Appendix
Lemma 3.5. Under Assumptions
3.7 & 3.2 for any 8 E Ryo, we have
P
sup (0) OEB
- S 5 21
P(Z)
where a (Z)) = n1 Em= ? and letting Emax 2x mam Amax be upper bound on (fo(x) fo- (x)b
18P2 mar
(Prnc E (az ( max 1, 1, (1+ 8emar Kn f2
}
2
log
(7)
n
We give a proof in Appendix B.4 Note that e + 0 as n + 0o. Next, we prove our main result, which says that quadratic neural networks can be functionally identified.
4

---

Under review as a conference paper at ICLR 2022
Theorem 3.6. Under Assumptions 3.7 & 3.2 we have
2
P P(2) Vo X. (f(2)(() - fo. (x))² _
2 1 1
a
Proof. By Lemma 3.3 and since vol(o(07)) = O, we have
Lp()(Z)) Lp(0)) = Lp(g(((2))) - Lp(g((0)) allg(0(()) -((0))
(8)
Next, by Lemma 3.5 and the fact that 0 minimizes i(0; 2), have
[z)) (0: Z) oZ) I(00: + o (Z) - (0*) +
(9)
with probability at least 8. Combining S and ( we have
2e
[lg(0(Z)) - IIo S
with probability at least - 6. Finally, by Lemma 3.4 we have
2K (vr E X)
(f)(2)(z) - fo- (z))² = (z()(()) f -(()))
with probability at least 8, as claimed.
As result, we provide robust generalization error bound for quadratic neural networks with potential distribution shifts.
Corollary 3.7. Under Assumptions 3.1 & 3.2
for any distribution q(x) with support Bz(0, Tmaa),
2K
P p(z)
Ly (0(2))
2 1 - 8.
Finally, we also prove that gradient descent can find the global minima of L(;; Z), which ensures that gradient descent can perform function identification practice; we give proof in Appendix B3.) Proposition 3.8. All local minima of L(;; Z) are also global minima.
QUADRATIC NEURAL BANDITS
key application of robust generalization bounds is to parametric bandits; this because, in bandit learning, the distribution of inputs æ used to estimate 0 as can differ from the distribution under which ja is used. Thus, generalization bounds based on notions such as Rademacher complexity cannot used. Unlike prior literature in bandits, we consider an over- weer-parameterized function that does admit unique solution; in contrast, recent work neural tangent kernel bandits (Zhoul et al. 2020) assumes that there is unique, identifiable solution. Note that this assumption cannot hold for quadratic neural networks because they are invariant to transformations such rotations.
We consider a standard linear bandit (Rusmevichientong & Tsitsiklis 2010 Abbasi- Yadkori et al.) 2011) with fixed horizon 7' N, but where the expected reward parameterized by quadratic neural network instead of a linear function. At each time step 4, the algorithm chooses among a continuum of actions Fe X' and receives reward
yt = So- + &. = 2
(10)
where 0* E ggdk is unknown parameter matrix, and & are bounded i.i.d. random variables. For simplicity, we assume that = =B,00, 1) is the unit ball. Then, our goal is to bound the regret
T R(T) Σ (Emeolvd 8²)
where
max fo- (x). y
xx
We make the following assumption, which says that o3t T has gap in its top eigenvalue:
5

---

Under review as a conference paper at ICLR 2022
Algorithm
Algorithm for Quadratic Neural Network Bandit
procedure QUADRATICNEURALBBAND
Initialize Z + Ø
Let m be in (11
for E (1......)) j do
Sample i..d. action 2, ~ p. where p is as (12)
Take action and obtain reward y/ in (T))
Update
Z U ((1r, yc}}
end for
Compute 0 = arg min, L(: 2). where L(0; 2) = m Di (0f((x))-000 - y2)2
Compute è = arg m x (x)
for t E {m + 1..,, T) do
Take action H, & and obtain reward yt in 10
end for end procedure
Let o² 0*0 and let A A2 2 ... 2 A be the eigenvalues of d², There exists a constant M E Rş0 such that λ X2 2 4/M.
Assumption 4.1.
This assumption ensures that the eigenvectors of ø² to be stable under perturbations. The eigenvectors of Φ² correspond the optimal action x² arg mxx fo- (z) since fo- (x) = "0 x; thus, this assumption ensures that if 8" then the optimal action = arg max,ex fa(x) satisfies & = x*.
Next, we describe algorithm, summarized in Algorithm| We consider an explore-then-cor thmm- strategy for simplicity, since it already achieves the asymptotically optimal regret rate (Rusmevichien tong & Tsitsiklis 2010); approach
similarly be applied more sophisticated algorithms such as UCB (Abbasi-Yadkori Yadkori et al., 2011) and Thompson sampling (Agrawal & Goyal 2013).
Our algorithm proceeds in two stages: (i) the exploration stage (for E {1, ..., m)), and (ii) the
exploitation stage (for E {m + 1,.,, T)) where
2/3
(35.M(cma + Emax) ) VeeTVoon Vlog(3 + 8pmax KT/Pk
P2
max
(11)
Pmax
In the exploration stage, we randomly choose actions 24 p. where
p(x) II Uniform zy:
[ Vd
1 1 d
Val
(12)
Note that ||=||e < 1 for x in the support of p. 2, E X. Following the discussion in Section 3 for this choice p. Assumption 3.2 holds for the dataset Z with = =4(42d²).
Next, we compute an estimate 0 of 0" based on the data Z collected so far, and compute the optimal action à assuming are the true parameters. Then, the exploitation stage, we always use action
The key challenge providing theoretical guarantees using traditional generalization bounds is handling the optimization problem over æ E X' used to compute á. Since ă is not sampled from the distribution p. traditional bounds do not provide any guarantees about the accuracy of fg (ă) compared fo- (ø). In contrast, Theorem 3.6 provides a uniform guarantee, and can therefore be used to bound the regret.
Theorem 4.2. Under Assumptions
3.7 3.2 & 4.1
the expected regret of Algorithm is
R(T) S Co + C1
72/3
(log
80mar KT max
1/3
where Co and C₁ do not depend on T' (see Appendix C
In particular, have R(T) = ((T give proof in Appendix C Note that this is worse than the usual Ö(VT) regret because Theorem 3.6 only admits 4/1 convergence rate.

---

Under review as a conference paper at ICLR 2022
5 TRANSFER LEARNING OF QUADRATIC NEURAL NETWORKS
So far, we have considered shifts in the covariate distribution but not in the label distribution. Now, consider transfer learning problem where there is additionally small shift in the labels. In particular, we assume we have proxy data Zp C X X y from the source domain of the form sp.. - Joz (ap,a) + Ep.i (for [n,), where OF θ are the proxy parameters and p(z,) is the source covariate distribution, along with gold data Zg X' y from the target domain of the form gg.1 = fey (zy,a) + &g.a (for i [RgD, where 03 E θ are the gold parameters and g(xz) is the target covariate distribution. We are interested in the setting np > Ml. and where 110 0g F S B is small.
We consider a two-stage estimator (Bastani 2020) that first computes an estimate of the proxy parameters 8, = arg mingee [(;; Zp), and then computes an estimate of the gold parameters in way that is constrained towards the proxy parameters, First, note that we have
[.)(0p)) S
2k
6 2
P PM(
2 1
where
1 1862 max
(Lens + 22ax)
de max 1, 80mna K Knp
+ log 4 g
6p
10p
max
where we have highlighted the differences from in 7 in red. Next, we make a technical assumption: Assumption 5.1. For some du E Ryu, miin (87) 2 du. where T (f) is the dth singular value of A.
Equivalently, the minimum eigenvalue of g g(0)) is positive; intuitively, this assumption ensures good estimate of 070%T implies good estimate of 07 (up to an orthogonal transformation). Then, letting
1
2cp
B B
be an expanded radius to account for error in our estimate of Ó,. we use the estimator
0g = arg min L(:: Zg)
where BB(6,, B) 10 θ ||0 -pple S B}.
B)
Note that we have assume is known; in practice, this constraint can be included as an additive regularization term. Intuitively, this formulation mirrors transfer learning algorithms based fine- tuning-i.e.. initializing the parameters to the proxy data 6, and then taking a small number of steps of stochastic gradient descent (SGD) the gold data 0g. In particular, SGD can be interpreted L2 regularization on the parameters (Ali et al. 2020). so fine-tuning Le-regularizes 0g towards 0p- for any g(z) with support on (2 xmax).
Theorem 5.2. Under Assumptions
3.1
3.2 &
5.1
2K²cg
P ((())
Lg
11-
a
where
18K2 + E2mx))
80 Kng
}}
B
max
1, 1
+ log
ng
(2
5
Thus, if B is small and np is large, fa, is accurate even if ng is small; we give a proof in Appendix D
GENERALIZATION BOUNDS FOR NEURAL MODULE NETWORKS
While function identification enables robust generalization, many data generating processes are too complex to be identifiable. Neural module networks designed to break complex prediction problems into smaller tasks that are individually easier to solve. These models take two kinds input: (i) a sequence of tokens (e.g., word embeddings) indicating the correct composition of modules,
7

---

Under review as a conference paper at ICLR 2022
and (ii) the input to the modules. Then, the model predicts the sequence of modules f1...fr based on w, and runs the modules in sequence to obtain output x' = for (.. (...(f),, (z))...)
We study conditions under which neural module networks can robustly generalize. Rather than study arbitrary distribution shifts, we consider two separate shifts:
Module inputs: We assume that the individual modules are identifiable; consequence, assume the shift to the module input z can be arbitrary.
Module composition: We consider shifts to the token sequence w. If the model mapping w to j1.--r is identifiable, then the entire model is identifiable. Instead, show that when this model is not identifiable, compositional structure can still aid generalization. Intuitively, we show that while small shifts in the compositional structure can cause large shifts in the distribution p(w), models that leverage the structure of p(w) can still generalize well.
In more detail, consider simplified neural module network f. which includes (i) set of neural modules {fq X + X' } and (ii) a parser z + [k]" where Z C IR", with model class E G. We assume each component of fz) is computed by separate quadratic neural network; we discuss the architecture of g below. Then, given an input E X' S R2 and w E W = z" the corresponding neural module network f X' W is defined by
[x, w) (Jrr o... [z))()) (-- (=))))))))))...)
where
jr--rr g(w).
We assume that g has compositional structure i.., for some ğ : [K] x 2 + [k]. have
g(w) = f1.-f
where
je = (((2),)--1) _ otherwise,
ift=0
where = 2T. Intuitively, w is sequence of word vectors; then, the current neural module je = J(-)) depends both on the current word vector 24 and the previous neural module Jt-1- First, we assume that the individual modules have been functionally identified.
We assume we have fully labeled data we can use to train the neural modules-ic.. i... for each input x and sequence w, have both the desired sequence f1...Jr of neural modules, as well as the entire execution Zo. T1,-.., TT, where To = x: and 2++1 = fj. (I:) otherwise. Thus, we can use supervised learning to train the neural modules in particular, we can construct labeled examples (Je-1, 24. je) used train the parser g. and labeled examples (Ie, 2x+1) to train the modules fj.- For simplicity, we assume we have a uniform lower bound n on the number of training examples for the parser and for each module. Then, have the following straightforward result:
Lemma 6.1. Under Assumptions
3.7
3.2
with probability at least - dkó, for each j [k].
2dK a
(z) S
Ef
(Vr X),
where î3 is the estimated module and f7 is the ground truth module.
This result follows straightforwarddly from Theorem (.6 along with union bound. In contrast, do not assume the parsing model robustly generalizes, but only on distribution. For the subsequent analysis, can any neural network models that satisfy the statement of Lemma 6.1
Lemma 6.2. Under Assumptions
3.7 & 3.2 with probability at least - 8, we have
PR(Za)) ((2,0) + 3" (2. s (R) (2log(2/0) Eg.
where (9) the Rademacher complexity of (including its loss function), where p(z,j) T-1 ET p), and where
1 (e(z.)) = - -1 - 11-1(2', dd) otherwise.
ift=1
Neural modules are often trained with only partial supervision (Andreas al. 2016); we leave an analysis of this strategy future work since our focus is understanding generalization rather than learning dynamics.
8

---

Under review as a conference paper at ICLR 2022
This result is standard Rademacher generalization bound (Bartlett & Mendelson 2002). Note that we have also assumed that the distribution over token sequences is structured, which is necessary for compositional implementation of g to generalize, even distribution. Intuitively, the distribution over (2 consists of both unigram model over the word vectors:
pl 21, m. 2T) II
p(z 2-11;;
t=1
where we define p(z1 z0) = 21 as well as a unigram model over neural modules:
T
p(i...jr 21, 21, 2T) = II 1 = - (E).
=
Next, we consider shifted distribution q(= 2²), which is close to p(z 2'). Assumption 6.3. We have | =') p(- =') S a.
Importantly, despite this assumption, the shift between the overall distributions p(21 ...,...2))) 2r) and ((21,..,,7)) 27) can still be large since it compounds exponentially across the steps E [T].
Proposition 6.4. There exist p and q that satisfy Assumption 6.3 but ||p- -/zzv 2(1
a 2/)
That is, even the single step probabilities 5(z =') and ((2 x²) have total variation (TV) distance bounded as in Assumption 66.3 the overall distributions p and q can have TV distance exponentially close to the maximum possible distance of 2 in T; we give a proof in Appendix E.
We show that neural module networks generalize since leverages the compositional structure of p. First, we show that under Assumption 6.3 the overall shift in the input distribution of ğ is bounded: Lemma 6.5. Under Assumptions 3.1
3.2 & 6.3 we have |ă - Ty S To, where p is defined in
Lemma 6. 2 and ğ is defined in Assumption 6.3
That is, while the shift can compound across steps I, it does so only linearly; we give a proof in Appendix F.2 Next, we show that as consequence, the error of ğ is bounded.
Lemma 6.6. Under Assumptions
3
3.2 6.3 and assuming that TP p() (((c,))0(22,)) # g" (2,11) S cg. Teg.
have that IP (P(())0(w) # gww)
We give a proof in Appendix E.3 Finally, we have our main result.
Theorem 6.7. Under Assumptions 3.7
3.2 & 6.3 with probability at least 1 - (dk + 1)8, we have
(q(w) [)) (x, w)ll Tey max(KT-1,, .13]
Teg T
We give proof in Appendix E.4 Intuitively, Theorem 6.7 says that the error of the neural module network is linear in T as long as K < 1. Note that even if there is no distribution shift, its error is
P Pp(w) [1f(z, w) f (x, < Tey max(k T-1 1 2 21-Te,,
by the same argument the proof of Theorem 6.7 The exponential dependence on K is unavoidable since K >> 1 says that the modules f, can expand the input, which leads to exponential blowup in the magnitude of the output function of T, which also makes the estimation error exponential in T. Thus, the only cost to the distribution shift from to is the additional error probability T2a.
7 CONCLUSION
We have presented a number of results demonstrating that over pere-aarmeetiizationn does not funda- mentally harm learning models that are robust to arbitrary distribution shifts. In particular, even though we can no longer identify the true parameters for quadratic neural networks, we show that we can identify the true function, thereby enabling to prove new results in bandits and transfer learning. Finally, we provide preliminary analysis extending these results to neural module networks to handle complex data generating processes. A limitation of our work is that it is specialized to quadratic neural networks; a key direction for future work is to explore how our results extend to other activation functions and deeper architectures.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Yasin Yadkor,, Dávid Pál, and Csaba Szepesvári. Improved algorithms for linear stochastic bandits. In NIPS, volume 11, pp. 2312 2220, 2011.
Shipra Agrawal and Navin Goyal, Thompson sampling for contextual bandits with linear payoff's. In International Conference on Machine Learning, pp. 127-135. PMLR, 2013.
Alnur Ali, Edgar Dobriban, and Ryan Tibshirani. The implicit regularization of stochastic gradient flow for least squares. In International Conference Machine Learning, pp- 233-2444 PMLR, 2020.
Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Dan Klein. Neural module networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 39-48, 2016.
Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Yi Zhang. Stronger generalization bounds for deep nets via a compression approach. In International Conference on Machine Learning, pp. 254-263. PMLR, 2018.
Francis Bach, Julien Mairal, and Jean Ponce. Convex sparse matrix factorizations. arXiv preprint MrXiv:0012.1169, 2008.
Peter I. Bartlett and Shahar Mendelson. Rademacher and gaussian complexitie: Risk bounds and structural results. Journal of Machine Learning Research, 3(Nov):463 482, 2002.
Hamsa Bastani. Predicting with proxies: Transfer learning in high dimension. Management Science, 2020.
Shai Ben-David, John Blitzer, Koby Crammer, and Fernando Pereira. Analysis of representations for domain adaptation. In Advances in neural information processing systems, pp. 137-144, 2007.
John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Wortman. Learning bounds for domain adaptation. In Advances in neural information processing systems, pp. 129-136, 2008.
Emmanuel J Candes and Yaniv Plan. Tight oracle inequalities for low-rank matrix recovery from a minimal number of noisy random measurements. IEEE Transactions Information Theory, 57 (4):2342-2359, 2011.
Jeremy Cohen, Elan Rosenfeld, and Zico Kolter. Certified adversarial robustness via randomized smoothing. International Conference on Machine Learning, pp. 1310-1320. PMLR, 2019.
Ilias Diakonikolas, Gautam Kamath, Daniel Kane, Jerry Li, Ankur Moitra, and Alistair Stewart. Robust estimators in high-dimensions without the computational intractability. SIAM Journal on Computing, 48(2):742-888 864, 2019.
Simon Du and Jason Lee. On the power of over peremaaramttiintooonn in neural networks with quadratic activation. In International Conference Machine Learning, pp. 1329-1338. PMLR, 2018.
Simon Du Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International Conference on Machine Learning, pp. 1675-1685 PMLR, 2019.
John Duchi and Hongseok Namkoong, Learning models with uniform performance via distributionally robust optimization. arXiv preprint arXiv:1810.0875: 2018.
Dylan Foster and Alexander Rakhlin. Beyond ucb: Optimal and efficient contextual bandits with regression oracles. In International Conference on Machine Learning, pp. 3199-3210. PMLR, 2020.
Ruiqi Gao, Tianle Cai, Haochuan Li, Cho-Jui Hsieh, Liwei Wang, and Jason D Lee. Convergence of adversarial training in overparametrized neural networks. Advances in Neural Information Processing Systems, 52:13029-1304p 2019.
10

---

Under review as a conference paper at ICLR 2022
Rong Ge, Chi Jin, and Yi Zheng. No spurious local minima in nonconvex low rank problems: A unified geometric analysis. In International Conference on Machine Learning, pp. 1232. PMLR. 2017a.
Rong Ge, Jason Lee, and Tengyu Ma. Learning one-hidden-layer neural networks with landscape design. arXiv preprint urXiv:/711.0050/ 20171
lan J Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. arXiv preprint arXiv:1412.6572, 2014.
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness common corruptions and perturbations. arXiv preprint VrXiv:1903.12261 2019.
Daniel Hsu, Sham M Kakade, and Percy Liang. Identifiability and unmixing of latent parse trees. arXiv preprint ar 1206.3137, 2012.
Andrew Ilyas, Shibani Santurkar, Dimitris Tsipras, Logan Engstrom, Brandon Tran, and Aleksander Madry, Adversarial examples are not bugs, they are features. arXiv preprint arXiv: 1905.0217,, 2019.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. arXiv preprint arXiv: DxXxv:/806.07522 2018.
Maryia Kabanava, Richard Kueng, Holger Rauhut, and UIrich Terstiege. Stable low-rank matrix recovery via null space properties. Information and Inference: A Journal of the IMA, 5(4):405-441, 2016.
Michael J Kearns, Umesh Virkumar Vazirani, and Umesh Vazirani. An introduction to computational learning theory. MIT press, 1994.
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Bal- subramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, al. Wilds: A benchmark of in-the-wild distribution shifts. ar preprint arXiv:2012.07421, 2020.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang, Algorithmic regularization in wer-ppaameeterized matrix sensing and neural networks with quadratic activations. In Conference On Learning Theory, pp. 2-47. PMLR, 2018.
Philip Long and Hanie Sedghi. Generalization bounds for deep convolutional neural networks. arXiv preprint arXiv: /905///000/000000 2019.
Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nathan Srebro. Exploring general- ization in deep learning. arXiv preprint arXiv: 1rX.v:1706.08947, 2017.
Aditi Raghunathan, Jacob Steinhardt, and Percy Liang. Certified defenses against adversarial examples. arXiv preprint 1rXiv:1801.09:0 .09444, 2018.
Marco Tulio Ribeiro, Tongshuang Wu, Carlos Guestrin, and Sameer Singh. Beyond accuracy: Behavioral testing of nlp models with checklist. arXiv preprint arXiv:2005.04118, 2020.
Laura Ruis, Jacob Andreas, Marco Baroni, Diane Bouchacourt, and Brenden M Lake. A benchmark for systematic generalization in grounded language understanding. arXiv preprint arXiv:2003.05161, 2020.
Paat Rusmevichientong and John N Tsitsiklis. Linearly parameterized bandits. Mathematics of Operations Research, 35(2):395-411, 2010.
Hidetoshi Shimodaira. Improving predictive inference under covariate shift by weighting the log-
likelihood function. Journal of statistical planning and inference, 90(2):227-244, 2000.
Mahdi Soltanolkotabi, Adel Javanmard, and Jason D Lee. Theoretical insights into the optimization landscape of paraterraeeeri shallow neural networks. IEEE Transactions on Information Theory, 65(2):742-769 2018
11

---

Under review as a conference paper at ICLR 2022
Jacob Steinhardt, Pang Wei Koh, and Percy Liang. Certified defenses for data poisoning attacks. arXiv preprint 1tth::1706.03929) 2017.
Christian Szegedy, Wojciech Zaremba, llya Sutskever, Joan Bruna, Dumitru Erhan, lan Goodfellow, and Rob Fergus. Intriguing properties of neural networks. arXiv preprint arXiv:1312.6199 2013.
Rohan Taori, Achal Dave, Vaishaal Shankar, Nicholas Carlini, Benjamin Recht, and Ludwig Schmidt. Measuring robustness to natural distribution shifts in image classification. arXiv preprint arXiv:2007.00644, 2020.
Dimitris Tsipras. Shibani Santurkar, Logan Engstrom, Alexander Turner, and Aleksander Madry. Robustness may be odds with accuracy. arXiv preprint arXiv:/808. 12/52, 2018.
Roman Vershynin. High-dimensional probability: An introduction with applications in data science, volume 47, Cambridge university press, 2018.
Martin Wainwright. High-dimensionus statistics: A non-asymptotic viewpoint. Book Draft
(Working Publication), 2016. URL htttpp://ww..sttt. wwr..t......e...........e.ee.eee www.:yedddwwwwinnedwdddd.
nachdiplome Stachddpllomm/hha__seeeeee 2010__211000 201.2011.p.ff
Junfeng Wen, Chun-Nam Yu, and Russell Greiner. Robust learning under uncertain test distributions: Relating covariate shift to model misspecificatiion In International Conference on Machine Learning, pp. 631-639. PMLR, 2014.
Kan Xu, Xuanyi Zhao, Hamsa Bastani, and Osbert Bastani. Group-sparse matrix factorization for transfer learning of word embeddings. In International Conference on Machine Learning, 2021.
Yi Yu, Tengyao Wang, and Richard Samworth. A useful variant of the davis-kahan theorem for statisticians. Biometrika, 10532) 2015.
Dongruo Zhou, Lihong Li, and Quanquan Gu. Neural contextual bandits with ucb-based exploration. In International Conference Machine Learning, pp. 11492-11502 PMLR, 2020.
12

---

Under review as a conference paper at ICLR 2022
A ADDITIONAL RELATED WORK
Low-rank matrix factorization. Our notion of functional identification for quadratic neural network related to the low-rank matrix factorization literature. However, they impose low-rank structure on the matrix to recover, and hence typically require extra conditions to identify the matrix-e,, the restricted isometry property (RIP) (Candes & Plan 2011 Ge al.) 2017a), or bounded & norm of noise vector (Kabanava et al. 2016). In contrast, we consider more general case and do not assume any underlying structure of the matrix; in particular, since our goal is to capture over- P parammterrttinno of neural networks, our matrix is usually decomposed as é 00 where 0 E Rd* and k 2 d (and á is not necessarily low-rank). Also, we study the prediction error of neural networks in the presence of distribution shifts, whereas the goal of the low-rank literature is to recover the true matrix.
Multi-armed bandits. Prior literature on parameterized bandits has considered a number of func- tional forms, ranging from linear (Abbasi- Yadkori et al. 2011 |Ruseevichientong & Tsitsiklis 2010) to neural tangent kernels (Zhou et al. (020). Most of this work makes a realizability assumption that the model family contains the true model implicitly, they consider model families where there is unique, identifiable true parameter. These assumptions are necessary precisely due to the fact that the test and training distributions are different; thus, much of the bandit literature has focused on proving parameter identification results to enable learning. In contrast, the identifiability assumption does not hold for quadratic neural networks because they are invariant to parameter transformations. To the best of our knowledge, we consider the first werr-parameterized bandit problem that considers model that not parameter-identffable we find that similar regret results hold as long as the function represented by the model can be identified. Separately, Foster& & Rakhlin (2020) makes general connection between online regression oracles and the regret bandit algorithm: however, their approach only provides good guarantees when the regression oracle returns model that generalizes off-distribution. Finally, recent work on UCB with neural tangent kernels (Zhou et al., (2020) provides general regret bounds, but their bound is only sublinear under conditions such as the true reward function having small RKHS norm (see Remark 4.8 in their paper), which amounts to assuming they can recover the true parameters.
Novelty. We briefly discuss the novelty of our results compared to existing work. First, the results in Section are novel. To the best of our knowledge, the proof strategy in our main result, Theorem 3.6 is novel, though we note that the preceding lemmas are based on standard arguments e.g. bounding the convexity of (Ø) (Lemmaa 3.3) and the Lipschitz constant (Lemma 3.40 of fø: also, Lemma 3.5 relies on a standard covering number argument. For our applications bandits and transfer learning. our key novel results are Lemma C.3 for bandits, which proves smoothed bounded response for quadratic neural networks, and Lemma D.I for transfer learning. Finally, to the best of our knowledge,
arguments in Section 6 are novel.
Slightly different from realizahility in PAC learning Kearns et al. zero true loss.
1994], which says there is a model with
13

---

Under review as a conference paper at ICLR 2022
B PROOFS FOR SECTION 3
B.1 PROOF OF MINIMUM EIGENVALUE FOR UNIFORM DISTRIBUTION
In this section, we prove the claim that Assumption 3.2 is an i.ddd random variable with distribution Uniform(az aa 1122, To this end, note that
holds for the covariate distribution where ay
Ep(x) (IT Ar) Bp(x)
242 24
i,==1
Bp(x)
A
Σ
2 A A5
+
izj
(F)
Σ
A
Σ
ور + Σ
D
80
144
72
$
2
Σ A
Σ An
+
Σ
80
144
144
72
it
2
|| 1
180
claimed.
B.2 PROOF OF LEMMA 3.3
We use the notation U : (()) : V to denote the matrix inner product (U, V (V)) for U, V E Rdddd, The Hessian V²)(0) can be viewed as de x de matrix. As everything here is bounded, we can exchange the expectation and differentiation. Therefore, the Hessian of our loss function has for any symmetric matrix A
A V Lo)) A = 2x,()) (xx Ax) 2a A||
where the last inequality uses Assumption
B.3 PROOF OF LEMMA 3.4
By our definition, for any di ø6 $ (fo(x)) - (x)) (x To - 0')x)21 S xmaxlle> - dlll.-
Given our quadratic loss function, we have
((fo(z) for (x))² (f(()) for (x))²1
< -x)
- for (x) + (x) - (x)11f(f6(() - Fa)
< 44 max 100 FF
Next, the loss satisfies Lp(0) - _p(o') S (l) Fo. (x))) - (For(z) fz))) I 2 6'l...
Finally, the empirical loss satisfies
a
|i(0; 2) L(0'; -(())) 2)
Σ (fs(ax)) - y1)² -v1
- (xx)))
(ff((a)) - (o.) - (fo - (2.(1)))
1=1
Σ lsal (Fo(x.)) - f (x:))
< (40 D + 25mmx 2 1160 Φ'llr,
claimed.
1
14

---

Under review as a conference paper at ICLR 2022
B.4 PROOF OF LEMMA 3.5
First, we have the following results:
Lemma B.1 (Covering Number of Euclidean Ball). For a Euclidean ball in [R with radius R with respect to Frobenius norm, there exists an e-net such that
2R
R1n2
[ S 1
Proof. This claim follows by direct application of Proposition 4.2.12 in Vershynin (0018).
Lemma B.2 (Hoeffding's s Inequality for Subgaussian Random Variables). Letting { be a set of independent o-subgaussian random variables, then for all t 2 O, we have
n
2nd
Pr
24
exp
71
Proof. See Proposition 2.1 of Wainwright (2016).
Now, we prove Lemma such that
3,5
Consider an e/(4K) &. Then, for any & E 4, there exists ø' E E
(((() 2) - (i(0') -((())) (p(d)) S 2K l S
Therefore, have
P p(Z)
sup [[(()); 2) Lp(o(0)) o(Z))
[(: Z) (2(0) oZ)) 2
PM)
sup
NEN
i(0; 2) (b) o()
< 2(2))
max
DEE
2
$ Σ PP (())
[|L(0: Z) (,(0) o(Z) 2 2
(13)
6E8
Now, defining
L(o; 2) and recalling that (2)= Σm= 62, then we have
1 - - -f (2:))
and 7(b:
1 Σ
(f for (z.)))..
n1 i=1
i=1
i(0: 2) i(0: 2) 27(6: 2) oo()
Thus, continuing from
13 we have
[i Z) - - - 2
P
p(Z)
3Ep
[[[0: 2) (p(() + 2 2) 2 =
<
P P(2)
deg
[[(0:Z) [,() {
+ PP [ii(0: 2) 2 )
P p()))
(14)
6.
3ap
For the first term in subgaussian; thus, by Lemma B.2 have
14p. note that (f) - Fo- (x)) $ G. so (fo(z) - fo- (x)) is Pa
PA(2 [[(;z)] Lp(d))
> E -
< 2
exp
ne²
(15)
1861
3E0
max
15

---

Under review as a conference paper at ICLR 2022
Next, for the second term in (44. note that (fo ( (xx))Edl S max Emax, so (fo(z:)) (x:)) PmaxSmax -uugauuusass thus, by Lemma B.2 we have
nE²
Σ BEE
pl)
[i(: 2)
2|E1 exp
(16)
18 c2 max Smax
Combining (15 & 116
continuing from (14
we have
Σ 3EE
(P(Z) [i 2) Ip(p)
+
PP()))
(7(0: 2)1 )
6
ne2
$ 4|8|
exp
1862
(L + S
max
d2
8chmaxx K
n
2 2
1
exp
(-1862 max (emxx + C2
TEE
dl log (1+ 80max K
eex
(17)
18P2
(C2
max
Smax
where for the first inequality, we have used max follows since by Lemma [II the covering number of the c-net of t satisfies
( ( <2 I P2 + S and the second inequality
20max
|8|
<
Finally, we choose c so that 17 is smaller than 5-in particular, letting
c <=1
1862
max
( +
<aax)
dd max
1, 1,
1+
86 Kn
}
log
n
max
then continuing 17
we have
n
d2 log 1 8dmart(
))
exp
S
1862 max (e2 + E
as claimed.
B.5 PROOF OF PROPOSITION B.8 3.8
L(0: Z) is twice differentiable and convex in é. Note that the minimization problem of i(; 2) is equivalent to that of (g((); Z). We consider two cases. First, consider the where 9 has rank d. The first order condition (i Z) is the same vi(g((); Z) 0, which gives
V ( Z) O.
(18)
As Ô is of full rank, there exists matrix Ô1 kxd such that 00 (e.g. 01 - ŌJ - We can right multiply the above equation by @t and obtain that
(00
vi.(0 2) 0.
As L(;; Z) is convex in the above implies : ==(()) is global minimum of L(: Z). Therefore. Ô is global minimum of f.(0; Z). Next, consider case where the rank of θ smaller than In this case, we follow the proof strategy in Proposition 4 in Bach et al. (2008); we provide here for completeness. In this case, Equation (18 still holds, which implics
0 vi Z)00T vi. Z)
(19)
The Hessian of L(g(0); Z) has
( A) 2(vi( Z), AAT) + ( + A A =
16

---

Under review as a conference paper at ICLR 2022
As R.R. is also local minimum for any orthogonal matrix R (i.e., RR = I), find i with the last column being 0 by right multiplying certain R. Then, consider any A with the first columns being 0 and the last column being any u. Rd, With this choice of A and 0, 0A 0. Therefore.
() 2u v ( Z)u. Since 0 is local minimum of (g)(;;; 2), holds that ()) 2 0, which implies
vi. 2) = 0.
(20)
Equation 19 and 20 together comprise the first order conditions of the convex minimization
problem minazo ((:: 2). Thus, 0 also a global minimum.
C PROOFS FOR SECTION
First, we provide the full statement of Theorem
4.2 7
(including constants).
Theorem C.1. The expected regret of Algorithm
is
80me KTT) 1/3
))
R(T) S Co + C1 TT2/3
log
3 +
max
where
2ir2+2 +2 24
64( mmxx)
Co (135M Emax
2nd2 12
3d
) 2213) 2a--1
(80) max K
e
23 -1
G 162d2 (M2 (fmax + Emax) )maax 1/3
Before we prove Theorem we first prove preliminary result establishing an analog of the smooth best arm response property Rusmevichientong & Tsitsiklis (010)) to our setting. First, we have the following useful result:
Lemma C.2. Let d. o' Igdxd be symmetric matrices, let x, E Rd be eigenvectors of Φ. d corresponding to their top eigenvalue, such that ||>||| = ||>1|| = and let A2 2 2 da be the eigenvalues of ø. Suppose that (x, 2 0. Then, we have
23/2 0'1l2
-
-
d A2
Proof. See Corollary of Yu et al. 2015
Next, let X: [Rdxd + x denote the subset of reward-maximizing arms for g()) = 000 iic.,
x arg max dox, aex
where the argmax returns the set of all optimal values. Then, we have the following analog of smooth best arm response:
Lemma C.3. For any Rdd there exists x(δ) and xt x(x²) such that
| | x < M|lø IF
Proof. First, note that 2, xx are eigenvectors of Φ 0* corresponding to their top eigenvalues, re- spectively. Next, note that if E x(x"), ) then we also have -r" x( ); thus, without loss of generality, we can assume that (x",x) 2 0. Also, note that ||æ||| = | I| = 1 since the optimizer maximizes the magnitude of æ. Thus, we have
23/2 0012
| x
<
Ml |r,
A1 X2
where the second inequality follows by by Lemma tion 4.1 claimed.
C.2 and the third inequality follows by Assump-
17

---

Under review as a conference paper at ICLR 2022
Now, we prove Theorem C.1
The cumulative regret R(T) of horizon of T has that
R(T) E
Σ (fo.(x) fo-(
m
T Σ (Jo-(x*) fo (1i))
E
Σ (So (x*) (x-() +
t=m+1
< 2mm
E
E
(0(0)
g(00 1,
f xxxxx1
(g(0), x x=x
1:
1=m+1
(21)
where 1 is an estimator that minimizes the empirical loss of the first samples, & E x(o(0)) maximizes the estimated expected reward fT(r). and (Ø, o') Eij=1 Diyoy is the matrix inner
product. Since & maximizer of fg(x)) (g))), xxx ) have (o(0), xx= -0 S 0. Thus, continuing from (T), we have
R(T)= 2m1mmax + E
(g(0) g(0"), xt xxxx" x²x²1)
[==++1
2mm + (T m)E
[lg( g(0 )1 |1100 x²xx
(22)
To bound the second term in 12 note that
| x x0 |[
|20 =x=T || + ||xx"" - x xx-x"'1
2 || < 2M]19(0) 9(0*)
where the last step follows by Lemma
Next, by Theorem 3.6 we have
2e ll9(0) 9(0'))l )1 s V
a
3/2 Panx
(Paex 22
(dmax max 1, log
86max Km
}
))
1/4
1
log
10m
L
8
with probability at least
8. Now, defining the event
2e (10(0) g(00))lr S
g
a
letting
== exp (-d max 1, (1 11
80max Km p2
}}
max
and continuing from 22 have
R(T) < 2m1
T E [lx(0) (1=))1=)118"") xtx" -x 1[))
4T Amxmax P(g
< 2m0maxx 2MM IE [|ls(0) g(0 |2
g) 4T Omax P(g))
/22 Caa log(3 80max KT
81 O'
< 2mmmax 270M
(80 km/ 82 P2 dz
(23)
The third term in inequality 23 is smaller than the second term when
1/(df 11/2-
40max
2
(24)
135M (Emax + Emax 2d (880 K/ K/f2 P )d2 Vlog(3 (3 + 86 MT KT/
mux
18

---

Under review as a conference paper at ICLR 2022
For choice of m satisfying 244. continuing from 23 we have
R(T)= S 2mmmax 540M (4max + Emax ) T log(3 + S0max KT P
(25)
Next, we choose m minimize the upper bound in (24) for sufficiently large - particular,
135M + 5 T 1og(3 + S0max KT L (
$
(26)
Omax
With this choice of m, we have
R(T) S 162(M22 (fmux + Emax)* (mmax) 32f (1og(3 + Schmax KT/ /p. (aax ))
Finally, note that 24 holds under the choice of m in 26 for T satisfying
64( (4mma)
24
T Vog + 86max KT k max)
2d2 (135M (Emax + Smax) 2dd 20d
34
(80max K/ ( ad
The claim follows.
D PROOF OF THEOREM 5.2
First, we have the following key result:
Lemma D.1. Let 0,0 RRdkk, and let é = 00 and Φ' = Assume that ||o 6'Ilr S m, and that (0) 2 do > O, where (f) is the minimum singular value of A (more precisely, the dth largest singular value). Then, there exist orthogonal matrices R, R' Bgak such that
||AR 8' R'R'1 5 7)
(27)
00
Proof. Consider the SVDs UEVT and 8' U'E'V' where U, U Bddd E, E' Rddd and V, V E R&xd, then, we have & U.-U and ø' = U'' Then, we claim that the choices R = VUT and R' = V'U"T satisfy (27). In particular, note that 0R = UEUT and O'R''R = U'UTU since V VT la since l 2 d. where Ia Rddd is the d-dimensional identity matrix. Thus, it suffices to show that
oo uoUUUTT UEU
(28)
To this end, note that
n >
1 -
(29)
where the last step, we have multiplied the expression inside the Frobenius norm by U on the left and by U on the right, using the fact that the Frobenius norm is invariant under multiplication by orthogonal matrices. Defining W = U U, note that
(WE), =
WaE Wi
(30)
1=1
(EW); Σ
E/ Wkj WojE%p
(31)
k=1
(W = Σ
Wi (E²) )kj = WoyE3
(32)
Am
(E"))a Σ
( ): WR5 W.
(33)
1=1
19

---

Under review as a conference paper at ICLR 2022
Then, continuing from 29 we have
r 2 IW - W
-
Σ
( -2,
i,=1
W ( -
2
Σ (F Edj
1,=1
WiIWE E'WIZ
UE -N- E'U'E
-U'EUU 112
where the first line, have used (32 & (33). the third line have used Enj 2 on- the fourth line we have used (30) & (31). and on the last line we have multiplied on by U" on the left
U the right, again using the fact that the Frohenius norm is invariant under multiplication by
orthogonal matrices. Thus, we have shown 28 so the claim follows.
We note here that our result provides an analog of Lemma 6 in Ge al. (017aa) for quadratic neural networks.
Now, we prove Theorem 5.2 First, by directly applying the arguments the proof of Theorem 3.6 we have
(lo(0,) - g(02))l S $/22
2cp
with probability at least 1 - 8/2. However, 0p itself may not be close to 65. Instead, applying Lemma D.1 with 8 8, and 0' 87. and with 17 /2,, a, there exists an orthogonal matrix Rp RRR that ""ligns" 0p with 0% yielding
2cp 0o V
10, R, S
where oo is the minimum singular value of 67 Now, let 0g = 87 Rp and note that this is global minimizer (i.e., g(O,) = g(0))), since Rd is orthogonal. Then, we have
||6 8pllr I 110, Rs
0g RR Rr
100 0p11r
1
22p
< 110g
071l +
σ
2Ep
< B
(34)
00
with probability at least 1- 8/2. In other words, an alternative global minimizer 0, exists within small Frobenius norm of our proxy estimator 8p. even t, is not close to 07.
Finally, the event that (34) holds, note that for B2(0p, B), we have the alternative upper bound [f(()) - Sog (x))1 < K Kllg(() g(0))ll KB,
where the first inequality holds by Lemma 3.4 thus, we can take lmax KB. Thus, on the that 34 holds, by Theorem 3.6. we have
2K2cg cg
8
PP(Z)
(g(gg) <
2 1
2
so the claim follows by a union bound.
n
20

---

Under review as a conference paper at ICLR 2022
E PROOFS FOR SECTION 6
E. PROOF OF PROPOSITION 6.4
Suppose that 2 {0, 1} is binary, 2n 0, and
{1
if
p(= |
0
otherwise.
In particular, since 20 0, p())= 1 (w wo) places all weight the sequence we 0...0.
Next, consider the shifted distribution
=
0(2 F(-1)
1 a/2
2=== 0
a/2 otherwise.
Note that (1p(- - q(- 1l I a, so Assumption
6.3 is satisfied. Note that
(ww))= II
6(0 (() = (1-a/2)²
t=1
As consequence, we have
lp glllv =
Σ (p(w) - [(w)
weW
= (p(wo) - q(wo))] +
Σ ((w)
wEW mom
(1- (1 -2/2222)) + (1 - (1 (1-(1)/)))
2(1 - (1 (1- (1]1=())/2))))
as claimed.
E.2
PROOF OF LEMMA 6.5
Note that
1ğe PelTTV
[ãc(2,g) - pe(
Σ / 1(j= = -d #(z 22)p--((2' d')''))
=l
lā(2
5(2 10) (z',5') dz'
10(= 1z -(= - $(2 2') 1/.
==1"
k
Σ a "1-1(=',5') + (=-1(=" p.- (1-1(2',
='=1"
< a + 1|ă:-1 1:-1/ITV:
Since ()(),)) po(2, j) for all 2 E 2 and j E [Rt], by induction, 1 - Pe ll S to. Thus, we have
l/q 5lTV
llā. Pell S Ta,
claimed. n
21

---

Under review as a conference paper at ICLR 2022
E.3
PROOF OF LEMMA 6.6
First, we prove the following lemma. Lemma E.1. We have
II 1 ((mn)r=1)) ) p(e,.... 2,)22...
[-1
pe( (0-1)
Σ
Proof. For the base case, we have
p2(22.51) = Σ /
k
1(i - p(22 ( (0)
1o=1
k
/
Σ
1(i (2)) p(22 I 2)) 1(j 0) p(2))de
j=
= (21, 22)21
as claimed. For the inductive case, we have
&
Pe(2))0-1)) = Σ
1(jr-1 = 9" (21-1. p(z0 22-1) 12-1 (1-
dt 2m1
(-1
Σ
II
1(j-11 = 9" (( 2id
as claimed.
Now, we prove Lemma 6.6
First, note that for each E [T], we have
(-1
((w)= = gg" (w))+ )
Pp() (g(w)e # (w)))
A
= -f /I(0(w), # (w))))
1-1
(0((w, = 9" (ww))) p(w)dw
k
1-1
Σ
# 10(())+
II =(())))='({{))))))
=
phi...-- I w) p(w)du
k
(
Σ
(g")))
II (((w), - -(()))))
1 = -" (((r II
Ao 1=1
p(w)dw
k Σ
((((zr)r-1))) f (zr.fr-1))
1 (((..tr-
1
1-1
II 1(jr 5" (3r1) p(w)dw
b
1-1
Σ
1(g(2r) (zr)trr))
II 1(jr = = ((r))rr))) p(w)dd
k
Σ
I 1(g(2r1 + ( fr-))r-)))
II 1(jr = x" (Er.)r-1)) (1)
[((2)))-1)) F 0¹ (2r)
Pr((t))
22

---

Under review as a conference paper at ICLR 2022
where the last step follows from Lemma E.1
Now, note that
(-1
)
P (((w)((()) # o (w))
P p(w)
^ [(w)) 0
(x(e.)) (3.
< pee(x.t)
$ Teg
claimed.
E.4 PROOF OF THEOREM 6.7
First, we show that P. Po(=.)) P((=,))((2,5) & (2, S + Ta. To this end, note that
Po Po (z
= P (plca)(((z,))) 5" (2) + P (),,))(()= + (=, PP (0(2,)) # (z,)))
cg Σ ((((z.)) # (z. (ā(2,)) j)l
5=1
1 + 1 - pl|TV
$ Eg + Ta.
Next, by Lemma with q in place of p and cg + Ta in place of cg. we have (o(w) [0w)) * g"((w]] s
Teg ao Then, assuming that j(w) 9 (w), we have
(I, w) w)lle
Il(frr f (x) (for o... fx)
110 .... Sjon o Ij, Sn o... o [z)(s)) - (ftr o... o o...o fimn o Sn Sn o.. o {z)(2)12
Σ T- kT- .1155,
sm o... o (z)(z) - (s. o sm o... (z)(2
t=1
Σ KT-
1=1
TTy max TTT 11.
The claim follows by union bound. C
23


Reviews:
[{"review_id": "yViVPIFejBz", "reviewer": "Reviewer_hpBY", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "1. Despite that the authors have already mentioned that the analysis is focused on quadratic neural networks, the analysis of function identification seems to be very limited to quadratic neural networks and hard to generalize to other activations. This somehow constrained the potential extension of the proposed framework. \n2. Identifying the functions are very different from identifying the parameters. Could the authors add a discussion to compare the pros and cons between them? Why function identification is a better definition besides the usage in over-parameterized networks? The authors mentioned that the analysis is valid for over-parameterized networks multiple times, but it is applicable simply because the formulation is very strong and avoids the parameters. \n3. The authors should point out which bounds are newly derived and if there are existing bounds with similar settings, a discussion should be added to compare them. \n4. The conclusion is too concise that it lacks a discussion on potential future works.", "summary_of_the_paper": "The author proposed the concept of function identification to address the limitation of parameter identification in the over-parameterized setting. The function identification states that the outputs of quadratic neural networks with the empirically minimized parameters and the true parameters have bounded difference with high probability. The reason to adopt the definition of function identification analysis is that it is hard ti make parameter identification with neural networks which usually are over-parameterized. The main theoretical result is Theorem 3.6, which is directly used in Corollary 3.7 to derive a robust generalization error bound, i.e., the loss function is evaluated on an arbitrary distribution. The robust generalization error bound is then used for three cases: (1) the upper bound of the regret of quadratic neural bandits, (2) bounds for transfer learning, and (3) generalization bounds for neural module networks.", "summary_of_the_review": "The consideration of function identification is an appealing tool for analysis since it only requests the outputs to be similar with high probability. The derivation then follows the analysis of quadratic neural networks. I think there are some discussion and comparisons needed (see bullet points above), and therefore I decided to hold my scores.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "oHHIFsE41v", "reviewer": "Reviewer_LNcd", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper investigates a very relevant and interesting topic of uncertainty estimation in overparameterized neural networks (here with one hidden layer network with quadratic activation function). The interesting and novel part of the paper is in attempt to give uncertainty estimation for predictors rather than parameters. The later is of course problematic in neural networks, since equivalent predictors can have different parameters (as acknowledged by the paper). \nThe paper starts of by bounding an excess risk through a fairly standard covering number argument. At this point the idea is to \"convert\" excess risk to a gap between predictor and the regression function. Perhaps, here comes the novel idea, namely observing that the risk is $\\\\alpha$-strongly convex in parameters of a neural network, and that the predictor is Lipschitz. Here strong convexity heavily relies on the quadratic nature of a neural network (one can see that by computing the Hessian). Interestingly, for considered covariate instance, $\\\\alpha = \\\\text{poly}(d)$ (where $d$ is a covariate dimension).\n\nThe resulting estimation bound is of order $\\\\text{poly}(d B) / (\\\\alpha \\\\sqrt{n})$, where the paper assumes that $\\\\|\\theta\\\\|_F, \\\\|\\theta^{\\\\star}\\\\|_F \\\\leq B$ such that $\\\\theta$, $\\\\theta^{\\\\star}$ are parameters of a fitted network and that of the regression function network.\nWhile dependence on $\\\\alpha$ (somewhat surprisingly) is not problematic, boundedness of $\\\\|\\theta\\\\|_F$ seems to be a bit problematic. When fitting overparameterized networks (especially till interpolation), $\\\\|\\theta\\\\|_F^2$ can scale linearly in $n$, which will make the bound vacuous. This is of course a pessimistic view, but to have an entire picture, we would need to have an algorithm-dependent bound on the norm.\n\nThis limitation later finds its way into the explore-and-commit-type algorithm, where one commits to the argmin predictor (which is an interpolant in the overparameterized case). Here, the length of the exploration stage $m$ is tuned based on $d, B$ among other parameters.\nThe resulting regret depends on the norm through constants $C_0$, $C_1$ (at this point I wouldn't call them constants). Can we ensure that they don't scale as $\\\\text{poly}(m)$?\n\nIt is likely that this issue can be mitigated by analyzing the norm (e.g. given by GD solution), and for instance using some arguments as in Arora et al. 2019, however I feel that right now the entire picture is not very clear.\n\nOn another note, what are the challenges in analyzing a UCB-style algorithm employing the confidence bound proved here?", "summary_of_the_paper": "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. \nThe uncertainty estimation developed here is later used to give an explore-then-commit type of bandit algorithm with $T^{\\\\frac{2}{3}}$ regret.", "summary_of_the_review": "The paper investigates a very relevant and interesting topic of uncertainty estimation in overparameterized neural networks. The interesting and novel part of the paper is in attempt to give uncertainty estimation for predictors rather than parameters. However, uncertainty estimation depends on the norm of a fitted network which is generally unclear in the interpolation setting. It is likely that this issue can be mitigated by analyzing the norm (e.g. given by GD solution -- right now no algorithm is considered for fitting the network), and for instance using some arguments as in Arora et al. 2019, however I feel that right now the entire picture is not very clear and the paper could benefit a lot from clarifying this.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "cHHtmad-tP", "reviewer": "Reviewer_QEKW", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths: I think the technical extension and its applications presented in the paper are conceptually interesting\n\nWeaknesses:\n1. The technical details of the paper (at least the main paper; I don't have access to the Appendix yet) is not very satisfactory: here are some examples\n - Lemma 6.2, the statement contains G (Rademacher complexity of G), but G is never defined\n - In section 6, it seems to me that the neural networks in neural network modules can only be quadratic NN (in order to apply Theorem 3.6). How realistic is this? And the authors should also state this fact in the paper\n - In Lemma 6.5, \\tilde q and \\tilde p are used in the statement and the discussion above, but were never defined; these definitions are important for the reader to understand the seemingly contradictory result of Proposition 6.4 and Lemma 6.5. \n\n2. The assumptions made in the paper seem to be a little too strong to me:\n - It assumes that the noise are independent of the data, which doesn't seem to be very common in robust learning; or in rebuttal, perhaps the authors should discuss/cite related work to justify this assumption\n - The single-layer quadratic NN seems to me a little over-simplified to be studied; or the authors should provide clear argument on why this is not the case\n\n3. The difficulty of the technical extension is not very obvious to me in the current paper: either the authors should clearly discuss and highlight these difficulties in the main paper or they should try to strengthen their technical result in a future version\n", "summary_of_the_paper": "Under the generative data paradigm (e.g., the hypothesis space contains the ground-truth), the paper extends bounds on test (generalization) error under data distribution shift from the common setup when parameter identification is assumed to the setup where the ground-truth parameter may not be identifiable (not unique). Then the paper shows how this technical extension can be applied to obtain test-error bounds under distributional shift in three scenarios: bandits represented by an over-parametrized quadratic NN; transfer learning with both covariate and label shift; neural network modules", "summary_of_the_review": "While I think the paper works on an interesting topic with several interesting applications, I feel technical quality of the current version of the paper lacking strength; however, I would not strongly argue for rejection if the authors can provide convincing rebuttal or are willing to improve on the technical quality of the paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper tackles an interesting problem: distribution shift generalization often requires parameter identification but this is not possible for over-parameterized neural networks. This paper shows for quadratic neural networks, it is possible to identify the function without identifying the parameter. 

This is an interesting result. However, reviewers raise concerns about the assumption and technical details. The meta-reviewer agrees with these concerns.
Decision: Reject