How Neural Networks Extrapolate: From Feedforward to Graph Neural Networks — part 8/11

Cleaned-text code points [107629, 123435). Bibliography omitted here; source.md includes it.

Let us consider the solution 𝜷 + ′ = 1 \bm{\beta}_{+}^{\prime}=1 and 𝜷 + = 𝜷 − = 𝜷 − ′ = 0 \bm{\beta}_{+}=\bm{\beta}_{-}=\bm{\beta}_{-}^{\prime}=0 . It is clear that this solution can fit the training data, and thus satisfies Eqn. 71 . Moreover, this solution is equivalent to the underlying reasoning function, maximum degree, g ⁡ ( G ) = max u ∈ G ⁡ 𝒉 u g(G)=\max_{u\in G}\bm{h}_{u} .

Hence, it remains to show that, given our training data, there exist λ i \lambda_{i} so that the remaining four constraints are satisfies for this solution. Let us rewrite these constraints as a linear systems where the variables are λ i \lambda_{i} ( 𝜷 + 𝜷 − 𝜷 + ′ 𝜷 − ′ ) = c ⋅ ∑ i = 1 n λ i ⋅ ( ∑ u ∈ G i 𝒉 u ⋅ 𝕀 ⁡ ( u = arg ⁡ max v ∈ G i ⁡ 𝒉 v ) ∑ u ∈ G i 𝒉 u ⋅ 𝕀 ⁡ ( u = arg ⁡ min v ∈ G i ⁡ 𝒉 v ) max u ∈ G i ⁡ 𝒉 u min u ∈ G i ⁡ 𝒉 u ) \displaystyle\begin{pmatrix}\bm{\beta}_{+}\\ \bm{\beta}_{-}\\ \bm{\beta}_{+}^{\prime}\\ \bm{\beta}_{-}^{\prime}\end{pmatrix}=c\cdot\sum\limits_{i=1}^{n}\lambda_{i}\cdot\begin{pmatrix}\sum_{u\in G_{i}}\bm{h}_{u}\cdot\mathbb{I}\left(u=\arg\max\limits_{v\in G_{i}}\bm{h}_{v}\right)\\ \sum_{u\in G_{i}}\bm{h}_{u}\cdot\mathbb{I}\left(u=\arg\min\limits_{v\in G_{i}}\bm{h}_{v}\right)\\ \max\limits_{u\in G_{i}}\bm{h}_{u}\\ \min\limits_{u\in G_{i}}\bm{h}_{u}\end{pmatrix} (73)

By standard theory of linear systems, there exist λ i \lambda_{i} to solve Eqn. 73 if there are at least four training data G i G_{i} whose following vectors linear independent ( ∑ u ∈ G i 𝒉 u ⋅ 𝕀 ⁡ ( u = arg ⁡ max v ∈ G i ⁡ 𝒉 v ) ∑ u ∈ G i 𝒉 u ⋅ 𝕀 ⁡ ( u = arg ⁡ min v ∈ G i ⁡ 𝒉 v ) max u ∈ G i ⁡ 𝒉 u min u ∈ G i ⁡ 𝒉 u ) = ( max u ∈ G i ⁡ 𝒉 u ⋅ N i max min u ∈ G i ⁡ 𝒉 u ⋅ N i min max u ∈ G i ⁡ 𝒉 u min u ∈ G i ⁡ 𝒉 u ) \displaystyle\begin{pmatrix}\sum_{u\in G_{i}}\bm{h}_{u}\cdot\mathbb{I}\left(u=\arg\max\limits_{v\in G_{i}}\bm{h}_{v}\right)\\ \sum_{u\in G_{i}}\bm{h}_{u}\cdot\mathbb{I}\left(u=\arg\min\limits_{v\in G_{i}}\bm{h}_{v}\right)\\ \max\limits_{u\in G_{i}}\bm{h}_{u}\\ \min\limits_{u\in G_{i}}\bm{h}_{u}\end{pmatrix}=\begin{pmatrix}\max\limits_{u\in G_{i}}\bm{h}_{u}\cdot N^{\max}_{i}\\ \min\limits_{u\in G_{i}}\bm{h}_{u}\cdot N^{\min}_{i}\\ \max\limits_{u\in G_{i}}\bm{h}_{u}\\ \min\limits_{u\in G_{i}}\bm{h}_{u}\end{pmatrix} (74)

Here, N i max N^{\max}_{i} denotes the number of nodes that achieve the maximum degree in the graph G i G_{i} , and N i min N^{\min}_{i} denotes the number of nodes that achieve the min degree in the graph G i G_{i} . By the assumption of our training data that there are at least four G i ∼ 𝒢 G_{i}\sim\mathcal{G} with linearly independent Eqn. 74 . Hence, our simple GNN learns the underlying function as desired.

This completes the proof.

### B.6 Proof of Lemma 2

Let W W denote the span of the feature maps of training data 𝒙 i \bm{x}_{i} , i.e. W = span ​ ( ϕ ⁡ ( 𝒙 1 ) , ϕ ⁡ ( 𝒙 2 ) , … , ϕ ⁡ ( 𝒙 n ) ) . \displaystyle W=\text{span}\left(\phi\left(\bm{x}_{1}\right),\phi\left(\bm{x}_{2}\right),...,\phi\left(\bm{x}_{n}\right)\right). Then we can decompose the coordinates of f NTK f_{\text{NTK}} in the RKHS space, 𝜷 NTK \bm{\beta}_{\text{NTK}} , into a vector 𝜷 0 \bm{\beta}_{0} for the component of f NTK f_{\text{NTK}} in the span of training data features W W , and a vector 𝜷 1 \bm{\beta}_{1} for the component in the orthogonal complement W ⊤ W^{\top} , i.e., 𝜷 NTK = 𝜷 0 + 𝜷 1 . \bm{\beta}_{\text{NTK}}=\bm{\beta}_{0}+\bm{\beta}_{1}. First, note that since f NTK f_{\text{NTK}} must be able to fit the training data (NTK is a universal kernel as we will discuss next), i.e., ϕ ​ ( 𝒙 i ) ⊤ ​ 𝜷 NTK = y i . \phi(\bm{x}_{i})^{\top}\bm{\beta}_{\text{NTK}}=y_{i}. Thus, we have ϕ ​ ( 𝒙 i ) ⊤ ​ 𝜷 0 = y i \phi(\bm{x}_{i})^{\top}\bm{\beta}_{0}=y_{i} . Then, 𝜷 0 \bm{\beta}_{0} is uniquely determined by the kernel regression solution with respect to the neural tangent kernel f NTK ( 𝒙 ) = ( ⟨ ϕ ( 𝒙 ) , ϕ ( 𝒙 1 ) ⟩ , … , ⟨ ϕ ( 𝒙 ) , ϕ ( 𝒙 n ) ⟩ ) ⋅ NTK train − 1 𝒀 , \displaystyle f_{\text{NTK}}(\bm{x})=\left(\bigl<\phi(\bm{x}),\phi(\bm{x}_{1})\bigr>,...,\bigl<\phi(\bm{x}),\phi(\bm{x}_{n})\bigr>\right)\cdot\text{NTK}_{\text{train}}^{-1}\bm{Y}, where NTK train \text{NTK}_{\text{train}} is the n × n n\times n kernel for training data, ⟨ ϕ ( 𝒙 ) , ϕ ( 𝒙 i ) ⟩ \bigl<\phi(\bm{x}),\phi(\bm{x}_{i})\bigr> is the kernel between test data 𝒙 \bm{x} and training data 𝒙 i \bm{x}_{i} , and 𝒀 \bm{Y} is training labels.

The kernel regression solution f NTK f_{\text{NTK}} is uniquely determined because the neural tangent kernel NTK train \text{NTK}_{\text{train}} is positive definite assuming no two training data are parallel, which can be enforced with a bias term ( Du et al., 2019c ) . In any case, the solution is a min-norm by pseudo-inverse.

Moreover, a unique kernel regression solution f NTK f_{\text{NTK}} that spans the training data features corresponds to a unique representation in the RKHS space 𝜷 0 \bm{\beta}_{0} .

Since 𝜷 0 \bm{\beta}_{0} and 𝜷 1 \bm{\beta}_{1} are orthogonal, we also have the following ‖ 𝜷 NTK ‖ 2 2 = ‖ 𝜷 0 + 𝜷 1 ‖ 2 2 = ‖ 𝜷 0 ‖ 2 2 + ‖ 𝜷 1 ‖ 2 2 . \displaystyle\|\bm{\beta}_{\text{NTK}}\|_{2}^{2}=\|\bm{\beta}_{0}+\bm{\beta}_{1}\|_{2}^{2}=\|\bm{\beta}_{0}\|_{2}^{2}+\|\bm{\beta}_{1}\|_{2}^{2}. This implies the norm of 𝜷 NTK \bm{\beta}_{\text{NTK}} is at least as large as the norm of any 𝜷 \bm{\beta} such that ϕ ​ ( 𝒙 i ) ⊤ ​ 𝜷 NTK = y i \phi(\bm{x}_{i})^{\top}\bm{\beta}_{\text{NTK}}=y_{i} . Moreover, observe that the solution to kernel regression Eqn. 9 is in the feature span of training data, given the kernel matrix for training data is full rank. f NTK ( 𝒙 ) = ( ⟨ ϕ ( 𝒙 ) , ϕ ( 𝒙 1 ) ⟩ , … , ⟨ ϕ ( 𝒙 ) , ϕ ( 𝒙 n ) ⟩ ) ⋅ NTK train − 1 𝒀 . \displaystyle f_{\text{NTK}}(\bm{x})=\left(\bigl<\phi(\bm{x}),\phi(\bm{x}_{1})\bigr>,...,\bigl<\phi(\bm{x}),\phi(\bm{x}_{n})\bigr>\right)\cdot\text{NTK}_{\text{train}}^{-1}\bm{Y}. Since 𝜷 1 \bm{\beta}_{1} is for the component of f NTK f_{\text{NTK}} in the orthogonal complement of training data feature span, we must have 𝜷 1 = 𝟎 \bm{\beta}_{1}=\bm{0} . It follows that 𝜷 NTK \bm{\beta}_{\text{NTK}} is equivalent to min 𝜷 ⁡ ‖ 𝜷 ‖ 2 \displaystyle\min_{\bm{\beta}}\|\bm{\beta}\|_{2} s.t. ϕ ( 𝒙 i ) ⊤ 𝜷 = y i , for i = 1 , … , n . \displaystyle\phi(\bm{x}_{i})^{\top}\bm{\beta}=y_{i},\;\;\;\text{for }i=1,...,n. as desired.

### B.7 Proof of Lemma 3

We first compute the neural tangent kernel NTK ​ ( 𝒙 , 𝒙 ′ ) \text{NTK}(\bm{x},\bm{x}^{\prime}) for a two-layer multi-layer perceptron (MLP) with ReLU activation function, and then show that it can be induced by the feature space ϕ ⁡ ( 𝒙 ) \phi(\bm{x}) specified in the lemma so that NTK ( 𝒙 , 𝒙 ′ ) = ⟨ ϕ ( 𝒙 ) , ϕ ( 𝒙 ′ ) ⟩ \text{NTK}(\bm{x},\bm{x}^{\prime})=\bigl<\phi(\bm{x}),\phi(\bm{x}^{\prime})\bigr> .

Recall that Jacot et al. (2018) have derived the general framework for computing the neural tangent kernel of a neural network with general architecture and activation function. This framework is also described in Arora et al. (2019b) ; Du et al. (2019b) , which, in addition, compute the exact kernel formula for convolutional networks and Graph Neural Networks, respectively. Following the framework in Jacot et al. (2018) and substituting the general activation function σ \sigma with ReLU gives the kernel formula for a two-layer MLP with ReLU activation. This has also been described in several previous works ( Du et al., 2019c ; Chizat et al., 2019 ; Bietti & Mairal, 2019 ) .

Below we describe the general framework in Jacot et al. (2018) and Arora et al. (2019b) . Let σ \sigma denote the activation function. The neural tangent kernel for an h h -layer multi-layer perceptron can be recursively defined via a dynamic programming process. Here, Σ ( i ) : ℝ d × ℝ d → ℝ \Sigma^{(i)}:\mathbb{R}^{d}\times\mathbb{R}^{d}\rightarrow\mathbb{R} for i = 0 ​ … ​ h i=0...h is the covariance for the i i -th layer. Σ ( 0 ) ​ ( 𝒙 , 𝒙 ′ ) = 𝒙 ⊤ ​ 𝒙 ′ , \displaystyle\Sigma^{(0)}(\bm{x},\bm{x}^{\prime})=\bm{x}^{\top}\bm{x}^{\prime}, ∧ ( i ) ( 𝒙 , 𝒙 ′ ) = ( Σ ( i − 1 ) ​ ( 𝒙 , 𝒙 ) Σ ( i − 1 ) ​ ( 𝒙 , 𝒙 ′ ) Σ ( i − 1 ) ​ ( 𝒙 ′ , 𝒙 ) Σ ( i − 1 ) ​ ( 𝒙 ′ , 𝒙 ′ ) ) , \displaystyle\bm{\wedge}^{(i)}(\bm{x},\bm{x}^{\prime})=\left(\begin{array}[]{cc}\Sigma^{(i-1)}(\bm{x},\bm{x})&\Sigma^{(i-1)}(\bm{x},\bm{x}^{\prime})\\ \Sigma^{(i-1)}(\bm{x}^{\prime},\bm{x})&\Sigma^{(i-1)}(\bm{x}^{\prime},\bm{x}^{\prime})\end{array}\right), Σ ( i ) ( 𝒙 , 𝒙 ′ ) = c ⋅ 𝔼 u , v ∼ 𝒩 ( 𝟎 , ∧ ( i ) ) [ σ ( u ) σ ( v ) ] . \displaystyle\Sigma^{(i)}(\bm{x},\bm{x}^{\prime})=c\cdot\mathop{\mathbb{E}}_{u,v\sim\mathcal{N}(\bm{0},\bm{\wedge}^{(i)})}\left[\sigma(u)\sigma(v)\right]. The derivative covariance is defined similarly: Σ ˙ ( i ) ( 𝒙 , 𝒙 ′ ) = c ⋅ 𝔼 u , v ∼ 𝒩 ( 𝟎 , ∧ ( i ) ) [ σ ˙ ( u ) σ ˙ ( v ) ] . \displaystyle\dot{\Sigma}^{(i)}(\bm{x},\bm{x}^{\prime})=c\cdot\mathop{\mathbb{E}}_{u,v\sim\mathcal{N}(\bm{0},\bm{\wedge}^{(i)})}\left[\dot{\sigma}(u)\dot{\sigma}(v)\right]. Then the neural tangent kernel for an h h -layer network is defined as NTK ( h − 1 ) ​ ( 𝒙 , 𝒙 ′ ) = ∑ i = 1 h ( Σ ( i − 1 ) ​ ( 𝒙 , 𝒙 ′ ) ⋅ ∏ k = i h Σ ˙ ( k ) ​ ( 𝒙 , 𝒙 ′ ) ) , \displaystyle\text{NTK}^{(h-1)}(\bm{x},\bm{x}^{\prime})=\sum\limits_{i=1}^{h}\left(\Sigma^{(i-1)}(\bm{x},\bm{x}^{\prime})\cdot\prod\limits_{k=i}^{h}\dot{\Sigma}^{(k)}(\bm{x},\bm{x}^{\prime})\right), where we let Σ ˙ ( h ) ​ ( 𝒙 , 𝒙 ′ ) = 1 \dot{\Sigma}^{(h)}(\bm{x},\bm{x}^{\prime})=1 for the convenience of notations.

We compute the explict NTK formula for a two-layer MLP with ReLU activation function by following this framework and substituting the general activation function with ReLU, i.e. σ ⁡ ( a ) = max ⁡ ( 0 , a ) = a ⋅ 𝕀 ⁡ ( a ≥ 0 ) \sigma(a)=\max(0,a)=a\cdot\mathbb{I}(a\geq 0) and σ ˙ ​ ( a ) = 𝕀 ​ ( a ≥ 0 ) \dot{\sigma}(a)=\mathbb{I}(a\geq 0) . NTK ( 1 ) ​ ( 𝒙 , 𝒙 ′ ) = \displaystyle\text{NTK}^{(1)}(\bm{x},\bm{x}^{\prime})=\; ∑ i = 1 2 ( Σ ( i − 1 ) ​ ( 𝒙 , 𝒙 ′ ) ⋅ ∏ k = i h Σ ˙ ( k ) ​ ( 𝒙 , 𝒙 ′ ) ) \displaystyle\sum\limits_{i=1}^{2}\left(\Sigma^{(i-1)}(\bm{x},\bm{x}^{\prime})\cdot\prod\limits_{k=i}^{h}\dot{\Sigma}^{(k)}(\bm{x},\bm{x}^{\prime})\right) = \displaystyle= Σ ( 0 ) ​ ( 𝒙 , 𝒙 ′ ) ⋅ Σ ˙ ( 1 ) ​ ( 𝒙 , 𝒙 ′ ) + Σ ( 1 ) ​ ( 𝒙 , 𝒙 ′ ) \displaystyle\;\Sigma^{(0)}(\bm{x},\bm{x}^{\prime})\cdot\dot{\Sigma}^{(1)}(\bm{x},\bm{x}^{\prime})+\Sigma^{(1)}(\bm{x},\bm{x}^{\prime}) So we can get the NTK via Σ ( 1 ) ​ ( 𝒙 , 𝒙 ′ ) \Sigma^{(1)}(\bm{x},\bm{x}^{\prime}) and Σ ˙ ( 1 ) ​ ( 𝒙 , 𝒙 ′ ) \dot{\Sigma}^{(1)}(\bm{x},\bm{x}^{\prime}) , Σ ( 0 ) ​ ( 𝒙 , 𝒙 ′ ) \Sigma^{(0)}(\bm{x},\bm{x}^{\prime}) . Precisely, Σ ( 0 ) ​ ( 𝒙 , 𝒙 ′ ) = 𝒙 ⊤ ​ 𝒙 ′ , \displaystyle\Sigma^{(0)}(\bm{x},\bm{x}^{\prime})=\bm{x}^{\top}\bm{x}^{\prime}, ∧ ( 1 ) ( 𝒙 , 𝒙 ′ ) = ( 𝒙 ⊤ ​ 𝒙 𝒙 ⊤ ​ 𝒙 ′ 𝒙 ′ ⊤ 𝒙 𝒙 ′ ⊤ 𝒙 ′ ) = ( 𝒙 𝒙 ′ ) ⋅ ( 𝒙 𝒙 ′ ) , \displaystyle\bm{\wedge}^{(1)}(\bm{x},\bm{x}^{\prime})=\left(\begin{array}[]{cc}\bm{x}^{\top}\bm{x}&\bm{x}^{\top}\bm{x}^{\prime}\\ \bm{x}^{\prime^{\top}}\bm{x}&\bm{x}^{\prime^{\top}}\bm{x}^{\prime}\end{array}\right)=\left(\begin{array}[]{c}\bm{x}\\ \bm{x}^{\prime}\end{array}\right)\cdot\left(\begin{array}[]{cc}\bm{x}&\bm{x}^{\prime}\end{array}\right), Σ ( 1 ) ( 𝒙 , 𝒙 ′ ) = c ⋅ 𝔼 u , v ∼ 𝒩 ( 𝟎 , ∧ ( 1 ) ) [ u ⋅ 𝕀 ( u ≥ 0 ) ⋅ v ⋅ 𝕀 ( v ≥ 0 ) ] . \displaystyle\Sigma^{(1)}(\bm{x},\bm{x}^{\prime})=c\cdot\mathop{\mathbb{E}}_{u,v\sim\mathcal{N}(\bm{0},\bm{\wedge}^{(1)})}\left[u\cdot\mathbb{I}(u\geq 0)\cdot v\cdot\mathbb{I}(v\geq 0)\right]. To sample from 𝒩 ( 𝟎 , ∧ ( 1 ) ) \mathcal{N}(\bm{0},\bm{\wedge}^{(1)}) , we let L L be a decomposition of ∧ ( 1 ) \bm{\wedge}^{(1)} , such that ∧ ( 1 ) = L L ⊤ \bm{\wedge}^{(1)}=LL^{\top} . Here, we can see that L = ( 𝒙 , 𝒙 ′ ) ⊤ L=(\bm{x},\bm{x}^{\prime})^{\top} . Thus, sampling from 𝒩 ( 𝟎 , ∧ ( 1 ) ) \mathcal{N}(\bm{0},\bm{\wedge}^{(1)}) is equivalent to first sampling 𝒘 ∼ 𝒩 ⁡ ( 𝟎 , 𝑰 ) \bm{w}\sim\mathcal{N}(\bm{0},\bm{I}) , and output L ​ 𝒘 = 𝒘 ⊤ ​ ( 𝒙 , 𝒙 ′ ) . L\bm{w}=\bm{w}^{\top}(\bm{x},\bm{x}^{\prime}). Then we have the equivalent sampling ( u , v ) = ( 𝒘 ⊤ ​ 𝒙 , 𝒘 ⊤ ​ 𝒙 ′ ) (u,v)=(\bm{w}^{\top}\bm{x},\bm{w}^{\top}\bm{x}^{\prime}) . It follows that Σ ( 1 ) ​ ( 𝒙 , 𝒙 ′ ) \displaystyle\Sigma^{(1)}(\bm{x},\bm{x}^{\prime}) = c ⋅ 𝔼 𝒘 ∼ 𝒩 ⁡ ( 𝟎 , 𝑰 ) [ 𝒘 ⊤ ​ 𝒙 ⋅ 𝕀 ⁡ ( 𝒘 ⊤ ​ 𝒙 ≥ 0 ) ⋅ 𝒘 ⊤ ​ 𝒙 ′ ⋅ 𝕀 ⁡ ( 𝒘 ⊤ ​ 𝒙 ′ ≥ 0 ) ] \displaystyle=c\cdot\mathop{\mathbb{E}}_{\bm{w}\sim\mathcal{N}(\bm{0},\bm{I})}\left[\bm{w}^{\top}\bm{x}\cdot\mathbb{I}\left(\bm{w}^{\top}\bm{x}\geq 0\right)\cdot\bm{w}^{\top}\bm{x}^{\prime}\cdot\mathbb{I}\left(\bm{w}^{\top}\bm{x}^{\prime}\geq 0\right)\right] It follows from the same reasoning that Σ ˙ ( 1 ) ​ ( 𝒙 , 𝒙 ′ ) = c ⋅ 𝔼 𝒘 ∼ 𝒩 ⁡ ( 𝟎 , 𝑰 ) [ 𝕀 ⁡ ( 𝒘 ⊤ ​ 𝒙 ≥ 0 ) ⋅ 𝕀 ⁡ ( 𝒘 ⊤ ​ 𝒙 ′ ≥ 0 ) ] . \displaystyle\dot{\Sigma}^{(1)}(\bm{x},\bm{x}^{\prime})=c\cdot\mathop{\mathbb{E}}_{\bm{w}\sim\mathcal{N}(\bm{0},\bm{I})}\left[\mathbb{I}\left(\bm{w}^{\top}\bm{x}\geq 0\right)\cdot\mathbb{I}\left(\bm{w}^{\top}\bm{x}^{\prime}\geq 0\right)\right]. The neural tangent kernel for a two-layer MLP with ReLU activation is then NTK ( 1 ) ​ ( 𝒙 , 𝒙 ′ ) = \displaystyle\text{NTK}^{(1)}(\bm{x},\bm{x}^{\prime})= Σ ( 0 ) ​ ( 𝒙 , 𝒙 ′ ) ⋅ Σ ˙ ( 1 ) ​ ( 𝒙 , 𝒙 ′ ) + Σ ( 1 ) ​ ( 𝒙 , 𝒙 ′ ) \displaystyle\;\Sigma^{(0)}(\bm{x},\bm{x}^{\prime})\cdot\dot{\Sigma}^{(1)}(\bm{x},\bm{x}^{\prime})+\Sigma^{(1)}(\bm{x},\bm{x}^{\prime}) = \displaystyle= c ⋅ 𝔼 𝒘 ∼ 𝒩 ⁡ ( 𝟎 , 𝑰 ) [ 𝒙 ⊤ ​ 𝒙 ′ ⋅ 𝕀 ⁡ ( 𝒘 ⊤ ​ 𝒙 ≥ 0 ) ⋅ 𝕀 ⁡ ( 𝒘 ⊤ ​ 𝒙 ′ ≥ 0 ) ] \displaystyle\;c\cdot\mathop{\mathbb{E}}_{\bm{w}\sim\mathcal{N}(\bm{0},\bm{I})}\left[\bm{x}^{\top}\bm{x}^{\prime}\cdot\mathbb{I}\left(\bm{w}^{\top}\bm{x}\geq 0\right)\cdot\mathbb{I}\left(\bm{w}^{\top}\bm{x}^{\prime}\geq 0\right)\right] + \displaystyle+ c ⋅ 𝔼 𝒘 ∼ 𝒩 ⁡ ( 𝟎 , 𝑰 ) [ 𝒘 ⊤ ​ 𝒙 ⋅ 𝕀 ⁡ ( 𝒘 ⊤ ​ 𝒙 ≥ 0 ) ⋅ 𝒘 ⊤ ​ 𝒙 ′ ⋅ 𝕀 ⁡ ( 𝒘 ⊤ ​ 𝒙 ′ ≥ 0 ) ] . \displaystyle\;c\cdot\mathop{\mathbb{E}}_{\bm{w}\sim\mathcal{N}(\bm{0},\bm{I})}\left[\bm{w}^{\top}\bm{x}\cdot\mathbb{I}\left(\bm{w}^{\top}\bm{x}\geq 0\right)\cdot\bm{w}^{\top}\bm{x}^{\prime}\cdot\mathbb{I}\left(\bm{w}^{\top}\bm{x}^{\prime}\geq 0\right)\right]. Next, we use the kernel formula to compute a feature map for a two-layer MLP with ReLU activation function. Recall that by definition a valid feature map must satisfy the following condition NTK ( 1 ) ( 𝒙 , 𝒙 ′ ) = ⟨ ϕ ( 𝒙 ) , ϕ ( 𝒙 ′ ) ⟩ \text{NTK}^{(1)}(\bm{x},\bm{x}^{\prime})=\bigl<\phi(\bm{x}),\phi(\bm{x}^{\prime})\bigr> It is easy to see that the way we represent our NTK formula makes it easy to find such a decomposition. The following infinite-dimensional feature map would satisfy the requirement because the inner product of ϕ ⁡ ( 𝒙 ) \phi(\bm{x}) and ϕ ⁡ ( 𝒙 ′ ) \phi(\bm{x}^{\prime}) for any 𝒙 \bm{x} , 𝒙 ′ \bm{x}^{\prime} would be equivalent to the expected value in NTK, after we integrate with respect to the density function of 𝒘 \bm{w} . ϕ ⁡ ( 𝒙 ) = c ′ ​ ( 𝒙 ⋅ 𝕀 ⁡ ( 𝒘 ( k ) ⊤ ​ 𝒙 ≥ 0 ) , 𝒘 ( k ) ⊤ ​ 𝒙 ⋅ 𝕀 ⁡ ( 𝒘 ( k ) ⊤ ​ 𝒙 ≥ 0 ) , … ) , \displaystyle\phi\left(\bm{x}\right)=c^{\prime}\left(\bm{x}\cdot\mathbb{I}\left(\bm{w}^{(k)^{\top}}\bm{x}\geq 0\right),\bm{w}^{(k)^{\top}}\bm{x}\cdot\mathbb{I}\left(\bm{w}^{(k)^{\top}}\bm{x}\geq 0\right),...\right), where 𝒘 ( k ) ∼ 𝒩 ⁡ ( 𝟎 , 𝑰 ) \bm{w}^{(k)}\sim\mathcal{N}(\bm{0},\bm{I}) , with k k going to infinity. c ′ c^{\prime} is a constant, and 𝕀 \mathbb{I} is the indicator function. Note that here the density of features of ϕ ⁡ ( 𝒙 ) \phi(\bm{x}) is determined by the density of 𝒘 \bm{w} , i.e. Gaussian.

## Appendix C Experimental Details

In this section, we describe the model, data and training details for reproducing our experiments. Our experiments support all of our theoretical claims and insights.

##### Overview.

We classify our experiments into the following major categories, each of which includes several ablation studies: 1) Learning tasks where the target functions are simple nonlinear functions in various dimensions and training/test distributions: quadratic, cosine, square root, and l1 norm functions, with MLPs with a wide range of hyper-parameters.