Going Beyond Linear RL: Sample Efficient Neural Function Approximation

Baihe Huang [email protected]. Peking University. Kaixuan Huang [email protected]. Princeton University. Sham M. Kakade [email protected]. University of Washington Jason D. Lee [email protected]. Princeton University. Qi Lei [email protected]. Princeton University Runzhe Wang [email protected]. Princeton University Jiaqi Yang [email protected]. Tsinghua University

Abstract

Deep Reinforcement Learning (RL) powered by neural net approximation of the Q function has had enormous empirical success. While the theory of RL has traditionally focused on linear function approximation (or eluder dimension) approaches, little is known about nonlinear RL with neural net approximations of the Q functions. This is the focus of this work, where we study function approximation with two-layer neural networks (considering both ReLU and polynomial activation functions). Our first result is a computationally and statistically efficient algorithm in the generative model setting under completeness for two-layer neural networks. Our second result considers this setting but under only realizability of the neural net function class. Here, assuming deterministic dynamics, the sample complexity scales linearly in the algebraic dimension. In all cases, our results significantly improve upon what can be attained with linear (or eluder dimension) methods.

1 Introduction

In reinforcement learning (RL), an agent aims to learn the optimal decision-making rule by interacting with an unknown environment [SB18]. Deep Reinforcement Learning, empowered by deep neural networks [LBH15, GBC16], has achieved tremendous success in various real-world applications, such as Go [SHM⁺16], Atari [MKS⁺13a], Dota2 [BBC⁺19], Texas Holdém poker [MSB⁺17], and autonomous driving [SSSS16]. Those modern RL applications are characterized by large state-action spaces, and the empirical success of deep RL corroborates the observation that deep neural networks can extrapolate well across state-action spaces [HIB⁺18, MKS⁺13b, LHP⁺15].

Although in practice non-linear function approximation scheme is prevalent, theoretical understandings of the sample complexity of RL mainly focus on tabular or linear function approximation settings [SLW⁺06, JOA10, AOM17, JAZBJ18, Rus19, ZB19, AYLSW19, JYWJ20a, JYWJ20b, WWDK21]. These results rely on finite state space or exact linear approximations. Recently, sample efficient algorithms under non-linear function approximation settings are proposed [WVR17, DJK⁺18, DKJ⁺19, DPWZ20, LCYW19, WCYW20, DYM21, ZLG20, YJW⁺20]. Those algorithms are based on Bellman rank [JKA⁺17], eluder dimension [RVR13b], neural tangent kernel [JGH18, AZLS19, DLL⁺19, ZCZG18], or sequential Rademacher complexity [RST15a, RST15b]. However, the understanding of how deep RL learns and generalizes in large state spaces is far from complete. Whereas the aforementioned works study function approximation structures that possess the nice properties of linear models, such as low information gain and low eluder dimensions, the highly non-linear nature of neural networks renders challenges on the applicability of existent analysis to deep RL. For one thing, recent wisdoms in deep learning theory cast doubt on the ability of neural tangent kernel and random features to model the actual neural networks. Indeed, the neural tangent kernel approximately reduces neural networks to random feature models, but the RKHS norm of neural networks is exponential [YS19]. Moreover, it remains unclear what neural networks models have low eluder dimensions. For example, recent works [DYM21, LKFS21] show that two layer neural networks have exponential eluder dimension in the dimension of features. Furthermore, [MPSL21] demonstrates hard instances where learning RL is exponentially hard even in the case that the target value function can be approximated as a neural network and the optimal policy is softmax linear. Thus, the mismatch between the empirical success of deep RL and its theoretical understanding remains significant, which yields the following important question:

What are the structural properties that allow sample-efficient algorithms for RL with neural network function approximation?

In this paper, we advance the understanding of the above question by displaying several structural properties that allow efficient RL algorithms with neural function approximations. We consider several value function approximation models that possess high information gain and high eluder dimension. Specifically, we study two structures, namely two-layer neural networks and structured polynomials (i.e. two-layer neural networks with polynomial activation functions), under two RL settings, namely RL with simulator model and online RL. In the simulator (generative model) setting [Kak03, SWW⁺18], the agent can simulate the MDP at any state-action pair. In online RL, the agent can only start at an initial state and interact with the MDP step by step. The goal in both settings is to find a near-optimal policy while minimizing the number of samples used.

We obtain the following results. For the simulator setting, we propose sample-efficient algorithms for RL with two-layer neural network function approximation. Under either policy completeness, Bellman completeness, or gap conditions, our method provably learns near-optimal policy with polynomial sample complexities. For online RL, we provide sample-efficient algorithms for RL with structured polynomial function approximation. When the transition is deterministic, we also present sample-efficient algorithms under only the realizability assumption [DKWY20a, WWK21]. Our main techniques are based on neural network recovery [ZSJ⁺17, JSA15, GLM18a], and algebraic geometry [Mil17, Sha13, BCR13, WX19].

1.1 Summary of our results

Our main results in different settings are summarized in Table 1. We consider two-layer neural networks $f(x)=\langle v,\sigma(Wx)\rangle$ (where $\sigma$ is ReLU activation) and rank $k$ polynomials (see Example 4.3).

Table 1: Baselines and our main results for the sample complexity to find an

\epsilon

-optimal policy.

	rank $k$ polynomial		Neural Net of Width $k$
	Sim. + Det. (R)	Onl. + Det. (R)	Sim. + Det. (R)	Sim. + Gap. (R)	Sim. + Stoch. (C)
Baseline	$O(d^{p})$	$O(d^{p})$	$O(d^{\mathrm{poly}(1/\epsilon)})$ (*)	$O(d^{\mathrm{poly}(1/\epsilon)})$	$O(d^{\mathrm{poly}(1/\epsilon)})$
Our results	${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}O(dk)}$	${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}O(dk)}$	${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\tilde{O}(\mathrm{poly}(d)\cdot\exp(k))}$	${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\tilde{O}(\mathrm{poly}(d,k))}$	${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\tilde{O}(\mathrm{poly}(d,k)/\epsilon^{2})}$

We make the following elaborations on Table 1.

•

For simplicity, we display only the dependence on the feature dimension $d$ , network width or polynomial rank $k$ , precision $\epsilon$ , and degree $p$ (of polynomials).
•

In the table Sim. denotes simulator model, Onl. denotes online RL, Det. denotes deterministic transitions, Stoch. denotes stochastic transitions, Gap. denotes gap condition, (R) denotes realizability assumption only, and (C) denotes completeness assumption (either policy complete or Bellman complete) together with realizability assumption.
•

We apply [DLMW20] for the deterministic transition baseline, and apply [DKWY20b] for the stochastic transition baseline. We are unaware of any methods that can directly learn MDP with neural network value function approximation¹¹1Prior work on neural function approximation has focused on neural tangent kernels, which would require $d^{\mathrm{poly}(1/\epsilon)}$ to approximate a two-layer network [GMMM21]. .
•

In polynomial case, the baseline first vectorizes the tensor $\begin{pmatrix}1\\ x\end{pmatrix}^{\otimes p}$ into a vector in $\mathbb{R}^{(d+1)^{p}}$ and then performs on this vector. In the neural network case, the baseline uses a polynomial of degree $1/\epsilon$ to approximate the neural network with precision $\epsilon$ and then vectorizes the polynomial into a vector in $\mathbb{R}^{d^{\mathrm{poly}(1/\epsilon)}}$ . The baseline method for realizable model (denoted by (*)) needs a further gap assumption of $\text{gap}\geq d^{\mathrm{poly}(1/\epsilon)}\epsilon$ to avoid the approximation error from escalating [DLMW20]; note for small $\epsilon$ this condition never holds but we include it in the table for the sake of comparison.
•

In rank $k$ polynomial case, our result $O(dk)$ in simulator model can be found in Theorem 4.7 and our result $O(dk)$ in online RL model can be found in Theorem 4.8. These results only require a realizability assumption. Efficient explorations are guaranteed by algebraic-geometric arguments. In neural network model, our result $\tilde{O}(\mathrm{poly}(d)\cdot\exp(k))$ in simulator model can be found in Theorem 3.4. This result also only relies on the realizability assumption. For stochastic transitions, our result $\tilde{O}(\mathrm{poly}(d,k)/\epsilon^{2})$ works for either policy complete or Bellman complete settings, as in Theorem 3.5 and Theorem 3.6 respectively. The $\tilde{O}(\mathrm{poly}(d,k))$ result for gap condition can be found in Theorem 3.8.

1.2 Related Work

Linear Function Approximation.

RL with linear function approximation has been widely studied under various settings, including linear MDP and linear mixture MDP [JYWJ20b, ZLKB20, YW20]. While these papers have proved efficient regret and sample complexity bounds, their analyses relied heavily on two techniques: they used the confidence ellipsoid to quantify the uncertainty, and they used the elliptical potential lemma to bound the total uncertainty [AYPS11]. These two techniques were integral to their analyses but are so restrictive that they generally do not extend to nonlinear cases.

Eluder Dimension.

[RVR13a, ORVR13] proposed eluder dimension, a complexity measure of the function space, and proved regret and sample complexity bounds that scaled with the eluder dimension, for bandits and reinforcement learning [WSY20a, JLM21]. They also showed that the eluder dimension is small in several settings, including generalized linear models and LQR. However, as shown in [DYM21], the eluder dimension could be exponentially large even with a single ReLU neuron, which suggested the eluder dimension would face difficulty in dealing with neural network cases. The eluder dimension is only known to give non-trivial bounds for linear function classes and monotone functions of linear function classes. For structured polynomial classes, the eluder dimension simply embeds into an ambient linear space of dimension $d^{p}$ , where $d$ is the dimension, and $p$ is the degree. This parallels the lower bounds in linearization / neural tangent kernel (NTK) works [WLLM19, GMMM19, AZL19], which show that linearization also incurs a similarly large penalty of $d^{p}$ sample complexity, and more advanced algorithm design is need to circumvent linearization[BL20, CBL⁺20, FLYZ20, WGL⁺19, GCL⁺19, NGL⁺19, GLM18b, MGW⁺20, HWLM20, WWL⁺20, DML21].

Bellman Rank and Completeness.

[JKA⁺17, SJK⁺19] studied RL with general function approximation. They used Bellman rank to measure the error of the function class under the Bellman operator and gave proved bounds in the term of it. Recently, [DKL⁺21a] propose bilinear rank and encompass more function approximation models. However, it is hard to bound either the Bellman rank or the bilinear rank for neural nets. Therefore, their results are not known to include the neural network approximation setting. Another line of work shows that exponential sample complexity is unavoidable even with good representations [DKWY20b, WAS20, WWK21], which implies the realizability assumption alone might be insufficient for function approximations.

Deterministic RL

Deterministic system is often the starting case in the study of sample-efficient algorithms, where the issue of exploration and exploitation trade-off is more clearly revealed since both the transition kernel and reward function are deterministic. The seminal work [WVR13] proposes a sample-efficient algorithm for Q-learning that works for a family of function classes. Recently, [DLMW20] studies the agnostic setting where the optimal Q-function can only be approximated by a function class with approximation error. The algorithm in [DLMW20] learns the optimal policy with the number of trajectories linear with the eluder dimension.

2 Preliminaries

An episodic Markov Decision Process (MDP) is defined by the tuple $\mathcal{M}=(\mathcal{S},\mathcal{A},H,\mathbb{P},r)$ where $\mathcal{S}$ is the state space, $\mathcal{A}$ is the action set, $H$ is the number of time steps in each episode, $\mathbb{P}$ is the transition kernel and $r$ is the reward function. In each episode the agent starts at a fixed initial state $s_{1}$ and at each time step $h\in[H]$ it takes action $a_{h}$ , receives reward $r_{h}(s_{h},a_{h})$ and transits to $s_{h+1}\sim\mathbb{P}(\cdot|s_{h},a_{h})$ .

A deterministic policy $\pi$ is a length- $H$ sequence of functions $\pi=\{\pi_{h}:\mathcal{S}\mapsto\mathcal{A}\}_{h=1}^{H}$ . Given a policy $\pi$ , we define the value function $V_{h}^{\pi}(s)$ as the expected sum of reward under policy $\pi$ starting from $s_{h}=s$ :

\displaystyle V_{h}^{\pi}(s):=\mathbb{E}\left[\sum_{t=h}^{H}r_{t}(s_{t},a_{t})|s_{h}=s\right]

and we define the Q function $Q_{h}^{\pi}(s,a)$ as the the expected sum of reward taking action $a$ in state $s_{h}=s$ and then following $\pi$ :

\displaystyle Q_{h}^{\pi}(s,a):=\mathbb{E}\left[\sum_{t=h}^{H}r_{t}(s_{t},a_{t})|s_{h}=s,a_{h}=a\right].

The Bellman operator ${\cal T}_{h}$ applied to Q-function $Q_{h+1}$ is defined as follow

\displaystyle{\cal T}_{h}(Q_{h+1})(s,a):=r_{h}(s,a)+\mathbb{E}_{s^{\prime}\sim\mathbb{P}(\cdot|s,a)}[\max_{a^{\prime}}Q_{h+1}(s^{\prime},a^{\prime})].

There exists an optimal policy $\pi^{\ast}$ that gives the optimal value function for all states, i.e. $V^{\pi^{\ast}}_{h}(s)=\sup_{\pi}V_{h}^{\pi}(s)$ for all $h\in[H]$ and $s\in\mathcal{S}$ . For notational simplicity we abbreviate $V^{\pi^{\ast}}$ as $V^{\ast}$ and correspondingly $Q^{\pi^{\ast}}$ as $Q^{\ast}$ . Therefore $Q^{\ast}$ satisfies the following Bellman optimality equations for all $s\in\mathcal{S}$ , $a\in\mathcal{A}$ and $h\in[H]$ :

\displaystyle Q^{\ast}_{h}(s,a)={\cal T}_{h}(Q^{\ast}_{h+1})(s,a).

The goal is to find a policy $\pi$ that is $\epsilon$ -optimal in the sense that $V_{1}^{\ast}(s_{1})-V_{1}^{\pi}(s_{1})\leq\epsilon$ , within a small number of samples. We consider two query models of interacting with MDP:

•

In the simulator model ([Kak03], [SWW⁺18]), the agent interacts with a black-box that simulates the MDP. At each time step $h\in[H]$ , the agent can start at a state-action pair $(s,a)$ and interact with the black box by executing some policy $\pi$ chosen by the agent.
•

In online RL, the agent can only start at the initial state and interact with the MDP by using a policy and observing the rewards and the next states. In each episode $k$ , the agent proposes a policy $\pi^{k}$ based on all history up to episode $k-1$ and executes $\pi^{k}$ to generate a single trajectory $\{s_{h}^{k},a_{h}^{k}\}_{h=1}^{H}$ with $a_{h}^{k}=\pi^{k}_{h}(s_{h}^{k})$ and $s_{h+1}^{k}\sim\mathbb{P}(\cdot|s_{h}^{k},a_{h}^{k})$ .

2.1 Function approximation

In reinforcement learning with value function approximation, the learner is given a function class ${\cal F}={\cal F}_{1}\times\cdots\times{\cal F}_{H}$ , where ${\cal F}_{h}\subset\{f:\mathcal{S}\times\mathcal{A}\mapsto[0,1]\}$ is a set of candidate functions to approximate $Q^{\ast}$ . The following assumption is commonly adopted in the literature [JYWJ20a, WSY20b, JLM21, DKL⁺21b].

Assumption 2.1 (Realizability).

$Q^{\ast}_{h}\in{\cal F}_{h}$ for all $h\in[H]$ .

The function approximation is equipped with feature mapping $\phi:\mathcal{S}\times\mathcal{A}\mapsto\{u\in\mathbb{R}^{d}:\|u\|_{2}\leq B_{\phi}\}$ that is known to the agent. We focus the continuous action setting (e.g. in control and robotics problems) and make the following regularity assumption on the feature function $\phi$ .

Assumption 2.2 (Bounded features).

Assume $\phi(s,a)\leq B_{\phi},\forall(s,a)\in{\cal S\times A}$ .

Notation

For any vector $x\in\mathbb{R}^{d}$ , let $x_{\max}:=\max_{i\in[d]}x_{i}$ and $x_{\min}:=\min_{i\in[d]}x_{i}$ . Let $s_{i}(\cdot)$ denote the $i$ -th singular value, $s_{\min}(\cdot)$ denotes the minimum eigenvalue and $s_{\max}(\cdot)$ denotes the maximum eigenvalue. The conditional number is defined by $\kappa(\cdot):=s_{\max}(\cdot)/s_{\min}(\cdot)$ . We use $\otimes$ to denote Kronecker product and $\circ$ to denote Hadamard product. For a given integer $H$ , we use $[H]$ to denote the set $\{1,2,\ldots,H\}$ . For a function $f:\mathfrak{X}\mapsto\mathfrak{Y}$ , we use $f^{-1}(y):=\{x\in\mathfrak{X}:f(x)=y\}$ to denote the preimage of $y\in\mathfrak{Y}$ . We use the shorthand $x\lesssim y$ ( $x\gtrsim y$ ) to indicate $x\leq O(y)$ ( $x\geq\Omega(y)$ ).

3 Neural Network Function Approximation

In this section we show sample-efficient algorithms with neural network function approximations. The function class of interest is given in the following definition. More general neural network class is discussed in Appendix A.5.

Definition 3.1 (Neural Network Function Class).

We use ${\cal F}_{NN}$ to denote the function class of $f(\phi(s,a)):\mathcal{S}\times\mathcal{A}\mapsto\mathbb{R}$ where $f(x)=\langle v,\sigma(Wx)\rangle:\mathbb{R}^{d}\mapsto\mathbb{R}$ is a two-layer neural network where $\sigma$ is ReLU, $\|W\|_{F}\leq B_{W}$ , $v\in\{\pm 1\}^{k}$ , $\prod_{i=1}^{k}s_{i}(W)/s_{\min}(W)\leq\lambda$ , $s_{\max}(W)/s_{\min}(W)\leq\kappa$ and $k\leq d$ . Here $\phi:\mathcal{A}\times\mathcal{S}\mapsto\mathbb{R}^{d}$ is a known feature map whose image contains a ball $\{u\in\mathbb{R}^{d}:\|u\|_{2}\leq\delta_{\phi}\}$ with $\delta_{\phi}\geq d\cdot\mathrm{polylog}(d)$ .²²2Here the $\delta_{\phi}$ is chosen only for simplicity. In general this assumption can be relaxed to that the image of $\phi$ contains an arbitrary dense ball near the origin, since one can always rescale the feature mapping in the neural function approximation.

We introduce the following completeness properties in the setting of value function approximations. Along with Assumption 2.1, they are commonly adopted in the literature .

Definition 3.2 (Policy complete).

Given MDP $\mathcal{M}=(\mathcal{S},\mathcal{A},\mathbb{P},r,H)$ , function class ${\cal F}_{h}:\mathcal{S}\times\mathcal{A}\mapsto\mathbb{R},h\in[H]$ is called policy complete if for all $\pi$ and $h\in[H]$ , $Q^{\pi}_{h}\in{\cal F}_{h}$ .

Definition 3.3 (Bellman complete).

Given MDP $\mathcal{M}=(\mathcal{S},\mathcal{A},\mathbb{P},r,H)$ , function class ${\cal F}_{h}:\mathcal{S}\times\mathcal{A}\mapsto\mathbb{R},h\in[H]$ is called Bellman complete if for all $h\in[H]$ and $Q_{h+1}\in{\cal F}_{h+1}$ , ${\cal T}_{h}(Q_{h+1})\in{\cal F}_{h}$ .

3.1 Warmup: Realizable $Q^{\ast}$ with deterministic transition

We start by considering the case when the transition kernel is deterministic. In this case only Assumption 2.1 is required for the expressiveness of neural network function approximations. Algorithm 1 learns optimal policy from time step $H$ to $1$ . Suppose we have learned policies $\pi_{h+1},\dots,\pi_{H}$ at level $h$ and they are exactly the optimal policies. We first explore features $\phi(s_{h}^{i},a_{h}^{i})$ over a standard Gaussian distribution, and if $\|\phi(s_{h}^{i},a_{h}^{i})\|_{2}\geq\delta_{\phi}$ then we simply skip this trial. Recall that $\delta_{\phi}\geq d\cdot\mathrm{poly}\log(d)$ , so with high probability (w.r.t $d$ ) almost all feature samples will be explored. We next construct an estimate $\widehat{Q}_{h}^{i}$ of $Q^{\ast}(s^{i}_{h},a^{i}_{h})$ by collecting cumulative rewards using $\pi_{h+1},\dots,\pi_{H}$ as the roll-out. Since the transition is deterministic, $\widehat{Q}_{h}^{i}=Q^{\ast}(s^{i}_{h},a^{i}_{h})$ for all explored samples $(s^{i}_{h},a^{i}_{h})$ . Recall that $Q^{\ast}_{h}$ is a two-layer neural network, we can now recover its parameters in Line 12 exactly by invoking techniques in neural network optimization (see, e.g. [JSA15, GLM17, ZSJ⁺17]). Details of this step can be found in Appendix A.5, where the method is mainly based on [ZSJ⁺17]. This means the reconstructed $\widehat{Q}_{h}$ in Line 13 is precisely $Q^{\ast}$ , and the algorithm can thus find optimal policy $\pi^{\ast}_{h}$ in the $h$ -th level.

Algorithm 1 Learning realizable

Q^{\ast}

with deterministic transition

1:for

h=H,\ldots 1

2: Sample

x^{i}_{h},i\in[n]

from standard Gaussian

\mathcal{N}(0,I_{d})

3: for

i\in[n]

4: if

\|x^{i}_{h}\|\leq\delta_{\phi}

then

5: Find

(s^{i}_{h},a^{i}_{h})\in\phi^{-1}(x^{i}_{h})

and locate the state

s^{i}_{h}

in the generative model

6: Pull action

a^{i}_{h}

and use

\pi_{h+1},\dots,\pi_{H}

as the roll-out to collect rewards

r_{h}^{(i)},\dots,r_{H}^{(i)}

7: Construct estimation

{\widehat{Q}}_{h}^{i}\leftarrow r_{h}^{(i)}+\cdots+r_{H}^{(i)}

8: else

9: Let

{\widehat{Q}}_{h}^{i}\leftarrow 0

10: end if

11: end for

12: Compute

(v_{h},W_{h})\leftarrow\textsc{NeuralNetRecovery}(\{(x^{i}_{h},{\widehat{Q}}_{h}^{i}):i\in[n]\})

13: Set

\widehat{Q}_{h}(s,a)\leftarrow v_{h}^{\top}\sigma(W_{h}\phi(s,a))

14: Let

\pi_{h}(s)\leftarrow\operatorname*{arg\,max}_{a\in\mathcal{S}}\leavevmode\nobreak\ \widehat{Q}_{h}(s,a)

15:end for

16:Return

\pi_{1},\dots,\pi_{H}

Theorem 3.4.

(Informal) If $n\geq d\cdot\mathrm{poly}(\kappa,k,\lambda,\log d,B_{W},B_{\phi},H)$ , then with high probability Algorithm 1 learns the optimal policy.

The formal statement and complete proof are deferred to the Appendix A.1. The main idea of exact neural network recovery can be summarized in the following. We first use method of moments to find a ‘rough’ parameter recovery. If this ‘rough’ recovery is sufficiently close to the true parameter, the empirical $\ell_{2}$ loss function is locally strongly convex and there is unique global minimum. Then we can apply gradient descent to find this global minimum which is exactly the true parameter.

3.2 Policy complete neural function approximation

Now we consider general stochastic transitions. Difficulties arise in this scenario due to noises in the estimation of Q-functions. In the presence of model misspecification, these noises cause estimation errors to amplify through levels and require samples to be exponential in $H$ . In this section, we show that neural network function approximation is still learnable, assuming the function class ${\cal F}_{NN}$ is policy complete with regard to MDP $\mathcal{M}$ . Thus for all $\pi\in\Pi$ , we can denote $Q_{h}^{\pi}(s,a)=\langle v^{\pi},\sigma(W^{\pi}\phi(s,a))\rangle$ .

Algorithm 2 Learn policy complete NN with simulator.

1:for

h=H,\ldots 1

2: Sample

x^{i}_{h},i\in[n]

from standard Gaussian

\mathcal{N}(0,I_{d})

3: for

i\in[n]

4: if

\|x^{i}_{h}\|\leq\delta_{\phi}

then

5: Find

(s^{i}_{h},a^{i}_{h})\in\phi^{-1}(x^{i}_{h})

and locate the state

s^{i}_{h}

in the generative model

6: Pull action

a^{i}_{h}

and use

\pi_{h+1},\dots,\pi_{H}

as the roll-out to collect rewards

r_{h}^{(i)},\dots,r_{H}^{(i)}

7: Construct unbiased estimation of

{Q}_{h}^{\pi_{h+1},\dots,\pi_{H}}(s^{i}_{h},a^{i}_{h})

{\widehat{Q}}_{h}^{i}\leftarrow r_{h}^{(i)}+\cdots+r_{H}^{(i)}

8: else

9: Let

{\widehat{Q}}_{h}^{i}\leftarrow 0

10: end if

11: end for

12: Compute

(v_{h},W_{h})\leftarrow\textsc{NeuralNetNoisyRecovery}(\{(x^{i}_{h},{\widehat{Q}}_{h}^{i}):i\in[n]\})

13: Set

\widehat{Q}_{h}(s,a)\leftarrow v_{h}^{\top}\sigma(W_{h}\phi(s,a))

14: Let

\pi_{h}(s)\leftarrow\operatorname*{arg\,max}_{a\in\mathcal{S}}\leavevmode\nobreak\ \widehat{Q}_{h}(s,a)

15:end for

16:Return

\pi_{1},\dots,\pi_{H}

Algorithm 2 learns policy from level $H,H-1,\dots,1$ . In level $h$ , the algorithm has learned policy $\pi_{h+1},\dots,\pi_{H}$ that is only sub-optimal by $(H-h)\epsilon/H$ . Then it explores features $\phi(s,a)$ from $\mathcal{N}(0,I_{d})$ . The algorithm then queries $(s,a)$ and uses learned policy $\pi_{h+1},\dots,\pi_{H}$ as roll out to collect an unbiased estimate of the Q-function ${Q}_{h}^{\pi_{h+1},\dots,\pi_{H}}(s,a)$ . Since ${Q}_{h}^{\pi_{h+1},\dots,\pi_{H}}(s,a)\in{\cal F}_{NN}$ is a two-layer neural network, it can then be recovered from samples. Details of this step can be found in Appendix A.5, where the methods are mainly based on [ZSJ⁺17]. The algorithm then reconstructs this Q-function and finds its optimal policy $\pi_{h}$ .

Theorem 3.5.

(Informal) Fix $\epsilon,t$ , if $n\geq\epsilon^{-2}\cdot d\cdot\mathrm{poly}(\kappa,k,B_{W},B_{\phi},H,\log(d/t))$ , then with probability at least $1-t$ Algorithm 2 returns an $\epsilon$ -optimal policy $\pi$ .

The formal statement and complete proof are deferred to Appendix A.2. Notice that unlike the case of Theorem 3.4, the sample complexity does not depend on $\lambda$ , thus avoiding the potential exponential dependence in $k$ .

The main idea of the proof is that at each time step a neural network surrogate of $Q^{\ast}$ can be constructed by the policy already learned. Suppose we have learned $\pi_{h+1},\dots,\pi_{H}$ in level $h$ , then from policy completeness $Q^{\pi_{h+1},\dots,\pi_{H}}_{h}$ belongs to ${\cal F}_{NN}$ and we can interact with the simulator to obtain its estimate $\widehat{Q}_{h}$ . If $\|\widehat{Q}_{h}-Q^{\pi_{h+1},\dots,\pi_{H}}_{h}\|_{\infty}$ is small, the optimistic planning based on $\widehat{Q}_{h}$ is not far from the optimal policy of $Q^{\pi_{h+1},\dots,\pi_{H}}_{h}$ . Therefore the errors can be decoupled into the errors in recovering $Q^{\pi_{h+1},\dots,\pi_{H}}_{h}$ and the suboptimality of $Q^{\pi_{h+1},\dots,\pi_{H}}_{h}$ , which depends on level $h+1$ . This reasoning can then be recursively performed to level $H$ , and thus we can bound the suboptimality of $\pi_{h}$ .

3.3 Bellman complete neural function approximation

In addition to policy completeness, we show that neural network function approximation can also learn efficiently under the setting where the function class ${\cal F}_{NN}$ is Bellman complete with regard to MDP $\mathcal{M}$ . Specifically, for $Q_{h+1}\in{\cal F}_{h+1}$ , there are $v^{Q_{h+1}}$ and $W^{Q_{h+1}}$ such that ${\cal T}_{h}(Q_{h+1})(s,a)=\langle v^{Q_{h+1}},\sigma(W^{Q_{h+1}}\phi(s,a))\rangle$ .

Algorithm 3 is similar to the algorithm in previous section. Suppose in level $h$ , the algorithm has constructed the Q-function $\widehat{Q}_{h+1}(s,a)=v_{h+1}^{\top}\sigma(W_{h+1}\phi(s,a))$ that is $(H-h)\epsilon/H$ -close to the optimal $Q^{\ast}_{h+1}$ . It then recovers weights $v_{h},W_{h}$ from ${\cal T}_{h}(\widehat{Q}_{h+1})(s,a)=\langle v^{\widehat{Q}_{h+1}},\sigma(W^{\widehat{Q}_{h+1}}\phi(s,a))\rangle$ , using unbiased estimates $r_{h}(s^{i}_{h},a^{i}_{h})+\widehat{V}_{h+1}(s_{h+1}^{i})$ . The Q-function $\widehat{Q}_{h}(s,a)=v_{h}^{\top}\sigma(W_{h}\phi(s,a))$ reconstructed from weights $v_{h},W_{h}$ is thus $(H-h+1)\epsilon/H$ -close to the $Q^{\ast}_{h}$ .

Algorithm 3 Learn Bellman complete NN with simulator.

1:for

h=H,\ldots 1

2: Sample

x^{i}_{h},i\in[n]

from standard Gaussian

\mathcal{N}(0,I_{d})

3: for

i\in[n]

4: if

\|x^{i}_{h}\|\leq\delta_{\phi}

then

5: Find

(s^{i}_{h},a^{i}_{h})\in\phi^{-1}(x^{i}_{h})

and locate the state

s^{i}_{h}

in the generative model

6: Pull action

a^{i}_{h}

and and observe

r_{h}(s^{i}_{h},a^{i}_{h}),s_{h+1}^{i}

7: Construct unbiased estimation of

{\cal T}_{h}(\widehat{Q}_{h+1})(s^{i}_{h},a^{i}_{h})

{\widehat{Q}}_{h}^{i}\leftarrow r_{h}(s^{i}_{h},a^{i}_{h})+\widehat{V}_{h+1}(s_{h+1}^{i})

8: else

9: Let

{\widehat{Q}}_{h}^{i}\leftarrow 0

10: end if

11: end for

12: Compute

(v_{h},W_{h})\leftarrow\textsc{NeuralNetNoisyRecovery}(\{(x^{i}_{h},{\widehat{Q}}_{h}^{i}):i\in[n]\})

13: Set

\widehat{Q}_{h}(s,a)\leftarrow v_{h}^{\top}\sigma(W_{h}\phi(s,a))

and

\widehat{V}_{h}\leftarrow\max_{a\in\mathcal{A}}\widehat{Q}_{h}(s,a)

14: Let

\pi_{h}(s)\leftarrow\operatorname*{arg\,max}_{a\in\mathcal{A}}\leavevmode\nobreak\ \widehat{Q}_{h}(s,a)

15:end for

16:Return

\pi_{1},\dots,\pi_{H}

Theorem 3.6.

Due to Bellman completeness, the error of estimation $\widehat{Q}_{h}$ can be controlled recursively. In fact, we can show $\|\widehat{Q}_{h}-Q^{\ast}(s,a)\|_{\infty}$ is small by induction. The formal statement and detailed proof are deferred to Appendix A.3. Similar to Theorem 3.5, the sample complexity does not explicitly depend on $\lambda$ , thus avoiding potentially exponential dependence in $k$ .

3.4 Realizable $Q^{\ast}$ with optimality gap

In this section we consider MDPs where there is a non-zero gap between the optimal policy and any other ones. This concept, known as optimality gap, is widely used in reinforcement learning and bandit literature [DKJ⁺19, DKWY20b, DLMW20].

Definition 3.7.

The optimality gap is defined as

\displaystyle\rho=\inf_{a:Q^{\ast}(s,a)\neq V^{\ast}(s)}V^{\ast}(s)-Q^{\ast}(s,a).

We show that in the presence positive optimality gap, there exists an algorithm that can learn the optimal policy with polynomial samples even without the completeness assumptions. Intuitively, this is because one only needs to recover the neural network up to precision $\rho/4$ in order to make sure the greedy policy is identical to the optimal one. The formal statement and proof are deferred to Appendix A.4.

Theorem 3.8.

(Informal) Fix $t\in(0,1)$ , if $n=\frac{d}{\rho^{2}}\cdot\mathrm{poly}(\kappa,k,B_{W},B_{\phi},H,\log(d/t))$ , then with probability at least $1-t$ there exists an algorithm that returns the optimal policy $\pi^{\ast}$ .

Remark 3.9.

In all aforementioned methods, there are two key components that allow efficient learning. First, the exploration is conducted in a way that guarantees an $\ell_{\infty}$ recovery of candidate functions. By $\ell_{\infty}$ recovery we mean the algorithm recovers a candidate Q-function in this class deviating from the target function $Q^{\ast}$ by at most $\epsilon$ uniformly for all state-action pairs in the domain of interest. This notion of learning guarantee has received study in active learning [Han14, KAH⁺17] and recently gain interest in contextual bandits [FAD⁺18]. Second, the agent constructs unbiased estimators of certain approximations to $Q^{\ast}$ that lie in the neural function approximation class. This allows the recovery error to decouple linearly across time steps, which is made possible in several well-posed MDP instances, such as deterministic MDPs, MDPs with completeness assumptions, and MDPs with gap conditions. In principle, we note that provably efficient RL algorithms with general function approximation is possible as long as the above two components are present. We will see in the next section another example of learning RL with highly non-convex function approximations, where the function class of interest, admissible polynomial families, also allows for exploration schemes to achieve $\ell_{\infty}$ recovery.

4 Polynomial Realizability

In this section, we study the sample complexity to learn deterministic MDPs under polynomial realizability. We identify sufficient and necessary conditions for efficiently learning the MDPs for two different settings — the generative model setting and the online RL setting. Specifically, we show that if the image of the feature map $\phi_{h}(s_{h},a_{h})$ satisfies some positive measure conditions, then by random exploring, we can identify the optimal policy with samples linear in the algebraic dimension of the underlying polynomial class. We also provide a lower bound example showing the separation between the two settings.

Next, we introduce the notion of Admissible Polynomial Families, which are the families of structured polynomials that enable efficient learning.

Definition 4.1 (Admissible Polynomial Families).

For $x\in\mathbb{R}^{d}$ , denote $\widetilde{x}=[1,x^{\top}]^{\top}$ . Let $\mathcal{X}:=\left\{\widetilde{x}^{\otimes p}:x\in\mathbb{R}^{d}\right\}$ . For any algebraic variety ${\mathcal{V}}$ , we define ${\cal F}_{{\mathcal{V}}}:=\{f_{\Theta}(x)=\left\langle\Theta,\widetilde{x}^{\otimes p}\right\rangle:\Theta\in\mathcal{V}\}$ as the polynomial family parameterized by $\Theta\in{\mathcal{V}}$ . We say ${\cal F}_{{\mathcal{V}}}$ is admissible³³3Admissible means the dimension of $\mathcal{X}$ decreases by one when there is an additional linear constraint $\langle\Theta,X\rangle=0$ w.r.t. $\mathcal{X}$ , if for any $\Theta\in{\mathcal{V}}$ , $\text{dim}(\mathcal{X}\cap\{X\in\mathcal{X}:\langle X,\Theta\rangle=0\rangle\})<\text{dim}(\mathcal{X})=d$ . We define the dimension $D$ of the family to be the dimension of ${\mathcal{V}}$ .

The following theorem shows that to learn an admissible polynomial family, the sample complexity only scales with the algebraic dimension of the polynomial family.

Theorem 4.2 ([HHK⁺21]).

Consider the polynomial family ${\cal F}_{{\mathcal{V}}}$ of dimension $D$ . For $n\geq 2D$ , there exists a Lebesgue-measure zero set $N\in\mathbb{R}^{d}\times\ldots\mathbb{R}^{d}$ , such that if $(x_{1},\cdots,x_{n})\notin N$ , for any $y_{i}$ , there is a unique $f$ (or no such $f$ ) to the system of equations $y_{i}=f(x_{i})$ for $f\in{\cal F}_{{\mathcal{V}}}$ .

We give two important examples of admissible polynomial families with low dimension.

Example 4.3.

(Low-rank Polynomial of rank $k$ ) The function $f\in\mathcal{F}_{{\mathcal{V}}}$ is a polynomial with $k$ terms, that is

F(x)=\sum_{i=1}^{k}\lambda_{i}\langle v_{i},x\rangle^{p_{i}},

where $p=\max\{p_{i}\}$ . The dimension of this family is upper bounded by $D\leq dk$ . Neural network with monomial/polynomial activation functions are low-rank polynomials.

Example 4.4.

The function $f\in\mathcal{F}_{\mathcal{V}}$ is of the form $f(x)=q(Ux)$ , where $U\in\mathbb{R}^{k\times d}$ and $q$ is a degree $p$ polynomial. The polynomial $q$ and matrix $U$ are unknown. The dimension of this family is upper bounded by $D\leq d(k+1)^{p}$ .

Next, we introduce the notion of positive measure.

Definition 4.5.

We say a measurable set $E\in\mathbb{R}^{d}$ is of positive measure if $\mu(E)>0$ , where $\mu$ is the standard Lebesgue measure on $\mathbb{R}^{d}$ .

If a measurable set $E$ satisfies $\mu(E)>0$ , then there exists a procedure to draw samples from $E$ , such that for any $N\subset\mathbb{R}^{d}$ of Lebesgue-measure zero, the probability that the sample falls in $N$ is zero. In fact, the sampling probability can be given by $\mathbb{P}_{x\in\mathcal{N}(0,I_{d})}(\cdot|x\in E)$ . The intuition behind its definition is that for all admissible polynomial families, the set of $(x_{1},\cdots,x_{n})$ with "redundant information" about learning the parameter $\Theta$ is of Lebesgue-measure zero. Therefore, a positive measure set allows you to query randomly and avoids getting coherent measurements.

Next two theorems identify the sufficent conditions for efficiently learning deterministic MDPs under polynomial realizability. Specifically, under online RL setting, we require the strong assumption that the set $\{\phi_{h}(s,a)|a\in{\mathcal{A}}\}$ is of positive measure for all $h\in[H]$ and all $s\in{\mathcal{S}}$ , while under generative model setting, we only require the union set $\bigcup_{s\in{\mathcal{S}}}\{\phi_{h}(s,a)|a\in{\mathcal{A}}\}$ to be of positive measure for all $h\in[H]$ . The algorithms for solving the both cases are summarized in Algorithms 4 and 5.

Assumption 4.6 (Polynomial Realizability).

For all $h\in[H]$ , $Q^{\ast}_{h}(s_{h},a_{h})$ , viewed as the function of $\phi_{h}(s_{h},a_{h})$ , lies in some admissible polynomial family ${\cal F}_{{\mathcal{V}}_{h}}$ with dimension bounded by $D$ .

Theorem 4.7.

For the generative model setting, assume that the set $\{\phi_{h}(s,a)\mid s\in{\mathcal{S}},a\in{\mathcal{A}}\}$ is of positive measure at any level $h$ . Under the polynomial realizability, Algorithm 4 almost surely learns the optimal policy $\pi^{\ast}$ with at most $N=2DH$ samples.

Theorem 4.8.

For the online RL setting, assume that $\{\phi_{h}(s,a)\mid a\in{\mathcal{A}}\}$ is of positive measure for every state $s$ at every level $h$ . Under polynomial realizability, within $T=2DH$ episodes, Algorithm 5 learns the optimal policy $\pi^{\ast}$ almost surely.

Algorithm 4 Dynamic programming under generative model settings

1:for

h=H,\cdots,1

2: Sample

2D

points

\{\phi_{h}(s_{h}^{(i)},a_{h}^{(i)})\}_{i=1}^{2D}

according to

\mathbb{P}_{x\in\mathcal{N}(0,I_{d})}(\cdot|x\in E_{h})

where

E_{h}=\{\phi_{h}(s,a)\mid s\in{\mathcal{S}},a\in{\mathcal{A}}\}

3: Query the generative model with state-action pair

(s_{h}^{(i)},a_{h}^{(i)})

at level

h

for

i=1,\dots,2D

, and observe the next state

\widetilde{s}_{h}^{(i)}

and reward

r_{h}^{(i)}

4: Solve for

Q^{*}_{h}

with the

2D

equations

Q^{*}_{h}(s_{h}^{(i)},a_{h}^{(i)})=r_{h}^{(i)}+V^{*}_{h+1}(\widetilde{s}_{h}^{(i)})

5: Set

\pi^{*}_{h}(s)=\operatorname*{arg\,max}_{a}Q^{*}_{h}(s,a)

and

V_{h}^{*}(s)=\max_{a}Q^{*}_{h}(s,a)

6:end for

7:Output

\pi^{*}

Algorithm 5 Dynamic programming under online RL settings

1:for

h=H,\cdots,1

2: Fix any action sequence

a_{1},\cdots,a_{h-1}

3: Play

a_{1},\cdots,a_{h-1}

for the first

h-1

levels and reach a state

s_{h}

. Sample

2D

points

\{\phi_{h}(s_{h},a_{h}^{(i)})\}_{i=1}^{2D}

according to

\mathbb{P}_{x\in\mathcal{N}(0,I_{d})}(\cdot|x\in E_{h})

where

E_{h}=\{\phi_{h}(s_{h},a)\mid a\in{\mathcal{A}}\}

4: Play

a_{h}^{(i)}

s_{h}

for

i=1,\dots,2D

, and observe the next state

\widetilde{s}_{h}^{(i)}

and reward

r_{h}^{(i)}

5: Solve for

Q^{*}_{h}

with the

2D

equations

Q^{*}_{h}(s_{h}^{(i)},a_{h}^{(i)})=r_{h}^{(i)}+V^{*}_{h+1}(\widetilde{s}_{h}^{(i)})

6: Set

\pi^{*}_{h}(s)=\operatorname*{arg\,max}_{a}Q^{*}_{h}(s,a)

and

V_{h}^{*}(s)=\max_{a}Q^{*}_{h}(s,a)

7:end for

8:Output

\pi^{*}

We remark that our Theorem 4.8 for learning MDPs under the online RL setting relies on a very strong assumption that allows the learner to explore randomly for any state. However, this assumption is necessary in some sense, as is suggested by our lower bound example in the next subsection.

4.1 Necessity of Generic Feature Maps in Online RL

In this section, we consider lower bounds for learning deterministic MDPs with polynomial realizable $Q^{*}$ under online RL setting. Our goal is to show that in the online setting the generic assumption on the feature maps $\phi_{h}(s,\cdot)$ is necessary. On the contrary, under the generative model setting one can efficiently learn the MDPs without such a strong assumption, since at every level $h$ the we can set the state arbitrarily.

MDP construction

We briefly introduce the intuition of our construction. Consider a family of MDPs with only two states ${\mathcal{S}}=\{S_{\text{good}},S_{\text{bad}}\}$ . we set the feature map $\phi_{h}$ such that, for the good state $S_{\text{good}}$ , it allows the learner to explore randomly, i.e., $\{\phi_{h}(S_{\text{good}},a)\mid a\in{\mathcal{A}}\}$ is of postive measure.

However, for the bad state $S_{\text{bad}}$ , all actions are mapped to some restricted set, which forbids random exploration, i.e., $\{\phi_{h}(S_{\text{bad}},a)\mid a\in{\mathcal{A}}\}$ is measure zero. This is illustrated in Figure 1.

Specifically, at least $\Omega(d^{p})$ actions are needed to identify the groud-truth polynomial of $Q_{h}^{*}$ for $S_{\text{bad}}$ , while $O(d)$ actions suffice for $S_{\text{good}}$ .

The transition $\mathbb{P}_{h}$ is constructed as $\mathbb{P}_{h}(s_{\text{bad}}|s,a)=1\text{ for all }s\in{\mathcal{S}},a\in{\mathcal{A}}$ , which means it is impossible for the online scenarios to reach the good state for $h>1$ .

Refer to caption — Figure 1: An illustration of the hard case for deterministic MDPs with polynomial realizable $Q^{*}$ . The image of the feature map $\phi_{h}$ at $S_{\text{good}}$ is of positive measure, while the image of $\phi_{h}$ at $S_{\text{bad}}$ is not. This makes it difficult to learn under the online RL setting.

Theorem 4.9.

There exists a family of MDPs satisfying Assumption 4.6, such that the set $\left\{\phi_{h}(s,a)\mid s\in{\mathcal{S}},a\in{\mathcal{A}}\right\}$ is of positive measure at any level $h$ , but for all $h$ there is some $s_{\mathrm{bad}}\in{\mathcal{S}}$ such that $\{\phi_{h}(s_{\mathrm{bad}},a)\mid a\in{\mathcal{A}}\}$ is measure zero. Under the online RL setting, any algorithm needs to play at least $\Omega(d^{p})$ episodes to identify the optimal policy. On the contrary, under the generative model setting, only $O(d)H$ samples are needed.

5 Conclusions

In this paper, we consider neural network and polynomial function approximations in the simulator and online settings. To our knowledge, this is the first paper that shows sample-efficient reinforcement learning is possible with neural net function approximation. Our results substantially improve upon what can be achieved with existing results that primarily rely on embedding neural networks into linear function classes. The analysis reveals that for function approximations that allows for efficient $\ell_{\infty}$ recovery, such as two layer neural networks and admissible polynomial families, reinforcement learning can be reduced to parameter recovery problems, as well-studied in theories for deep learning, phase retrieval, and etc. Our method can also be potentially extended to handle three-layer and deeper neural networks, with advanced tools in [FKR19, FVD19].

Our results for polynomial activation require deterministic transitions, since we cannot handle how noise propagates in solving polynomial equations. We leave to future work an in-depth study of the stability of roots of polynomial systems with noise, which is a fundamental mathematical problem and even unsolved for homogeneous polynomials. In particular, noisy tensor decomposition approaches combined with zeroth-order optimization may allow for stochastic transitions [HHK⁺21].

In the online RL setting, we can only show efficient learning under a very strong yet necessary assumption on the feature mapping. We leave to future work identifying more realistic and natural conditions which permit efficient learning in the online RL setting.

Finally, in future work, we hope to consider deep neural networks where parameter recovery or $\ell_{\infty}$ error is unattainable, and deep reinforcement learning with representation learning [YHLD20, DHK⁺20].

Acknowledgment

JDL acknowledges support of the ARO under MURI Award W911NF-11-1-0303, the Sloan Research Fellowship, NSF CCF 2002272, and an ONR Young Investigator Award. QL is supported by NSF 2030859 and the Computing Research Association for the CIFellows Project. SK acknowledges funding from the NSF Award CCF-1703574 and the ONR award N00014-18-1-2247. BH is supported by the Elite Undergraduate Training Program of School of Mathematical Sciences at Peking University.

References

[AOM17] Mohammad Gheshlaghi Azar, Ian Osband, and Rémi Munos. Minimax regret bounds for reinforcement learning. In Proceedings of the 34th International Conference on Machine Learning, pages 263–272, 2017.
[AYLSW19] Yasin Abbasi-Yadkori, Nevena Lazic, Csaba Szepesvari, and Gellert Weisz. Exploration-enhanced POLITEX. arXiv preprint arXiv:1908.10479, 2019.
[AYPS11] Yasin Abbasi-Yadkori, Dávid Pál, and Csaba Szepesvári. Improved algorithms for linear stochastic bandits. In Advances in Neural Information Processing Systems, pages 2312–2320, 2011.
[AZL19] Zeyuan Allen-Zhu and Yuanzhi Li. What can resnet learn efficiently, going beyond kernels? arXiv preprint arXiv:1905.10337, 2019.
[AZLS19] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-parameterization. In International Conference on Machine Learning (ICML), pages 242–252, 2019.
[BBC⁺19] Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemyslaw Debiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, and Chris Hesse. Dota 2 with large scale deep reinforcement learning. arXiv preprint arXiv:1912.06680, 2019.
[BCR13] Jacek Bochnak, Michel Coste, and Marie-Françoise Roy. Real algebraic geometry, volume 36. Springer Science & Business Media, 2013.
[BL20] Yu Bai and Jason D. Lee. Beyond linearization: On quadratic and higher-order approximation of wide neural networks. In International Conference on Learning Representations, 2020.
[CBL⁺20] Minshuo Chen, Yu Bai, Jason D Lee, Tuo Zhao, Huan Wang, Caiming Xiong, and Richard Socher. Towards understanding hierarchical learning: Benefits of neural representations. Neural Information Processing Systems (NeurIPS), 2020.
[DHK⁺20] Simon S Du, Wei Hu, Sham M Kakade, Jason D Lee, and Qi Lei. Few-shot learning via learning the representation, provably. arXiv preprint arXiv:2002.09434, 2020.
[DJK⁺18] Christoph Dann, Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E. Schapire. On oracle-efficient PAC-RL with rich observations. In Advances in Neural Information Processing Systems, 2018.
[DKJ⁺19] Simon S Du, Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudik, and John Langford. Provably efficient RL with rich observations via latent state decoding. In International Conference on Machine Learning, pages 1665–1674, 2019.
[DKL⁺21a] Simon S. Du, Sham M. Kakade, Jason D. Lee, Shachar Lovett, Gaurav Mahajan, Wen Sun, and Ruosong Wang. Bilinear classes: A structural framework for provable generalization in rl. arXiv preprint arXiv:2103.10897, 2021.
[DKL⁺21b] Simon S Du, Sham M Kakade, Jason D Lee, Shachar Lovett, Gaurav Mahajan, Wen Sun, and Ruosong Wang. Bilinear classes: A structural framework for provable generalization in RL. arXiv preprint arXiv:2103.10897, 2021.
[DKWY20a] Simon S. Du, Sham M. Kakade, Ruosong Wang, and Lin F. Yang. Is a good representation sufficient for sample efficient reinforcement learning? In International Conference on Learning Representations, 2020.
[DKWY20b] Simon S Du, Sham M Kakade, Ruosong Wang, and Lin F Yang. Is a good representation sufficient for sample efficient reinforcement learning? In International Conference on Learning Representations, 2020.
[DLL⁺19] 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 (ICML), pages 1675–1685. PMLR, 2019.
[DLMW20] Simon S Du, Jason D Lee, Gaurav Mahajan, and Ruosong Wang. Agnostic Q-learning with function approximation in deterministic systems: Tight bounds on approximation error and sample complexity. Advances in Neural Information Processing Systems, 2020.
[DML21] Alex Damian, Tengyu Ma, and Jason Lee. Label noise sgd provably prefers flat global minimizers. arXiv preprint arXiv:2106.06530, 2021.
[DPWZ20] Kefan Dong, Jian Peng, Yining Wang, and Yuan Zhou. $\sqrt{n}$ -regret for learning in Markov decision processes with function approximation and low Bellman rank. In Conference on Learning Theory, pages 1554–1557. PMLR, 2020.
[DYM21] Kefan Dong, Jiaqi Yang, and Tengyu Ma. Provable model-based nonlinear bandit and reinforcement learning: Shelve optimism, embrace virtual curvature. arXiv preprint arXiv:2102.04168, 2021.
[FAD⁺18] Dylan J. Foster, Alekh Agarwal, Miroslav Dudík, Haipeng Luo, and Robert E. Schapire. Practical contextual bandits with regression oracles. Proceedings of the 35th International Conference on Machine Learning, 2018.
[FKR19] Massimo Fornasier, Timo Klock, and Michael Rauchensteiner. Robust and resource efficient identification of two hidden layer neural networks. arXiv preprint arXiv:1907.00485, 2019.
[FLYZ20] Cong Fang, Jason D Lee, Pengkun Yang, and Tong Zhang. Modeling from features: a mean-field framework for over-parameterized deep neural networks. arXiv preprint arXiv:2007.01452, 2020.
[FVD19] Massimo Fornasier, Jan Vybíral, and Ingrid Daubechies. Robust and resource efficient identification of shallow neural networks by fewest samples. arXiv preprint arXiv:1804.01592, 2019.
[GBC16] Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep learning. MIT press, 2016.
[GCL⁺19] Ruiqi Gao, Tianle Cai, Haochuan Li, Liwei Wang, Cho-Jui Hsieh, and Jason D Lee. Convergence of adversarial training in overparametrized networks. Neural Information Processing Systems (NeurIPS), 2019.
[GLM17] Rong Ge, Jason D. Lee, and Tengyu Ma. Learning one-hidden-layer neural networks with landscape design. arXiv preprint arXiv:1711.00501, 2017.
[GLM18a] Rong Ge, Jason D. Lee, and Tengyu Ma. Learning one-hidden-layer neural networks with landscape design. In International Conference on Learning Representations (ICLR), 2018.
[GLM18b] Rong Ge, Jason D Lee, and Tengyu Ma. Learning one-hidden-layer neural networks with landscape design. International Conference on Learning Representations (ICLR), 2018.
[GMMM19] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Linearized two-layers neural networks in high dimension. arXiv preprint arXiv:1904.12191, 2019.
[GMMM21] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Linearized two-layers neural networks in high dimension. The Annals of Statistics, 49(2):1029–1054, 2021.
[Han14] Steve Hanneke. Active learning for cost-sensitive classification. Foundations and Trends^® in Machine Learning, 7(2-3):131–309, 2014.
[HHK⁺21] Baihe Huang, Kaixuan Huang, Sham M Kakade, Jason D Lee, Qi Lei, Runzhe Wang, and Jiaqi Yang. Optimal gradient-based algorithms for non-concave bandit optimization. arXiv preprint arXiv:2107.04518, 2021.
[HIB⁺18] Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. Deep reinforcement learning that matters. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
[HWLM20] Jeff Z. HaoChen, Colin Wei, Jason D. Lee, and Tengyu Ma. Shape matters: Understanding the implicit bias of the noise covariance. arXiv preprint arXiv:2006.08680, 2020.
[JAZBJ18] Chi Jin, Zeyuan Allen-Zhu, Sebastien Bubeck, and Michael I Jordan. Is Q-learning provably efficient? In Advances in Neural Information Processing Systems, pages 4863–4873, 2018.
[JGH18] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in neural information processing systems (NeurIPS), pages 8571–8580, 2018.
[JKA⁺17] Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E Schapire. Contextual decision processes with low Bellman rank are PAC-learnable. In Proceedings of the 34th International Conference on Machine Learning, pages 1704–1713, 2017.
[JLM21] Chi Jin, Qinghua Liu, and Sobhan Miryoosefi. Bellman eluder dimension: New rich classes of rl problems, and sample-efficient algorithms. arXiv preprint arXiv:2102.00875, 2021.
[JOA10] Thomas Jaksch, Ronald Ortner, and Peter Auer. Near-optimal regret bounds for reinforcement learning. Journal of Machine Learning Research, 11(Apr):1563–1600, 2010.
[JSA15] Majid Janzamin, Hanie Sedghi, and Anankumar Anima. Beating the perils of nonconvexity: Guaranteed training of neural networks using tensor methods. arXiv preprint arXiv:1506.08473, 2015.
[JYWJ20a] Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael I Jordan. Provably efficient reinforcement learning with linear function approximation. In Conference on Learning Theory, pages 2137–2143. PMLR, 2020.
[JYWJ20b] Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael I Jordan. Provably efficient reinforcement learning with linear function approximation. In Conference on Learning Theory, pages 2137–2143, 2020.
[KAH⁺17] Akshay Krishnamurthy, Alekh Agarwal, Tzu-Kuo Huang, Hal Daume III, and John Langford. Active learning for cost-sensitive classification. Proceedings of the 34th International Conference on Machine Learning, 2017.
[Kak03] Sham Machandranath Kakade. On the sample complexity of reinforcement learning. PhD thesis, University of College London, 2003.
[KCL15] Volodymyr Kuleshov, Arun Chaganty, and Percy Liang. Tensor factorization via matrix factorization. In Proceedings of the Eighteenth International Conference on Artificial Intelligence and Statistics (AISTATS), page 507–516, 2015.
[LBH15] Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. Nature, 521(7553):436–444, 2015.
[LCYW19] Boyi Liu, Qi Cai, Zhuoran Yang, and Zhaoran Wang. Neural proximal/trust region policy optimization attains globally optimal policy. arXiv preprint arXiv:1906.10306, 2019.
[LHP⁺15] Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv:1509.02971, 2015.
[LKFS21] Gene Li, Pritish Kamath, Dylan J. Foster, and Nathan Srebro. Eluder dimension and generalized rank. arXiv preprint arXiv:2104.06970, 2021.
[MGW⁺20] Edward Moroshko, Suriya Gunasekar, Blake Woodworth, Jason D Lee, Nathan Srebro, and Daniel Soudry. Implicit bias in deep linear classification: Initialization scale vs training accuracy. Neural Information Processing Systems (NeurIPS), 2020.
[Mil17] James S. Milne. Algebraic geometry (v6.02), 2017. Available at www.jmilne.org/math/.
[MKS⁺13a] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, DaanWierstra, and Martin Riedmiller. Playing atari with deep reinforcement learningep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013.
[MKS⁺13b] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing Atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013.
[MPSL21] Dhruv Malik, Aldo Pacchiano, Vishwak Srinivasan, and Yuanzhi Li. Sample efficient reinforcement learning in continuous state spaces: A perspective beyond linearity. arXiv preprint arXiv:2106.07814, 2021.
[MSB⁺17] Matej Moravčík, Martin Schmid, Neil Burch, Viliam Lisỳ, Dustin Morrill, Nolan Bard, Trevor Davis, Kevin Waugh, Michael Johanson, and Michael Bowling. Deepstack: Expert-level artificial intelligence in heads-up no-limit poker. Science, 356(6337):508–513, 2017.
[NGL⁺19] Mor Shpigel Nacson, Suriya Gunasekar, Jason Lee, Nathan Srebro, and Daniel Soudry. Lexicographic and depth-sensitive margins in homogeneous and non-homogeneous deep models. In International Conference on Machine Learning, pages 4683–4692. PMLR, 2019.
[ORVR13] Ian Osband, Daniel Russo, and Benjamin Van Roy. (more) efficient reinforcement learning via posterior sampling. In Advances in Neural Information Processing Systems, pages 3003–3011, 2013.
[RST15a] Alexander Rakhlin, Karthik Sridharan, and Ambuj Tewari. Online learning via sequential complexities. Journal of Machine Learning Research, 16(6):155–186, 2015.
[RST15b] Alexander Rakhlin, Karthik Sridharan, and Ambuj Tewari. Sequential complexities and uniform martingale laws of large numbers. Probability Theory and Related Fields, 161(1-2):111–153, 2015.
[Rus19] Daniel Russo. Worst-case regret bounds for exploration via randomized value functions. In Advances in Neural Information Processing Systems, pages 14433–14443, 2019.
[RVR13a] Dan Russo and Benjamin Van Roy. Eluder dimension and the sample complexity of optimistic exploration. In Advances in Neural Information Processing Systems, pages 2256–2264, 2013.
[RVR13b] Daniel Russo and Benjamin Van Roy. Eluder dimension and the sample complexity of optimistic exploration. In Advances in Neural Information Processing Systems, 2013.
[SB18] Richard S Sutton and Andrew G. Barto. Reinforcement learning: An introduction. MIT press, 2018.
[Sha13] Igor R Shafarevich. Basic Algebraic Geometry 1: Varieties in Projective Space. Springer Science & Business Media, 2013.
[SHM⁺16] David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, and Marc Lanctot. Mastering the game of Go with deep neural networks and tree search. Nature, 529(7587):484, 2016.
[SJK⁺19] Wen Sun, Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, and John Langford. Model-based RL in contextual decision processes: PAC bounds and exponential improvements over model-free approaches. In Conference on Learning Theory, pages 2898–2933, 2019.
[SLW⁺06] Alexander L Strehl, Lihong Li, Eric Wiewiora, John Langford, and Michael L Littman. PAC model-free reinforcement learning. In Proceedings of the 23rd international conference on Machine learning, pages 881–888. ACM, 2006.
[SSSS16] Shai Shalev-Shwartz, Shaked Shammah, and Amnon Shashua. Safe, multi-agent, reinforcement learning for autonomous driving. arXiv preprint arXiv:1610.03295, 2016.
[SWW⁺18] Aaron Sidford, Mengdi Wang, Xian Wu, Lin F Yang, and Yinyu Ye. Near-optimal time and sample complexities for solving discounted markov decision process with a generative model. arXiv preprint arXiv:1806.01492, 2018.
[WAS20] Gellert Weisz, Philip Amortila, and Csaba Szepesvári. Exponential lower bounds for planning in mdps with linearly-realizable optimal action-value functions. arXiv preprint arXiv:2010.01374, 2020.
[WCYW20] Lingxiao Wang, Qi Cai, Zhuoran Yang, and Zhaoran Wang. Neural policy gradient methods: Global optimality and rates of convergence. arXiv preprint arXiv:1909.01150, 2020.
[WGL⁺19] Blake Woodworth, Suriya Genesekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Kernel and deep regimes in overparametrized models. In Conference on Learning Theory (COLT), 2019.
[WLLM19] Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets vs their induced kernel. In Advances in Neural Information Processing Systems, pages 9709–9721, 2019.
[WSY20a] Ruosong Wang, Ruslan Salakhutdinov, and Lin F Yang. Provably efficient reinforcement learning with general value function approximation. arXiv preprint arXiv:2005.10804, 2020.
[WSY20b] Ruosong Wang, Ruslan Salakhutdinov, and Lin F Yang. Provably efficient reinforcement learning with general value function approximation. Advances in Neural Information Processing Systems, 2020.
[WVR13] Zheng Wen and Benjamin Van Roy. Efficient exploration and value function generalization in deterministic systems. In Advances in Neural Information Processing Systems, pages 3021–3029, 2013.
[WVR17] Zheng Wen and Benjamin Van Roy. Efficient reinforcement learning in deterministic systems with value function generalization. Math. Oper. Res., 42(3):762–782, 2017.
[WWDK21] Yining Wang, Ruosong Wang, Simon S. Du, and Akshay Krishnamurthy. Optimism in reinforcement learning with generalized linear function approximation. In International Conference on Learning Representations, 2021.
[WWK21] Yuanhao Wang, Ruosong Wang, and Sham M. Kakade. An exponential lower bound for linearly-realizable mdps with constant suboptimality gap. arXiv preprint arXiv:2103.12690, 2021.
[WWL⁺20] Xiang Wang, Chenwei Wu, Jason D Lee, Tengyu Ma, and Rong Ge. Beyond lazy training for over-parameterized tensor decomposition. Neural Information Processing Systems (NeurIPS), 2020.
[WX19] Yang Wang and Zhiqiang Xu. Generalized phase retrieval: measurement number, matrix recovery and beyond. Applied and Computational Harmonic Analysis, 47(2):423–446, 2019.
[YHLD20] Jiaqi Yang, Wei Hu, Jason D Lee, and Simon S Du. Provable benefits of representation learning in linear bandits. arXiv preprint arXiv:2010.06531, 2020.
[YJW⁺20] Zhuoran Yang, Chi Jin, Zhaoran Wang, Mengdi Wang, and Michael Jordan. Provably efficient reinforcement learning with kernel and neural function approximations. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 13903–13916. Curran Associates, Inc., 2020.
[YS19] Gilad Yehudai and Ohad Shamir. On the power and limitations of random features for understanding neural networks. arXiv preprint arXiv:1904.00687, 2019.
[YW20] Lin F Yang and Mengdi Wang. Reinforcement leaning in feature space: Matrix bandit, kernels, and regret bound. International Conference on Machine Learning, 2020.
[ZB19] Andrea Zanette and Emma Brunskill. Tighter problem-dependent regret bounds in reinforcement learning without domain knowledge using value function bounds. In International Conference on Machine Learning, pages 7304–7312, 2019.
[ZCZG18] Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Stochastic gradient descent optimizes over-parameterized deep relu networks. arXiv preprint arXiv:1811.08888, 2018.
[ZLG20] Dongruo Zhou, Lihong Li, and Quanquan Gu. Neural contextual bandits with ucb-based exploration. Proceedings of the 37th International Conference on Machine Learning, 2020.
[ZLKB20] Andrea Zanette, Alessandro Lazaric, Mykel Kochenderfer, and Emma Brunskill. Learning near optimal policies with low inherent Bellman error. In International Conference on Machine Learning, 2020.
[ZSJ⁺17] Kai Zhong, Zhao Song, Prateek Jain, Peter L. Bartlett, and Inderjit S. Dhillon. Recovery guarantees for one-hidden-layer neural networks. Proceedings of the Thirty-fourth International Conference on Machine Learning (ICML), 70, 2017.

Appendix A Omitted Proofs in Section 3

A.1 Proof of Section 3.1

Theorem A.1 (Formal statement of Theorem 3.4).

Consider MDP $\mathcal{M}$ where the transition is deterministic. Assume the function class in Definition 3.1 satisfies Assumption 2.1 and Assumption 2.2. For any $t\in(0,1)$ , if $d\geq\Omega(\log(B_{W}/\lambda))$ and $n\geq d\cdot\mathrm{poly}(\kappa,k,\lambda,B_{W},B_{\phi},H,\log(d/t))$ , then with probability at least $1-t$ Algorithm 1 returns the optimal policy $\pi^{\ast}$ .

Proof.

Use $\pi^{\ast}_{1},\dots,\pi^{\ast}_{H}$ to denote the global optimal policy. We prove that Algorithm 1 learns $\pi_{h}^{\ast}$ from $h=H$ to $h=1$ .

At level $H$ , the query obtains exact $Q_{H}^{\ast}(s,a)$ . Therefore by Theorem A.15, $\widehat{Q}_{H}=Q_{H}^{\ast}$ and thus the optimal planning finds $\pi_{H}=\pi_{H}^{\ast}$ . Suppose we have learned $\pi_{h+1}^{\ast},\dots,\pi_{H}^{\ast}$ at level $h$ . Due to deterministic transition, the query obtains exact $Q_{h}^{\ast}(s,a)$ . Therefore by Theorem A.15, $\widehat{Q}_{h}=Q_{h}^{\ast}$ and thus the optimal planning finds $\pi_{h}=\pi_{h}^{\ast}$ . Recursively applying this process to $h=1$ , we complete the proof. ∎

A.2 Proof of Section 3.2

Theorem A.2 (Formal statement of Theorem 3.5).

Assume the function class in Definition 3.1 satisfies Assumption 2.1, Assumption 2.2 and is policy complete. For any $\epsilon>0$ and $t\in(0,1)$ such that $d\geq\Omega(\log(B_{W}B_{\phi}/\epsilon))$ , if $n\geq\epsilon^{-2}\cdot d\cdot\mathrm{poly}(\kappa,k,B_{W},B_{\phi},H,\log(d/t))$ , then with probability at least $1-t$ Algorithm 2 returns a policy $\pi$ such that $V^{\ast}-V^{\pi}\leq\epsilon$ .

Proof.

Use $\pi^{\ast}_{1},\dots,\pi^{\ast}_{H}$ to denote the global optimal policy. We prove for all $s\in\mathcal{S}$ ,

\displaystyle V^{\pi^{\ast}_{h},\pi^{\ast}_{h+1},\dots,\pi^{\ast}_{H}}_{h}(s)-V^{\pi_{h},\pi_{h+1},\dots,\pi_{H}}_{h}(s)\leq\frac{(H-h+1)\epsilon}{H}.

At level $H$ , let $e_{H}(s^{i}_{H},a^{i}_{H})=r_{H}(s^{i}_{H},a^{i}_{H})-Q^{\ast}_{H}(s^{i}_{H},a^{i}_{H})$ , then $e_{H}(s^{i}_{H},a^{i}_{H})=0$ . From Theorem A.13, we have $\widehat{Q}_{H}(s,a):=v_{H}^{\top}\sigma(W_{H}\phi(s,a))$ satisfies $|\widehat{Q}_{H}(s,a)-Q^{\ast}_{H}(s,a)|\leq\frac{\epsilon}{2H}$ for all $s\in\mathcal{S},a\in\mathcal{A}$ . Therefore for all $s\in\mathcal{S}$ ,

	$\displaystyle V_{H}^{\ast}(s)-V^{\pi_{H}}_{H}(s)=$	$\displaystyle\leavevmode\nobreak\ \mathbb{E}_{a\sim\pi^{\ast}_{H}}[Q^{\ast}_{H}(s,a)]-\mathbb{E}_{a\sim\pi^{\ast}_{H}}[\widehat{Q}_{H}(s,a)]$
		$\displaystyle\leavevmode\nobreak\ +\mathbb{E}_{a\sim\pi^{\ast}_{H}}[\widehat{Q}_{H}(s,a)]-\mathbb{E}_{a\sim\pi_{H}}[\widehat{Q}_{H}(s,a)]$
		$\displaystyle\leavevmode\nobreak\ +\mathbb{E}_{a\sim\pi_{H}}[\widehat{Q}_{H}(s,a)]-\mathbb{E}_{a\sim\pi_{H}}[Q^{\ast}_{H}(s,a)]$
	$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ \frac{\epsilon}{H}$

where in the second step we used $\mathbb{E}_{a\sim\pi^{\ast}_{H}}[\widehat{Q}_{H}(s,a)]\leq\mathbb{E}_{a\sim\pi_{H}}[\widehat{Q}_{H}(s,a)]$ by optimality of $\pi_{H}$ and $|\widehat{Q}_{H}(s,a)-Q^{\ast}_{H}(s,a)|\leq\frac{\epsilon}{2H}$ .

Suppose we have learned policies $\pi_{h+1},\dots,\pi_{H}$ , we use $\widetilde{\pi}_{h}$ to denote the optimal policy of $Q^{\pi_{h+1},\dots,\pi_{H}}_{h}(s,a)$ . Let

e_{h}(s^{i}_{h},a^{i}_{h})={\widehat{Q}}_{h}^{i}-Q_{h}^{\pi_{h+1},\dots,\pi_{H}}(s^{i}_{h},a^{i}_{h})

then $e_{h}(s^{i}_{h},a^{i}_{h})$ is zero mean $H^{2}$ sub-Gaussian (notice that ${\widehat{Q}}_{h}^{i}$ is unbiased estimate of $Q_{h}^{\pi_{h+1},\dots,\pi_{H}}(s^{i}_{h},a^{i}_{h})$ , and ${\widehat{Q}}_{h}^{i}\leq O(H)$ ). From Theorem A.13, we have $\widehat{Q}_{h}(s,a)=v_{h}^{\top}\sigma(W_{h}\phi(s,a))$ satisfies $|\widehat{Q}_{h}(s,a)-Q_{h}^{\pi_{h+1},\dots,\pi_{H}}(s,a)|\leq\frac{\epsilon}{2H}$ for all $s\in\mathcal{S},a\in\mathcal{A}$ . Therefore for all $s\in\mathcal{S}$ ,

		$\displaystyle\leavevmode\nobreak\ V^{\widetilde{\pi}_{h},\pi_{h+1},\dots,\pi_{H}}_{h}(s)-V^{\pi_{h},\pi_{h+1},\dots,\pi_{H}}_{h}(s)$
	$\displaystyle=$	$\displaystyle\leavevmode\nobreak\ \mathbb{E}_{a\sim\widetilde{\pi}_{h}}[Q_{h}^{\pi_{h+1},\dots,\pi_{H}}(s,a)]-\mathbb{E}_{a\sim\widetilde{\pi}_{h}}[\widehat{Q}_{h}(s,a)]$
		$\displaystyle\leavevmode\nobreak\ +\mathbb{E}_{a\sim\widetilde{\pi}_{h}}[\widehat{Q}_{h}(s,a)]-\mathbb{E}_{a\sim\pi_{h}}[\widehat{Q}_{h}(s,a)]$
		$\displaystyle\leavevmode\nobreak\ +\mathbb{E}_{a\sim\pi_{h}}[\widehat{Q}_{h}(s,a)]-\mathbb{E}_{a\sim\pi_{h}}[Q_{h}^{\pi_{h},\pi_{h+1},\dots,\pi_{H}}(s,a)]$
	$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ \frac{\epsilon}{H}$

where in the second step we used $\mathbb{E}_{a\sim\widetilde{\pi}_{h}}[\widehat{Q}_{h}(s,a)]\leq\mathbb{E}_{a\sim\pi_{h}}[\widehat{Q}_{h}(s,a)]$ by optimality of $\pi_{h}$ and $|\widehat{Q}_{h}(s,a)-Q_{h}^{\pi_{h+1},\dots,\pi_{H}}(s,a)|\leq\frac{\epsilon}{2H}$ .

It thus follows that

	$\displaystyle V^{\pi^{\ast}_{h},\pi^{\ast}_{h+1},\dots,\pi^{\ast}_{H}}_{h}(s)-V^{\pi_{h},\pi_{h+1},\dots,\pi_{H}}_{h}(s)=$	$\displaystyle\leavevmode\nobreak\ V^{\pi^{\ast}_{h},\pi^{\ast}_{h+1},\dots,\pi^{\ast}_{H}}_{h}(s)-V^{\pi^{\ast}_{h},\pi_{h+1},\dots,\pi_{H}}_{h}(s)$
		$\displaystyle\leavevmode\nobreak\ +V^{\pi^{\ast}_{h},\pi_{h+1},\dots,\pi_{H}}_{h}(s)-V^{\widetilde{\pi}_{h},\pi_{h+1},\dots,\pi_{H}}_{h}(s)$
		$\displaystyle\leavevmode\nobreak\ +V^{\widetilde{\pi}_{h},\pi_{h+1},\dots,\pi_{H}}_{h}(s)-V^{\pi_{h},\pi_{h+1},\dots,\pi_{H}}_{h}(s)$
	$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ V^{\pi^{\ast}_{h},\pi^{\ast}_{h+1},\dots,\pi^{\ast}_{H}}_{h}(s)-V^{\pi^{\ast}_{h},\pi_{h+1},\dots,\pi_{H}}_{h}(s)+\frac{\epsilon}{H}$
	$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ \cdots$
	$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ \frac{(H-h+1)\epsilon}{H}.$

where in the second step we use $V^{\pi^{\ast}_{h},\pi_{h+1},\dots,\pi_{H}}_{h}(s)\leq V^{\widetilde{\pi}_{h},\pi_{h+1},\dots,\pi_{H}}_{h}(s)$ from optimality of $\widetilde{\pi}_{h}$ . Repeating this argument to $h=1$ completes the proof ∎

A.3 Proof of Section 3.3

Theorem A.3 (Formal statement of Theorem 3.6).

Assume the function class in Definition 3.1 satisfies Assumption 2.1, Assumption 2.2, and is Bellman complete. For any $\epsilon>0$ and $t\in(0,1)$ such that $d\geq\Omega(\log(B_{W}B_{\phi}/\epsilon))$ , if $n\geq\epsilon^{-2}\cdot d\cdot\mathrm{poly}(\kappa,k,B_{W},B_{\phi},H,\log(d/t))$ , then with probability at least $1-t$ Algorithm 3 returns a policy $\pi$ such that $V^{\ast}-V^{\pi}\leq\epsilon$ .

Proof.

Use $\pi^{\ast}_{1},\dots,\pi^{\ast}_{H}$ to denote the global optimal policy. We prove

\displaystyle|\widehat{Q}_{h}(s,a)-Q^{\ast}_{h}(s,a)|\leq\frac{(H-h+1)\epsilon}{H}

(1)

for all $s\in\mathcal{S},a\in\mathcal{A}$ .

At level $H$ , let

e_{H}(s^{i}_{H},a^{i}_{H})=r_{H}(s^{i}_{H},a^{i}_{H})-Q^{\ast}_{H}(s^{i}_{H},a^{i}_{H})

then $e_{H}(s^{i}_{H},a^{i}_{H})=0$ . From Theorem A.13, we have $\widehat{Q}_{H}(s,a):=v_{H}^{\top}\sigma(W_{H}\phi(s,a))$ satisfies $|\widehat{Q}_{H}(s,a)-Q^{\ast}_{H}(s,a)|\leq\frac{\epsilon}{H}$ for all $s\in\mathcal{S},a\in\mathcal{A}$ .

Suppose we have learned $\widehat{Q}_{h+1}(s,a)$ with $|\widehat{Q}_{h+1}(s,a)-Q^{\ast}_{h+1}(s,a)|\leq\frac{(H-h)\epsilon}{H}$ . At level $h$ , let

e_{h}(s^{i}_{h},a^{i}_{h})=r_{h}(s^{i}_{h},a^{i}_{h})+\widehat{V}_{h+1}(s_{h+1}^{i})-{\cal T}_{h}(\widehat{Q}_{h+1})(s^{i}_{h},a^{i}_{h})

then $e_{h}(s^{i}_{h},a^{i}_{h})$ is zero mean $H^{2}$ sub-Gaussian (notice that $r_{h}(s^{i}_{h},a^{i}_{h})+\widehat{V}_{h+1}(s_{h+1}^{i})$ is unbiased estimate of ${\cal T}_{h}(\widehat{Q}_{h+1})(s^{i}_{h},a^{i}_{h})$ , and $r_{h}(s^{i}_{h},a^{i}_{h})+\widehat{V}_{h+1}(s_{h+1}^{i})\leq O(H)$ ). From Theorem A.13, we have $\widehat{Q}_{h}(s,a):=v_{h}^{\top}\sigma(W_{h}\phi(s,a))$ satisfies $|\widehat{Q}_{h}(s,a)-{\cal T}_{h}(\widehat{Q}_{h+1})(s^{i}_{h},a^{i}_{h})|\leq\frac{\epsilon}{H}$ for all $s\in\mathcal{S},a\in\mathcal{A}$ . Therefore

	$\displaystyle\|\widehat{Q}_{h}(s,a)-Q^{\ast}_{h}(s,a)\|\leq$	$\displaystyle\leavevmode\nobreak\ \|\widehat{Q}_{h}(s,a)-{\cal T}_{h}(\widehat{Q}_{h+1})(s,a)\|+\|{\cal T}_{h}(\widehat{Q}_{h+1})(s,a)-Q^{\ast}_{h}(s,a)\|$
	$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ \frac{\epsilon}{H}+\max_{s\in\mathcal{S},a\in\mathcal{A}}\|\widehat{Q}_{h+1}(s,a)-Q^{\ast}_{h+1}(s,a)\|$
	$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ \frac{(H-h+1)\epsilon}{H}$

holds for all $s\in\mathcal{S},a\in\mathcal{A}$ .

It thus follows that for all $s_{1}\in\mathcal{S}$ ,

	$\displaystyle V^{\pi^{\ast}_{1},\dots,\pi^{\ast}_{H}}_{h}(s_{1})-V^{\pi_{1},\dots,\pi_{H}}_{h}(s_{1})=$	$\displaystyle\leavevmode\nobreak\ \mathbb{E}_{a\sim\pi^{\ast}_{1}}[Q_{1}^{\ast}(s_{1},a)]-\mathbb{E}_{a\sim\pi_{1}}[Q_{1}^{\pi_{2},\dots,\pi_{H}}(s_{1},a)]$
	$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ \mathbb{E}_{a\sim\pi^{\ast}_{1}}[\widehat{Q}_{1}(s_{1},a)]-\mathbb{E}_{a\sim\pi_{1}}[Q_{1}^{\pi_{2},\dots,\pi_{H}}(s_{1},a)]+\epsilon$
	$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ \mathbb{E}_{a\sim\pi_{1}}[\widehat{Q}_{1}(s_{1},a)-Q_{1}^{\pi_{2},\dots,\pi_{H}}(s_{1},a)]+\epsilon$
	$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ \mathbb{E}_{a\sim\pi_{1}}[Q_{1}^{\ast}(s_{1},a)-Q_{1}^{\pi_{2},\dots,\pi_{H}}(s_{1},a)]+2\epsilon$
	$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ \mathbb{E}_{a\sim\pi_{1}}\mathbb{E}_{s_{2}\sim\mathbb{P}(\cdot\|s,a)}[V_{2}^{\pi^{\ast}_{2},\dots,\pi^{\ast}_{H}}(s_{2})-V_{2}^{\pi_{2},\dots,\pi_{H}}(s_{2})]+2\epsilon$
	$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ \cdots$
	$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ 2H\epsilon$

where the first step comes from definition of value function; the second step comes from Eq. (1); the third step comes from optimality of $\pi_{1}$ ; the fourth step comes from Eq. (1); the fifth step comes from Bellman equation. The proof is complete by rescaling $\epsilon\leftarrow\epsilon/H$ . ∎

A.4 Proof of Section 3.4

With gap condition, either Algorithm 2 or Algorithm 3 will work as long as we select $\epsilon\approx\rho$ . The following displays an adaption from Algorithm 2.

Algorithm 6 Learning realizable

Q^{\ast}

with optimality gap

1:for

h=H,\ldots 1

2: Sample

x^{i}_{h},i\in[n]

from standard Gaussian

\mathcal{N}(0,I_{d})

3: for

i\in[n]

4: if

\|x^{i}_{h}\|\leq\delta_{\phi}

then

5: Find

(s^{i}_{h},a^{i}_{h})\in\phi^{-1}(x^{i}_{h})

and locate the state

s^{i}_{h}

in the generative model

6: Pull action

a^{i}_{h}

and use

\pi_{h+1},\dots,\pi_{H}

as the roll-out to collect rewards

r_{h}^{(i)},\dots,r_{H}^{(i)}

7: Construct unbiased estimation of

{Q}_{h}^{\pi_{h+1},\dots,\pi_{H}}(s^{i}_{h},a^{i}_{h})

{\widehat{Q}}_{h}^{i}\leftarrow r_{h}^{(i)}+\cdots+r_{H}^{(i)}

8: else

9: Let

{\widehat{Q}}_{h}^{i}\leftarrow 0

10: end if

11: end for

12: Compute

(v_{h},W_{h})\leftarrow\textsc{NeuralNetNoisyRecovery}(\{(x^{i}_{h},{\widehat{Q}}_{h}^{i}):i\in[n]\})

13: Set

\widehat{Q}_{h}(s,a)\leftarrow v_{h}^{\top}\sigma(W_{h}\phi(s,a))

14: Let

\pi_{h}(s)\leftarrow\operatorname*{arg\,max}_{a\in\mathcal{S}}\leavevmode\nobreak\ \widehat{Q}_{h}(s,a)

15:end for

16:Return

\pi_{1},\dots,\pi_{H}

Theorem A.4 (Formal statement of Theorem 3.8).

Assume the function class in Definition 3.1 satisfies Assumption 2.1 and Assumption 2.2. Suppose $\rho>0$ and $d\geq\Omega(\log(B_{W}B_{\phi}/\rho))$ , for any $t\in(0,1)$ , if $n=\frac{d}{\rho^{2}}\cdot\mathrm{poly}(\kappa,k,B_{W},B_{\phi},H,\log(d/t))$ , then with probability at least $1-t$ Algorithm 6 returns the optimal policy $\pi^{\ast}$ .

Proof.

Use $\pi^{\ast}_{1},\dots,\pi^{\ast}_{H}$ to denote the global optimal policy. Similar to Theorem A.1, we prove that Algorithm 6 learns $\pi_{h}^{\ast}$ from $h=H$ to $h=1$ .

At level $H$ , the algorithm uses $n=\frac{d}{\rho^{2}}\cdot\mathrm{poly}(\kappa,k,\log d,B_{W},B_{\phi},H,\log(1/t))$ trajectories to obtain $\widehat{Q}_{H}$ such that $|\widehat{Q}_{H}(s,a)-Q^{\ast}(s,a)|\leq\rho/4$ by Theorem A.13. Therefore

		$\displaystyle\leavevmode\nobreak\ V^{\ast}_{H}(s)-Q^{\ast}_{H}(s,\pi_{H}(s))$
	$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ Q^{\ast}_{H}(s,\pi^{\ast}_{H}(s))-Q^{\ast}_{H}(s,\pi_{H}(s))$
	$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ Q^{\ast}_{H}(s,\pi^{\ast}_{H}(s))-\widehat{Q}_{H}(s,\pi^{\ast}_{H}(s))+\widehat{Q}_{H}(s,\pi^{\ast}_{H}(s))-\widehat{Q}_{H}(s,\pi_{H}(s))$
		$\displaystyle\leavevmode\nobreak\ +\widehat{Q}_{H}(s,\pi_{H}(s))-Q^{\ast}_{H}(s,\pi_{H}(s))$
	$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ \rho/2$

where the third inequality uses the optimality of $\pi_{H}(s)$ under $\widehat{Q}_{H}$ . Thus Definition 3.7 gives $\pi_{H}(s)=\pi^{\ast}_{H}(s)$ . Suppose we have learned $\pi_{h+1}^{\ast},\dots,\pi_{H}^{\ast}$ at level $h$ . We apply the same argument to derive $\pi_{h}=\pi^{\ast}_{h}$ . Recursively applying this process to $h=1$ , we complete the proof. ∎

A.5 Neural network recovery

This section considers recovering neural network $\langle v,\sigma(Wx)\rangle$ from the following two models, where $B=\Omega(d\cdot\mathrm{poly}\log(d))$ .

•

Noisy samples from

\displaystyle x\sim\mathcal{N}(0,I_{d}),\leavevmode\nobreak\ \leavevmode\nobreak\ y=(\langle v,\sigma(Wx)\rangle+\xi)\cdot\mathbbm{1}(\|x\|\leq B)

(2)

where $\xi$ is $\vartheta$ sub-Gaussian noise.

•

Noiseless samples from

\displaystyle x\sim\mathcal{N}(0,I_{d}),\leavevmode\nobreak\ \leavevmode\nobreak\ y=(\langle v,\sigma(Wx)\rangle)\cdot\mathbbm{1}(\|x\|\leq B)

(3)

Recovering neural network has received comprehensive study in deep learning theory [JSA15, ZSJ⁺17, GLM17]. The analysis in this section is mainly based on the method of moments in [ZSJ⁺17]. However, notice that the above learning tasks are different from those considered in [ZSJ⁺17], due to the presence of noise and the truncated signals. Therefore, additional considerations must be made in the analysis.

We consider more general homogeneous activation functions, specified by the assumptions that follow. Since the activation function is homogeneous, we assume $v_{i}\in\{\pm 1\}$ in the following without loss of generality.

Assumption A.5 (Property 3.1 of [ZSJ⁺17]).

Assume $\sigma^{\prime}(x)$ is nonnegative and homogeneously bounded, i.e. $0\leq\sigma^{\prime}(x)\leq L_{1}|x|^{p}$ for some constants $L_{1}>0$ and $p\geq 0$ .

Definition A.6 (Part of property 3.2 of [ZSJ⁺17]).

Define $\rho(z):=\min\{\beta_{0}(z)-\alpha_{0}^{2}(z)-\alpha_{1}^{2}(z),\beta_{2}(z)-\alpha_{1}^{2}(z)-\alpha_{2}^{2}(z),\alpha_{0}(z)\alpha_{2}(z)-\alpha_{1}^{2}(z)\}$ , where $\alpha_{q}(z):=\mathbb{E}_{x\sim\mathcal{N}(0,1)}[\sigma^{\prime}(zx)x^{q}],q\in\{0,1,2\}$ , and $\beta_{q}(z):=\mathbb{E}_{x\sim\mathcal{N}(0,1)}[(\sigma^{\prime})^{2}(zx)x^{q}]$ for $q\in\{0,2\}$ .

Assumption A.7 (Part of property 3.2 of [ZSJ⁺17]).

The first derivative $\sigma^{\prime}(z)$ satisfies that, for all $z>0$ , we have $\rho(z)>0$ .

Assumption A.8 (Property 3.3 of [ZSJ⁺17]).

The second derivative $\sigma^{\prime\prime}(x)$ is either (a) globally bounded or (b) $\sigma^{\prime\prime}(x)=0$ except for finite points.

Notice that ReLU, squared ReLU, leaky ReLU, and polynomial activation function functions all satisfies the above assumption. We make the following assumption on the dimension of feature vectors, which corresponds to how features can extract information about neural networks from noisy samples. The dimension only has to be greater than a logarithmic term in $1/\epsilon$ and the norm of parameters.

Assumption A.9 (Rich feature).

Assume $d\geq\Omega(\log(B_{W}/\epsilon))$ .

First we introduce a notation from [ZSJ⁺17].

Definition A.10.

Define outer product $\tilde{\otimes}$ as follows. For a vector $v\in\mathbb{R}^{d}$ and an identity matrix $I\in\mathbb{R}^{d\times d}$ ,

v\tilde{\otimes}I=\sum_{j=1}^{d}[v\otimes e_{j}\otimes e_{j}+e_{j}\otimes v\otimes e_{j}+e_{j}\otimes e_{j}\otimes v].

For a symmetric rank- $r$ matrix $M=\sum_{i=1}^{r}s_{i}v_{i}v_{i}^{\top}$ and an identity matrix $I\in\mathbb{R}^{d\times d}$ ,

M\tilde{\otimes}I=\sum_{i=1}^{r}s_{i}\sum_{j=1}^{d}\sum_{l=1}^{6}A_{l,i,j}

where $A_{1,i,j}=v_{i}\otimes v_{i}\otimes e_{j}\otimes e_{j}$ , $A_{2,i,j}=v_{i}\otimes e_{j}\otimes v_{i}\otimes e_{j}$ , $A_{3,i,j}=e_{j}\otimes v_{i}\otimes v_{i}\otimes e_{j}$ , $A_{4,i,j}=v_{i}\otimes e_{j}\otimes e_{j}\otimes v_{i}$ , $A_{5,i,j}=e_{j}\otimes v_{i}\otimes e_{j}\otimes v_{i}$ , $A_{6,i,j}=e_{j}\otimes e_{j}\otimes v_{i}\otimes v_{i}$ .

Now we define some moments.

Definition A.11.

Define $M_{1},M_{2},M_{3},M_{4},m_{1,i},m_{2,i},m_{3,i},m_{4,i}$ as follows:

	$\displaystyle M_{1}:=$	$\displaystyle\leavevmode\nobreak\ \mathbb{E}[y\cdot x]$
	$\displaystyle M_{2}:=$	$\displaystyle\leavevmode\nobreak\ \mathbb{E}[y\cdot(x\otimes x-I)]$
	$\displaystyle M_{3}:=$	$\displaystyle\leavevmode\nobreak\ \mathbb{E}[y\cdot(x^{\otimes 3}-x\tilde{\otimes}I)]$
	$\displaystyle M_{4}:=$	$\displaystyle\leavevmode\nobreak\ \mathbb{E}[y\cdot(x^{\otimes 4}-(x\otimes x)\tilde{\otimes}I+I\tilde{\otimes}I)]$
	$\displaystyle\gamma_{j}(x):=$	$\displaystyle\leavevmode\nobreak\ \mathbb{E}_{z\sim\mathcal{N}(0,1)}[\sigma(x\cdot z)z^{j}],\forall j\in 0,1,2,3,4$
	$\displaystyle m_{1,i}:=$	$\displaystyle\leavevmode\nobreak\ \gamma_{1}(\\|w_{i}\\|)$
	$\displaystyle m_{2,i}:=$	$\displaystyle\leavevmode\nobreak\ \gamma_{2}(\\|w_{i}\\|)-\gamma_{0}(\\|w_{i}\\|)$
	$\displaystyle m_{3,i}:=$	$\displaystyle\leavevmode\nobreak\ \gamma_{3}(\\|w_{i}\\|)-3\gamma_{1}(\\|w_{i}\\|)$
	$\displaystyle m_{4,i}:=$	$\displaystyle\leavevmode\nobreak\ \gamma_{4}(\\|w_{i}\\|)+3\gamma_{0}(\\|w_{i}\\|)-6\gamma_{2}(\\|w_{i}\\|)$

The above expectations are all with respect to $x\sim\mathcal{N}(0,I_{d})$ and $y=\langle v,\sigma(Wx)\rangle$ .

Assumption A.12 (Assumption 5.3 of [ZSJ⁺17]).

Assume the activation function satisfies the followings:

•

If $M_{i}\neq 0$ , then $m_{j,i}\neq 0$ for all $i\in[k]$ .
•

At least one of $M_{3}$ and $M_{4}$ is not zero.
•

If $M_{1}=M_{3}=0$ , then $\sigma(z)$ is an even function.
•

If $M_{2}=M_{4}=0$ , then $\sigma(z)$ is an odd function.

Now we state the theoretical result that recovers neural networks from noisy data.

Theorem A.13 (Neural network recovery from noisy data).

Let the activation function $\sigma$ satisfies Assumption A.5 and Assumption A.12. Let $\kappa$ be the condition number of $W$ . Given $n$ samples from Eq. (2). For any $t\in(0,1)$ and $\epsilon\in(0,1)$ such that Assumption A.9 holds, if

n\geq\epsilon^{-2}\cdot d\cdot\mathrm{poly}(\kappa,k,\vartheta,\log(d/t))

then there exists an algorithm that takes $\tilde{O}(nkd)$ time and returns a matrix $\widehat{W}\in\mathbb{R}^{k\times d}$ and a vector $\widehat{v}\in\{\pm 1\}^{k}$ such that with probability at least $1-t$ ,

\displaystyle\|\widehat{W}-W\|_{F}\leq\epsilon\cdot\mathrm{poly}(k,\kappa)\cdot\|W\|_{F},\text{ and }\widehat{v}=v.

The algorithm and proof are shown in Appendix A.5.1. By Assumption A.5, the following corollary is therefore straightforward.

Corollary A.14.

In the same setting as Theorem A.13. For any $t\in(0,1)$ and suppose $\|W\|_{F}\leq B_{W}$ and Assumption A.9 holds. Given $n$ samples from Eq. (2). If

n\geq\epsilon^{-2}\cdot d\cdot\mathrm{poly}(\kappa,k,\log(d/t),B_{W},B_{\phi},\vartheta)

then there exists an algorithm that takes $\tilde{O}(nkd)$ time and outputs a matrix $\widehat{W}\in\mathbb{R}^{k\times d}$ and a vector $\widehat{v}\in\{\pm 1\}^{k}$ such that with probability at least $1-t$ , for all $\|x\|_{2}\leq B_{\phi}$

\displaystyle|\langle\widehat{v},\sigma(\widehat{W}x)\rangle-\langle v,\sigma(Wx)\rangle|\leq\epsilon.

In particular, when $B_{\phi}=O(d\cdot\mathrm{poly}\log d)$ the following sample complexity suffices

\displaystyle n\geq\epsilon^{-2}\cdot d^{O(1+p)}\cdot\mathrm{poly}(\kappa,k,\log(d/t),B_{W},\vartheta).

Now we state the theoretical result that precisely recovers neural networks from noiseless data. The proof and method are shown in Appendix A.5.2.

Theorem A.15 (Exact neural network recovery from noiseless data).

Let the activation function satisfies Assumption A.5 and Assumption A.12, Assumption A.7 and Assumption A.8(b). Given $n$ samples from Eq. (3). For any $t\in(0,1)$ , suppose $d\geq\Omega(\log(B_{W}/\lambda))$ and

\displaystyle n\geq d\cdot\mathrm{poly}(\kappa,k,\lambda,\log(d/t)),

then there exists an algorithm that output exact $W$ and $v$ with probability at least $1-t$ .

A.5.1 Recover neural networks from noisy data

In this section we prove Theorem A.13. Denote $W=[w_{1},\cdots,w_{k}]^{\top}$ where $w_{i}\in\mathbb{R}^{d}$ and $\bar{w_{i}}=w_{i}/\|w_{i}\|_{2}$ .

Definition A.16.

Given a vector $\alpha\in\mathbb{R}^{d}$ . Define $P_{2}:=M_{j_{2}}(I,I,\alpha,\cdots,\alpha)$ where $j_{2}=\min\{j\geq 2:M_{j}\neq 0\}$ and $P_{3}:=M_{j_{3}}(I,I,I,\alpha,\cdots,\alpha)$ where $j_{3}=\min\{j\geq 3:M_{j}\neq 0\}$ .

The method of moments is presented in Algorithm 7. Here we sketch its ideas, and refer readers to [ZSJ⁺17] for thorough explanations. There are three main steps. In the first step, it computes the span of the rows of $W$ . By power method, Line 7 finds the top- $k$ eigenvalues of $CI+\widehat{P}_{2}$ and $CI-\widehat{P}_{2}$ . It then picks the largest $k$ eigenvalues from $CI+\widehat{P}_{2}$ and $CI-\widehat{P}_{2}$ , by invoking TopK in Line 15. Finally it orthogonalizes the corresponding eigenvectors in Line 19 and finds an orthogonal matrix $V$ in the subspace spanned by $\{\bar{w}_{1},\dots,\bar{w}_{k}\}$ .

In the second step, the algorithm forms third order tensor $R_{3}=P_{s}(V,V,V)\in\mathbb{R}^{k\times k\times k}$ and use the robust tensor decomposition method in [KCL15] to find $\widehat{u}$ that approximates $s_{i}V^{\top}\bar{w}_{i}$ with unknown signs $s_{i}$ . In the third step, the algorithm determines $s$ , $v$ and $w_{i},i\in[k]$ . Since the activation function is homogeneous, we assume $v_{i}\in\{\pm 1\}$ and $m_{j,i}=c_{j}\|w_{i}\|^{p+1}$ for universal constants $c_{j}$ without loss of generality. For illustration, we define $Q_{1}$ and $Q_{2}$ as follows.

	$\displaystyle Q_{1}=M_{l_{1}}(I,\underbrace{\alpha,\cdots,\alpha}_{(l_{1}-1)\leavevmode\nobreak\ \alpha\text{'s}})=\sum_{i=1}^{k}v_{i}c_{l_{1}}\\|w_{i}\\|^{p+1}(\alpha^{\top}\overline{w}_{i})^{{l_{1}-1}}\overline{w}_{i},$		(4)
	$\displaystyle Q_{2}=M_{l_{2}}(V,V,\underbrace{\alpha,\cdots,\alpha}_{(l_{2}-2)\leavevmode\nobreak\ \alpha\text{'s}})=\sum_{i=1}^{k}v_{i}c_{l_{2}}\\|w_{i}\\|^{p+1}(\alpha^{\top}\overline{w}_{i})^{{l_{2}-2}}(V^{\top}\overline{w}_{i})(V^{\top}\overline{w}_{i})^{\top},$		(5)

where $l_{1}\geq 1$ such that $M_{l_{1}}\neq 0$ and $l_{2}\geq 2$ such that $M_{l_{2}}\neq 0$ are specified later. Then the solutions of the following linear systems

\displaystyle z^{*}=\operatorname*{arg\,min}_{z\in\mathbb{R}^{k}}\left\|\sum_{i=1}^{k}z_{i}s_{i}\overline{w}_{i}-Q_{1}\right\|,\leavevmode\nobreak\ \leavevmode\nobreak\ r^{*}=\operatorname*{arg\,min}_{r\in\mathbb{R}^{k}}\left\|\sum_{i=1}^{k}r_{i}V^{\top}\overline{w}_{i}(V^{\top}\overline{w}_{i})^{\top}-Q_{2}\right\|_{F}.

(6)

are the followings

\displaystyle z^{*}_{i}

\displaystyle=v_{i}s_{i}^{l_{1}}c_{l_{1}}\|w_{i}\|^{p+1}(\alpha^{\top}s_{i}\overline{w}_{i})^{{l_{1}-1}},\leavevmode\nobreak\ \leavevmode\nobreak\ r_{i}=v_{i}s_{i}^{l_{2}}c_{l_{2}}\|w_{i}\|^{p+1}(\alpha^{\top}s_{i}\overline{w}_{i})^{{l_{2}-2}}.

When $c_{l_{1}}$ and $c_{l_{2}}$ do not have the same sign, we can recover $v_{i}$ and $s_{i}$ by $v_{i}=\operatorname{sign}({r}^{\ast}_{i}c_{l_{2}}),s_{i}=\operatorname{sign}(v_{i}{z}^{\ast}_{i}c_{l_{1}})$ , and recover $w_{i}$ by

\displaystyle{w}_{i}=\left(\left|\frac{z^{\ast}_{i}}{c_{l_{1}}(\alpha^{\top}s_{i}\overline{w}_{i})^{l_{1}-1})}\right|\right)^{1/(p+1)}\overline{w}_{i}.

In Algorithm 7, we use $V\widehat{u}_{i}$ to approximate $s_{i}\overline{w}_{i}$ , and use moment estimators $\widehat{Q}_{1}$ and $\widehat{Q}_{2}$ to approximate $Q_{1}$ and $Q_{1}$ . Then the solutions $\widehat{z},\widehat{r}$ to the optimization problems in Line 29 should approximate $z^{*}$ and $r^{*}$ , due to robustness for solving linear systems. As such, the outputs $\widetilde{v},\widetilde{W}$ approximately recover the true model parameters.

Algorithm 7 Using method of moments to recover neural network parameters

1:procedure NeuralNetNoisyRecovery(

S=\{(x_{i},y_{i}):i\in[n]\}

)

2: Choose

\alpha

to be a random unit vector

3: Partition

S

into

S_{1},S_{2},S_{3},S_{4}

of equal size

\widehat{P}_{2}\leftarrow\mathbb{E}_{S_{1}}[P_{2}]

C\leftarrow 3\|P_{2}\|

T\leftarrow\log(1/\epsilon)

5: Choose

\widehat{V}_{1}^{(0)},\widehat{V}_{1}^{(0)}\in\mathbb{R}^{d\times k}

to be random matrices

\triangleright

Estimate subspace

V

6: for

t=1,\dots,T

\widehat{V}_{1}^{(t)}\leftarrow\mathrm{QR}(C\widehat{V}_{1}^{(t-1)}+\widehat{P}_{2}\widehat{V}_{1}^{(t-1)}),\widehat{V}_{2}^{(t)}\leftarrow\mathrm{QR}(C\widehat{V}_{2}^{(t-1)}-\widehat{P}_{2}\widehat{V}_{2}^{(t-1)})

8: end for

9: for j = 1,2 do

10:

\widehat{V}_{1}^{(T)}\leftarrow[\widehat{V}_{j,1},\cdots,\widehat{V}_{j,k}]

11: for

i\in[k]

12:

\lambda_{j,i}\leftarrow|\widehat{V}_{j,i}\widehat{P}_{2}\widehat{V}_{j,i}|

13: end for

14: end for

15:

\pi_{1},\pi_{2},k_{1},k_{2}\leftarrow\textsc{TopK}(\lambda,k)

16: for j = 1,2 do

17:

V_{j}\leftarrow[\widehat{V}_{j,\pi_{j}(1)},\cdots,\widehat{V}_{j,\pi_{j}(k_{j})}]

18: end for

19:

\tilde{V}_{2}\leftarrow\textsc{QR}((I-V_{1}V_{1}^{\top})V_{2})

V\leftarrow[V_{1},\tilde{V}_{2}]

20:

\widehat{R}_{3}\leftarrow\mathbb{E}_{S_{2}}[P_{3}(V,V,V)]

\{\widehat{u}_{i}\}_{i\in[k]}\leftarrow\textsc{TensorDecomposition}(\widehat{R}_{3})

\triangleright

Learn

s_{i}V^{\top}\bar{w}_{i}

21: if

M_{1}=M_{3}=0

then

22:

l_{1},l_{2}=\min\{j\in\{2,4\}:M_{j}\neq 0\}

23: else if

M_{2}=M_{4}=0

then

24:

l_{1}\leftarrow\min\{j\in\{1,3\}:M_{j}\neq 0\}

l_{2}\leftarrow 3

25: else

26:

l_{1}\leftarrow\min\{j\in\{1,3\}:M_{j}\neq 0\}

l_{2}=\min\{j\in\{2,4\}:M_{j}\neq 0\}

27: end if

28:

\widehat{Q}_{1}\leftarrow\mathbb{E}_{S_{3}}[Q_{1}]

\widehat{Q}_{2}\leftarrow\mathbb{E}_{S_{4}}[Q_{2}]

29:

\widehat{z}\leftarrow\operatorname*{arg\,min}_{z}\|\sum_{i=1}^{k}z_{i}V\widehat{u}_{i}-\widehat{Q}_{1}\|

\widehat{r}\leftarrow\operatorname*{arg\,min}_{r}\|\sum_{i=1}^{k}r_{i}\widehat{u}_{i}\widehat{u}_{i}-\widehat{Q}_{2}\|_{F}

30: for

i=1,\dots,k

\triangleright

Learn parameters

v,W

31:

{\widehat{v}}_{i}\leftarrow\mathrm{sign}(\widehat{r}_{i}c_{l_{2}})

{\widehat{s}}_{i}\leftarrow\mathrm{sign}({\widehat{v}}_{i}\widehat{z}_{i}c_{l_{1}})

32:

{\widehat{w}}_{i}\leftarrow{\widehat{s}}_{i}(|\frac{\widehat{z}_{i}}{c_{l_{1}}(\alpha^{\top}V\widehat{u}_{i})^{l_{1}-1})}|)^{1/(p+1)}V\widehat{u}_{i}

33: end for

34:

{\widehat{W}}\leftarrow[{\widehat{w}}_{i},\cdots,{\widehat{w}}_{k}]

35: Return

({\widehat{v}},{\widehat{W}})

36:end procedure

Since Algorithm 7 carries out the same computation as [ZSJ⁺17], the computational complexity is the same. The difference of sample complexity comes from the noise $\xi$ in the model and the truncation of standard Gaussian. The proof entails bounding the error in estimating $P_{2}$ in Line 4, $R_{3}$ in Line 20 and $Q_{1},Q_{2}$ in Line 28. In the following, unless further specified, the expectations are all with respect to $x\sim\mathcal{N}(0,I_{d})$ and $y\sim(\langle v,\sigma(Wx)\rangle+\xi)\cdot\mathbbm{1}(\|x\|\leq B)$ .

Lemma A.17.

Let $\widehat{P}_{2}$ be computed in Line 4 of Algorithm 7 and $P_{2}$ defined in Definition A.16. Suppose $m_{0}=\min_{i\in[k]}\{|m_{j_{2},i}|^{2}(\overline{w}_{i}^{\top}\alpha)^{2(j_{2}-2)}\}$ and

\displaystyle|S|\gtrsim d\cdot\mathrm{poly}(\kappa,\vartheta,\log(d/t))/(\epsilon^{2}m_{0})

then with probability at least $1-t$ ,

\displaystyle\|P_{2}-\widehat{P}_{2}\|\lesssim\epsilon\sum_{i=1}^{k}|v_{i}m_{j_{2},i}(\bar{w_{i}}^{\top}\alpha)^{j_{2}-2}|+\epsilon.

Proof.

It suffices to bound $\|M_{2}-\widehat{M}_{2}\|$ , $\|M_{3}(I,I,\alpha)-\widehat{M}_{3}(I,I,\alpha)\|$ and $\|M_{4}(I,I,\alpha,\alpha)-\widehat{M}_{4}(I,I,\alpha,\alpha)\|$ . The main strategy is to bound all relevant moment terms and to invoke Claim D.6.

Specifically, we show that with probability at least $1-t/4$ ,

\displaystyle\|M_{2}-\widehat{M}_{2}\|\lesssim\epsilon\sum_{i=1}^{k}|v_{i}m_{2,i}|+\epsilon.

(7)

\displaystyle\|M_{3}(I,I,\alpha)-\widehat{M}_{3}(I,I,\alpha)\|\lesssim\epsilon\sum_{i=1}^{k}|v_{i}m_{3,i}(\overline{w}_{i}^{\top}\alpha)|+\epsilon.

(8)

\displaystyle\|M_{4}(I,I,\alpha,\alpha)-\widehat{M}_{4}(I,I,\alpha,\alpha)\|\lesssim\epsilon\sum_{i=1}^{k}|v_{i}m_{4,i}|(\overline{w}_{i}^{\top}\alpha)^{2}+\epsilon.

(9)

Recall that for sample $(x_{j},y_{j})\in S$ , $y_{j}=\sum_{i=1}^{k}v_{i}\sigma(w_{i}^{\top}x_{j})+\xi_{j}$ where $\xi_{j}$ is independent of $x_{j}$ . Consider each component $i\in[k]$ . Define $C_{i}(x_{j}),B_{i}(x_{j})\in\mathbb{R}^{d\times d}$ as follows:

	$\displaystyle B_{i}(x_{j})=$	$\displaystyle\leavevmode\nobreak\ (\sigma(w_{i}^{\top}x_{j})+\xi_{j})\cdot(x_{j}^{\otimes 4}-(x_{j}\otimes x_{j})\tilde{\otimes}I+I\tilde{\otimes}I)(I,I,\alpha,\alpha)$
	$\displaystyle=$	$\displaystyle\leavevmode\nobreak\ (\sigma(w_{i}^{\top}x_{j})+\xi_{j})\cdot((x^{\top}\alpha)^{2}x^{\otimes 2}-(\alpha^{\top}x)^{2}I-2(\alpha^{\top}x)(x\alpha^{\top}+\alpha x^{\top})-xx^{\top}+2\alpha\alpha^{\top}+I),$

and $C_{i}(x_{j})=\mathbbm{1}({\|x_{j}\|\leq B})\cdot B_{i}(x_{j})$ . Then from Claim D.5 we have $\mathbb{E}[B_{i}(x_{j})]=m_{4,i}(\overline{w}_{i}^{\top}\alpha)^{2}\overline{w}_{i}\overline{w}_{i}^{\top}$ . We calculate

		$\displaystyle\leavevmode\nobreak\ \sigma(w_{i}^{\top}x_{j})\cdot(x_{j}^{\otimes 4}-(x_{j}\otimes x_{j})\tilde{\otimes}I+I\tilde{\otimes}I)(I,I,\alpha,\alpha)$
	$\displaystyle\lesssim$	$\displaystyle(\|w_{i}^{\top}x_{j}\|^{p+1}+\|\phi(0)\|)\cdot((x_{j}^{\top}\alpha)^{2}\\|x_{j}\\|^{2}+1+\\|x_{j}\\|^{2}+(\alpha^{\top}x_{j})^{2})$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ \|w_{i}\|^{p+1}\cdot\|x_{j}\|^{p+5},$

By Assumption A.5, using Claim D.1 and $B\geq d\cdot\mathrm{poly}\log(d)$ we have

	$\displaystyle\\|\mathbb{E}[C_{i}(x_{j})]-m_{4,i}(\overline{w}_{i}^{\top}\alpha)^{2}\overline{w}_{i}\overline{w}_{i}^{\top}\\|\lesssim$	$\displaystyle\leavevmode\nobreak\ \mathbb{E}[\mathbbm{1}_{\\|x_{j}\\|\geq B}\|w_{i}\|^{p+1}\cdot\|x_{j}\|^{p+5}]$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ (\\|w_{i}\\|d)^{p+5}\cdot e^{-\Omega(d\log d)}$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ \epsilon.$

Also, $\frac{1}{2}|m_{4,i}|(\overline{w}_{i}^{\top}\alpha)^{2}\leq\|\mathbb{E}[C_{i}(x_{j})]\|\leq 2|m_{4,i}|(\overline{w}_{i}^{\top}\alpha)^{2}$ .

For any constant $t\in(0,1)$ , we have with probability $1-t/4$ ,

	$\displaystyle\\|C_{i}(x_{j})\\|\lesssim$	$\displaystyle(\|w_{i}^{\top}x_{j}\|^{p+1}+\|\phi(0)\|+\|\xi_{j}\|)\cdot((x_{j}^{\top}\alpha)^{2}\\|x_{j}\\|^{2}+1+\\|x_{j}\\|^{2}+(\alpha^{\top}x_{j})^{2})$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ (\\|w_{i}\\|^{p+1}+\|\phi(0)\|+\vartheta)\cdot d\cdot\mathrm{poly}(\log(d/t))$

where the first step comes from Assumption A.5 and the second step comes from Claim D.2 and Claim D.3.

Using Claim D.4, we have

	$\displaystyle\left\\|\mathbb{E}[C_{i}(x_{j})^{2}]\right\\|\lesssim$	$\displaystyle\leavevmode\nobreak\ \left(\mathbb{E}[(\phi(w_{i}^{\top}x_{j})+\xi_{j})^{4}]\right)^{1/2}\left(\mathbb{E}[(x_{j}^{\top}\alpha)^{8}]\right)^{1/2}\left(\mathbb{E}[\\|x_{j}\\|^{4}]\right)^{1/2}$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ (\\|w_{i}\\|^{p+1}+\|\phi(0)\|+\vartheta)^{2}d.$

Furthermore we have,

\displaystyle\max_{\|a\|=1}\left(\mathbb{E}\left[(a^{\top}C_{i}(x_{j})a)^{2}\right]\right)^{1/2}\lesssim\left(\mathbb{E}\left[(\phi(w_{i}^{\top}x_{j})+\xi_{j})^{4}\right]\right)^{1/4}\lesssim\|w_{i}\|^{p+1}+|\phi(0)|+\vartheta.

Then by Claim D.6, with probability at least $1-t$ ,

		$\displaystyle\leavevmode\nobreak\ \left\\|m_{4,i}(\overline{w}_{i}^{\top}\alpha)^{2}\overline{w}_{i}\overline{w}_{i}^{\top}-\frac{1}{\|S\|}\sum_{x_{j}\in S}C_{i}(x_{j})\right\\|$
	$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ \left\\|m_{4,i}(\overline{w}_{i}^{\top}\alpha)^{2}\overline{w}_{i}\overline{w}_{i}^{\top}-\mathbb{E}[C_{i}(x_{j})]\right\\|+\left\\|\mathbb{E}[C_{i}(x_{j})]-\frac{1}{\|S\|}\sum_{x_{j}\in S}C_{i}(x_{j})\right\\|$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ \epsilon\|m_{4,i}\|(\overline{w}_{i}^{\top}\alpha)^{2}+\epsilon.$

Summing up all components $i\in[k]$ , we proved Eq. (9). Eq. (7) and Eq. (8) can be shown similarly.

∎

Lemma A.18.

Let $V\in\mathbb{R}^{d\times k}$ be an orthogonal matrix. Let $\widehat{R}_{3}$ be computed in Line 20 of Algorithm 7 and $R_{3}=P_{3}(V,V,V)$ . Suppose

m_{0}=\min_{i\in[k]}\{|m_{j_{3},i}|^{2}(\overline{w}_{i}^{\top}\alpha)^{2(j_{3}-3)}\}

and

\displaystyle|S|\gtrsim d\cdot\mathrm{poly}(\kappa,\vartheta,\log(d/t))/(\epsilon^{2}m_{0})

then with probability at least $1-t$ ,

\displaystyle\|R_{3}-\widehat{R}_{3}\|\lesssim\epsilon\sum_{i=1}^{k}|v_{i}m_{j_{3},i}(\bar{w_{i}}^{\top}\alpha)^{j_{3}-3}|+\epsilon.

Proof.

From the definition of $R_{3}$ , it suffices to bound $\|M_{3}(V,V,V)-\widehat{M}_{3}(V,V,V)\|$ and $\|M_{4}(V,V,V,\alpha)-\widehat{M}_{4}(V,V,V,\alpha)\|$ . The proof is similar to the previous one.

Specifically, we show that with probability at least $1-t/4$ ,

\displaystyle\|M_{3}(V,V,V)-\widehat{M}_{3}(V,V,V)\|\lesssim\epsilon\sum_{i=1}^{k}|v_{i}m_{3,i}|+\epsilon.

(10)

\displaystyle\|M_{4}(V,V,V,\alpha)-\widehat{M}_{4}(V,V,V,\alpha)\|\lesssim\epsilon\sum_{i=1}^{k}|v_{i}m_{4,i}(\overline{w}_{i}^{\top}\alpha)|+\epsilon.

(11)

Recall that for sample $(x_{j},y_{j})\in S$ , $y_{j}=\sum_{i=1}^{k}v_{i}\sigma(w_{i}^{\top}x_{j})+\xi_{j}$ where $\xi_{j}$ is independent of $x_{j}$ . Consider each component $i\in[k]$ . Define $T_{i}(x_{j}),S_{i}(x_{j})\in\mathbb{R}^{k\times k\times k}$ :

	$\displaystyle T_{i}(x_{j})=$	$\displaystyle\leavevmode\nobreak\ (\sigma(w_{i}^{\top}x_{j})+\xi_{j})$
		$\displaystyle\leavevmode\nobreak\ \cdot\left(x_{i}^{\top}\alpha\cdot v(x)^{\otimes 3}-(V^{\top}\alpha)\tilde{\otimes}(v(x)\otimes v(x))-\alpha^{\top}x\cdot v(x)\tilde{\otimes}I+(V^{\top}\alpha)\tilde{\otimes}I\right),$
	$\displaystyle S_{i}(x_{j})=$	$\displaystyle\leavevmode\nobreak\ \mathbbm{1}({\\|x_{j}\\|\leq B})\cdot T_{i}(x_{j})$

where $v(x)=V^{\top}x$ . Flatten $T_{i}(x_{j})$ along the first dimension to obtain $B_{i}(x_{j})\in\mathbb{R}^{k\times k^{2}}$ , flatten $S_{i}(x_{j})$ along the first dimension to obtain $C_{i}(x_{j})\in\mathbb{R}^{k\times k^{2}}$ .

From Claim D.7, $\mathbb{E}[B_{i}(x_{j})]=m_{4,i}(\alpha^{\top}\overline{w}_{i})(V^{\top}\overline{w}_{i})\text{vec}((V^{\top}\overline{w}_{i})(V^{\top}\overline{w}_{i})^{\top})^{\top}$ . Therefore we have,

\displaystyle\left\|\mathbb{E}[B_{i}(x)]\right\|=|m_{4,i}(\alpha^{\top}\overline{w}_{i})|\cdot\|V^{\top}\overline{w}_{i}\|^{3}.

We calculate

	$\displaystyle\\|\mathbb{E}_{\xi_{j}}[B_{i}(x_{j})]\\|\lesssim$	$\displaystyle(\|w_{i}^{\top}x_{j}\|^{p+1}+\|\phi(0)\|)\cdot((x_{j}^{\top}\alpha)^{2}\\|V^{\top}x_{j}\\|^{3}$
		$\displaystyle\leavevmode\nobreak\ +3\\|V^{\top}x_{j}\\|^{3}+3\|x_{j}^{\top}\alpha\|\\|V^{\top}x_{j}\\|\sqrt{k}+3\\|V^{\top}\alpha\\|\sqrt{k})$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ \sqrt{k}\cdot\\|w_{i}\\|^{p+1}\\|x_{j}\\|^{p+6}.$

By Assumption A.5, using Claim D.1 and $B\geq d\cdot\mathrm{poly}\log(d)$ ,

		$\displaystyle\leavevmode\nobreak\ \\|\mathbb{E}[C_{i}(x_{j})]-m_{4,i}(\alpha^{\top}\overline{w}_{i})(V^{\top}\overline{w}_{i})\text{vec}((V^{\top}\overline{w}_{i})(V^{\top}\overline{w}_{i})^{\top})^{\top}\\|$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ \mathbb{E}[\mathbbm{1}_{\\|x_{j}\\|\leq B}\sqrt{k}\\|w_{i}\\|^{p+1}\\|x_{j}\\|^{p+6}]$
	$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ \epsilon.$

For any constant $t\in(0,1)$ , we have with probability $1-t$ ,

	$\displaystyle\\|C_{i}(x_{j})\\|\lesssim$	$\displaystyle(\|w_{i}^{\top}x_{j}\|^{p+1}+\|\phi(0)\|+\|\xi_{j}\|)\cdot((x_{j}^{\top}\alpha)^{2}\\|V^{\top}x_{j}\\|^{3}$
		$\displaystyle\leavevmode\nobreak\ +3\\|V^{\top}x_{j}\\|^{3}+3\|x_{j}^{\top}\alpha\|\\|V^{\top}x_{j}\\|\sqrt{k}+3\\|V^{\top}\alpha\\|\sqrt{k})$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ (\\|w_{i}\\|^{p+1}+\|\phi(0)\|+\vartheta)k^{3/2}\mathrm{poly}(\log(d/t))$

where the first step comes from Assumption A.5 and the second step comes from Claim D.2 and Claim D.3.

Using Claim D.4, we have

	$\displaystyle\left\\|\mathbb{E}[C_{i}(x_{j})C_{i}(x_{j})^{\top}]\right\\|\lesssim$	$\displaystyle\leavevmode\nobreak\ \left(\mathbb{E}\left[(\phi(w_{i}^{\top}x_{j})+\xi_{j})^{4}\right]\right)^{1/2}\left(\mathbb{E}\left[(\alpha^{\top}x_{j})^{4}\right]\right)^{1/2}\left(\mathbb{E}\left[\\|V^{\top}x_{j}\\|^{6}\right]\right)^{1/2}$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ (\\|w_{i}\\|^{p+1}+\|\phi(0)\|+\vartheta)^{2}k^{3/2}.$

and

		$\displaystyle\leavevmode\nobreak\ \left\\|\mathbb{E}[C_{i}(x_{j})^{\top}C_{i}(x_{j})]\right\\|$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ \left(\mathbb{E}[(\phi(w_{i}^{\top}x_{j})+\xi_{j})^{4}]\right)^{1/2}\left(\mathbb{E}[(\alpha^{\top}x_{j})^{4}]\right)^{1/2}\left(\mathbb{E}[\\|V^{\top}x_{j}\\|^{4}]\right)^{1/2}$
		$\displaystyle\leavevmode\nobreak\ \cdot\left(\max_{\\|A\\|_{F}=1}\mathbb{E}\left[\langle A,(V^{\top}x_{j})(V^{\top}x_{j})^{\top}\rangle^{4}\right]\right)^{1/2}$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ (\\|w_{i}\\|^{p+1}+\|\phi(0)\|+\vartheta)^{2}k^{2}.$

Furthermore we have,

		$\displaystyle\leavevmode\nobreak\ \max_{\\|a\\|=\\|b\\|=1}\left(\mathbb{E}\left[(a^{\top}C_{i}(x_{j})b)^{2}\right]\right)^{1/2}$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ \left(\mathbb{E}[(\phi(w_{i}^{\top}x_{j})+\xi_{j})^{4}]\right)^{1/4}\left(\mathbb{E}\left[(\alpha^{\top}x_{j})^{4}\right]\right)^{1/4}\max_{\\|a\\|=1}\left(\mathbb{E}\left[(a^{\top}V^{\top}x_{j})^{4}\right]\right)^{1/2}$
		$\displaystyle\leavevmode\nobreak\ \cdot\max_{\\|A\\|_{F}=1}\left(\mathbb{E}\left[\langle A,(V^{\top}x_{j})(V^{\top}x_{j})^{\top}\rangle^{4}\right]\right)^{1/2}$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ (\\|w_{i}\\|^{p+1}+\|\phi(0)\|+\vartheta)k.$

Then by Claim D.6, with probability at least $1-t$ ,

		$\displaystyle\leavevmode\nobreak\ \left\\|m_{4,i}(\alpha^{\top}\overline{w}_{i})(V^{\top}\overline{w}_{i})\text{vec}((V^{\top}\overline{w}_{i})(V^{\top}\overline{w}_{i})^{\top})^{\top}-\frac{1}{\|S\|}\sum_{x_{j}\in S}C_{i}(x_{j})\right\\|$
	$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ \left\\|m_{4,i}(\alpha^{\top}\overline{w}_{i})(V^{\top}\overline{w}_{i})\text{vec}((V^{\top}\overline{w}_{i})(V^{\top}\overline{w}_{i})^{\top})^{\top}-\mathbb{E}[C_{i}(x_{j})]\right\\|$
		$\displaystyle\leavevmode\nobreak\ +\left\\|\mathbb{E}[C_{i}(x_{j})]-\frac{1}{\|S\|}\sum_{x_{j}\in S}C_{i}(x_{j})\right\\|$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ \epsilon\|v_{i}m_{4,i}(\overline{w}_{i}^{\top}\alpha)\|+\epsilon.$

Summing up all neurons $i\in[k]$ , we proved Eq. (11). Eq. (10) can be shown similarly.

∎

Lemma A.19.

Let $\widehat{Q}_{1}$ and $\widehat{Q}_{2}$ be computed in Line 28 of Algorithm 7. Let $Q_{1}$ be defined by Eq. 4 and $Q_{2}$ be defined by Eq. 5. Suppose

m_{0}=\min_{i\in[k]}\{|m_{j_{1},i}|^{2}(\overline{w}_{i}^{\top}\alpha)^{2(j_{1}-1)},|m_{j_{2},i}|^{2}(\overline{w}_{i}^{\top}\alpha)^{2(j_{2}-2)}\}

and

\displaystyle|S|\gtrsim d\cdot\mathrm{poly}(\kappa,\vartheta,\log(d/t))/(\epsilon^{2}m_{0})

then with probability at least $1-t$ ,

	$\displaystyle\\|Q_{1}-\widehat{Q}_{1}\\|\lesssim$	$\displaystyle\leavevmode\nobreak\ \epsilon\sum_{i=1}^{k}\|v_{i}m_{j_{1},i}(\bar{w_{i}}^{\top}\alpha)^{j_{1}-1}\|+\epsilon,$
	$\displaystyle\\|Q_{2}-\widehat{Q}_{2}\\|\lesssim$	$\displaystyle\leavevmode\nobreak\ \epsilon\sum_{i=1}^{k}\|v_{i}m_{j_{2},i}(\bar{w_{i}}^{\top}\alpha)^{j_{2}-2}\|+\epsilon.$

Proof.

Recall the expression of $Q_{1}$ and $Q_{2}$ ,

	$\displaystyle Q_{1}=M_{l_{1}}(I,\underbrace{\alpha,\cdots,\alpha}_{(j_{1}-1)\leavevmode\nobreak\ \alpha\text{'s}})=\sum_{i=1}^{k}v_{i}c_{j_{1}}\\|w_{i}\\|^{p+1}(\alpha^{\top}\overline{w}_{i})^{{j_{1}-1}}\overline{w}_{i},$
	$\displaystyle Q_{2}=M_{j_{2}}(V,V,\underbrace{\alpha,\cdots,\alpha}_{(j_{2}-2)\leavevmode\nobreak\ \alpha\text{'s}})=\sum_{i=1}^{k}v_{i}c_{j_{2}}\\|w_{i}\\|^{p+1}(\alpha^{\top}\overline{w}_{i})^{{j_{2}-2}}(V^{\top}\overline{w}_{i})(V^{\top}\overline{w}_{i})^{\top}.$

The proof is essentially similar to Lemma A.17 and Lemma A.18. ∎

We also use the following Lemmata from [KCL15, ZSJ⁺17].

Lemma A.20 (Adapted from Theorem 3 of [KCL15]).

Given a tensor $\widehat{T}=\sum_{i=1}^{k}\pi_{i}u_{i}^{\otimes 3}+\epsilon R\in\mathbb{R}^{d\times d\times d}$ . Assume incoherence $u_{i}^{\top}u_{j}\leq\mu$ . Let $L_{0}:=(\frac{50}{1-\mu^{2}})^{2}$ and $L\geq L_{0}\log(15d(k-1)/t)^{2}$ . Then there exists an algorithm such that, with probability at least $1-t$ , for every $u_{i}$ , the algorithm returns a $\tilde{u}_{i}$ such that

\displaystyle\|\tilde{u}_{i}-u_{i}\|_{2}\leq O\left(\frac{\sqrt{\|\pi\|_{1}\pi_{\max}}}{\pi_{\min}^{2}}\cdot\frac{\|V^{\top}\|_{2}^{2}}{1-\mu^{2}}\cdot(1+C(t))\right)\epsilon+o(\epsilon),

where $C(t):=\log(kd/t)\sqrt{d/L}$ and $V$ is the inverse of the full-rank extension of $(u_{1}\dots u_{k})$ with unit-norm columns.

Lemma A.21 (Adapted from Lemma E.6 of [ZSJ⁺17]).

Let $P_{2}$ be defined as in Definition A.16 and $\widehat{P}_{2}$ be its empirical version calculated in Line 4 of Algorithm 7. Let $U\in\mathbb{R}^{d\times k}$ be the orthogonal column span of $W\in\mathbb{R}^{d\times k}$ . Assume $\|\widehat{P}_{2}-P_{2}\|\leq s_{k}(P_{2})/10$ . Let $C$ be a large enough positive number such that $C>2\|P_{2}\|$ . Then after $T=O(\log(1/\epsilon))$ iterations, the $V\in\mathbb{R}^{d\times k}$ computed in Algorithm 7 will satisfy

\displaystyle\|UU^{\top}-VV^{\top}\|\lesssim\|\widehat{P}_{2}-P_{2}\|/s_{k}(P_{2})+\epsilon,

which implies

\displaystyle\|(I-VV^{\top})w_{i}\|\lesssim(\|\widehat{P}_{2}-P_{2}\|/s_{k}(P_{2})+\epsilon)\|w_{i}\|.

Lemma A.22 (Adapted from Lemma E.13 in [ZSJ⁺17]).

Let $U\in\mathbb{R}^{d\times k}$ be the orthogonal column span of $W^{*}$ . Let $V\in\mathbb{R}^{d\times k}$ denote an orthogonal matrix satisfying that $\|VV^{\top}-UU^{\top}\|\leq\widehat{\delta}_{2}\lesssim 1/(\kappa^{2}\sqrt{k})$ . For each $i\in[k]$ , let $\widehat{u}_{i}$ denote the vector satisfying $\|s_{i}\widehat{u}_{i}-V^{\top}\overline{w}_{i}\|\leq\widehat{\delta}_{3}\lesssim 1/(\kappa^{2}\sqrt{k})$ . Let $Q_{1}$ be defined as in Eq. (4) and $\widehat{Q}_{1}$ be the empirical version of $Q_{1}$ such that $\|Q_{1}-\widehat{Q}_{1}\|\leq\widehat{\delta}_{4}\|Q_{1}\|\leq\frac{1}{4}\|Q_{1}\|$ .

Let $z^{*}\in\mathbb{R}^{k}$ and $\widehat{z}\in\mathbb{R}^{k}$ be defined as in Eq. (6) and Line 29. Then

\displaystyle|\widehat{z}_{i}-z_{i}^{*}|\leq(\kappa^{4}k^{3/2}(\widehat{\delta}_{2}+\widehat{\delta}_{3})+\kappa^{2}k^{1/2}\widehat{\delta}_{4})\|z^{*}\|_{1}.

Lemma A.23 (Adapted from Lemma E.14 in [ZSJ⁺17]).

Let $U\in\mathbb{R}^{d\times k}$ be the orthogonal column span of $W^{*}$ and $V$ be an orthogonal matrix satisfying that $\|VV^{\top}-UU^{\top}\|\leq\widehat{\delta}_{2}\lesssim 1/(\kappa\sqrt{k})$ . For each $i\in[k]$ , let $\widehat{u}_{i}$ denote the vector satisfying $\|s_{i}\widehat{u}_{i}-V^{\top}\overline{w}_{i}^{*}\|\leq\widehat{\delta}_{3}\lesssim 1/(\sqrt{k}\kappa^{3})$ .

Let $Q_{2}$ be defined as in Eq. (5) and $\widehat{Q}_{2}$ be the empirical version of $Q_{2}$ such that $\|Q_{2}-\widehat{Q}_{2}\|_{F}\leq\widehat{\delta}_{4}\|Q_{2}\|_{F}\leq\frac{1}{4}\|Q_{2}\|_{F}$ . Let $r^{*}\in\mathbb{R}^{k}$ and $\widehat{r}\in\mathbb{R}^{k}$ be defined as in Eq. (6) and Line 29. Then

\displaystyle|\widehat{r}_{i}-r_{i}^{*}|\leq(k^{3}\kappa^{8}\widehat{\delta}_{3}+\kappa^{2}k^{2}\widehat{\delta}_{4})\|r^{*}\|.

Now we are in the position of proving Theorem A.13.

Proof.

Consider Algorithm 7. First, by Lemma A.21 and Lemma A.17, we have

$\displaystyle\\|VV^{\top}\overline{w}_{i}-\overline{w}_{i}\\|\leq$	$\displaystyle\leavevmode\nobreak\ (\\|\widehat{P}_{2}-P_{2}\\|/s_{k}(P_{2})+\epsilon)$
$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ (\mathrm{poly}(k,\kappa)\\|\widehat{P}_{2}-P_{2}\\|+\epsilon)$
$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ \mathrm{poly}(k,\kappa)\epsilon.$	(12)

Next, combining Lemma A.20 and Lemma A.18, we have

\displaystyle\|V^{\top}\overline{w}_{i}-s_{i}\widehat{u}_{i}\|\leq\mathrm{poly}(k,\kappa)\|\widehat{R}_{3}-R_{3}\|\leq\epsilon\mathrm{poly}(k,\kappa).

(13)

It thus follows that

$\displaystyle\\|\overline{w}_{i}-s_{i}V\widehat{u}_{i}\\|\leq$	$\displaystyle\leavevmode\nobreak\ \\|VV^{\top}\overline{w}_{i}-\overline{w}_{i}\\|+\\|VV^{\top}\overline{w}_{i}-Vs_{i}\widehat{u}_{i}\\|$
$\displaystyle=$	$\displaystyle\leavevmode\nobreak\ \\|VV^{\top}\overline{w}_{i}-\overline{w}_{i}\\|+\\|V^{\top}\overline{w}_{i}-s_{i}\widehat{u}_{i}\\|$
$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ \epsilon\mathrm{poly}(k,\kappa),$	(14)

where the first step applies triangle inequality and the last step uses Eq. (A.5.1) and Eq. (13).

We proceed to bound the error in $\widehat{r}$ and $\widehat{z}$ . We have,

	$\displaystyle\|\widehat{r}_{i}-r_{i}^{*}\|\lesssim$	$\displaystyle\leavevmode\nobreak\ \mathrm{poly}(k,\kappa)(\\|Q_{2}-\widehat{Q}_{2}\\|+\\|s_{i}\widehat{u}_{i}-V^{\top}\overline{w}_{i}\\|)\cdot\\|r^{*}\\|$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ \epsilon\mathrm{poly}(k,\kappa)\cdot\\|r^{*}\\|$		(15)

where the first step comes from Lemma A.23 and the second step comes from Lemma A.19 and Eq. (A.5.1). Furthermore,

	$\displaystyle\|\widehat{z}_{i}-z_{i}^{*}\|\lesssim$	$\displaystyle\leavevmode\nobreak\ \mathrm{poly}(k,\kappa)\cdot(\\|Q_{1}-\widehat{Q}_{1}\\|+\\|s_{i}\widehat{u}_{i}-V^{\top}\overline{w}_{i}\\|+\\|VV^{\top}-UU^{\top}\\|)\cdot\\|z^{*}\\|_{1}$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ \epsilon\mathrm{poly}(k,\kappa)\\|z^{*}\\|_{1},$		(16)

where the first step comes from Lemma A.22 and the second step comes from combining Lemma A.19, Lemma A.21, and Eq. (A.5.1). Finally, combining Eq. (A.5.1), Eq. (A.5.1) and Eq. (A.5.1), the output in Line 32 satisfies $\|{\widehat{w}}_{i}-w_{i}\|_{F}\leq\epsilon\mathrm{poly}(k,\kappa)\cdot\|w_{i}\|_{F}$ . Since $v_{i}$ are discrete values, they are exactly recovered.

∎

A.5.2 Exact recovery of neural networks from noiseless data

In this section we prove Theorem A.15. Similar to Appendix A.5.1, denote $W=[w_{1},\cdots,w_{k}]^{\top}$ where $w_{i}\in\mathbb{R}^{d}$ and $\widehat{W}=[\widehat{w}_{1},\cdots,\widehat{w}_{k}]^{\top}$ . We use $\cal D$ to denote the distribution of $x\sim\mathcal{N}(0,I_{d})$ and $y=\langle v,\sigma(Wx)\rangle$ . We define the empirical loss for explored features and the population loss as follows,

	$\displaystyle L_{n}(\widehat{W})=$	$\displaystyle\leavevmode\nobreak\ \frac{1}{2n}\sum_{(x,y)\in S^{(1)}}\left(\sum_{i=1}^{k}v_{i}\sigma(\widehat{w}_{i}^{\top}x_{i})-y_{i}\right)^{2},$		(17)
	$\displaystyle L(\widehat{W})=$	$\displaystyle\leavevmode\nobreak\ \frac{1}{2}{\mathbb{E}_{\cal D}}\left[\left(\sum_{i=1}^{k}v_{i}\sigma(\widehat{w}_{i}^{\top}x)-y\right)^{2}\right].$		(18)

Algorithm 8 Using method of moments and gradient descent to recover neural network parameters

1:procedure NeuralNetRecovery(

S=\{(x_{i},y_{i}):i\in[n]\}

)

2: Let

S^{(1)}\leftarrow\{(x,y)\in S:\|x\|\leq B\}

3: Compute

(v,\widehat{W})\leftarrow\textsc{NeuralNetNoisyRecovery}(S^{(1)})

4: Find

W^{(1)}

as the global minimum of

L_{n}(\cdot)

by gradient descent initialized at

\widehat{W}

, where

L_{n}(\cdot)

is defined in Eq. (17).

5: Return

(v,W^{(1)})

6:end procedure

Definition A.24.

Let $s_{i}$ be the $i$ -th singular value of $W$ , $\lambda:=\prod_{i=1}^{k}(s_{i}/s_{k})$ . Let $\tau=(3s_{1}/2)^{4p}/\min_{z\in[s_{k}/2,3s_{1}/2]}\{\rho^{2}(z)\}$ .

We use the follow results adapted from [ZSJ⁺17]. The only difference is that the rewards are potentially truncated if $\|x\|\geq B$ , and due to $B=d\cdot\mathrm{poly}\log(d)$ we can bound its difference between standard Gaussian in the same way as Appendix A.5.1.

Lemma A.25 (Concentration, adapted from Lemma D.11 in [ZSJ⁺17]).

Let samples size $n\geq\epsilon^{-2}d\tau\mathrm{poly}(\log(d/t))$ , then with probability at least $1-t$ ,

\displaystyle\|\nabla^{2}L_{n}(W)-\nabla^{2}L(W)\|\lesssim ks_{1}^{2p}\epsilon+\mathrm{poly}(B_{W},d)e^{-\Omega(d)}.

Lemma A.26 (Adapted from Lemma D.16 in [ZSJ⁺17]).

Assume activation $\sigma(\cdot)$ satisfies Assumption A.8 and Assumption A.7. Then for any $t\in(0,1)$ , if $n\geq d\cdot\mathrm{poly}(\log(d/t))$ , with probability at least $1-t$ , for any $\widehat{W}$ (which is not necessarily to be independent of samples) satisfying $\|W-\widehat{W}\|\leq s_{k}/4$ , we have

\displaystyle\|\nabla^{2}L_{n}(\widehat{W})-\nabla^{2}L_{n}(W)\|\leq ks_{1}^{p}\|W-\widehat{W}\|d^{(p+1)/2}.

Now we prove Theorem A.15.

Proof.

The exact recovery consists of first finding (exact) $v$ and (approximate) $\widehat{W}$ close enough to $W$ by tensor method (Appendix A.5.1), and then minimizing the empirical loss $L_{n}(\cdot)$ . We will prove that $L_{n}(\cdot)$ is locally strongly convex, thus we find the precise $W$ .

From Lemma D.3 from [ZSJ⁺17] we know:

\Omega(\rho(s_{k})/\lambda)I\preceq\nabla^{2}L(W)\preceq O(ks_{1}^{2p})I.

(19)

Combining Lemma A.25, $d\geq\log(B_{W}/\lambda)$ , and $n\geq\frac{k^{2}\lambda^{2}s_{1}^{4p}}{\rho^{2}(s_{k})}d\tau\mathrm{poly}(\log(d/t))$ , we know $\nabla^{2}L_{n}(W)$ must be positive definite.

Next we uniformly bound Lipschitzness of $\nabla^{2}L_{n}$ . From Lemma A.26 there exists a universal constant $c$ , such that for all $\widehat{W}$ that satisfies $\|W-\widehat{W}\|\leq cks_{1}^{2p}/(ks_{1}^{p}d^{(p+1)/2})=cs_{1}^{p}d^{-(p+1)/2}$ , $\nabla L_{n}^{2}(\widehat{W})\gtrsim ks_{1}^{2p}$ holds uniformly. So there is a unique miminizer of $L_{n}$ in this region.

Notice $L_{n}(W)=0$ , therefore we can find $W$ by directly minimizing the empirical loss as long as we find any $\widehat{W}$ in this region. This can be achieved by tensor method in Appendix A.5.1. We thus complete the proof. ∎

Appendix B Omitted Proofs in Section 4

For the proofs of Theorem 4.2, Example 4.3, and Example 4.4, we refer the readers to [HHK⁺21].

Lemma B.1.

Consider the polynomial family ${\cal F}_{{\mathcal{V}}}$ of dimension $D$ . Assume that $n>2D$ . For any $E\in\mathbb{R}^{d}$ that is of positive measure, by sampling $n$ samples $\{x_{i}\}$ i.i.d. from $\mathbb{P}_{x\in\mathcal{N}(0,I_{d})}(\cdot|x\in E)$ and observing the noiseless feedbacks $y_{i}=f^{*}(x_{i})$ , one can almost surely uniquely determine the $f^{*}$ by solving the system of equations $y_{i}=f(x_{i})$ , $i=1,\dots,n$ , for $f\in{\cal F}_{{\mathcal{V}}}$ .

Proof.

By Theorem 4.2, there exists a set $N\in\mathbb{R}^{d}\times\ldots\mathbb{R}^{d}$ of Lebesgue measure zero, such that if $(x_{1},\cdots,x_{n})\notin N$ , one can uniquely determine the $f^{*}$ by the observations on the $n$ samples. Therefore, we only need to show that with probability 1, the sampling procedure returns $(x_{1},\dots,x_{n})\notin N$ . This is because

	$\displaystyle\mathbb{P}(x_{1},\dots,x_{n}\in N)$	$\displaystyle=\mathbb{P}_{x_{i}\in{\mathcal{N}}(0,I_{d})}((x_{1},\dots,x_{n})\in N\mid x_{1},\dots,x_{n}\in E)$
		$\displaystyle=\frac{\mathbb{P}_{x_{i}\in{\mathcal{N}}(0,I_{d})}((x_{1},\dots,x_{n})\in N\cap(E\times\dots\times E))}{\mathbb{P}_{x_{i}\in{\mathcal{N}}(0,I_{d})}((x_{1},\dots,x_{n})\in(E\times\dots\times E))}$
		$\displaystyle=\frac{0}{[\mathbb{P}_{x_{1}\in{\mathcal{N}}(0,I_{d})}(x_{1}\in E)]^{n}}$
		$\displaystyle=0.$

∎

By Lemma B.1 above, it is not hard to see that Algorithms 4 and 5 work.

Appendix C Omitted Constructions and Proofs in Subsection 4.1

Construction of the Reward Functions

The following construction of the polynomial hard case is adopted from [HHK⁺21].

Let $d$ be the dimension of the feature space. Let $e_{i}$ denotes the $i$ -th standard orthonormal basis of $\mathbb{R}^{d}$ , i.e., $e_{i}$ has only one $1$ at the $i$ -th entry and $0$ ’s for other entries. Let $p$ denote the highest order of the polynomial. We assume $d\gg p$ . We use $\Lambda$ to denote a subset of the $p$ -th multi-indices

\Lambda=\{(\alpha_{1},\dots,\alpha_{p})|1\leq\alpha_{1}\leq\dots\leq\alpha_{p}\leq d\}.

For an $\alpha=(\alpha_{1},\dots,\alpha_{p})\in\Lambda$ , denote $M_{\alpha}=e_{\alpha_{1}}\otimes\dots\otimes e_{\alpha_{p}}$ , $x_{\alpha}=e_{\alpha_{1}}+\dots+e_{\alpha_{p}}$ .

The model space $\mathcal{M}$ is a subset of rank-1 $p$ -th order tensors, which is defined as $\mathcal{M}=\{M_{\alpha}|\alpha\in\Lambda\}$ . We define two subsets of feature space $\mathcal{F}_{0}$ and $\mathcal{F}$ as $\mathcal{F}_{0}=\{x_{\alpha}|\alpha\in\Lambda\},\ \mathcal{F}=\mathrm{conv}(\mathcal{F}_{0})$ . For $M_{\alpha}\in\mathcal{M}$ , $x\in\mathcal{F}$ , define $r(M_{\alpha},x)$ as $r(M_{\alpha},x)=\langle M_{\alpha},x^{\otimes p}\rangle=\prod_{i=1}^{p}\langle e_{\alpha_{i}},x\rangle$ . We assume that for each level $h$ , there is a $M^{(h)}=M_{\alpha^{(h)}}\in\mathcal{M}$ , and the noiseless reward is $r_{h}(s,a)=r(M^{(h)},\phi_{h}(s,a))$ .

We have the following properties.

Proposition C.1 ([HHK⁺21]).

For $M_{\alpha}\in\mathcal{M}$ and $x_{\alpha^{\prime}}\in\mathcal{F}_{0}$ , we have

r(M_{\alpha},x_{\alpha^{\prime}})=\mathbb{I}_{\{\alpha=\alpha^{\prime}\}}.

Proposition C.2.

For $M_{\alpha}\in\mathcal{M}$ , we have

\max_{x\in\mathcal{F}}r(M_{\alpha},x)=1.

proof of Proposition C.2.

For all $x\in\mathcal{F}$ , since $\mathcal{F}=\mathrm{conv}(\mathcal{F}_{0})$ , we can write

x=\sum_{\alpha\in\Lambda}p_{\alpha}(e_{\alpha_{1}}+\dots+e_{\alpha_{p}}),

where $\sum_{\alpha\in\Lambda}p_{\alpha}=1$ and $p_{\alpha}\geq 0$ . Therefore,

r(M_{\alpha^{\prime}},x)=\prod_{i=1}^{p}\langle e_{\alpha^{\prime}_{i}},x\rangle.

Plug in the expression of $x$ , we have

	$\displaystyle\langle e_{\alpha^{\prime}_{i}},x\rangle$	$\displaystyle=\sum_{\alpha}p_{\alpha}\langle e_{\alpha^{\prime}_{i}},e_{\alpha_{1}}+\dots+e_{\alpha_{p}}\rangle$
		$\displaystyle=\sum_{\alpha}p_{\alpha}\mathbb{I}_{\{\alpha^{\prime}_{i}\in\alpha\}}$
		$\displaystyle\leq\sum_{\alpha}p_{\alpha}=1.$

Therefore,

	$\displaystyle r(M_{\alpha^{\prime}},x)$	$\displaystyle=\prod_{i=1}^{p}\langle e_{\alpha^{\prime}_{i}},x\rangle$
		$\displaystyle=\Big{(}\sum_{\alpha}p_{\alpha}\mathbb{I}_{\{e_{\alpha^{\prime}_{1}}\in\alpha\}}\Big{)}\cdots\Big{(}\sum_{\alpha}p_{\alpha}\mathbb{I}_{\{e_{\alpha^{\prime}_{p}}\in\alpha\}}\Big{)}$
		$\displaystyle\leq 1.$

Finally, since $r(M_{\alpha^{\prime}},x_{\alpha^{\prime}})=1$ , we have $\max_{x\in\mathcal{F}}r(M_{\alpha},x)=1$ . ∎

MDP constructions

Consider a family of MDPs with only two states ${\mathcal{S}}=\{S_{\text{good}},S_{\text{bad}}\}$ . The action set ${\mathcal{A}}$ is set to be $\mathcal{F}$ . Let $f$ be a mapping from $\mathcal{F}$ to $\mathcal{F}_{0}$ such that $f$ is identity when restricted to $\mathcal{F}_{0}$ . For all level $h\in[H]$ , we define the feature map $\phi_{h}:S\times A\to\mathcal{F}$ to be

\phi_{h}(s,a)=\left\{\begin{array}[]{ll}a&\text{ if }s=S_{\text{good}},\\ f(a)&\text{ if }s=S_{\text{bad}}.\end{array}\right.

Given an unknown sequence of indices $\alpha^{(1)},\dots,\alpha^{(H)}$ , the reward function at level $h$ is $r_{h}(s,a)=r(M_{\alpha^{(h)}},\phi_{h}(s,a))$ . Specifically, we have

r_{h}(S_{\text{good}},a)=r(M_{\alpha^{(h)}},a),\ r_{h}(S_{\text{bad}},a)=r(M_{\alpha^{(h)}},f(a)).

The transition $P_{h}$ is constructed as

P_{h}(S_{\text{bad}}|s,a)=1\text{ for all }s\in{\mathcal{S}},a\in{\mathcal{A}}.

This construction means it is impossible for the online scenarios to reach the good state for $h>1$ .

The next proposition shows that $Q^{*}_{h}$ is polynomial realizable and falls into the case of Example 4.4.

Proposition C.3.

We have for all $h\in[H]$ and $s\in{\mathcal{S}}$ , $a\in{\mathcal{A}}$ , $V^{*}_{h}(s)=H-h+1$ and $Q^{*}_{h}(s,a)=r_{h}(s,a)+H-h+1$ . Furthermore, $Q^{*}_{h}(s,a)$ , viewed as the function of $\phi_{h}(s,a)$ , is a polynomial of the form $q_{h}(U_{h}\phi_{h}(s,a))$ for some degree- $p$ polynomial $q_{h}$ and $U_{h}\in\mathbb{R}^{p\times d}$ .

proof of Proposition C.3.

First notice that by Proposition C.2, for all $h\in[H]$ and $s\in{\mathcal{S}}$ , we have

\max_{a\in{\mathcal{A}}}r_{h}(s,a)=1.

Therefore, by induction, suppose we have proved for all $s^{\prime}$ , $V^{*}_{h+1}(s^{\prime})=H-h$ , then we have

	$\displaystyle V^{*}_{h}(s)$	$\displaystyle=\max_{a\in{\mathcal{A}}}Q^{*}_{h}(s,a)$
		$\displaystyle=\max_{a\in{\mathcal{A}}}\{r_{h}(s,a)+\mathbb{E}_{s^{\prime}\sim P_{h}(\cdot\|s,a)}[V^{*}_{h+1}(s^{\prime})]\}$
		$\displaystyle=1+H-h.$

Then we have $Q^{*}_{h}(s,a)=r_{h}(s,a)+H-h+1$ .

Furthermore, we have

	$\displaystyle Q^{*}_{h}(s,a)$	$\displaystyle=r_{h}(s,a)+H-h+1$
		$\displaystyle=r(M_{\alpha^{(h)}},\phi_{h}(s,a))+H-h+1$
		$\displaystyle=\prod_{i=1}^{p}\langle e_{\alpha_{i}^{(h)}},\phi_{h}(s,a)\rangle+H-h+1$
		$\displaystyle=q_{h}(U_{h}\phi_{h}(s,a)),$

where $q_{h}(x_{1},\dots,x_{p})=x_{1}x_{2}\cdots x_{p}+(H-h+1)$ and $U_{h}\in\mathbb{R}^{p\times d}$ is a matrix with $e_{\alpha_{i}^{(h)}}$ as the $i$ -th row. ∎

Theorem C.4.

Under the online RL setting, any algorithm needs to play at least $({d\choose p}-1)=\Omega(d^{p})$ episodes to identify $\alpha^{(2)},\dots,\alpha^{(H)}$ and thus to identify the optimal policy.

proof of Theorem C.4.

Under the online RL setting, any algorithm enters and remains in $S_{\text{bad}}$ for $h>1$ . When $s_{h}=S_{\text{bad}}$ , no matter what $a_{h}$ the algorithm chooses, we have $\phi_{h}(s_{h},a_{h})=f(a_{h})\in\mathcal{F}_{0}$ . Notice that for any $M_{\alpha^{(h)}}\in\mathcal{M}$ and any $x_{\alpha}\in\mathcal{F}_{0}$ , we have $r(M_{\alpha^{(h)}},x_{\alpha})=\mathbb{I}_{\{\alpha=\alpha^{(h)}\}}$ as Proposition C.1 suggests. Hence, we need to play $({d\choose p}-1)$ times at level $h$ in the worst case to find out $\alpha^{(h)}$ . The argument holds for all $h=2,3,\dots,H$ . ∎

Theorem C.5.

Under the generative model setting, by querying $2d(p+1)^{p}H=O(dH)$ samples, we can almost surely identify $\alpha^{(1)}$ , $\alpha^{(2)},\dots,\alpha^{(H)}$ and thus identify the optimal policy.

proof of Theorem C.5.

By Proposition C.3, we know that $Q^{*}_{h}(s,a)$ , viewed as the function of $\phi_{h}(s,a)$ , falls into the case of Example 4.4 with $k=p$ .

Next, notice that for all $h\in[H]$ , $\{\phi_{h}(s,a)\mid s\in{\mathcal{S}},a\in{\mathcal{A}}\}={\mathcal{F}}$ . Although ${\mathcal{F}}$ is not of positive measure, we can actually know the value of $Q^{*}_{h}$ when $\phi_{h}(s,a)$ is in $\mathrm{conv}({\mathcal{F}},\mathbf{0})$ since the reward is $p$ -homogenous. Specifically, for every feature of the form $c\cdot\phi_{h}(s,a)$ , where $0\leq c\leq 1$ and $\phi_{h}(s,a)\in{\mathcal{F}}$ , the reward is $c^{p}$ times the reward of $(s,a)$ . Therefore, to get the reward at $c\cdot\phi_{h}(s,a)$ , we only need to query the generative model at $(s,a)$ of level $h$ , and then multiply the reward by $c^{p}$ .

Notice that $\mathrm{conv}({\mathcal{F}},\mathbf{0})$ is of positive Lebesgue measure. By Theorem 4.7, we know that only $2d(p+1)^{p}H=O(dH)$ samples are needed to determine the optimal policy almost surely.

∎

Appendix D Technical claims

Claim D.1.

Let $\chi^{2}(d)$ denote $\chi^{2}$ -distribution with degree of freedom $d$ . For any $t>0$ we have,

\displaystyle\Pr_{z\sim\chi^{2}(d)}(z\geq d+2t+2\sqrt{dt})\leq e^{-t}

We use the following facts from [ZSJ⁺17].

Claim D.2.

Given a fixed vector $z\in\mathbb{R}^{d}$ , for any $C\geq 1$ and $n\geq 1$ , we have

\displaystyle\underset{x\sim{\cal N}(0,I_{d})}{\Pr}[|\langle x,z\rangle|^{2}\leq 5C\|z\|^{2}\log n]\geq 1-1/(nd^{C}).

Claim D.3.

For any $C\geq 1$ and $n\geq 1$ , we have

\displaystyle\underset{x\sim{\cal N}(0,I_{d})}{\Pr}[\|x\|^{2}\leq 5Cd\log n]\geq 1-1/(nd^{C}).

Claim D.4.

Let $a,b,c\geq 0$ be three constants, let $u,v,w\in\mathbb{R}^{d}$ be three vectors, we have

\displaystyle\underset{x\sim{\cal N}(0,I_{d})}{\mathbb{E}}\left[|u^{\top}x|^{a}|v^{\top}x|^{b}|w^{\top}x|^{c}\right]\approx\|u\|^{a}\|v\|^{b}\|w\|^{c}.

Claim D.5.

Let $M_{j},j\in[4]$ be defined in Definition A.11. For each $j\in[4]$ , $M_{j}=\sum_{i=1}^{k}v_{i}m_{j,i}\overline{w}_{i}^{\otimes j}$ .

Claim D.6.

Let ${\cal B}$ denote a distribution over $\mathbb{R}^{d_{1}\times d_{2}}$ . Let $d=d_{1}+d_{2}$ . Let $B_{1},B_{2},\cdots B_{n}$ be i.i.d. random matrices sampled from ${\cal B}$ . Let $\overline{B}=\mathbb{E}_{B\sim{\cal B}}[B]$ and $\widehat{B}=\frac{1}{n}\sum_{i=1}^{n}B_{i}$ . For parameters $m\geq 0,\gamma\in(0,1),\nu>0,L>0$ , if the distribution ${\cal B}$ satisfies the following four properties,

	$\displaystyle\mathrm{({1})}\quad$	$\displaystyle\quad\underset{B\sim{\cal B}}{\Pr}\left[\left\\|B\right\\|\leq m\right]\geq 1-\gamma;$
	$\displaystyle\mathrm{({2})}\quad$	$\displaystyle\quad\left\\|\underset{B\sim{\cal B}}{\mathbb{E}}[B]\right\\|>0;$
	$\displaystyle\mathrm{({3})}\quad$	$\displaystyle\quad\max\left(\left\\|\underset{B\sim{\cal B}}{\mathbb{E}}[BB^{\top}]\right\\|,\left\\|\underset{B\sim{\cal B}}{\mathbb{E}}[B^{\top}B]\right\\|\right)\leq\nu;$
	$\displaystyle\mathrm{({4})}\quad$	$\displaystyle\quad\max_{\\|a\\|=\\|b\\|=1}\left(\underset{B\sim{\cal B}}{\mathbb{E}}\left[\left(a^{\top}Bb\right)^{2}\right]\right)^{1/2}\leq L.$

Then we have for any $0<\epsilon<1$ and $t\geq 1$ , if

\displaystyle n\geq(18t\log d)\cdot(\nu+\|\overline{B}\|^{2}+m\|\overline{B}\|\epsilon)/(\epsilon^{2}\|\overline{B}\|^{2})\quad\text{\leavevmode\nobreak\ and\leavevmode\nobreak\ }\quad\gamma\leq(\epsilon\|\overline{B}\|/(2L))^{2}

with probability at least $1-1/d^{2t}-n\gamma$ ,

\|\widehat{B}-\overline{B}\|\leq\epsilon\|\overline{B}\|.

Claim D.7.

Let $P_{2}$ and $P_{3}$ be defined in Definition A.16. Then

P_{2}=\sum_{i=1}^{k}v_{i}m_{j_{2},i}(\alpha^{\top}\overline{w}_{i})^{{j_{2}-2}}\overline{w}_{i}^{\otimes 2}

and

P_{3}=\sum_{i=1}^{k}v_{i}m_{j_{3},i}(\alpha^{\top}\overline{w}_{i})^{{j_{3}-3}}\overline{w}_{i}^{\otimes 3}.

	$\displaystyle\\|C_{i}(x_{j})\\|\lesssim$	$\displaystyle(\|w_{i}^{\top}x_{j}\|^{p+1}+\|\phi(0)\|+\|\xi_{j}\|)\cdot((x_{j}^{\top}\alpha)^{2}\\|x_{j}\\|^{2}+1+\\|x_{j}\\|^{2}+(\alpha^{\top}x_{j})^{2})$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ (\\|w_{i}\\|^{p+1}+\|\phi(0)\|+\vartheta)\cdot d\cdot\mathrm{poly}(\log(d/t))$

	$\displaystyle\left\\|\mathbb{E}[C_{i}(x_{j})^{2}]\right\\|\lesssim$	$\displaystyle\leavevmode\nobreak\ \left(\mathbb{E}[(\phi(w_{i}^{\top}x_{j})+\xi_{j})^{4}]\right)^{1/2}\left(\mathbb{E}[(x_{j}^{\top}\alpha)^{8}]\right)^{1/2}\left(\mathbb{E}[\\|x_{j}\\|^{4}]\right)^{1/2}$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ (\\|w_{i}\\|^{p+1}+\|\phi(0)\|+\vartheta)^{2}d.$

		$\displaystyle\leavevmode\nobreak\ \left\\|m_{4,i}(\overline{w}_{i}^{\top}\alpha)^{2}\overline{w}_{i}\overline{w}_{i}^{\top}-\frac{1}{\|S\|}\sum_{x_{j}\in S}C_{i}(x_{j})\right\\|$
	$\displaystyle\leq$	$\displaystyle\leavevmode\nobreak\ \left\\|m_{4,i}(\overline{w}_{i}^{\top}\alpha)^{2}\overline{w}_{i}\overline{w}_{i}^{\top}-\mathbb{E}[C_{i}(x_{j})]\right\\|+\left\\|\mathbb{E}[C_{i}(x_{j})]-\frac{1}{\|S\|}\sum_{x_{j}\in S}C_{i}(x_{j})\right\\|$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ \epsilon\|m_{4,i}\|(\overline{w}_{i}^{\top}\alpha)^{2}+\epsilon.$

	$\displaystyle\\|\mathbb{E}_{\xi_{j}}[B_{i}(x_{j})]\\|\lesssim$	$\displaystyle(\|w_{i}^{\top}x_{j}\|^{p+1}+\|\phi(0)\|)\cdot((x_{j}^{\top}\alpha)^{2}\\|V^{\top}x_{j}\\|^{3}$
		$\displaystyle\leavevmode\nobreak\ +3\\|V^{\top}x_{j}\\|^{3}+3\|x_{j}^{\top}\alpha\|\\|V^{\top}x_{j}\\|\sqrt{k}+3\\|V^{\top}\alpha\\|\sqrt{k})$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ \sqrt{k}\cdot\\|w_{i}\\|^{p+1}\\|x_{j}\\|^{p+6}.$

	$\displaystyle\\|C_{i}(x_{j})\\|\lesssim$	$\displaystyle(\|w_{i}^{\top}x_{j}\|^{p+1}+\|\phi(0)\|+\|\xi_{j}\|)\cdot((x_{j}^{\top}\alpha)^{2}\\|V^{\top}x_{j}\\|^{3}$
		$\displaystyle\leavevmode\nobreak\ +3\\|V^{\top}x_{j}\\|^{3}+3\|x_{j}^{\top}\alpha\|\\|V^{\top}x_{j}\\|\sqrt{k}+3\\|V^{\top}\alpha\\|\sqrt{k})$
	$\displaystyle\lesssim$	$\displaystyle\leavevmode\nobreak\ (\\|w_{i}\\|^{p+1}+\|\phi(0)\|+\vartheta)k^{3/2}\mathrm{poly}(\log(d/t))$

Going Beyond Linear RL: Sample Efficient Neural Function Approximation

Abstract

1 Introduction

1.1 Summary of our results

1.2 Related Work

Linear Function Approximation.

Eluder Dimension.

Bellman Rank and Completeness.

Deterministic RL

2 Preliminaries

2.1 Function approximation

Assumption 2.1 (Realizability).

Assumption 2.2 (Bounded features).

Notation

3 Neural Network Function Approximation

Definition 3.1 (Neural Network Function Class).

Definition 3.2 (Policy complete).

Definition 3.3 (Bellman complete).

3.1 Warmup: Realizable Q∗Q^{\ast} with deterministic transition

Theorem 3.4.

3.2 Policy complete neural function approximation

Theorem 3.5.

3.3 Bellman complete neural function approximation

Theorem 3.6.

3.4 Realizable Q∗Q^{\ast} with optimality gap

Definition 3.7.

Theorem 3.8.

Remark 3.9.

4 Polynomial Realizability

Definition 4.1 (Admissible Polynomial Families).

Theorem 4.2 ([HHK+21]).

Example 4.3.

Example 4.4.

Definition 4.5.

Assumption 4.6 (Polynomial Realizability).

Theorem 4.7.

Theorem 4.8.

4.1 Necessity of Generic Feature Maps in Online RL

MDP construction

Theorem 4.9.

5 Conclusions

Acknowledgment

References

Appendix A Omitted Proofs in Section 3

A.1 Proof of Section 3.1

Theorem A.1 (Formal statement of Theorem 3.4).

Proof.

A.2 Proof of Section 3.2

Theorem A.2 (Formal statement of Theorem 3.5).

Proof.

A.3 Proof of Section 3.3

Theorem A.3 (Formal statement of Theorem 3.6).

Proof.

A.4 Proof of Section 3.4

Theorem A.4 (Formal statement of Theorem 3.8).

Proof.

A.5 Neural network recovery

Assumption A.5 (Property 3.1 of [ZSJ+17]).

Definition A.6 (Part of property 3.2 of [ZSJ+17]).

Assumption A.7 (Part of property 3.2 of [ZSJ+17]).

Assumption A.8 (Property 3.3 of [ZSJ+17]).

Assumption A.9 (Rich feature).

Definition A.10.

Definition A.11.

Assumption A.12 (Assumption 5.3 of [ZSJ+17]).

Theorem A.13 (Neural network recovery from noisy data).

Corollary A.14.

Theorem A.15 (Exact neural network recovery from noiseless data).

A.5.1 Recover neural networks from noisy data

Definition A.16.

Lemma A.17.

Proof.

Lemma A.18.

Proof.

Lemma A.19.

Proof.

Lemma A.20 (Adapted from Theorem 3 of [KCL15]).

Lemma A.21 (Adapted from Lemma E.6 of [ZSJ+17]).

Lemma A.22 (Adapted from Lemma E.13 in [ZSJ+17]).

Lemma A.23 (Adapted from Lemma E.14 in [ZSJ+17]).

3.1 Warmup: Realizable $Q^{\ast}$ with deterministic transition

3.4 Realizable $Q^{\ast}$ with optimality gap

Theorem 4.2 ([HHK⁺21]).

Assumption A.5 (Property 3.1 of [ZSJ⁺17]).

Definition A.6 (Part of property 3.2 of [ZSJ⁺17]).

Assumption A.7 (Part of property 3.2 of [ZSJ⁺17]).

Assumption A.8 (Property 3.3 of [ZSJ⁺17]).

Assumption A.12 (Assumption 5.3 of [ZSJ⁺17]).

Lemma A.21 (Adapted from Lemma E.6 of [ZSJ⁺17]).

Lemma A.22 (Adapted from Lemma E.13 in [ZSJ⁺17]).

Lemma A.23 (Adapted from Lemma E.14 in [ZSJ⁺17]).

Lemma A.25 (Concentration, adapted from Lemma D.11 in [ZSJ⁺17]).

Lemma A.26 (Adapted from Lemma D.16 in [ZSJ⁺17]).

Proposition C.1 ([HHK⁺21]).