Year: 2021

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 over-parameterized, 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 over-parameterized 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.

heading: 1 INTRODUCTION
text: 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 cannot 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.

heading: 2 PROBLEM FORMULATION
text: 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.

heading: 3 FUNCTION IDENTIFICATION OF QUADRATIC NEURAL NETWORKS
text: Traditional statistical bounds on parameter identification can provide guarantees for arbitrary covariate shift. In particular, suppose we have a bound of the form
Pp(Z) [ ‖θ̂(Z)− θ∗‖2 ≤ ] ≥ 1− δ, (5)
and assume that the model family fθ is K-Lipschitz continuous in θ; then, we have
Lq(θ̂(Z)) ≤ K2 · ‖θ̂(Z)− θ∗‖22 ≤ K2 2 (6)
with probability at least 1− δ according to p(Z). In particular, this bound holds for any covariate distribution q. Our goal is to extend these techniques to quadratic neural networks, which are over-parameterized so we cannot identify the true parameters θ∗—i.e., (5) does not hold.
1Note that in (3), the loss Lp omits the label errors ξ; including it would result in an additive constant to Lp. This choice ensures that the optimal parameters have zero loss—i.e., Lq(θ∗) = 0 for any q.

heading: 3.1 QUADRATIC NEURAL NETWORKS
text: We consider a neural network fθ, where θ ∈ Rd×k, with a single hidden layer with k neurons—i.e., fθ(x) = ∑k j=1 aj · σ(〈θj , x〉). We consider the over-parameterization case where k can be much larger than d. Following prior work (Du & Lee, 2018), we assume that fθ has quadratic activations and output weights equal to one—i.e., σ(z) = z2 and aj = 1 for each j ∈ [k], so we have
fθ(x) = k∑ j=1 〈θj , x〉2.
We assume the true (training) loss is the mean-squared error Lp(θ) = Ep(x)[(fθ(x)− fθ∗(x))2], and we consider a model trained using an empirical estimate of this loss on the training dataset:
θ̂(Z) = arg min θ∈Θ
L̂(θ;Z) where L̂(θ;Z) = 1
n n∑ i=1 (fθ(xi)− yi)2.
Now, our goal is to obtain a bound of the form (6); to this end, we assume the following: Assumption 3.1. ‖x‖2 ≤ xmax and ‖θ‖F ≤ θmax. Assumption 3.2. There exists α ∈ R>0 such that Ep(x)[(x>∆x)2] ≥ α‖∆‖2F for any symmetric ∆ ∈ Rd×d.
Our second assumption is standard; in particular, it is closely related to the assumption in linear regression that the minimum eigenvalue of the covariance matrix is lower bounded—i.e., Σ = Ep(x)[xx>] 0. As an example, when x is i.i.d. uniform in each component, e.g., p(x) =∏d i=1 Uniform(xi; [−1/2, 1/2]), then we can take α = 1/180; we give a proof in Appendix B.1.

heading: 3.2 ROBUST GENERALIZATION
text: Our approach leverages the fact that fθ(x) = x>(θθ>)x; thus, fθ resembles a matrix factorization model. Recent work has leveraged this connection to translate matrix factorization theory to quadratic neural networks (Du & Lee, 2018). We let g(θ) = θθ> and f̃φ(x) = x>φx, where φ ∈ Φ ⊆ Rd×d, in which case fθ(x) = f̃g(θ)(x); in addition, we define L̃p(φ) = Ep(x)[(f̃φ(x)− f̃φ∗(x))2], where φ∗ = g(θ∗), and ˆ̃L(φ;Z) = n−1 ∑n i=1(f̃φ(xi) − yi)2, so Lp(θ) = L̃p(g(θ)) and L̂(θ;Z) =
ˆ̃L(g(θ);Z). Finally, we assume that ‖φ‖F ≤ φmax; in general, we have φmax ≤ θ2max 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 L̃p(φ) is 2α-strongly convex in φ.
We give a proof in Appendix B.2. Our next lemma says that our model family is Lipschitz in φ.
Lemma 3.4. Under Assumptions 3.1 & 3.2, f̃φ and L̃ are K-Lipschitz in φ, where K = 4φmaxx4max.
We give a proof in Appendix B.3. Our final lemma says that our estimate of the loss function is a uniformly good approximation of the true loss. Lemma 3.5. Under Assumptions 3.1 & 3.2, for any δ ∈ R>0, we have
Pp(Z) [ sup θ∈Θ |Lp(θ)− L̂(θ;Z)− σ(Z)| ≤ ] ≥ 1− δ,
where σ(Z) = n−1 ∑n i=1 ξ 2 i , and letting `max = 2x 2 maxφmax be an upper bound on |fθ(x)− fθ∗(x)|,
=
√ 18`2max(` 2 max + ξ 2 max)
n
( d2 max { 1, log ( 1 + 8φmaxKn
`2max
)} + log 2
δ
) . (7)
We give a proof in Appendix B.4. Note that → 0 as n→∞. Next, we prove our main result, which says that quadratic neural networks can be functionally identified.
Theorem 3.6. Under Assumptions 3.1 & 3.2, we have Pp(Z) [ ∀x ∈ X . (fθ̂(Z)(x)− fθ∗(x)) 2 ≤ 2K 2
α
] ≥ 1− δ.
Proof. By Lemma 3.3, and since∇φL̃(g(θ∗)) = 0, we have
Lp(θ̂(Z))− Lp(θ∗) = L̃p(g(θ̂(Z)))− L̃p(g(θ∗)) ≥ α‖g(θ̂(Z))− g(θ∗)‖2F . (8)
Next, by Lemma 3.5 and the fact that θ̂ minimizes L̂(θ;Z), we have
Lp(θ̂(Z)) ≤ L̂(θ̂;Z) + − σ(Z) ≤ L̂(θ∗;Z) + − σ(Z) ≤ Lp(θ∗) + 2 (9) with probability at least 1− δ. Combining (8) and (9), we have
‖g(θ̂(Z))− g(θ∗)‖F ≤ √ 2
α
with probability at least 1− δ. Finally, by Lemma 3.4, we have
(fθ̂(Z)(x)− fθ∗(x)) 2 = (f̃g(θ̂(Z))(x)− f̃g(θ∗)(x)) 2 ≤ K2‖g(θ̂(Z))− g(θ∗)‖22 ≤ 2K2
α (∀x ∈ X )
with probability at least 1− δ, as claimed.
As a result, we provide a 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 on B2(0, xmax),
Pp(Z) [ Lq(θ̂(Z)) ≤ 2K2
α
] ≥ 1− δ.
Finally, we also prove that gradient descent can find the global minima of L̂(θ;Z), which ensures that gradient descent can perform function identification in practice; we give a proof in Appendix B.5.
Proposition 3.8. All local minima of L̂(θ;Z) are also global minima.

heading: 4 QUADRATIC NEURAL BANDITS
text: A key application of robust generalization bounds is to parametric bandits; this is because, in bandit learning, the distribution of inputs x used to estimate θ̂ ≈ θ∗ can differ from the distribution under which fθ̂ is used. Thus, generalization bounds based on notions such as Rademacher complexity cannot be used. Unlike prior literature in bandits, we consider an over-parameterized function that does not admit a unique solution; in contrast, recent work on neural tangent kernel bandits (Zhou et al., 2020) assumes that there is a unique, identifiable solution. Note that this assumption cannot hold for quadratic neural networks because they are invariant to transformations such as rotations.
We consider a standard linear bandit (Rusmevichientong & Tsitsiklis, 2010; Abbasi-Yadkori et al., 2011) with a fixed horizon T ∈ N, but where the expected reward is parameterized by a quadratic neural network instead of a linear function. At each time step t, the algorithm chooses among a continuum of actions xt ∈ X , and receives a reward
yt = fθ∗(xt) + ξt = k∑ j=1 〈θ∗j , xt〉2 + ξt, (10)
where θ∗ ∈ Rd×k is an unknown parameter matrix, and ξt are bounded i.i.d. random variables. For simplicity, we assume that X = B2(0, 1) is the unit ball. Then, our goal is to bound the regret
R(T ) = T∑ t=1 (Ep(ξt)[yt]− y ∗) where y∗ = max x∈X fθ∗(x).
We make the following assumption, which says that φ∗ = θ∗θ∗> has a gap in its top eigenvalue:
Algorithm 1 Explore-Then-Commit Algorithm for Quadratic Neural Network Bandit procedure QUADRATICNEURALBANDIT
Initialize Z ← ∅ Let m be as in (11) for t ∈ {1, ...,m} do
Sample i.i.d. action xt ∼ p, where p is as in (12) Take action xt and obtain reward yt as in (10) Update Z ← Z ∪ {(xt, yt)}
end for Compute θ̂ = arg minθ L̂(θ;Z), where L̂(θ;Z) = m −1∑m i=1(fθ(xi)− yi)2 Compute x̂ = arg maxx∈X fθ̂(x) for t ∈ {m+ 1, ..., T} do
Take action xt = x̂ and obtain reward yt as in (10) end for
end procedure
Assumption 4.1. Let φ∗ = θ∗θ∗>, and let λ1 ≥ λ2 ≥ ... ≥ λd be the eigenvalues of φ∗. There exists a constant M ∈ R>0 such that λ1 − λ2 ≥ 4/M .
This assumption ensures that the eigenvectors of φ∗ to be stable under perturbations. The eigenvectors of φ∗ correspond to the optimal action x∗ = arg maxx∈X fθ∗(x) since fθ∗(x) = x
>φ∗x; thus, this assumption ensures that if θ̂ ≈ θ∗, then the optimal action x̂ = arg maxx∈X fθ̂(x) satisfies x̂ ≈ x ∗.
Next, we describe our algorithm, summarized in Algorithm 1. We consider an explore-then-commit strategy for simplicity, since it already achieves the asymptotically optimal regret rate (Rusmevichientong & Tsitsiklis, 2010); our approach can similarly be applied to more sophisticated algorithms such as UCB (Abbasi-Yadkori et al., 2011) and Thompson sampling (Agrawal & Goyal, 2013). Our algorithm proceeds in two stages: (i) the exploration stage (for t ∈ {1, ...,m}), and (ii) the exploitation stage (for t ∈ {m+ 1, ..., T}), where
m =  ( 135M(`max + ξmax) 2d3T √ log(3 + 8φmaxKT/`2max) φmax )2/3 . (11) In the exploration stage, we randomly choose actions xt ∼ p, where
p(x) = d∏ i=1 Uniform ( xi; [ − 1√ d , 1√ d ]) . (12)
Note that ‖x‖2 ≤ 1 for x in the support of p, so xt ∈ X . Following the discussion in Section 3, for this choice of p, Assumption 3.2 holds for the dataset Z with α = 4/(45d2).
Next, we compute an estimate θ̂ of θ∗ based on the data Z collected so far, and compute the optimal action x̂ assuming θ̂ are the true parameters. Then, in the exploitation stage, we always use action x̂.
The key challenge providing theoretical guarantees using traditional generalization bounds is handling the optimization problem over x ∈ X used to compute x̂. Since x̂ is not sampled from the distribution p, traditional bounds do not provide any guarantees about the accuracy of fθ̂(x̂) compared to fθ∗(x̂). In contrast, Theorem 3.6 provides a uniform guarantee, and can therefore be used to bound the regret.
Theorem 4.2. Under Assumptions 3.1, 3.2 & 4.1, the expected regret of Algorithm 1 is
R(T ) ≤ C0 + C1 · T 2/3 ( log ( 3 + 8φmaxKT
`2max
))1/3 ,
where C0 and C1 do not depend on T (see Appendix C).
In particular, we have R(T ) = Õ(T 2/3); we give a proof in Appendix C. Note that this is worse than the usual Õ( √ T ) regret because Theorem 3.6 only admits a n1/4 convergence rate.

heading: 5 TRANSFER LEARNING OF QUADRATIC NEURAL NETWORKS
text: So far, we have considered shifts in the covariate distribution but not in the label distribution. Now, we consider a transfer learning problem where there is additionally a small shift in the labels. In particular, we assume we have proxy data Zp ⊆ X × Y from the source domain of the form yp,i = fθ∗p (xp,i) + ξp,i (for i ∈ [np]), where θ ∗ p ∈ Θ are the proxy parameters and p(xp) is the source covariate distribution, along with gold data Zg ⊆ X × Y from the target domain of the form yg,i = fθ∗g (xg,i) + ξg,i (for i ∈ [ng]), where θ ∗ g ∈ Θ are the gold parameters and q(xg) is the target covariate distribution. We are interested in the setting np ng , and where ‖θ∗g − θ∗p‖F ≤ B is small. We consider a two-stage estimator (Bastani, 2020) that first computes an estimate of the proxy parameters θ̂p = arg minθ∈Θ L̂(θ;Zp), and then computes an estimate of the gold parameters in a way that is constrained towards the proxy parameters. First, note that we have
Pp(Z) [ Lq(θ̂p) ≤
2K2 p α
] ≥ 1− δ
2 ,
where
p =
√ 18`2max(` 2 max + ξ 2 max)
np
( d2 max { 1, log ( 1 +
8φmaxKnp `2max
)} + log 4
δ
) ,
where we have highlighted the differences from in (7) in red. Next, we make a technical assumption: Assumption 5.1. For some σ0 ∈ R>0, σmin(θ∗p) ≥ σ0, where σmin(θ) is the dth singular value of θ.
Equivalently, the minimum eigenvalue of g(θ∗p) is positive; intuitively, this assumption ensures a good estimate of θ∗pθ ∗> p implies a good estimate of θ ∗ p (up to an orthogonal transformation). Then, letting
B̂ = B + 1
σ0
√ 2 p α
be an expanded radius to account for error in our estimate of θ̂p, we use the estimator
θ̂g = arg min θ∈B2(θ̂p,B̂)
L̂(θ;Zg) where B2(θ̂p, B̂) = {θ ∈ Θ | ‖θ − θ̂p‖F ≤ B̂}.
Note that we have assume B̂ is known; in practice, this constraint can be included as an additive regularization term. Intuitively, this formulation mirrors transfer learning algorithms based on finetuning—i.e., initializing the parameters to the proxy data θp and then taking a small number of steps of stochastic gradient descent (SGD) on the gold data θg . In particular, SGD can be interpreted as L2 regularization on the parameters (Ali et al., 2020), so fine-tuning L2-regularizes θ̂g towards θp. Theorem 5.2. Under Assumptions 3.1, 3.2 & 5.1, for any q(x) with support on B2(0, xmax),
Pp(Z) [ Lq(θ̂g) ≤
2K2 g α
] ≥ 1− δ,
where
g = B̂ ·
√ 18K2(K2B̂2 + ξ2max)
ng
( d2 max { 1, log ( 1 +
8φmaxKng `2max
)} + log 4
δ
) .
Thus, if B is small and np is large, fθ̂g is accurate even if ng is small; we give a proof in Appendix D.

heading: 6 GENERALIZATION BOUNDS FOR NEURAL MODULE NETWORKS
text: While function identification enables robust generalization, many data generating processes are too complex to be identifiable. Neural module networks are designed to break complex prediction problems into smaller tasks that are individually easier to solve. These models take two kinds of input: (i) a sequence of tokens w (e.g., word embeddings) indicating the correct composition of modules,
and (ii) the input x to the modules. Then, the model predicts the sequence of modules j1...jT based on w, and runs the modules in sequence to obtain output x′ = fjT (...(fj1(x))...).
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; as a consequence, we assume the shift to the module input x can be arbitrary.
• Module composition: We consider shifts to the token sequence w. If the model mapping w to j1...jT is identifiable, then the entire model is identifiable. Instead, we 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 a simplified neural module network f , which includes (i) a set of neural modules {fj : X → X}kj=1, and (ii) a parser g : ZT → [k]T , where Z ⊆ Rr, with model class g ∈ G. We assume each component of fj(x) is computed by a separate quadratic neural network; we discuss the architecture of g below. Then, given an input x ∈ X ⊆ Rd and w ∈ W = ZT , the corresponding neural module network f : X ×W → X is defined by
f(x,w) = (fjT ◦ ... ◦ fj1)(x) = fjT (...(fj1(x))...) where j1...jT = g(w).
We assume that g has compositional structure—i.e., for some g̃ : [k]×Z → [k], we have g(w) = j1...jT where jt = {
0 if t = 0 g̃(zt, jt−1) otherwise,
where w = z1...zT . Intuitively, w is a sequence of word vectors; then, the current neural module jt = g̃(zt, jt−1) depends both on the current word vector zt 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—i.e., for each input x and sequence w, we have both the desired sequence j1...jT of neural modules, as well as the entire execution x0, x1, ..., xT , where x0 = x and xt+1 = fjt(xt) otherwise. Thus, we can use supervised learning to train the neural modules;2 in particular, we can construct labeled examples (jt−1, zt, jt) used to train the parser g̃, and labeled examples (xt, xt+1) to train the modules fjt . 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, we have the following straightforward result:
Lemma 6.1. Under Assumptions 3.1 & 3.2, with probability at least 1− dkδ, for each j ∈ [k], ‖f̂j(x)− f∗j (x)‖2 ≤ √ 2dK2
α =: f (∀x ∈ X ),
where f̂j is the estimated module and f∗j is the ground truth module.
This result follows straightforwardly from Theorem 3.6 along with a union bound. In contrast, we do not assume the parsing model robustly generalizes, but only on distribution. For the subsequent analysis, we can use any neural network models that satisfy the statement of Lemma 6.1.
Lemma 6.2. Under Assumptions 3.1 & 3.2, with probability at least 1− δ, we have
Pp̃(z,j) [ ˆ̃g(z, j) 6= g̃∗(z, j) ] ≤ 4Rn(G) + √ 2 log(2/δ)
n =: g,
where Rn(G) is the Rademacher complexity of G (including its loss function), where p(z, j) = T−1 ∑T t=1 pt(z, j), and where
p̃t(z, j) = { 1(j = 0) · p̃(z) if t = 1∑k j′=1 ∫ 1(j = g̃∗(z′, j′)) · p̃(z | z′) · p̃t−1(z′, j′)dz′ otherwise.
2Neural modules are often trained with only partial supervision (Andreas et al., 2016); we leave an analysis of this strategy to future work since our focus is on understanding generalization rather than learning dynamics.
This result is a 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 our compositional implementation of g to generalize, even on distribution. Intuitively, the distribution over (z, j) consists of both a unigram model over the word vectors:
p(z1, ..., zT ) = T∏ t=1 p̃(zt | zt−1),
where we define p̃(z1 | z0) = p̃(z1), as well as a unigram model over neural modules:
p(j1...jT | z1, ..., zT ) = T∏ t=1 1(jt = g̃ ∗(zt, jt−1)).
Next, we consider a shifted distribution q̃(z | z′), which is close to p̃(z | z′). Assumption 6.3. We have ‖q̃(· | z′)− p̃(· | z′)‖TV ≤ α.
Importantly, despite this assumption, the shift between the overall distributions p(z1, ..., zT ) and q(z1, ..., zT ) can still be large since it compounds exponentially across the steps t ∈ [T ]. Proposition 6.4. There exist p and q that satisfy Assumption 6.3, but ‖p−q‖TV = 2(1− (1−α/2)T ).
That is, even if the single step probabilities p̃(z | z′) and q̃(z | z′) have total variation (TV) distance bounded as in Assumption 6.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.1.
We show that neural module networks generalize since ĝ leverages the compositional structure of p. First, we show that under Assumption 6.3, the overall shift in the input distribution of ˆ̃g is bounded: Lemma 6.5. Under Assumptions 3.1, 3.2 & 6.3, we have ‖q̃ − p̃‖TV ≤ Tα, where p̃ is defined in Lemma 6.2 and q̃ is defined in Assumption 6.3.
That is, while the shift can compound across steps t, it does so only linearly; we give a proof in Appendix E.2. Next, we show that as a consequence, the error of ĝ is bounded.
Lemma 6.6. Under Assumptions 3.1, 3.2 & 6.3, and assuming that Pp(z,j)[ˆ̃g(z, j) 6= g̃∗(z, j)] ≤ g , we have that Pp(w)[ĝ(w) 6= g∗(w)] ≤ T g .
We give a proof in Appendix E.3. Finally, we have our main result. Theorem 6.7. Under Assumptions 3.1, 3.2 & 6.3, with probability at least 1− (dk + 1)δ, we have
Pq(w) [ ‖f̂(x,w)− f∗(x,w)‖2 ≤ T f ·max{KT−1, 1} ] ≥ 1− T g − T 2α.
We give a 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
Pp(w) [ ‖f̂(x,w)− f∗(x,w)‖2 ≤ T f ·max{KT−1, 1} ] ≥ 1− T g,
by the same argument as the proof of Theorem 6.7. The exponential dependence on K is unavoidable since K > 1 says that the modules fj can expand the input, which leads to exponential blowup in the magnitude of the output as a function of T , which also makes the estimation error exponential in T . Thus, the only cost to the distribution shift from p to q is the additional error probability T 2α.

heading: 7 CONCLUSION
text: We have presented a number of results demonstrating that over-parameterization does not fundamentally 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 us 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.

heading: A ADDITIONAL RELATED WORK
text: Low-rank matrix factorization. Our notion of functional identification for quadratic neural network is related to the low-rank matrix factorization literature. However, they impose a low-rank structure on the matrix to recover, and hence typically require extra conditions to identify the matrix—e.g., the restricted isometry property (RIP) (Candes & Plan, 2011; Ge et al., 2017a), or bounded `2 norm of noise vector (Kabanava et al., 2016). In contrast, we consider a more general case and do not assume any underlying structure of the matrix; in particular, since our goal is to capture over-parameterization of neural networks, our matrix is usually decomposed as φ = θθ>, where θ ∈ Rd×k and k ≥ 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 functional forms, ranging from linear (Abbasi-Yadkori et al., 2011; Rusmevichientong & Tsitsiklis, 2010) to neural tangent kernels (Zhou et al., 2020). Most of this work makes a realizability assumption that the model family contains the true model;3 implicitly, they consider model families where there is a 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 over-parameterized bandit problem that considers a model that is not parameter-identifiable; we find that similar regret results hold as long as the function represented by the model can be identified. Separately, Foster & Rakhlin (2020) makes a general connection between online regression oracles and the regret of a bandit algorithm; however, their approach only provides good guarantees when the regression oracle returns a 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 3 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 L̃p(φ) (Lemma 3.3) and the Lipschitz constant (Lemma 3.4) of f̃φ; also, Lemma 3.5 relies on a standard covering number argument. For our applications to 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.1 for transfer learning. Finally, to the best of our knowledge, our arguments in Section 6 are novel.
3Slightly different from realizability in PAC learning (Kearns et al., 1994), which says there is a model with zero true loss.

heading: B PROOFS FOR SECTION 3

heading: B.1 PROOF OF MINIMUM EIGENVALUE FOR UNIFORM DISTRIBUTION
text: In this section, we prove the claim that Assumption 3.2 holds for the covariate distribution where xi is an i.i.d. random variable with distribution Uniform(xi; [−1/2, 1/2]). To this end, note that
Ep(x)[(x>∆x)2] = Ep(x)   d∑ i,j=1 xixj∆ij 2 
= Ep(x) ∑ i x4i∆ 2 ii + ∑ i 6=j x2ix 2 j∆ii∆jj + 2 ∑ i 6=j x2ix 2 j∆ 2 ij  = 1
80 ∑ i ∆2ii + 1 144 ∑ i 6=j ∆ii∆jj + 1 72 ∑ i 6=j ∆2ij
=
( 1
80 − 1 144 )∑ i ∆2ii + 1 144 (∑ i ∆ii )2 + 1 72 ∑ i6=j ∆2ij
≥ 1 180 ‖∆‖2F ,
as claimed.

heading: B.2 PROOF OF LEMMA 3.3
text: We use the notation U : ∇2f(φ) : V to denote the matrix inner product 〈U,∇2f(φ)(V )〉 for U, V ∈ Rd×d. The Hessian ∇2f(φ) can be viewed as a d2 × d2 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 ∆
∆ : ∇2L̃p(φ) : ∆ = 2Ep(x)[(x>∆x)2] ≥ 2α‖∆‖2F , where the last inequality uses Assumption 3.2.

heading: B.3 PROOF OF LEMMA 3.4
text: By our definition, for any φ, φ′ ∈ Φ, |f̃φ(x)− f̃φ′(x)| = |(x>(φ− φ′)x)2| ≤ x2max‖φ− φ′‖F . Given our quadratic loss function, we have |(f̃φ(x)− f̃φ∗(x))2 − (f̃φ′(x)− f̃φ∗(x))2| ≤ |f̃φ(x)− f̃φ∗(x) + f̃φ′(x)− f̃φ∗(x)||f̃φ(x)− f̃φ′(x)| ≤ 4φmaxx4max‖φ− φ′‖F . Next, the true loss satisfies |L̃p(φ)− L̃p(φ′)| ≤ Ep(x)[|(f̃φ(x)− f̃φ∗(x))2 − (f̃φ′(x)− f̃φ∗(x))2|] ≤ 4φmaxx4max‖φ− φ′‖F . Finally, the empirical loss satisfies
| ˆ̃L(φ;Z)− ˆ̃L(φ′;Z)| = ∣∣∣∣∣ 1n n∑ i=1 [(f̃φ(xi)− yi)2 − (f̃φ′(xi)− yi)2] ∣∣∣∣∣ ≤ 1 n n∑ i=1 |(f̃φ(xi)− f̃φ∗(xi))2 − (f̃φ′(xi)− f̃φ∗(xi))2|
+ 1
n n∑ i=1 |ξi| · |f̃φ(xi)− f̃φ′(xi)|
≤ (4φmaxx4max + 2ξmaxx2max)‖φ− φ′‖F , as claimed.

heading: B.4 PROOF OF LEMMA 3.5
text: First, we have the following results: Lemma B.1 (Covering Number of Euclidean Ball). For a Euclidean ball in Rn1×n2 with radius R with respect to Frobenius norm, there exists an -net E such that
|E| ≤ ( 1 + 2R )n1n2 .
Proof. This claim follows by a direct application of Proposition 4.2.12 in Vershynin (2018).
Lemma B.2 (Hoeffding’s Inequality for Subgaussian Random Variables). Letting {zi}ni=1 be a set of independent σ-subgaussian random variables, then for all t ≥ 0, we have
Pr
[ 1
n n∑ i=1 zi ≥ t
] ≤ exp ( −2nt 2
σ2
) .
Proof. See Proposition 2.1 of Wainwright (2016).
Now, we prove Lemma 3.5. Consider an /(4K)-net E . Then, for any φ ∈ Φ, there exists φ′ ∈ E such that
|( ˆ̃L(φ;Z)− L̃p(φ))− ( ˆ̃L(φ′;Z)− L̃p(φ′))| ≤ 2K‖φ− φ′‖F ≤
2 .
Therefore, we have
Pp(Z) [ sup θ |L̂(g(θ);Z)− Lp(g(θ))− σ(Z)| ≥ ] = Pp(Z) [ sup φ∈Φ | ˆ̃L(φ;Z)− L̃p(φ)− σ(Z)| ≥ ]
≤ Pp(Z) [ max φ∈E | ˆ̃L(φ;Z)− L̃p(φ)− σ(Z)| ≥ 2 ] ≤ ∑ φ∈E Pp(Z) [ | ˆ̃L(φ;Z)− L̃p(φ)− σ(Z)| ≥ 2 ] . (13)
Now, defining
¯̃L(φ;Z) = 1
n n∑ i=1 (f̃φ(xi)− f̃φ∗(xi))2 and η̃(φ;Z) = 1 n n∑ i=1 (f̃φ(xi)− f̃φ∗(xi))ξi,
and recalling that σ(Z) = n−1 ∑n i=1 ξ 2 i , then we have
ˆ̃L(φ;Z) = ¯̃L(φ;Z) + 2η̃(φ;Z) + σ(Z).
Thus, continuing from (13), we have∑ φ∈E Pp(Z) [ | ˆ̃L(φ;Z)− L̃p(φ)− σ(Z)| ≥ 2 ] ≤ ∑ φ∈E Pp(Z) [ | ¯̃L(φ;Z)− L̃p(φ)|+ 2|η̃(φ;Z)| ≥ 2
] ≤ ∑ φ∈E ( Pp(Z) [ | ¯̃L(φ;Z)− L̃p(φ)| ≥ 6 ] + Pp(Z) [ |η̃(φ;Z)| ≥ 6 ]) . (14)
For the first term in (14), note that |(f̃φ(x) − f̃φ∗(x))2| ≤ `2max, so (f̃φ(x) − f̃φ∗(x))2 is `2maxsubgaussian; thus, by Lemma B.2, we have∑
φ∈E
Pp(Z) [ | ¯̃L(φ;Z)− L̃p(φ)| ≥
6
] ≤ 2|E| · exp ( − n 2
18`4max
) . (15)
Next, for the second term in (14), note that |(f̃φ(xi)−f̃φ∗(xi))ξi| ≤ `maxξmax, so (f̃φ(xi)−f̃φ∗(xi))ξi is `maxξmax-subgaussian; thus, by Lemma B.2, we have∑
φ∈E
Pp(Z) [ |η̃(φ;Z)| ≥
6
] ≤ 2|E| · exp ( − n 2
18`2maxξ 2 max
) . (16)
Combining (15) & (16), continuing from (14), we have∑ φ∈E ( Pp(Z) [ | ¯̃L(φ;Z)− L̃p(φ)| ≥ 6 ] + Pp(Z) [ |η̃(φ;Z)| ≥ 6 ]) ≤ 4|E| · exp ( − n 2
18`2max(` 2 max + ξ 2 max) ) ≤ 2 ( 1 + 8φmaxK )d2 · exp ( − n 2
18`2max(` 2 max + ξ 2 max) ) = 2 exp ( − n 2
18`2max(` 2 max + ξ 2 max)
+ d2 log ( 1 + 8φmaxK )) , (17)
where for the first inequality, we have used max{`2max, ξ2max} ≤ `2max + ξ2max, and the second inequality follows since by Lemma B.1, the covering number of the -net E of Φ satisfies
|E| ≤ ( 1 + 2φmax )d2 .
Finally, we choose so that (17) is smaller than δ—in particular, letting
=
√ 18`2max(` 2 max + ξ 2 max)
n
( d2 max { 1, log ( 1 + 8φmaxKn
`2max
)} + log 2
δ
) .
then continuing (17), we have
2 exp ( − n 2
18`2max(` 2 max + ξ 2 max)
+ d2 log ( 1 + 8φmaxK )) ≤ δ,
as claimed.

heading: B.5 PROOF OF PROPOSITION 3.8
text: ˆ̃L(φ;Z) is twice differentiable and convex in φ. Note that the minimization problem of L̂(θ;Z) is equivalent to that of ˆ̃L(g(θ̂);Z). We consider two cases. First, consider the case where θ̂ has rank d. The first order condition ∇L̂(θ;Z) = 0 is the same as∇ ˆ̃L(g(θ̂);Z) = 0, which gives
∇ ˆ̃L(φ̂;Z)θ̂ = 0. (18)
As θ̂ is of full row rank, there exists a matrix θ̂† ∈ Rk×d such that θ̂θ̂† = I (e.g. θ̂† = θ̂>(θ̂θ̂>)−1). We can right multiply the above equation by θ̂† and obtain that
∇ ˆ̃L(φ̂;Z) = 0.
As ˆ̃L(φ;Z) is convex in φ, the above implies φ̂ = g(θ̂) is a global minimum of ˆ̃L(φ;Z). Therefore, θ̂ is a global minimum of L̂(θ;Z). Next, consider the case where the rank of θ̂ is smaller than d. 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 implies
0 = ∇ ˆ̃L(φ̂;Z)θ̂θ̂> = ∇ ˆ̃L(φ̂;Z)φ̂. (19)
The Hessian of ˆ̃L(g(θ̂);Z) has
∇2 ˆ̃L(g(θ̂);Z)(∆,∆) = 2〈∇ ˆ̃L(φ̂;Z),∆∆>〉+∇2 ˆ̃L(φ̂;Z)(θ̂∆> + ∆θ̂>, θ̂∆> + ∆θ̂>).
As θ̂R is also a local minimum for any orthogonal matrix R (i.e., RR> = R>R = I), we can find a θ̂ with the last column being 0 by right multiplying certain R. Then, consider any ∆ with the first k− 1 columns being 0 and the last column being any u ∈ Rd. With this choice of ∆ and θ̂, θ̂∆> = 0. Therefore,
∇2 ˆ̃L(g(θ̂);Z)(∆,∆) = 2u>∇ ˆ̃L(φ̂;Z)u.
Since θ̂ is a local minimum of ˆ̃L(g(θ̂);Z), it holds that∇2 ˆ̃L(g(θ̂);Z)(∆,∆) ≥ 0, which implies
∇ ˆ̃L(φ̂;Z) 0. (20) Equation (19) and (20) together comprise the first order conditions of the convex minimization problem minφ 0 ˆ̃L(φ;Z). Thus, θ̂ is also a global minimum.

heading: C PROOFS FOR SECTION 4
text: First, we provide the full statement of Theorem 4.2 (including constants). Theorem C.1. The expected regret of Algorithm 1 is
R(T ) ≤ C0 + C1 · T 2/3 ( log ( 3 + 8φmaxKT
`2max
))1/3 ,
where
C0 = 64(φmax)
2d2+2 2d2−1
(135M(`max + ξmax)2d3) 2d2+2 2d2−1 (8φmaxK/`2max) 3d2 2d2−1
,
C1 = 162d 2(M2(`max + ξmax) 4φmax) 1/3.
Before we prove Theorem C.1, we first prove a preliminary result establishing an analog of the smooth best arm response property Rusmevichientong & Tsitsiklis (2010) to our setting. First, we have the following useful result: Lemma C.2. Let φ, φ′ ∈ Rd×d be symmetric matrices, let x, x′ ∈ Rd be eigenvectors of φ, φ′ corresponding to their top eigenvalue, such that ‖x‖2 = ‖x′‖2 = 1, and let λ1 ≥ λ2 ≥ ... ≥ λd be the eigenvalues of φ′. Suppose that 〈x, x′〉 ≥ 0. Then, we have
‖x− x′‖2 ≤ 23/2‖φ− φ′‖2
λ1 − λ2 .
Proof. See Corollary 3 of Yu et al. (2015).
Next, let χ : Rd×d → 2X denote the subset of reward-maximizing arms for g(θ) = θθ>—i.e., χ(φ) = arg max
x∈X x>φx,
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 φ ∈ Rd×d, there exists x ∈ χ(φ) and x∗ ∈ χ(φ∗) such that
‖x− x∗‖2 ≤M‖φ− φ∗‖F .
Proof. First, note that x, x∗ are eigenvectors of φ, φ∗ corresponding to their top eigenvalues, respectively. Next, note that if x∗ ∈ χ(φ∗), then we also have −x∗ ∈ χ(φ∗); thus, without loss of generality, we can assume that 〈x∗, x〉 ≥ 0. Also, note that ‖x‖2 = ‖x∗‖2 = 1 since the optimizer maximizes the magnitude of x. Thus, we have
‖x− x∗‖2 ≤ 23/2‖φ− φ∗‖2
λ1 − λ2 ≤M‖φ− φ∗‖F ,
where the second inequality follows by by Lemma C.2, and the third inequality follows by Assumption 4.1, as claimed.
Now, we prove Theorem C.1. The cumulative regret R(T ) of a horizon of T has that
R(T ) = E [ T∑ t=1 (fθ∗(x ∗)− fθ∗(xt)) ]
= E [ m∑ t=1 (fθ∗(x ∗)− fθ∗(xt)) + T∑ t=m+1 (fθ∗(x ∗)− fθ∗(xt)) ]
≤ 2mφmax + E
[ T∑
t=m+1
〈g(θ̂)− g(θ∗), x̂x̂> − x∗x∗>〉+ T∑
t=m+1
〈g(θ̂), x∗x∗> − x̂x̂>〉 ] ,
(21)
where θ̂ is an estimator that minimizes the empirical loss of the first m samples, x̂ ∈ χ(g(θ̂)) maximizes the estimated expected reward fθ̂(x), and 〈φ, φ ′〉 = ∑d i,j=1 φijφ ′ ij is the matrix inner product. Since x̂ is a maximizer of fθ̂(x) = 〈g(θ̂), xx >〉, we have 〈g(θ̂), x∗x∗> − x̂x̂>〉 ≤ 0. Thus, continuing from (21), we have
R(T ) ≤ 2mφmax + E
[ T∑
t=m+1
〈g(θ̂)− g(θ∗), x̂x̂> − x∗x∗>〉 ] ≤ 2mφmax + (T −m)E [ ‖g(θ̂)− g(θ∗)‖F ‖x̂x̂> − x∗x∗>‖F ] . (22)
To bound the second term in (22), note that
‖x̂x̂> − x∗x∗>‖F ≤ ‖x̂x̂> − x̂x∗>‖F + ‖x̂x∗> − x∗x∗>‖F ≤ 2M‖g(θ̂)− g(θ∗)‖F ,
where the last step follows by Lemma C.3. Next, by Theorem 3.6, we have
‖g(θ̂)− g(θ∗)‖F ≤ √ 2
α
= d
( 453`2max(` 2 max + ξ 2 max)
10m
( d2 max { 1, log ( 1 + 8φmaxKm
`2max
)} + log 2
δ ))1/4 with probability at least 1− δ. Now, defining the event
G = { ‖g(θ̂)− g(θ∗)‖F ≤ √ 2
α
} ,
letting
δ = 2 exp ( −d2 max { 1, log ( 1 + 8φmaxKm
`2max
)}) ,
and continuing from (22), we have R(T ) ≤ 2mφmax + T · E [ ‖g(θ̂)− g(θ∗)‖F ‖x̂x̂> − x∗x∗>‖F1(G) ] + 4Tφmax · P(Gc)
≤ 2mφmax + 2MT · E [ ‖g(θ̂)− g(θ∗)‖2F ∣∣∣ G]+ 4Tφmax · P(Gc) ≤ 2mφmax + 270M(`max + ξmax)2d3T √ log(3 + 8φmaxKT/`2max)
m + 8Tφmax (8φmaxKm/`2max) d2 .
(23)
The third term in inequality (23) is smaller than the second term when
m ≥
( 4φmax
135M(`max + ξmax)2d3(8φmaxK/`2max) d2 √ log(3 + 8φmaxKT/`2max)
)1/(d2−1/2) . (24)
For a choice of m satisfying (24), continuing from (23), we have R(T ) ≤ 2mφmax + 540M(`max + ξmax)2d3T √ log(3 + 8φmaxKT/`2max)
m . (25)
Next, we choose m to minimize the upper bound in (24) for sufficiently large T—in particular,
m =  ( 135M(`max + ξmax) 2d3T √ log(3 + 8φmaxKT/`2max) φmax ) 2 3  . (26) With this choice of m, we have
R(T ) ≤ 162(M2(`max + ξmax)4φmax) 1 3 d2T 2/3(log(3 + 8φmaxKT/` 2 max)) 1 3 .
Finally, note that (24) holds under the choice of m in (26) for T satisfying
T √ log(3 + 8φmaxKT/`2max) ≥ 64(φmax) 2d2+2 2d2−1
(135M(`max + ξmax)2d3) 2d2+2 2d2−1 (8φmaxK/`2max) 3d2 2d2−1
.
The claim follows.

heading: D PROOF OF THEOREM 5.2
text: First, we have the following key result:
Lemma D.1. Let θ, θ′ ∈ Rd×k, and let φ = θθ> and φ′ = θ′θ′>. Assume that ‖φ− φ′‖F ≤ η, and that σmin(θ) ≥ σ0 > 0, where σmin(θ) is the minimum singular value of θ (more precisely, the dth largest singular value). Then, there exist orthogonal matrices R,R′ ∈ Rk×k such that
‖θR− θ′R′‖F ≤ η
σ0 . (27)
Proof. Consider the SVDs θ = UΣV > and θ′ = U ′Σ′V ′>, where U,U ′ ∈ Rd×d, Σ,Σ′ ∈ Rd×d, and V, V ′ ∈ Rk×d; then, we have φ = UΣ2U> and φ′ = U ′Σ′2U ′>. Then, we claim that the choices R = V U> andR′ = V ′U ′> satisfy (27). In particular, note that θR = UΣU> and θ′R′ = U ′Σ′U ′>, since V >V = V ′>V ′ = Id since k ≥ d, where Id ∈ Rd×d is the d-dimensional identity matrix. Thus, it suffices to show that
σ0‖UΣU> − U ′Σ′U ′>‖F ≤ η. (28)
To this end, note that
η ≥ ‖φ− φ′‖F = ‖UΣ2U> − U ′Σ′2U ′>‖F = ‖U ′>UΣ2 − Σ′2U ′>U‖F , (29)
where in 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
(WΣ)ij = d∑ k=1 WikΣkj = WijΣjj (30)
(ΣW )ij = d∑ k=1 Σ′ikWkj = WijΣ ′ ii (31)
(WΣ2)ij = d∑ k=1 Wik(Σ 2)kj = WijΣ 2 jj (32)
(Σ2W )ij = d∑ k=1 (Σ′2)ikWkj = WijΣ ′2 ii . (33)
Then, continuing from (29), we have
η2 ≥ ‖WΣ2 − Σ′2W‖2F = d∑
i,j=1
W 2ij(Σ 2 jj − Σ′2ii)2
= d∑ i,j=1 W 2ij(Σjj − Σ′ii)2(Σjj + Σ′ii)2 ≥ d∑
i,j=1
W 2ij(Σjj − Σ′ii)2σ20
= σ20‖WΣ− Σ′W‖2F = σ20‖U ′>UΣ− Σ′U ′>U‖2F = σ20‖UΣU> − U ′Σ′U ′>‖2F ,
where on the first line, we have used (32) & (33), on the third line we have used Σjj ≥ σ0, on the fourth line we have used (30) & (31), and on the last line we have multiplied on by U ′ on the left U> on the right, again using the fact that the Frobenius 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 et al. (2017a) for quadratic neural networks.
Now, we prove Theorem 5.2. First, by directly applying the arguments in the proof of Theorem 3.6, we have ‖g(θ̂p)− g(θ∗p)‖F ≤ √
2 p α
with probability at least 1 − δ/2. However, θ̂p itself may not be close to θ∗p. Instead, applying Lemma D.1 with θ = θ̂p and θ′ = θ∗p, and with η = √ 2 p/α, there exists an orthogonal matrix Rp = R ′R> that “aligns” θ̂p with θ∗p , yielding
‖θ̂p − θ∗pRp‖F ≤ 1
σ0
√ 2 p α ,
where σ0 is the minimum singular value of θ∗p. Now, let θ̃g = θ ∗ gRp, and note that this is a global minimizer (i.e., g(θ̃g) = g(θ∗g)), since Rp is orthogonal. Then, we have
‖θ̃g − θ̂p‖F ≤ ‖θ∗gRp − θ∗pRp‖F + ‖θ∗pRp − θ̂p‖F
≤ ‖θ∗g − θ∗p‖F + 1
σ0
√ 2 p α
≤ B + 1 σ0 √ 2 p α
(34)
with probability at least 1− δ/2. In other words, an alternative global minimizer θ̃g exists within a small Frobenius norm of our proxy estimator θ̂p, even if θ̂p is not close to θ∗p .
Finally, on the event that (34) holds, note that for θ ∈ B2(θ̂p, B̂), we have the alternative upper bound
|fθ(x)− fθ∗g (x))| ≤ K‖g(θ)− g(θ ∗ g)‖F ≤ KB̂,
where the first inequality holds by Lemma 3.4; thus, we can take `max = KB̂. Thus, on the event that (34) holds, by Theorem 3.6, we have
Pp(Z) [ Lq(θ̂g) ≤
2K2 g α
] ≥ 1− δ
2 ,
so the claim follows by a union bound.

heading: E PROOFS FOR SECTION 6

heading: E.1 PROOF OF PROPOSITION 6.4
text: Suppose that zt ∈ {0, 1} is binary, z0 = 0, and p̃(z | z′) = { 1 if z = z′
0 otherwise.
In particular, since z0 = 0, p(w) = 1(w = w0) places all weight on the zero sequence w0 = 0...0. Next, consider the shifted distribution
q̃(zt | zt−1) =  1 if z = z′ = 1 1− α/2 if z = z′ = 0 α/2 otherwise.
Note that ‖p̃(· | z′)− q̃(· | z′)‖TV ≤ α, so Assumption 6.3 is satisfied. Note that
q(w0) = T∏ t=1 q̃(0 | 0) = (1− α/2)T .
As a consequence, we have ‖p− q‖TV = ∑ w∈W |p(w)− q(w)|
= |p(w0)− q(w0)|+ ∑
w∈W\{w0}
q(w)
= (1− (1− α/2)T ) + (1− (1− α/2)T ) = 2(1− (1− α/2)T ),
as claimed.

heading: E.2 PROOF OF LEMMA 6.5
text: Note that
‖q̃t − p̃t‖TV
= k∑ j=1 ∫ |q̃t(z, j)− p̃t(z, j)|dz
= k∑ j=1 k∑ j′=1 ∫ 1(j = g̃∗(z′, j′)) · |q̃(z | z′)q̃t−1(z′, j′)− p̃(z | z′)p̃t−1(z′, j′)|dz′dz
= k∑ j′=1 ∫ |q̃(z | z′)q̃t−1(z′, j′)− p̃(z | z′)p̃t−1(z′, j′)|dz′dz ≤ k∑
j′=1
∫ |q̃(z | z′)− p̃(z | z′)| · q̃t−1(z′, j′) + p̃(z | z′) · |q̃t−1(z′, j′)− p̃t−1(z′, j′)|dz′dz
≤ k∑
j′=1
∫ α · q̃t−1(z′, j′) + |q̃t−1(z′, j′)− p̃t−1(z′, j′)|dz′
≤ α+ ‖q̃t−1 − p̃t−1‖TV. Since q0(z, j) = p0(z, j) for all z ∈ Z and j ∈ [k], by induction, ‖q̃t − p̃t‖TV ≤ tα. Thus, we have
‖q̃ − p̃‖TV ≤ 1
T T∑ t=1 ‖q̃t − p̃t‖ ≤ Tα,
as claimed.

heading: E.3 PROOF OF LEMMA 6.6
text: First, we prove the following lemma. Lemma E.1. We have
p̃t(zt, jt−1) = k∑ j1,...,jt−2 ∫ (t−1∏ τ=1 1(jτ = g̃ ∗(zτ , jτ−1)) ) · p(z1, ..., zt)dz1...dzt−1.
Proof. For the base case, we have
p̃2(z2, j1) = k∑ j0=1 ∫ 1(j1 = g̃ ∗(z1, j0)) · p̃(z2 | z1) · p̃1(z1, j0)dz1
= k∑ j0=1 ∫ 1(j1 = g̃ ∗(z1, j0)) · p̃(z2 | z1) · 1(j = 0) · p̃(z1)dz1
= ∫ 1(j1 = g̃ ∗(z1, j0)) · p(z1, z2)dz1,
as claimed. For the inductive case, we have
p̃t(zt, jt−1) = k∑ jt−2=1 ∫ 1(jt−1 = g̃ ∗(zt−1, jt−2)) · p̃(zt | zt−1) · p̃t−1(zt−1, jt−2)dzt−1
= k∑ j1,...,jt−2=1 ∫ (t−1∏ τ=1 1(jτ−1 = g̃ ∗(zτ−1, jτ−2)) ) · p(z1, ..., zt)dz1...dzt−1.
as claimed.
Now, we prove Lemma 6.6. First, note that for each t ∈ [T ], we have
Pp(w) [ (ĝ(w)t 6= g∗(w)t) ∧ ( t−1∧ τ=1 ĝ(w)τ = g ∗(w)τ )]
= ∫ 1(ĝ(w)t 6= g∗(w)t) · ( t−1∏ τ=1 1(ĝ(w)τ = g ∗(w)τ ) ) · p(w)dw
= k∑ j1,...,jt−1=1 ∫ 1(ĝ(w)t 6= g∗(w)t) · ( t−1∏ τ=1 1(ĝ(w)τ = g ∗(w)τ ) ) · p(j1...jt−1 | w) · p(w)dw = k∑
j1,...,jt−1=1
∫ 1(ĝ(w)t 6= g∗(w)t) · · ( t−1∏ τ=1 1(ĝ(w)τ = g ∗(w)τ ) ) · ( t−1∏ τ=1 1(jτ = g̃ ∗(zτ , jτ−1)) ) · p(w)dw
= k∑ j1,...,jt−1=1 ∫ 1(ˆ̃g(zτ , jτ−1) 6= g̃∗(zτ , jτ−1)) · ( t−1∏ τ=1 1(ˆ̃g(zτ , jτ−1) = g̃ ∗(zτ , jτ−1)) )
· ( t−1∏ τ=1 1(jτ = g̃ ∗(zτ , jτ−1)) ) · p(w)dw
≤ k∑
j1,...,jt−1=1
∫ 1(ˆ̃g(zτ , jτ−1) 6= g̃∗(zτ , jτ−1)) · ( t−1∏ τ=1 1(jτ = g̃ ∗(zτ , jτ−1)) ) · p(w)dw
= k∑ j1,...,jt−1=1 ∫ 1(ˆ̃g(zτ , jτ−1) 6= g̃∗(zτ , jτ−1)) · ( t−1∏ τ=1 1(jτ = g̃ ∗(zτ , jτ−1)) ) · p(z1, ..., zt)dz1...dzt
= Ppt(z,j) [ ˆ̃g(zτ , jτ−1) 6= g̃∗(zτ , jτ−1) ] ,
where the last step follows from Lemma E.1. Now, note that
Pp(w)[ĝ(w) 6= g∗(w)] = T∑ t=1 Pp(w)
[ (ĝ(w)t 6= g∗(w)t) ∧ ( t−1∧ τ=1 ĝ(w)τ = g ∗(w)τ )]
≤ T∑ t=1 Pp̃t(z,j) [ ˆ̃g(z, j) 6= g̃∗(z, j) ] ≤ T g,
as claimed.

heading: E.4 PROOF OF THEOREM 6.7
text: First, we show that Pq(z,j)[ˆ̃g(z, j) 6= g̃∗(z, j)] ≤ g + Tα. To this end, note that
Pq(z,j)[ˆ̃g(z, j) 6= g̃∗(z, j)] = Pp(z,j)[ˆ̃g(z, j) 6= g̃∗(z, j)] + Pq(z,j)[ˆ̃g(z, j) 6= g̃∗(z, j)]− Pp(z,j)[ˆ̃g(z, j) 6= g̃∗(z, j)]
≤ g + k∑ j=1 ∫ 1(ˆ̃g(z, j) 6= g̃∗(z, j)) · |q̃(z, j)− p̃(z, j)|dz
≤ g + ‖q̃ − p̃‖TV ≤ g + Tα.
Next, by Lemma 6.6 with q in place of p and g+Tα in place of g , we have Pq(w)[ĝ(w) 6= g∗(w)] ≤ T g + T 2α. Then, assuming that ĝ(w) = g∗(w), we have
‖f∗(x,w)− f̂(x,w)‖2 = ‖(f∗jT ◦ ... ◦ f ∗ j1)(x)− (f̂jT ◦ ... ◦ f̂j1)(x)‖2
≤ T∑ t=1 ‖(f∗jT ◦ ... ◦ f ∗ jt+1 ◦ f ∗ jt ◦ f̂jt−1 ◦ ... ◦ f̂j1)(x)− (f ∗ jT ◦ ... ◦ f ∗ jt+1 ◦ f̂jt ◦ f̂jt−1 ◦ ... ◦ f̂j1)(x)‖2
≤ T∑ t=1 KT−t · ‖(f∗jt ◦ f̂jt−1 ◦ ... ◦ f̂j1)(x)− (f̂jt ◦ f̂jt−1 ◦ ... ◦ f̂j1)(x)‖2
≤ T∑ t=1 KT−t f ≤ T f ·max{KT−1, 1}.
The claim follows by a union bound.

Review summaries:
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.

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.