How Neural Networks Extrapolate: From Feedforward to Graph Neural Networks — part 7/11
Cleaned-text code points [91985, 107629). Bibliography omitted here; source.md includes it.
First there must exist 𝑺 0 = { 𝒔 i } i = 1 2 d ∈ 𝒮 \bm{S}_{0}=\{\bm{s}_{i}\}_{i=1}^{2d}\in\mathcal{S} of size 2 d 2d , e.g., orthogonal basis and their opposite vectors. Observe that if we scale any 𝒔 i \bm{s}_{i} by k > 0 k>0 , the resulting dataset is still in 𝒮 \mathcal{S} by the definition of 𝒮 \mathcal{S} . We denote the set of datasets where we are allowed to scale elements of 𝑺 0 \bm{S}_{0} by 𝒮 0 \mathcal{S}_{0} . It follows that ℙ ( 𝒟 ∗ ( 𝑿 2 d ) > ϵ ) \displaystyle\mathbb{P}\left(\mathcal{D}^{\ast}(\bm{X}_{2d})>\epsilon\right) = ℙ ( min 𝑺 ∈ 𝒮 d ( 𝑿 2 d , 𝑺 ) > ϵ ) \displaystyle=\mathbb{P}\left(\min\limits_{\bm{S}\in\mathcal{S}}d\left(\bm{X}_{2d},\bm{S}\right)>\epsilon\right) ≤ ℙ ( min 𝑺 ∈ 𝒮 0 d ( 𝑿 2 d , 𝑺 ) > ϵ ) \displaystyle\leq\mathbb{P}\left(\min\limits_{\bm{S}\in\mathcal{S}_{0}}d\left(\bm{X}_{2d},\bm{S}\right)>\epsilon\right) = ℙ ( min 𝑺 ∈ 𝒮 0 min ∑ i = 1 n σ ℓ ( σ ( 𝑿 2 d , 𝑺 ) i ) > ϵ ) \displaystyle=\mathbb{P}\left(\min\limits_{\bm{S}\in\mathcal{S}_{0}}\min\limits_{\sigma}\sum\limits_{i=1}^{n}\ell\left(\sigma\left(\bm{X}_{2d},\bm{S}\right)_{i}\right)>\epsilon\right) = 1 − ℙ ( min 𝑺 ∈ 𝒮 0 min ∑ i = 1 n σ ℓ ( σ ( 𝑿 2 d , 𝑺 ) i ) ≤ ϵ ) \displaystyle=1-\mathbb{P}\left(\min\limits_{\bm{S}\in\mathcal{S}_{0}}\min\limits_{\sigma}\sum\limits_{i=1}^{n}\ell\left(\sigma\left(\bm{X}_{2d},\bm{S}\right)_{i}\right)\leq\epsilon\right) ≤ 1 − ℙ ( min 𝑺 ∈ 𝒮 0 min σ max i = 1 n ℓ ( σ ( 𝑿 2 d , 𝑺 ) i ) ≤ ϵ ) \displaystyle\leq 1-\mathbb{P}\left(\min\limits_{\bm{S}\in\mathcal{S}_{0}}\min\limits_{\sigma}\max\limits_{i=1}^{n}\ell\left(\sigma\left(\bm{X}_{2d},\bm{S}\right)_{i}\right)\leq\epsilon\right) ≤ δ < 1 \displaystyle\leq\delta<1 where we denote the bound of ℙ ( 𝒟 ∗ ( 𝑿 2 d ) > ϵ ) \mathbb{P}\left(\mathcal{D}^{\ast}(\bm{X}_{2d})>\epsilon\right) by δ < 1 \delta<1 , and the last step follows from ℙ ( min 𝑺 ∈ 𝒮 0 min σ max i = 1 n ℓ ( σ ( 𝑿 2 d , 𝑺 ) i ) ≤ ϵ ) > 0 \displaystyle\mathbb{P}\left(\min\limits_{\bm{S}\in\mathcal{S}_{0}}\min\limits_{\sigma}\max\limits_{i=1}^{n}\ell\left(\sigma\left(\bm{X}_{2d},\bm{S}\right)_{i}\right)\leq\epsilon\right)>0 which further follows from the fact that for any 𝒔 i ∈ 𝒮 0 \bm{s}_{i}\in\mathcal{S}_{0} , by the assumption on training distribution, we can always find k > 0 k>0 so that k 𝒔 i ∈ 𝒮 ∗ k\bm{s}_{i}\in\mathcal{S}^{\ast} , a connected set in the support of training distribution. By the connectivity of support 𝒮 ∗ \mathcal{S}^{\ast} , k 𝒔 i k\bm{s}_{i} cannot be an isolated point in 𝒮 ∗ \mathcal{S}^{\ast} , so for any ϵ > 0 \epsilon>0 , we must have ∫ ‖ 𝒙 − k 𝒔 i ‖ ≤ ϵ , 𝒙 ∈ 𝒮 ∗ f 𝑿 ( 𝒙 ) 𝑑 𝒙 > 0 \displaystyle\int\limits_{\|\bm{x}-k\bm{s}_{i}\|\leq\epsilon,\bm{x}\in\mathcal{S}^{\ast}}f_{\bm{X}}(\bm{x})\mathrm{d}\bm{x}>0
Hence, we can now apply Eqn. 62 to bound 𝒟 ∗ ( 𝑿 n ) \mathcal{D}^{\ast}(\bm{X}_{n}) . Given any ϵ > 0 \epsilon>0 , we have ℙ ( 𝒟 ∗ ( 𝑿 n ) > ϵ ) \displaystyle\mathbb{P}\left(\mathcal{D}^{\ast}(\bm{X}_{n})>\epsilon\right) = 1 − ℙ ( 𝒟 ∗ ( 𝑿 n ) ≤ ϵ ) \displaystyle=1-\mathbb{P}\left(\mathcal{D}^{\ast}(\bm{X}_{n})\leq\epsilon\right) ≤ 1 − ℙ ( min k = 1 ⌊ n / 2 d ⌋ 𝒟 ∗ ( { 𝒙 j } j = ( k − 1 ) ∗ 2 d + 1 k ∗ 2 d ) ≤ ϵ ) \displaystyle\leq 1-\mathbb{P}\left(\min\limits_{k=1}^{\lfloor n/2d\rfloor}\mathcal{D}^{\ast}\left(\{\bm{x}_{j}\}_{j=(k-1)*2d+1}^{k*2d}\right)\leq\epsilon\right) ≤ 1 − ( 1 − ∏ k = 1 ⌊ n / 2 d ⌋ ℙ ( 𝒟 ∗ ( { 𝒙 j } j = ( k − 1 ) ∗ 2 d + 1 k ∗ 2 d ) > ϵ ) ) \displaystyle\leq 1-\left(1-\prod\limits_{k=1}^{\lfloor n/2d\rfloor}\mathbb{P}\left(\mathcal{D}^{\ast}\left(\{\bm{x}_{j}\}_{j=(k-1)*2d+1}^{k*2d}\right)>\epsilon\right)\right) = ∏ k = 1 ⌊ n / 2 d ⌋ ℙ ( 𝒟 ∗ ( { 𝒙 j } j = ( k − 1 ) ∗ 2 d + 1 k ∗ 2 d ) > ϵ ) \displaystyle=\prod\limits_{k=1}^{\lfloor n/2d\rfloor}\mathbb{P}\left(\mathcal{D}^{\ast}\left(\{\bm{x}_{j}\}_{j=(k-1)*2d+1}^{k*2d}\right)>\epsilon\right) ≤ δ ⌊ n / 2 d ⌋ \displaystyle\leq\delta^{\lfloor n/2d\rfloor} Here δ < 1 \delta<1 . This implies 𝒟 ∗ ( 𝑿 n ) ⟶ p 0 \mathcal{D}^{\ast}(\bm{X}_{n})\stackrel{{\scriptstyle p}}{{\longrightarrow}}0 , i.e., lim n → ∞ ℙ ( 𝒟 ∗ ( 𝑿 n ) > ϵ ) = 0 ∀ ϵ > 0 \displaystyle\lim\limits_{n\rightarrow\infty}\mathbb{P}\left(\mathcal{D}^{\ast}(\bm{X}_{n})>\epsilon\right)=0\;\;\;\forall\epsilon>0 (63) Eqn. 63 says as the number of training samples n → ∞ n\rightarrow\infty , our training set will converge in probability to a dataset that satisfies the requirement for exact extrapolation.
##### Asymptotic convergence of predictions.
Let NTK ( 𝒙 , 𝒙 ′ ) : ℝ d × ℝ d → ℝ \text{NTK}(\bm{x},\bm{x}^{\prime}):\mathbb{R}^{d}\times\mathbb{R}^{d}\rightarrow\mathbb{R} denote the neural tangent kernel for a two-layer ReLU MLP. It is easy to see that if 𝒙 → 𝒙 ∗ \bm{x}\rightarrow\bm{x}^{\ast} , then NTK ( 𝒙 , ⋅ ) → NTK ( 𝒙 ∗ , ⋅ ) \text{NTK}(\bm{x},\cdot)\rightarrow\text{NTK}(\bm{x}^{\ast},\cdot) ( Arora et al. (2019b) ). Let NTK train \text{NTK}_{\text{train}} denote the n × n n\times n kernel matrix for training data.
We have shown that our training set converges to a perfect data set that satisfies conditions of exact extrapolation. Moreover, note that our training set will only have a finite number of (not increase with n n ) 𝒙 i \bm{x}_{i} that are not precisely the same as those in a perfect dataset. This is because a perfect data only contains a finite number of key points and the other points can be replaced by any other points while still being a perfect data set. Thus, we have NTK train → N ∗ \text{NTK}_{\text{train}}\rightarrow N^{\ast} , where N ∗ N^{\ast} is the n × n n\times n NTK matrix for some perfect data set.
Because neural tangent kernel is positive definite, we have NTK train − 1 → N ∗ − 1 \text{NTK}_{\text{train}}^{-1}\rightarrow N^{\ast^{-1}} . Recall that for any 𝒙 ∈ ℝ d \bm{x}\in\mathbb{R}^{d} , the prediction of NTK is f NTK ( 𝒙 ) = ( NTK ( 𝒙 , 𝒙 1 ) , … , NTK ( 𝒙 , 𝒙 n ) ) ⋅ NTK train − 1 𝒀 , \displaystyle f_{\text{NTK}}(\bm{x})=(\text{NTK}(\bm{x},\bm{x}_{1}),...,\text{NTK}(\bm{x},\bm{x}_{n}))\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, NTK ( 𝒙 , 𝒙 i ) \text{NTK}(\bm{x},\bm{x}_{i}) is the kernel value between test data 𝒙 \bm{x} and training data 𝒙 i \bm{x}_{i} , and 𝒀 \bm{Y} is training labels.
Similarly, we have ( NTK ( 𝒙 , 𝒙 1 ) , … , NTK ( 𝒙 , 𝒙 n ) ) → ( NTK ( 𝒙 , 𝒙 1 ∗ ) , … , NTK ( 𝒙 , 𝒙 n ∗ ) ) (\text{NTK}(\bm{x},\bm{x}_{1}),...,\text{NTK}(\bm{x},\bm{x}_{n}))\rightarrow(\text{NTK}(\bm{x},\bm{x}^{\ast}_{1}),...,\text{NTK}(\bm{x},\bm{x}^{\ast}_{n})) , where x i ∗ x^{\ast}_{i} is a perfect data set that our training set converges to. Combining this with NTK train − 1 → N ∗ − 1 \text{NTK}_{\text{train}}^{-1}\rightarrow N^{\ast^{-1}} gives f NTK ⟶ p f NTK ∗ = g , \displaystyle f_{\text{NTK}}\stackrel{{\scriptstyle p}}{{\longrightarrow}}f^{\ast}_{\text{NTK}}=g, where f NTK f_{\text{NTK}} is the function learned using our training set, and f NTK ∗ f^{\ast}_{\text{NTK}} is that learned using a perfect data set, which is equal to the true underlying function g g . This completes the proof.
### B.4 Proof of Corollary 1
In order for GNN with linear aggregations h u ( k ) \displaystyle h_{u}^{(k)} = ∑ v ∈ 𝒩 ( u ) MLP ( k ) ( h u ( k ) , h v ( k ) , 𝒙 ( u , v ) ) , \displaystyle=\sum\limits_{v\in\mathcal{N}(u)}\text{MLP}^{(k)}\Big(h_{u}^{(k)},h_{v}^{(k)},\bm{x}_{(u,v)}\Big), h G \displaystyle h_{G} = MLP ( K + 1 ) ( ∑ u ∈ G h u ( K ) ) , \displaystyle=\text{MLP}^{(K+1)}\Big(\sum_{u\in G}h_{u}^{(K)}\Big), to extrapolate in the maximum degree task, it must be able to simulate the underlying function h G = max ∑ v ∈ 𝒩 ( u ) u ∈ G 1 h_{G}=\max_{u\in G}\sum_{v\in\mathcal{N}(u)}1 Because the max function cannot be decomposed as the composition of piece-wise linear functions, the MLP ( K + 1 ) \text{MLP}^{(K+1)} module in GNN must learn a function that is not piece-wise linear over domains outside the training data range. Since Theorem 1 proves for two-layer overparameterized MLPs, here we also assume MLP ( K + 1 ) \text{MLP}^{(K+1)} is a two-layer overparameterized MLP, although the result can be extended to more layers. It then follows from Theorem 1 that for any input and label (and thus gradient), MLP ( K + 1 ) \text{MLP}^{(K+1)} will converge to linear functions along directions from the origin. Hence, there are always domains where the GNN cannot learn a correct target function.
### B.5 Proof of Theorem 3
Our proof applies the similar proof techniques for Lemma 1 and 2 to Graph Neural Networks (GNNs). This is essentially an analysis of Graph Neural Tangent Kernel (GNTK), i.e., neural tangent kernel of GNNs.
We first define the simple GNN architecture we will be analyzing, and then present the GNTK for this architecture. Suppose G = ( V , E ) G=(V,E) is an input graph without edge feature, and 𝒙 u ∈ ℝ d \bm{x}_{u}\in\mathbb{R}^{d} is the node feature of any node u ∈ V u\in V . Let us consider the simple one-layer GNN whose input is G G and output is h G h_{G} h G = W ( 2 ) max ∑ v ∈ 𝒩 ( u ) u ∈ G W ( 1 ) 𝒙 v \displaystyle h_{G}=W^{(2)}\max\limits_{u\in G}\sum\limits_{v\in\mathcal{N}(u)}W^{(1)}\bm{x}_{v} (64) Note that our analysis can be extended to other variants of GNNs, e.g., with non-empty edge features, ReLU activation, different neighbor aggregation and graph-level pooling architectures. We analyze this GNN for simplicity of exposition.
Next, let us calculate the feature map of the neural tangent kernel for this GNN. Recall from Section A that consider a graph neural network f ( 𝜽 , G ) : 𝒢 → ℝ f(\bm{\theta},G):\mathcal{G}\rightarrow\mathbb{R} where 𝜽 ∈ ℝ m \bm{\theta}\in\mathbb{R}^{m} is the parameters in the network and G ∈ 𝒢 G\in\mathcal{G} is the input graph. Then the neural tangent kernel is 𝑯 i j = ⟨ ∂ f ( 𝜽 , G i ) ∂ 𝜽 , ∂ f ( 𝜽 , G j ) ∂ 𝜽 ⟩ , \displaystyle\bm{H}_{ij}=\biggl<\frac{\partial f(\bm{\theta},G_{i})}{\partial\bm{\theta}},\frac{\partial f(\bm{\theta},G_{j})}{\partial\bm{\theta}}\biggr>, where 𝜽 \bm{\theta} are the infinite-dimensional parameters. Hence, the gradients with respect to all parameters give a natural feature map. Let us denote, for any node u u , the degree of u u by 𝒉 u = ∑ v ∈ 𝒩 ( u ) 𝒙 v \displaystyle\bm{h}_{u}=\sum_{v\in\mathcal{N}(u)}\bm{x}_{v} (65) It then follows from simple computation of derivative that the following is a feature map of the GNTK for Eqn. 64 ϕ ( G ) = c ⋅ ( max u ∈ G ( 𝒘 ( k ) ⊤ 𝒉 u ) , ∑ u ∈ G 𝕀 ( u = arg max v ∈ G 𝒘 ( k ) ⊤ 𝒉 v ) ⋅ 𝒉 u , … ) , \displaystyle\phi(G)=c\cdot\left(\max\limits_{u\in G}\left(\bm{w}^{(k)^{\top}}\bm{h}_{u}\right),\sum\limits_{u\in G}\mathbb{I}\left(u=\arg\max\limits_{v\in G}\bm{w}^{(k)^{\top}}\bm{h}_{v}\right)\cdot\bm{h}_{u},...\right), (66) where 𝒘 ( k ) ∼ 𝒩 ( 𝟎 , 𝑰 ) \bm{w}^{(k)}\sim\mathcal{N}(\bm{0},\bm{I}) , with k k going to infinity. c c is a constant, and 𝕀 \mathbb{I} is the indicator function.
Next, given training data { ( G i , y i } i = 1 n \{(G_{i},y_{i}\}_{i=1}^{n} , let us analyze the function learned by GNN through the min-norm solution in the GNTK feature space. The same proof technique is also used in Lemma 1 and 2 .
Recall the assumption that all graphs have uniform node feature, i.e., the learning task only considers graph structure, but not node feature. We assume 𝒙 v = 1 \bm{x}_{v}=1 without loss of generality. Observe that in this case, there are two directions, positive or negative, for one-dimensional Gaussian distribution. Hence, we can simplify our analysis by combining the effect of linear coefficients for 𝒘 \bm{w} in the same direction as in Lemma 1 and 2 .
Similarly, for any 𝒘 \bm{w} , let us define 𝜷 ^ 𝒘 ∈ ℝ \hat{\bm{\beta}}_{\bm{w}}\in\mathbb{R} as the linear coefficient corresponding to ∑ u ∈ G 𝕀 ( u = arg max v ∈ G 𝒘 ⊤ 𝒉 v ) ⋅ 𝒉 u \sum\limits_{u\in G}\mathbb{I}\left(u=\arg\max\limits_{v\in G}\bm{w}^{\top}\bm{h}_{v}\right)\cdot\bm{h}_{u} in RKHS space, and denote by 𝜷 ^ 𝒘 ′ ∈ ℝ \hat{\bm{\beta}}^{\prime}_{\bm{w}}\in\mathbb{R} the weight for max u ∈ G ( 𝒘 ⊤ 𝒉 u ) \max\limits_{u\in G}\left(\bm{w}^{\top}\bm{h}_{u}\right) . Similarly, we can combine the effect of all 𝜷 ^ \hat{\bm{\beta}} in the same direction as in Lemma 1 and 2 . We define the combined effect with 𝜷 𝒘 \bm{\beta}_{\bm{w}} and 𝜷 𝒘 ′ \bm{\beta}_{\bm{w}}^{\prime} . This allows us to reason about 𝒘 \bm{w} with two directions, + + and − - .
Recall that the underlying reasoning function, maximum degree, is g ( G ) = max u ∈ G 𝒉 u . \displaystyle g(G)=\max_{u\in G}\bm{h}_{u}.
We formulate the constrained optimization problem, i.e., min-norm solution in GNTK feature space that fits all training data, as min 𝜷 ^ , 𝜷 ^ ′ ∫ 𝜷 ^ 𝒘 2 + 𝜷 ^ 𝒘 ′ 2 d ℙ ( 𝒘 ) \displaystyle\min\limits_{\hat{\bm{\beta}},\hat{\bm{\beta}}^{\prime}}\int\hat{\bm{\beta}}_{\bm{w}}^{2}+\hat{\bm{\beta}}_{\bm{w}}^{\prime^{2}}\mathrm{d}\mathbb{P}(\bm{w}) s . t . \displaystyle s.t.\; ∫ ∑ u ∈ G i 𝕀 ( u = arg max v ∈ G 𝒘 ⋅ 𝒉 v ) ⋅ 𝜷 ^ 𝒘 ⋅ 𝒉 u + max u ∈ G i ( 𝒘 ⋅ 𝒉 u ) ⋅ 𝜷 ^ 𝒘 ′ 𝑑 ℙ ( 𝒘 ) = max u ∈ G i 𝒉 u ∀ i ∈ [ n ] , \displaystyle\int\sum\limits_{u\in G_{i}}\mathbb{I}\left(u=\arg\max\limits_{v\in G}\bm{w}\cdot\bm{h}_{v}\right)\cdot\hat{\bm{\beta}}_{\bm{w}}\cdot\bm{h}_{u}+\max\limits_{u\in G_{i}}\left(\bm{w}\cdot\bm{h}_{u}\right)\cdot\hat{\bm{\beta}}_{\bm{w}}^{\prime}\mathrm{d}\mathbb{P}(\bm{w})=\max\limits_{u\in G_{i}}\bm{h}_{u}\quad\forall i\in[n], where G i G_{i} is the i-th training graph and 𝒘 ∼ 𝒩 ( 0 , 1 ) \bm{w}\sim\mathcal{N}(0,1) . By combining the effect of 𝜷 ^ \hat{\bm{\beta}} , and taking the derivative of the Lagrange for the constrained optimization problem and setting to zero, we get the global optimum solution satisfy the following constraints. 𝜷 + \displaystyle\bm{\beta}_{+} = c ⋅ ∑ i = 1 n λ i ⋅ ∑ u ∈ G i 𝒉 u ⋅ 𝕀 ( u = arg max v ∈ G i 𝒉 v ) \displaystyle=c\cdot\sum\limits_{i=1}^{n}\lambda_{i}\cdot\sum_{u\in G_{i}}\bm{h}_{u}\cdot\mathbb{I}\left(u=\arg\max\limits_{v\in G_{i}}\bm{h}_{v}\right) (67) 𝜷 − \displaystyle\bm{\beta}_{-} = c ⋅ ∑ i = 1 n λ i ⋅ ∑ u ∈ G i 𝒉 u ⋅ 𝕀 ( u = arg min v ∈ G i 𝒉 v ) \displaystyle=c\cdot\sum\limits_{i=1}^{n}\lambda_{i}\cdot\sum_{u\in G_{i}}\bm{h}_{u}\cdot\mathbb{I}\left(u=\arg\min\limits_{v\in G_{i}}\bm{h}_{v}\right) (68) 𝜷 + ′ \displaystyle\bm{\beta}_{+}^{\prime} = c ⋅ ∑ i = 1 n λ i ⋅ max u ∈ G i 𝒉 u \displaystyle=c\cdot\sum_{i=1}^{n}\lambda_{i}\cdot\max\limits_{u\in G_{i}}\bm{h}_{u} (69) 𝜷 − ′ \displaystyle\bm{\beta}_{-}^{\prime} = c ⋅ ∑ i = 1 n λ i ⋅ min u ∈ G i 𝒉 u \displaystyle=c\cdot\sum_{i=1}^{n}\lambda_{i}\cdot\min\limits_{u\in G_{i}}\bm{h}_{u} (70) max u ∈ G i 𝒉 u \displaystyle\max\limits_{u\in G_{i}}\bm{h}_{u} = 𝜷 + ⋅ ∑ u ∈ G i 𝕀 ( u = arg max v ∈ G i 𝒉 v ) ⋅ 𝒉 u + 𝜷 + ′ ⋅ max u ∈ G i 𝒉 u \displaystyle=\bm{\beta}_{+}\cdot\sum\limits_{u\in G_{i}}\mathbb{I}\left(u=\arg\max\limits_{v\in G_{i}}\bm{h}_{v}\right)\cdot\bm{h}_{u}+\bm{\beta}_{+}^{\prime}\cdot\max\limits_{u\in G_{i}}\bm{h}_{u} (71) + 𝜷 − ⋅ ∑ u ∈ G i 𝕀 ( u = arg min v ∈ G i 𝒉 v ) ⋅ 𝒉 u + 𝜷 − ′ ⋅ min u ∈ G i 𝒉 u ∀ i ∈ [ n ] \displaystyle+\bm{\beta}_{-}\cdot\sum\limits_{u\in G_{i}}\mathbb{I}\left(u=\arg\min\limits_{v\in G_{i}}\bm{h}_{v}\right)\cdot\bm{h}_{u}+\bm{\beta}_{-}^{\prime}\cdot\min\limits_{u\in G_{i}}\bm{h}_{u}\quad\forall i\in[n] (72) where c c is some constant, λ i \lambda_{i} are the Lagrange parameters. Note that here we used the fact that there are two directions + 1 +1 and − 1 -1 . This enables the simplification of Lagrange derivative. For a similar step-by-step derivation of Lagrange, refer to the proof of Lemma 1 .