When narrower is better: the narrow width limit of Bayesian parallel branching neural networks

Thank you for your thoughtful feedback and critiques! We want to say first that although we feel that our work is a stand-alone story on GNNs, the results on the narrow width limit can be generalized to other architectures provided there are independent branches and trained in an offline fashion (more details below).

We are currently preparing new results on the narrow width limit for a 2-layer residual MLP architecture but please let us know at your earliest convenience if there are further concerns regarding our response below concurrently to best utilize the interactive period.

Responses to weakness:

1. Although the paper studies graph neural networks, our results on the narrow width limit can be generalized to other architectures with independent branches. We define a more general branching architecture with one hidden layer as follows: $y = \sum_{l=0}^{L-1} y_l = \sum_{l=0}^{L-1} a_l \phi_l(W_l x)$, where the final readout is a sum of $L$ branch readouts $y_l$'s and $W_l$'s are the independent hidden layer weights and $\phi_l$'s are different activation functions or convolutions. To further clarify the connection between the branching networks and common architectures we provide a dictionary below:

--GNN: $\phi_l(X W_l ) = A^l X W_l$. This is when the different branches represent different number of convolutions on nodes as in our BPB-GNN setup;

--CNN: $\phi_l$'s are different patches with the convolution filter;

--Transformers: $\phi_l(X W_l ) = A_l X W_l$ for a linear one hidden-layer attention network, with $A_l$'s now representing the $L$ attention heads;

--Residual-MLP: $\phi_l$'s can be different activation functions. eg. when $\phi_0$ is the identity function and $\phi_1$ is ReLu, this models the 2-layer residual MLP with one residual block.

The key insight from our theory is that the overall kernel undergoes a kernel renormalization with different $u_l$'s, ie. $K = \sum_l \frac{1}{L}u_l K_l$ as in Eq. 11 and all results regarding the narrow width limit follow from this provided the GP kernels $K_l$'s are sufficiently different for different branches $\phi_l$. Thus our result is general enough for the architectures mentioned above. To be scientifically honest, we started deriving the theory for branching-GNNs and this is the simplest test bed for the narrow width limit since the theory is exact for linear networks and the branches can be simply chosen to be different number of convolutions $A_l$ to resemble the residual-GNN architecture. Since this is the first paper that demonstrates such narrow width effect, the choice of architectures is not comprehensive; however, we will provide results on the residual-MLP shortly. Further work can be done for this dictionary.

2. Thank you for the catch! We redefined the bracket notation in the revised draft (see Eq 13 and Appendix A.1 on notations of statistical theory). Please clarify if there are other instances where we lack rigor.

3. As we mentioned in Appendix B.3, we used the standard Hamiltonian Monte Carlo package Numpyro to sample from the posterior distribution. Since we did not engineeringly optimize the sampling strategy, the experiment is limited by the number of nodes and width in the dataset. Citeseer and Pubmed ended to be too big for us to sample properly. However, as we mentioned before, we are preparing results on residual-MLP that also demonstrates the narrow width effect.

4. The PAC-Bayes approach relies on computing the norms of learned weights and does not have a decomposition of the generalization error into bias and variance. We revised the statement and provided more discussion on the relation to PAC-Bayes approach. Please see the related work section again.

5. We included this missing reference in the related works, thank you!

We added a new section (Appendix B) on the residual-MLP architecture which demonstrates the narrow width effect is more general. Please see our revised draft and see if it addresses your concerns. Thank you!

We thank the reviewer for questions and critiques of the paper. Weakness:

1. Although the paper considers the linear model, it is the first one to our knowledge that considers renormalization of kernels for GNN, which is already non-trivial and shows the robust learning of branches phenomenon we presented in the paper. The traditional line of work for NTK and NNGP considers the infinite width limit where the width $N\to \infty$ but sample size $P$ stays finite, which is discussed in the related work section. We also added a new section (appendix B) on residual MLPs which uses ReLu activation for one branch and linear activation for the other branch to show that the results are generalizable.

2. Thank you for pointing out the conjecture! We reformulated and proved it as a theorem in the main text and also prove a stronger version of it in appendix B. Please see the revised draft.

3. We try to organize the results such that each derivation in the main text conveys essential information, as the kernel renormalization theory is rather uncommon for the reader. Could you kindly point out if there is any result that is non-rigorous? We also cross-check our theory with experiments for validity.

1. $A^l$ is the $l$th power of $A$, which we make clearer in the revised draft. Thank you for the catch!
2. We added more details on the derivation in the appendix. In particular Eq. 26 uses the Gaussian trick to transform the original partition function to one that is linear in the readout weight $a_l$'s. The essential part of the derivation is Eq. 30 by integrating out the hidden layer weight $W_l$'s exactly. This actually also addresses your previous concern regarding the problem of linear networks: we showed that it is already non-trivial to perform this integration and it is not tractable if the activation function is non-linear (which we briefly discuss in the newly added appendix B). Please let us know if there are additional concerns!
3. We intentionally wrote in this form as the two terms with the kernel also scales with $P$. Therefore, the energy term scales with $\alpha N$ and the entropy term with $N$. As $\alpha \to 0$, the energy term can be ignored and thus the RHS of Eq. 12 is 0. We wrote more details in the appendix.

Thanks for the authors' response. I have some follow-up questions:

• Can you be more specific about what Gaussian trick you used and how you derived the current Eq. 22? Do the calculations only hold for random Gaussian parameters?
• Are the parameters $\Theta$ randomly initialized as Gaussian? This should be made clear in the paper. Does the theory hold for training or just for random initialization?
• "The two terms with the kernel also scale with $P$. Therefore, the energy term scales with $\alpha$ and the entropy term with $N$". This should be included in the derivation/explanation and made rigorous. The first term of kernel $Y^\top (K+TI)^{-1}Y \leq \frac{||Y||^2}{\lambda_0(K) + T}$ actually depends on the order of $\lambda_0(K)$, where $\lambda_0(K)$ is the smallest eigenvalue. If $\lambda_0(K)=O(1)$, then this term is $O(P)$. But if $\lambda_0(K)=O(P)$, then this term is $O(1)$ and does not scale with $P$. The order of $\lambda_0(K)$ depends on the distribution of the data and the weight matrices.

We thank the reviewer for further questions. We want to say that your feedback on the mathematical derivations and the equipartition conjecture really helped us to improve the revised draft, now with the main result as a theorem and with more details on partition function calculations. Thank you!

Regarding the questions:

1. By the Gaussian trick, we mean the identity $\int d^P t \exp(-\frac{1}{2}t^T A t-it^T x) = \exp(-\frac{1}{2}x^T A^{-1}x -\frac{1}{2}\log\det A)$, which is a Fourier transform representation of the multivariate Gaussian in terms of the $P$ dimensional vector $x^{\mu}$. In deriving Eq. 22, we use $x^{\mu}=\frac{1}{\sqrt{NL}}\sum_{l,i}a_{i,l}h_{i,l}^{\mu}-Y^{\mu}$ and $A=TI$ to represent the univariate Gaussian in each dimension the original partition function, and then insert $t^{\mu}$'s as the Fourier transform auxiliary variables. We use the same procedure in Eq. 24-26 for introducing $t^{\mu}$'s and Eq.34 in integrating out $t^{\mu}$'s.

Since the integration relies on the Fourier transform of Gaussians, they only hold for random Gaussian priors. However, we think there might be some confusion regarding our framework so we will explain in the following points.

2. Our setup is a Bayesian problem where the posterior distribution is given by Eq. 4. The parameters $\Theta$ has a random Gaussian prior as well as the likelihood term that depends on the squared loss function as detailed in our paper explaining Eq. 4. So strictly speaking, there is no "training" of the network, instead the parameters are drawn from the posterior distribution in both the theory and HMC sampling. One perspective of this paper is simply to take the Bayesian setup for granted, which is already an interesting study as many works exist in simply studying Bayesian networks [1][2][3][4].

However, as we said in the paper, this posterior distribution is the Boltzmann equilibrium distribution of the Langevin dynamics. More specifically, the gradient steps of the Langevin dynamics \begin{aligned} \Delta w &= -\eta (\nabla_{w}L+\gamma w) + \sqrt{2T\eta}\xi \ &= -\eta \nabla_{w}(L+\frac{\gamma}{2} |w|^2) + \sqrt{2T\eta}\xi \end{aligned} where $w$ is the concatenated weight vector combining all parameters in the network, $L$ is the squared loss, $\gamma$ is the decay rate, $\xi_i \sim \mathcal{N}(0,1)$ is unit white noise and $T$ is the temperature that represents the strength of stochastic noise. Taking $\gamma = \frac{T}{\sigma^2}$, the gradient updates converge exactly to the posterior distribution Eq. 4 in the equilibrium. Therefore, you could view the Bayesian distribution as the equilibrium of networks trained with Langevin dynamics, with weight decay that gives rise to the prior term. So regardless of initialization, such Langevin dynamics always converge to the equilibrium distribution. Furthermore, at near 0 temperature, the Langevin dynamics is simply GD with proper weight decay. There are works that show that early stopping is effectively a $L_2$ regularization, so that without the weight decay the network still exhibits similar behavior provided they are initialized with random Gaussian with variance $\sigma^2$ [5][6].

[1] https://arxiv.org/abs/2111.00034 [2] https://www.nature.com/articles/s41467-021-23103-1 [3] https://arxiv.org/abs/1711.00165 [4] https://www.pnas.org/doi/abs/10.1073/pnas.2301345120 [5] https://journals.aps.org/prx/abstract/10.1103/PhysRevX.11.031059 [6] https://www.sciencedirect.com/science/article/pii/S0893608020303117?via%3Dihub

3. Regarding the GP limit of Eq. 10 and 12, since the hidden layer weights $W$'s are all integrated out, the kernel $K=\sum_l \frac{1}{L}u_l K_l$ does not depend on the weight matrices but simply is a sum of the NNGP kernels $K_l$'s that are the input kernels with different powers of node convolutions. By Mercer's theorem, $K(x,x') = \sum_{i}\lambda_i e_i(x) e_i(x')$ can be written as a eigen decomposition in terms of the eigenfunctions $e_i(x)$, with eigenvalues $\lambda_i$.Therefore, as $P \to \infty$, this the largest eigenvalue of the Gram matrix $K^{\mu,\nu}$ converges to the largest eigenvalue in the Mercer's decomposition $\lambda_0$, which is of $O(1)$. Thus the term $Y^T (K+TI)^{-1}Y$ is of order $O(P)$ from the $\|Y\|^2$ term. The GP limit is also verified with the student-teacher experiments in the paper.

We really appreciate the reviewer's efforts in delving into the mathematical details and we will add the above points to the revised draft to make the mathematical derivations more transparent. Since it is past the revision deadline, we will make the final draft with these points in mind.

Please let us know if there are further questions/concerns and we are happy to provide more details!

Thanks for the authors's response.

• For the first point, these are necessary details to understand the paper and the proofs. Please include them in the revised paper. It is not mentioned anywhere in the submission that Fourier transformation is used. The $i$ in Eq 22 is confusing as it repeats with the summation index.

• Thanks for the authors' explanation. So the results hold for Langevin dynamics/GD trained networks in some sense if the equilibrium/equivalence holds.

• Regarding the order of $\lambda_0$, the existence of Mercer's decomposition and convergence to it does not mean the smallest eigenvalue $\lambda_0$ is $O(1)$. For example, for a simple linear kernel $<X, X>$, when the entries of $X \in \mathbb{R}^{N_0\times P}$ is independent standard Gaussian, then the singular value of $X$ is of order $\sqrt{P} - \sqrt{N_0}$ (suppose $P > N_0$) and $\lambda_0 = O(P)$ [1]. But I think it might not matter here because the second term in Eq 10 is $\sum_i^P \log(\lambda_i + T)$ and is at least order $O(P)$ if $T >0$.

I'm raising my score to 6. But I hope the authors can make the proof more rigorous and the paper more readable.

[1] Vershynin, Roman. "Introduction to the non-asymptotic analysis of random matrices." arXiv preprint arXiv:1011.3027 (2010).

Regarding the last question, in our case the $P<N_0$ so the kernel $K$ is full-rank. For a linear kernel with $X$ as Gaussian matrix, $K$ is the Wishart matrix and thus $Y^T K^{-1}Y = P (Y^T/P) (K/P)^{-1}(Y/P)$, where the eigenvalues of $K/P$ tends to the Marchenko-Pastur distribution with degree $\alpha_0 = P/N_0$ as $P,N_0 \to \infty$. Therefore, the term is still of $O(P)$ as the eigenspectrum is of $O(1)$. However, you are correct that the eigenspectrum of $K$ still depends on input statistics, but we expect this scaling still holds when we write the term as above.

We will ensure the points above and your other concerns are elaborated on in the final draft. Thank you for the helpful feedback and critiques!

Thank you for your thoughtful review and efforts in going through our paper! We want to say first that your intuition is totally correct and our results on the narrow width limit can be generalized to other architectures provided there are independent branches and trained in an offline fashion (more details below).

We are currently preparing new results on the narrow width limit for a 2-layer residual MLP architecture but please let us know at your earliest convenience if there are further concerns regarding our response below concurrently to best utilize the interactive period.

Responses to weakness:

1. Although the paper studies graph neural networks, our results on the narrow width limit can be generalized to other architectures with independent branches. We define a more general branching architecture with one hidden layer as follows: $y = \sum_{l=0}^{L-1} y_l = \sum_{l=0}^{L-1} a_l \phi_l(W_l x)$, where the final readout is a sum of $L$ branch readouts $y_l$'s and $W_l$'s are the independent hidden layer weights and $\phi_l$'s are different activation functions or convolutions. To further clarify the connection between the branching networks and common architectures we provide a dictionary below:

--GNN: $\phi_l(X W_l ) = A^l X W_l$. This is when the different branches represent different number of convolutions on nodes as in our BPB-GNN setup;

--CNN: $\phi_l$'s are different patches with the convolution filter;

--Transformers: $\phi_l(X W_l ) = A_l X W_l$ for a linear one hidden-layer attention network, with $A_l$'s now representing the $L$ attention heads;

--Residual-MLP: $\phi_l$'s can be different activation functions. eg. when $\phi_0$ is the identity function and $\phi_1$ is ReLu, this models the 2-layer residual MLP with one residual block.

The key insight from our theory is that the overall kernel undergoes a kernel renormalization with different $u_l$'s, ie. $K = \sum_l \frac{1}{L}u_l K_l$ as in Eq. 11 and all results regarding the narrow width limit follow from this provided the GP kernels $K_l$'s are sufficiently different for different branches $\phi_l$. Thus our result is general enough for the architectures mentioned above. To be scientifically honest, we started deriving the theory for branching-GNNs and this is the simplest test bed for the narrow width limit since the theory is exact for linear networks and the branches can be simply chosen to be different number of convolutions $A_l$ to resemble the residual-GNN architecture. Since this is the first paper that demonstrates such narrow width effect, the choice of architectures is not comprehensive; however, we will provide results on the residual-MLP shortly. Further work can be done by studying the architectures mentioned in the dictionary.

Finally, to address your point regarding the distinction of $L$ and $N$, the $L$ corresponds to the number of residual branches in residual-MLP, the number of different convolution branches in GNN, the number of patches in CNN and the number of heads in transformers. Therefore, our results apply to the architectures that have different independent "branches", which is not a bad approximation of real-world architectures that have residual connections or different schemes of convolutions. However, in the case of a valinna MLP that lacks residual connections, our results do not apply. This is thus the significance of our result that says something special about those branching-like networks.

2. As we wrote in the paper, the narrow width limit result only shows that the bias term decreases at narrower width; however there is a trade-off to the increase of variance at narrow width.Practically speaking, since the bias term corresponds to the prediction of the network averaged over random initialization, one method is to use an ensemble of networks that averages over different random seeds to arrive at our results, where the variance term tend to 0 for the ensemble of networks. Your intuition is very interesting regarding the task complexity and regularization, this is also our best guess so far, ie. in the overparametrized regime, narrow width helps with regularization for the average behavior and this is only true in the so-called "lazy" learning regime when the task is simple and in an offline fashion. It might be true that most real-world tasks are in the variance-dominated regime and narrow width can only be beneficial in the ensembled networks. Intriguingly, as shown in Figure 8f of [1], the ensembled network does have a narrow width effect on Resnet-18 trained offline and our results provide a convincing explanation for this!

3. We intentionally chose the same color for the same branches to illustrate the point that the teacher and student readout norms for the same branch coincides at narrow width, both theoretically and experiementally. Thank you for pointing out the readability issue which will be fixed in revision.

[1] https://arxiv.org/pdf/2305.18411

We added a new section (Appendix B) on the residual-MLP architecture which demonstrates the narrow width effect is more general. In addition, we added more discussion regarding your concerns on the variance-dominated scenarios and regularization of the narrow width in the discussion section. Please see our revised draft and see if it addresses your concerns. Thank you!

Thank you very much for your thoughtful reviews and concerns! We are currently in the process of revising the draft, with a new section on the narrow width limit on residual-MLP networks. In the meantime, please let us know if you have further concerns regarding our response below.

1. Thank you for pointing out the additional references! Indeed [1] studies the generalization performance across width; it is however mostly concerned with the muP regime which is the feature learning regime with online learning and our work focuses on the so called "lazy" regime in offline learning. Intriguingly, as shown in Figure 8f of [1], the ensembled network does have a narrow width effect on Resnet-18 trained offline and our results provide a convincing explanation for this (ie. the ensemble of network averaged over random initialization is exactly the bias term in our work). Indeed our bayesian framework is analogous to [4], which uses Bayesian GNN for node classifications; whereas [4] focuses on practical scaling up Bayesian-GNN on datasets with large number of nodes, our work emphasizes the theoretical results of the narrow width limit. [2] uses mutual information of representations for bounding the generalization, although relevant, the exact relation to our work is to be explored. Correct me if wrong, but [3] mainly develops an algorithm to divide up the graph for more efficient training and it is hard to see the relation to our current work.

2. In terms of the trade-off between expressivity and learnibility/generalization, indeed we hypothesize that in the highly overparametrized case, a narrower width serves as a regularization to make the network generalize better on unseen nodes. The narrow network lacks expressive power to compute certain algorithms as mentioned in [5]. We are not familiar with the TM perspective, but it might be related to the stable TM [6] for better learnibility. The work really stems from theory of infinitely wide networks and we think that the exact relationship to other perspectives can be further explored in future works.

We will mention all the points above related to [1]-[6] in the revised draft shortly and we hope it will answer your concerns.

3. Regarding your question on details of the experiment, we provided experiment details in section B of the paper. We used Hamiltonian Monte Carlo with a warmup period to sample the posterior distribution Eq.4, which is different from the usual SGD training. The Bayesian network at 0 temperature roughly approximates an ensemble of networks trained with GD with different random inializations, as we mentioned in the paper. All details on hyperparameters (hidden layer width $N$,regularization noise $\sigma_w$, number of branches $L$) are specified in our figures (they are key variables for our results, which says how generalization error depends on the hyperparameters). Specifically, the dimension of weight $W_l$ is $N \times N_0$, and the number of readout neurons $a_l$ is exactly the hidden layer width $N$, which we think are clearly stated in the paper. The depth of the network is fixed to be 1 hidden layer, as we mentioned in the beginning of the model setup. As the training is different from the usual SGD (it is really a statistical MCMC sampling from the posterior distribution), we do not provide the p-values etc as it already represents predictions from the ensemble of trained networks (if you will, the variance term is exactly this variation of networks trained with different initialization). Hope this clarifies!

4. Our results are true for both homophilous and heterophilous graphs, since if the kernels of different number of convolutions can be sufficiently different in both cases. Therefore the narrow width limit still holds. We chose homophilous CSBM model in our paper.

We added a new section (Appendix B) on the residual-MLP architecture which demonstrates the narrow width effect is more general. We also included your provided sources in the discussion section. Please see our revised draft. Thank you!