When narrower is better: the narrow width limit of Bayesian parallel branching neural networks
Authors: Zechen Zhang
Haim Sompolinsky
Venue: ICLR 2025 Poster (2025)
Keywords: Bayesian Networks
Gaussian Process
Kernel Renormalization
Graph Neural Networks
Residual Network
Theory of Generalization
Primary area: learning theory
TL;DR: We use statistical learning theory to understand parallel graph neural networks, and find a narrow width limit where networks generalize better with small width.

Abstract: The infinite width limit of random neural networks is known to result in Neural Networks as Gaussian Process (NNGP) (Lee et al. (2018)), characterized by task-independent kernels. It is widely accepted that larger network widths contribute to improved generalization (Park et al. (2019)). However, this work challenges this notion by investigating the narrow width limit of the Bayesian Parallel Branching
Neural Network (BPB-NN), an architecture that resembles neural networks with residual blocks. We demonstrate that when the width of a BPB-NN is significantly smaller compared to the number of training examples, each branch exhibits more robust learning due to a symmetry breaking of branches in kernel renormalization. Surprisingly, the performance of a BPB-NN in the narrow width limit is generally superior to or comparable to that achieved in the wide width limit in bias-limited scenarios. Furthermore, the readout norms of each branch in the narrow width limit are mostly independent of the architectural hyperparameters but generally reflective of the nature of the data. We demonstrate such phenomenon primarily in the branching graph neural networks, where each branch represents a different order of convolutions of the graph; we also extend the results to other more general architectures such as the residual-MLP and demonstrate that the narrow width effect is a general feature of the branching networks. Our results characterize a newly defined narrow-width regime for parallel branching networks in general.

Published as a conference paper at ICLR 2025
WHEN NARROWER IS BETTER: THE NARROW WIDTH LIMIT OF BAYESIAN PARALLEL BRANCHING NEURAL NETWORKS
Zechen Zhang
Haim Sompolinsky
Department of Physics, Harvard University Center for Brain Science, Harvard University zechen.zhanger harvard. edu
Center for Brain Science, Harvard University
Kempner Institute
Edmond and Lily Safra Center for Brain Sciences,
Hebrew University of Jerusalem
harvard edu
ABSTRACT
The infinite width limit of random neural networks is known to result in Neural Networks Gaussian Process (NNGP) (Lee al. (2018)), characterized by task- independent kernels. It widely accepted that larger network widths contribute to improved generalization (Park al. (2119). However, this work challenges this notion by investigating the narrow width limit of the Bayesian Parallel Branching Neural Network (BPB-NN), an architecture that resembles neural networks with residual blocks. We demonstrate that when the width of a BPB-NN is significantly smaller compared to the number of training examples, each branch exhibits robust learning due to symmetry breaking of branches in kernel renormalization. Surprisingly, the performance superior to or comparable to that achieved in the wide width limit in bias-limited scenarios. Furthermore, the readout norms of each branch in the narrow width limit are mostly independent of the architectural hyperparameters but generally reflective of nature of the data. We demonstrate such phenomenon primarily for branching graph neural networks, where each branch represents a different order of convolutions of the graph; we also extend the results to other more general architectures such the residual-MLP and show that the narrow width effect is a general feature of the branching networks. Our results characterize a newly defined narthw-width regime for parallel branching networks in general.
BPB-NN the narrow width limit is generally
INTRODUCTION
The study of neural network architectures has seen substantial growth, particularly in understanding how network width impacts learning and generalization. In general, wider networks are believed to perform better (Allen-Zhu et al. (2019;; Jacot et al. (2018); Gao al.] (2024)). However, this work challenges the prevailing assumption by exploring the narrow width limit of Bayesian Paral- lel Branching Neural Network (BPB-NN), an architecture inspired by neural networks with resid- ual blocks (He et al. (2016); Chen al. (2020a; 2022)). We show thcoretically and empirically that narrow-width BPB-NNs can perform better than their wider counterparts due to the symmetry- breaking effect in kernel renormalization, in bias- limited scenarios. We present detailed analysis of BPB-NNs primarily by exploring the Bayesian Parallel Branching Graph Convolution Networks (BPB-GCN), while extending the results to other more general architectures such the residual- MLP in the appendix.
Contributions
1. We introduce a novel yet simple graph neural architecture with parallel independent branches, and derive the exact generalization error for node regression in the statistical limit as the sample size P - oo and network width N - oo, with their ratio a finite number P/N, in the over-parameterized regime.

---

Published as a conference paper at ICLR 2025
2. We show that in the Bayesian setting the bias will decrease and saturate at narrow hidden layer width, a surprising phenomenon due to kernel renormalization. We demonstrate that this can be understood where each student's 's branch is learning the teacher's branch.
robust learning effect of each branch in the student-teacher task,
3. We demonstrate this narrow-width limit in the fal-world dataset Cora and understand each branch's 's importance as a nature of the dataset.
4. We further show that this narrow-width effect is a general feature of Bayesian parallel branching neural networks (BPB-NNs), with the residual-MI.P architecture as example.
2 RELATED WORKS
Infinitely wide neural networks: Our work follows a long tradition of mathematical analysis of infinitely-wide neural networks (Neal (2012); Jacot et al. (2018); Lee et al. (2018); Bahri et al.
22024), ), resulting in NTK or NNGP kernels. Recently, such analysis has been extended to structured
neural networks, including GCNs (Du al. (2019) Walker Glocker 2019); (Hung et al. (2021), ) However, they do not provide an analysis of feature learning in which the kernel depends on the tasks.
Kernel renormalization and feature learning: There has been progress in understanding simple
MLPs in the feature-learning regime as the shape of the kernel changes with task or time (Li & Sompolinsky (202T); Atanasov et al. (202I); Avidan et al. (2023); Waang & Jacot (2023)). We develop such understanding in waseh networks.
Theoretical analysis of GCN:
There is long line of works that theoretically analyze the ex-
pressiveness (uu et al. (2018):
; Geerts & Reutter (2022) and generalization performance (Tang &
Liu (2023): Garg al. 2020) Aminian et al. (2024)) of GCN. However. is challenging calcu-
late the dependence of generalization eтrors on tasks. In particular, the PAC-Bayes approach Liao ct al. (2020); Ju et al. (203)) results generalization bounds that are too large and that can be only computed with norms of learned weights. To our knowledge, our work is first to decompose the generalization error into bias and variance a priori (not dependent on learned weights) for linear GCNs with residual-like structures. The architecture closest to linear BPB-Gr CPB-CCN is the linearly decoupled GCN proposed by Cong et al.] (202T); however, the overall readout vector is shared for all branches, which will not result in kernel renormalization for different branches.
BPB-GCN
3.1 PARALLEL BRANCHING GCN ARCHITECTURE
We are motivated to study the parallel branching networks as they resemble residual blocks in com- monly used architectures (Kipf & Welling (016)) He et al. (2016): [Chnn et al.) (20205)) and easy to study analytically with our Bayesian framework. We primarily focus on the graph setting (BPB-GCN) in the main sections, and discuss the more general case in appendix B
Given graph G = (A, where A is the adjacency matrix and X the node feature matrix, the final readout for node y is scalar f θ) that depends on the graph and network parameters O. The parallel branching GCN is an ensemble of GCN branches, where each branch operates independently with no weight sharing. In this work, we focus on the simple setup of branches made of linear GCN with one hidden layer, but with different number of convolutions AT on the input node features (Figure (aa)) In this way, parallel branching GCN is analogous to GCN with residual connections, for which the final node readout can also be thought of ensemble of convolution layers (Veit et al. (2016)). Concretely, the overall readout ("; +) for node μ is a sum of L branch
In this paper, the convolution operation is normalized as A = ( + 1)2-1/2 where Á is the original Adjacency matrix and D is the degree matrix. We also use feature standardization after convolution for each branch normalize the input
2

---

Published as a conference paper at ICLR 2025
Student Norma Branch BPB-ONN
Rendoct
Frequany
Ilidden
layer
Teput layer
500000
Branch
ADD Seather
Squared
(a)
(b)
Figure I: Overview of the main takeaway: BPB-GCN learns robust representations for each branch narrow width. (a) The parallel branching GCN architecture, with 2 branches. The independent branches have non-sharing weights and produce the final output f as a sum of branch-level readouts fi. (b) Student and teacher readout norms squared for wide and narrow student BPB-GCN networks. The student network with width N is trained with the teacher network's 's output. Histograms corre- spond to the samples from Hamiltonian Monte Carlo simulations and solid lines correspond to the order parameters calculated theoretically. o; = dw = 1. At N 4, the HMC samples of branch
||a]ll"
readout norms squared (orange and red histograms) for the student network NN o. concentrate at their respective theoretical values un and overlap with the teacher's readout norms squared
|A:l" of (orange and red dashed lines) for corresponding branches. At N 1024 the samples for N
the student network (blue and green histograms) concentrate at their respective theoretical values but remain far from the teacher's s values, instead approaching the GP limit o (blue dashed line).
readouts f!::
L-1
f²(G:0) = Σ ff (G: = ay00,a011
(1)
1=0
where
N
No
(C = VL
1
(
W ( 1j
Σ
(Al
(2)
VN
No
1=1
In matrix notation,
F= VLNNO
(3)
Note that when L the network reduces exactly a 2-layer residual GCN (Chen al. (2020c)
Here No is the input dimension, N the width of the hidden layer, W(O the weight of the hidden layer for branch A' the (th power of the adjacency matrix and ( the final readout weight for branch I. We will consider only the linear activation function for this paper and provide a summary of notations in Appendix A.I
3.2 BAYESIAN NODE REGRESSION
We consider a Bayesian semi-supervised node regression problem, for which the posterior probabil- ity distribution for the weight parameters is given by
1 1 P P(O) = -((; = z Z
1
exp( 2T 1=1
O),
(4)
202
where the first term in the exponent corresponds to the likelihood term induced by learning P node labels yy with squared loss and the second term corresponds to the Gaussian prior with variance
3

---

Published as a conference paper at ICLR 2025
T.. 0 fe dO is the normalization constant. This Bayesian setup is well motivated, as the Langevin dynamics trained with energy potential and temperature T that results in this equilibrium posterior distribution shares a lot in common with the Gradient Descent (Avidan et al. (2033 Naveh et al. (02T)) and Stochastic GD optimizers (Mignacco & Urbani (002); Wu et al. 2020
In fact, Li & Sompolinsky (202T) shows empirically that the Bayesian equilibrium sta- tistical distribution of the usual gradient descent with early stopping opúmization and with random initializations at temperature in DNNs, where the Le regularization strength of corresponds the
Gaussian initialization variance.
We are interested in understanding the weight and predictor statistics of each branch and how they contribute to the overall generalization performance of the network. In the following theoretical derivations, our working regime is the overparameterizing high-dimensional limit (Li & Som- polinsky 2021 Montanari & Subag (2023); Bordelon & Pehlevan (2022); [Howard et al. (2024)): P,N,No 30, f a finite and the capacity dop LNg
As we will show later, this limit practically true even with P. N not large (our smallest is 4). We will also use near-O temperature, in which case the training error will be near O.
-
3.3 KERNEL RENORMALIZATION AND ORDER PARAMETERS
The normalization factor, or the partition function, Z J E(0)/T dO carries all the information to calculate the predictor statistics and the generalization dependence on network hyperparameters N, L, o. Using Eq. 24 we can integrate out the readout weights aj 's first, resulting in
z z=
dW dWe-H(ny -HW))
(5)
with effective Hamiltonian H(W in terms of the hidden layer weights for all branches
H(W)
1
Σ TWWT W
y +TK(W))TT-->
1 log det (K(W) + +TI,
(6)
20
where
K(W) 1 E 22 (H(W)
the P x P kernel matrix dependent on the observed I' nodes with node features H₁ XWWi and we denote |p as the matrix restricting to the elements generated by the training nodes.
As shown in Appendix A.2 we can further integrate out the Wi's and get the partition function 7 f Dre--Ht H(a) described by final effective Hamiltonian independent of weights
(u)) = S(u) E(u),
(8)
where we call S(u) the entropic term
S(u) S(u)= = Σ
N log ui + Σ
N
2 U1
and E(u) the energetic term Na E(u)
No log det
E
T
Ki + TI) +
L uKK + TI),
(10)
2P
2P
where K; e No Therefore, the final effective Hamiltonian has the overall kernel
[''XX A'Il the (P P) input node feature kernel for branch 1.
K K= == =
(11)
where wy's are order parameters which are the minimum of the effective Hamiltonian Eq. that satisfy the saddle point equations:
41 N(1 = -
= - Tre
(12)
where = YF -
(K
TT) -1 and Tr, W
4

---

Published as a conference paper at ICLR 2025
GP kernel renormalized kernel: Observe that as P/N - O. the entropic term domi- nates, and thus u = su for all branches from Eq. 12 This is the usual Gaussian Process (GP) limit, or the fffmite-wwdthh limit. In this case, the kernel K2 for each branch is not changed by the training data, and each branch has the same contribution in terms its strength.
However, when depend on the training data; therefore, we have feature learning in each branch with kernel renor- malization. turns out that u, is exactly the statistical average of readout norm squared over the posterior distribution (Appendix A.4). ie.
P/N is large, there is correction the GP prediction, and uy's in general
2 = 1)/N
(13)
for each branch I, where we use the bracket notation to represent the expectation value over the posterior distribution. Therefore, branches in general become more and more different as the width of the hidden layer decreases. We call this phenomenon symmelry breaking, which is discussed in Section 4.2
3.4 PREDICTOR STATISTICS AND GENERALIZATION
Under the theoretical framework, we obtain analytically (Appendix A.3 the mean predictor (J"(G))
and variance (G) for single test node Bq and Eq 55 calculate the generalization performance, defined by the MSE on test nodes
respectively. We use this to
Eg = 1 Σ (((G))
Bias + Variance,
(14)
where
Bias + (1r())) - Variance 1 F=1
2 E (8f2)
(15)
v=1
Note that our definition of bias and variance a statistical average over the posterior weight distri- bution, which is slightly different from the usual definition from GD.
THE NARROW WIDTH LIMIT
As we discussed briefly in Section 3.3 the kernel becomes highly renormalized at narrow width. In fact, in the extreme scenario as a P/N - oo, the energetic term in the Hamiltonian com- pletely dominates, and we would expect that the generalization performance saturates as the order parameters in the energetic terms become independent of width N. Therefore, just as infinitely wide networks correspond to the GP limit, we propose that there exists narrow width limit when the network width is extremely small compared to the number of training samples.
4.1 ROBUST LEARNING OF BRANCHES: THE EQUIPARTITION CONJECTURE
What happens in the narrow width limit? In the following, we demonstrate that each branch will learn robustly at narrow width.
The equipartition theorem given by
Consider a student-teacher network setup, where the teacher network
(') - E (TG Σ Σ-AL Σ (.
(16)
W ~ ~N(0,, of) and ai _N((088)) where 83 is the variance assigned to the readout weight for the teacher branch i and N, is the width of the hidden layer. Similarly, the student network is given by the same architecture, with layer width N and learns from P node labels from the teacher
Y, f% (G, O*) in the Bayesian regression setup of Eq. with prior variance o² We state that
T = 0 temperature, in the limit N₂ + so and P/N oc, there exists solution to the saddle
point equations 12
My of of 0
(17)
5

---

Published as a conference paper at ICLR 2025
Proof:
At the narrow width limit a + oo, the saddle equation 12
becomes ry = Tr. Now we calculate
Ni ΣE NoL Ci
Y N.L
(18)
by law of large numbers, this quantity concentrates at its expectation value
& Ki(ot),
y lim E.-,W (YY*)= No0
Σ
(A (A'XX" A') ki(o))
(19)
No
where Ki(0² rrepresens the GP input kernel with prior variance o². Therefore, becomes
r = Y U4 Ki
K K-
Tr(K 44 y.y- Tri Σ (u,K((oW)/))
K
(20)
Thus there exists solution that satisfies the saddle point euations
un = 0/
(21)
Furthermore, by Eq 13 show that the student branch norms learn exactly the teacher branch We call this equipartition theorem, as the mean-squared readout and the variance A.4) have to exactly balance each other, which contribute to the energy term in
norms, ie. (af) m a7 of
the Hamiltonian. Although the theorem is proven for the BPB-GCN architecture, we show that the
result general in Appendix B.1
4.2 TEACHER EXPERIMENT ON ROBUST BRANCH LEARNING
We demonstrate this robust learning phenomenon and provide first evidence of the equipartition conjecture with the student-teacher experiment setup introduced in the previous section. We use the contextual stochastic block model (CSBM) (Deshpande et al. (2018)) to generate the graph, where the adjacency matrix is given by a stochastic block model with two blocks, and the node feature is generated with latent vectors corresponding to the two blocks (Appendix C.I). Both student and teacher network has L = 2 branches. We calculate the order parameters uy 's for the student network with the saddle point equations, as well perform Hamiltonian Monte Carlo (HMC) sampling (Ap- pendix C from the posterior distribution of the student's 's network weights to generate a distribution of student readout norms squared of. We compare the values of the student branch readout norms to their corresponding teacher branches for the student network with either narrow or wide width. As shown in Figure [bb, an extremely narrow student network (N 4) learns the teacher's branch readout norms very robustly. while much wider network (N 1024) fails to learn the teachers' norms and approaches the GP (N- -> oo) limit where the two branches are indistinguishable.
Symmetry breaking and convergence of branches: We perform theoretical calculations and HMC sampling for the student branch squared readout norms as we vary the student network width N and the prior regularization strength ow. Figure 2 a) shows the statistical average of the student branch squared readout norms, ie. {||ax112))))))))))))) 12) as function of the network width N, where the branch norms split the network width gets smaller, which we call symmetry breaking. The sym- metry breaking of branch norms from the GP limit to the narrow width limit accompanies the con- vergence to learning teacher's norms at narrow width for different a., shown in Figure (b)()), supporting Eq. 62
Narrow-to-wide width transition: We can also determine the generalization properties of the student readout lables at both the branch level fi (Appendix A.3) and overall level usinn Eq. [415 We perform theoretical calculations and HMC sampling of generalization errors as a function of the network width N and regularization strength dw. with results shown in Figure 3 and 4 At narrow width, we expect individual branch ft to learn the corresponding teacher's branch output ff inde- pendently, causing the bias to increase with network width. This observed for both branches, with a transition from the narrow-width regime to the GP regime. The regularization strength o., controls the transition window, with larger o 's leading sharper transitions, This aligns with our analysis

---

Published as a conference paper at ICLR 2025
(6)
Symmetry Breaking of Brunches
[b)
Brunch Student Readout Norms
(6)
Branch Student Reudout Norms
Cranch Tierry HlymL.n
Camelina Pu
Buaeh Campling bouct
Gerplin Purtt
L30
Janoiig
Squerd
1.4
L23
sarsiing Ianun
L.0
1.5
6115
Figure 2: Statistical average of student readout norms squared as function of network width from
theory and HMC sampling, for student-teacher tasks described in Section 4.2 (a): <[|ax]] o function of network width N for fixed ww. The branch norms break the GP symmetry it goes to the narrow width limit. (b)(c): Branch and branch readout norm squared respectively for range of w regularization values. The student branch norms with different regularization strengths all converge to the same teacher readout norm values at narrow width.
/NV as
Branch Squared
Branch Variance vs.
3226
Samples (0
Hee samples Mu-D3)
0.24
Theory 1ou-0 8)
Theery (0u-0)))
Squared Bias
0.22
Samples (0
Squared Bins
HBC samples Muueaa)
9:20
samples (0p-1.00
HeC samples Mu-1.0)
Theory (ou-1.2)
Theory -1,21
HMC Samples 20
HNC Sampler 10,)
0.0
10
30
100
110
107
10
300
107
30a
10'
101
Branch Squared Bias vs. N
Branch Variance vs. N
Theory (0g-0.5)
Thenry (pp00.3
2000
HMC Samples (0y+0.5)
0.07
1991 Samples (0,)35)
Squared Bias
8.500
Theory (0y-0.8)
Squared Bins
0.00
Theory (ou-0.0
Theory (0g-2.00
Theory -0
8.05
HMC Samples 10pe1
0.04
19wc Sumples 10.0.00
Theory (mv-1122)
Theary -12)
0.40
0.03
0:2
2005
0.0
0.08
Figure 3: Student network squared bias and variance for individual branches as a function of network width N and regularization strength ow- The mean and variance of branch readout f for node is calculated in A.3 and the bias and variance for branch I can be infered similarly as Eq. 15 Generalization values are normalized over the average true readout labels.
of the entropic and energetic contributions, where the larger aw amplifies the distinction between the two terms. In contrast, the variance decreases with network width for small o resulting in a trade-off between the contributions of bias and variance to overall generalization performance, as shown clearly with the graph of network generalization vs. N with aw 0.5 in Figure
5 BPB-GCN ON CORA
We also perform experiments on the Cora benchmark dataset McCallum et al. (200)) by training the BPB-GCN with binary node regression, for a range of T,, N, on values (Appendix 2 for de-
7

---

Published as a conference paper at ICLR 2025
Overall Squared vs.
Overall Variance vs.
Generalization
Theiry 3.8
0.24
Tincy Nu-rda dal
Theeyy 1a1
Sanules
saxulee (c-5.
IBCC (4-1
330
6.18
Thotry (1
Teory/buel
Conmitiatton
Trooy (1-6.21
Squaral Riee
Symdio
Camd
0.14
Vickene
6.12
Figure 4: Student network generalization performance as a function of network width N and regu- larization strength ow Generalization is normalized over the average true readout labels.
Squared Blos
Squared Bise vs. L (N-4)
Thorry iI pempliggngg
Haar gntegge 2)
Thesey
1867 Campling 42.0
Squared Blaa
023
30
0322
Accumec vs. N (L--6)
Accuracy L.(N-4)
0.60
155
085
6.00
0.0
0.92
Figure 5: Cora generalization performance network width N and branch number L, for various regularization strength o 's. The accuracy is computed by turning the mean predictor from HMC samples into a class label using its sign.
tails). We observe a similar narrow-tor -orrww width transition for the bias term. As shown in Figure 5 the bias increases with network width, transitioning to the GP regime, and we observe the trend ex- tending to a potential narrow width limit Additionally, it is demonstrated that using more branches that involve higher-order convolutions improves performance.
Convergence of branch readout norms at narrow width An interesting aspect of the BPB-GF BCB-GCN network is that the branch readout norms converge at the narrow width for different hyperparameters au and L., reflecting the natural branch importance for the task.
PIn this case, the narrow width limit is hard to demonstrate as the transition window is below realistic minimum of network width.
8

---

Published as a conference paper at ICLR 2025
Cora Branch Squared Readout Norms
Besnch
Branch
Branch
Branch
Demch 4
Domch 5
New Hhmiii sourdon
N
Figure 6: Cora experiment: statistical average of squared readout norms <Ila (o /N for each branch / as function of the network width N, and regularization strength d.n The BPB-GCN has = 6 branches.
for 2
Branch
1,00
1,03
1.00
142
0.75
0.5
1075
1000
0.0
1000
3000
{ 1000
2003
2003
022
3000
2000
600
100
1045
Index
Komel fo Branch 4
Kome for Branch
Kemee for Branc's 6
Kannel for Branch
0.75
0.73
0.00
0.2
2007 2006
2013 20003
3230 2003
1001 2003 Noce
Inde:
Node 17200
Figure 7: Kernel K2 22 A' X' A for each branch I shown for the first branches Cora dataset, sorted by node labels. A is the normalized adjacency matrix of the Cora graph, Xo the node feature matrix and No the node feature dimension. Initialization variance of 1 and total node number n 2708.
No
As shown in Figure 6 the BPB-GCN with branches L = 6 robustly learns the readout norms at narrow width independently of ow" consistent with the student-teacher results. This suggests that we can recast the data as generated from a ground-truth teacher network even for real-world datasets. The last branch of the BPB-GCN network has larger contribution, reflecting the presence of higher-order convolutions in the Cora dataset. From a kernel perspective, increasing branches better distinguish the nodes, shown in Figure 7 This could explain the selective turn-off of intermediate branches and the increased contribution of the last branch. Our results also indicate that there is no oversmoothing problem narrow width: as increases, individual branches can still learn robustly.
Furthermore, the first two branches are learned most robustly narrow width, as shown in Figure 8, where the branch norms converge for the first two branches even for BPB-GCM with different L. This suggests that the branch importance, as reflected by the norms learned at narrow width, indicates the contribution of the bare data and the first convolution layer.
DISCUSSION
The findings presented in this paper reveal that BPB-GCNs exhibit unique characteristics in the nar- row width limit. Unlike the infinite-width limit, where neural networks behave as Gaussian Process (GP) with task-independent kernels, narrow-width BPB-GCNs undergo significant kemel renormal- ization. This renormalization leads to breaking of the symmetry between the branches, resulting in more robust and differentiated learning. Our experiments demonstrate that narrow-width BPB- GCNs can retain and, in some cases, improve generalization performance compared to their wider counterparts, particularly in bias-limited scenarios where the regularization effects dominate. There-

---

Published as a conference paper at ICLR 2025
Branch Squared Readout Norms vs. Branch Number at Narrow Width
lonping loeding
Vvnry
Squared Noems
Number
Figure 8: Branch Importance vs. branch number on Cora. Legends represent different total number of branches L. The branch importance is defined as the statistical average of branch squared readout norms <[|ar]] [[)) V at the narrow width limit; here we take the empirical branch norm values at N and fixed o 1.4.
fore, there a trade-off between the expressive power of GCN Xu et al. (2018): Loukas (2019) and stability Stogin et al. (2020). Additionally, the observed independence readout norms from archi- tectural hyperparameters suggests that nartow-width BPB-GCNs can capture the intrinsic properties of the data more effectively. These insights suggest potential new strategies for optimizing neural network architectures in practical applications, challenging the traditional emphasis increasing network width for better performance. Although our work is focused on GCN architectures. our findings are more general for family of branching networks as shown in Appendix
B and in par-
ticular the Bayesian theory with kernel renormalization is shown to be successful for CNN Aiudi al. (2023) well as transformers Tiberi et al. (2024). Intriguingly, Vyas et al. (2024) finds that an ensemble of ResNets averaged over random initializations accidentally shows this narrow width
effect, for which our findings provide an explanation.
Limitations: We demonstrate that the bias robustly decreases narrow width for the branching networks due to kernel renormalization, which defies the common belief that wider networks are better FAllen-Zhrr al. (2019) Lee al. (2018)). However, we do find that for real-world scenar- ios such as BPB-GCN variance-limited. In principle, training an ensemble branching networks averaged over random initialization helps to reduce the variance significantly such that the narrow-width effect is more pronounced as in Vyas et al. (0044). Furthermore, since our Bayesian setup corresponds to an en- semble of fully trained premmtrraetrieee networks in an offline fashion, the learning is still in the so-called lazy regime with minimal feature learning (Atanasov et al.) (2022), y Karkada (2024)) We hypothesize that narrower width serves regularizer in this regime which might break down as the network is trained with online learning and approaches the rich regime (Yang et al. (2022) with more feature learning. In that case, it might be possible that the bias does not decrease the narrow width.
Cora, and residual-MI.P on Cifar10, the generalization error is usually
CONCLUSION
In conclusion, this paper introduces and investigates the concept of narrow width limits Bayesian branching networks. Our results indicate that branching networks with significantly narrower widths can achieve better or competitive performance, contrary to common beliefs. This attributed to fective symmetry breaking and kernel renormalization in the nirtow-width limit, which leads to
10

---

Published as a conference paper at ICLR 2025
robust learning. Our theoretical analysis, supported by empirical evidence, establishes derstanding of how network architecture influences learning outcomes. This work provides a novel perspective the infinite width limit of neural networks and suggests further research directions in understanding narrow width as a regularization effect.
ACKNOWLEDGEMENT
We acknowledge support of the Swartz Foundation, Kempner Institute for the Study of Natu- ral and Artificial Intelligence at Harvard University and the Gatsby Charitable Foundation. This material is partially based upon work supported by the Center for Brains, Minds and Machines (CBMM), funded by NSF STC award CCF-1231211 We have benefitted from helpful discussions with Alexander van Meegen, Lorenzo Tiberi and Qianyi Li.
REFERENCES
Aiudi, Pacelli, Alessandro Vezzani, Raffaella Burioni, and Pietro Rotondo. Local kernel renor- malization as a mechanism for feature learning in overparametrized convolutional neural net- works. arXiv preprint arXiv:2307.1 11077, 2023.
Zeyuan Allen-Zhie Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparame- terized neural networks, going beyond two layers. Advances in neural information processing systems, 32, 2019.
Gholamali Aminian, Yixuan He, Gesine Reinert, Łukasz Szpruch, and Samuel N Cohen. General- ization error of graph neural networks in the mean-field regime. arXiv preprint 27X 2024.
Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan. Neural networks as kernel learners: The silent alignment effect. arXiv preprint 1rXiv:2111.00034, 2021.
Alexander Atanasov, Blake Bordelon, Sabarish Sainathan, and Cengiz Pehlevan. The on- set of variance-limited behavior for networks in the lazy and rich regimes. arXiv preprint arXiv:2212.12147, 2022.
Yehonatan Avidan, Qianyi Li, and Haim Sompolinsky. Connecting ntk and nngp: A unified the- oretical framework for neural network learning dynamics in the kernel regime. arXiv preprint arXiv:2309.04522, 2023.
Yasaman Bahri, Boris Hanin, Antonin Brossollet, Vittorio Erba, Christian Keup. Rosalba Pacelli, and James B. Simon. Les houches lectures on deep learning at large and infinite width, 2024.
Michael Betancourt. A conceptual introduction to hamiltonian monte carlo. arXiv preprint
arXiv:1701.02434, 2017.
Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks. Advances in Neural Information Processing Systems, 35:32240-32256. 2022.
Rohitash Chandra, Ayush Bhagat, Manavendra Maharana, and Pavel N Krivitsky. Bayesian graph convolutional neural networks via tempered meme. IEEE Access, 9: 9:130353-13033 2021.
Lei Chen, Le Wu, Richang Hong, Kun Zhang, and Meng Wang. Revisiting graph based collaborative filtering: A linear residual graph convolutional network approach. Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 27-34, 2020a.
Ming Chen, Zhewei Wei, Zengfeng Huang, Bolin Ding, and Yaliang Li. Simple and deep graph convolutional networks. International conference on machine learning, pp. 1725-1735 PMLR, 2020b.
11

---

Published as a conference paper at ICLR 2025
Ming Chen, Zhewei Wei, Zengfeng Huang. Bolin Ding, and Yaliang Li. Simple and deep graph convolutional networks. In Hal Daumé III and Aarti Singh (eds.). Proceedings of the 37th In- ternational Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 1725-1735 PMLR, 13-18 Jul 2020c. URL. https: httpp::/pprocee:: ngs. hgg.mlr.. mmlr 2Teas//1113/chenwv..1
Rong Chen, Li Guanghui, and Chenglong Dai. Feature fusion via deep residual graph convolutional network for hyperspectral image classification. IEEE Geoscience and Remote Sensing Letters, 2911-5,002..
Youngmin Cho and Lawrence Saul. Kernel methods for deep learning. Advances in neural informa- tion processing systems, 22, 2009.
Weilin Cong, Morteza Ramezani, and Mehrdad Mahdavi. On provable benefits of depth in training graph convolutional networks. Advances in Neural Information Processing Systems, 34:9936- 9949, 2021.
Yash Deshpande, Subhabrata Sen, Andrea Montanari, and Elchanan Mossel. Contextual stochastic block models. Advances in Neural Information Processing Systems, 31, 2018.
Simon Du, Kangcheng Hou, Russ R Salakhutdinov, Barnabas Poczos, Ruosong Wang. and Keyulu Xu. Graph neural tangent kernel: Fusing graph neural networks with graph kemels. Advances in neural information processing systems, 32, 2019.
Tianxiang Gao, Xiaokai Huo, Hailiang Liu, and Hongyang Gao. Wide neural networks as gaussian processes: Lessons from deep equilibrium models. Advances in Neural Information Processing Systems, 36, 2024.
Vikas Garg, Stefanie Jegelka, and Tommi Jaakkola. Generalization and representational limits of graph neural networks. In International Conference on Machine Learning, pp. 3419-3430. PMLR, 2020.
Floris Geerts and Juan L Reutter. Expressiveness and approximation properties of graph neural networks. arXiv preprint 2XXx::2204.04661, 2022.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Jessica Howard, Ro Jefferson, Anindita Maiti, and Zohar Ringel. Wilsonian renormalization of neural network gaussian processes. arXiv preprint 2405::200000000 2024.
Wei Huang, Yayong Li, Weitao Du, Jie Yin, Richard Yi Da Xu, Ling Chen, and Miao Zhang. Towards deepening graph neural networks: gatk-based optimization perspective. arXiv preprint arXhv:2103.03113. 2021.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kemnl:: Convergence and gen- eralization in neural networks. Advances in neural information processing systems, 31, 2018.
Haotian Ju, Dongyue Li, Aneesh Sharma, and Hongyang Zhang. Generalization in graph neural networks: Improved pac-bayesian bounds graph diffusion. In International Conference on Artificial Intelligence and Statistics, pp. 6314-6341. PMLR, 2023.
Dhruva Karkada. The lazy (ntk) and rich
mu p)regimes aggnttetuuorioll.at arXiv preprint arXiv:2404.19711 2024.
Thomas N Kipf and Max Welling, Semi-supervised classification with graph convolutional networks.
arXiv preprint arXiv: 2016.
Alex Krizhevsky et al. Learning multiple layers of features from tiny images. 2009.
Jachoon Lee, Yasaman Bahri, Roman Novak, Samuel S. Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks as gaussian processes. 2018.
12

---

Published as a conference paper at ICLR 2025
Qianyi Li and Haim Sompolinsky. Statistical mechanics of deep linear neural networks: The back- propagating kernel renormalization. Phys. Rev. X, 11:031059, Sep 2021. doi: 10.1105/PhysRevx. 11.031059. httphtts:: httpss///....kka......./////..2.
Renjie Liao, Raquel Urtasun, and Richard Zemel. A pac-bayesian approach to generalization bounds for graph neural networks, arXiv preprint arXiv:2012.07690, 2020.
Andreas Loukas. What graph neural networks cannot learn: depth vs width. arXiv preprint arXiv: /90.003999, 2019.
Andrew K. McCallum, Kamal Nigam, Jason Rennie, Automating the construction of inter- net portals with machine learning. Information Retrieval, 3(2):127-163, 2000. doi: 10.1023/A:
1009953814988.
Francesca Mignacco and Pierfrancesco Urbani.
The effective noise of stochastic gradient descent. Journal Statistical Mechanics: Theory and Experiment, 2022(8):083405, 2022.
Andrea Montanari and Eliran Subag. Solving overparametrized systems of random equations: model and algorithms for approximate solutions. 2023.
Gadi Naveh, Oded Ben David, Haim Sompolinsky, and Zohar Ringel. Predicting the outputs of finite deep neural networks trained with noisy gradients. Phys. Rev. 104:064301, Dec 2021. doi: 10. 1103/PhysRevE. 10..06.004..641111 URL 104.064301
aps wwwp://////oooo////// http://b///111111/////aaaaae.e/.
Radford M Neal. Bayesian learning for neural networks, volume 118. Springer Science & Business Media, 2012.
Daniel Park, Jaseha Sohl-Dickstein, Quoe Le, and Samuel Smith. The effect of network width stochastic gradient descent and generalization: an empirical study, In International Conference on Machine Learning, pp. 5042-5051, PMLR, 2019.
John Stogin, Ankur Mali, and C Lee Giles. A provably stable neural network turing machine. arXiv preprint arXiv:2006.0360 03651, 2020.
Huayi Tang and Yong Liu. Towards understanding generalization of graph neural networks. In An- dreas Krause, Emma Brunskill. Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 33674-33719. PMLR, 23-29 Jul 2023. URL wtw.ggsm......srr...
html
Lorenzo Tiberi. Francesca Mignacco, Kazuki Irie, and Haim Sompolinsky. Dissecting the interplay of
attention paths in statistical mechanics theory of transformers. arXiv preprint arXiv:2405. 15926, 2024.
Andreas Veit, Michael Wilber, and Serge Belongie. Residual networks behave like ensembles of relatively shallow networks. Advances in neural information processing systems, 29, 2016.
Nikhil Vyas, Alexander Atanasov, Blake Bordelon, Depen Morwani, Sabarish Sainathan, and Cengiz Pehlevan. Feature-learning networks are consistent across widths at realistic scales. Advances in Neural Information Processing Systems, 36, 2024.
Ian Walker and Ben Glocker. Graph convolutional gaussian processes. In International Conference on Machine Learning, pp. 6495-6504. PMLR, 2019.
Zihan Wang and Arthur Jacot. Implicit bias of sgd in L22) reggulaulr linear dnns: One-way jumps from high to low rank. arXiv preprint 2.X.::2305.1603.. 2023.
Jingfeng Wu, Wenqing Hu, Haoyi Xiong, Jun Huan, Vladimir Braverman, and Zhanxing Zhu. On the noisy gradient descent that generalizes as sgd. In International Conference on Machine Learning, pp. 10367-10376. PMLR, 2020.
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural net- works? arXiv preprint arXiv:/810.0000 2018.
13

---

Published as a conference paper at ICLR 2025
Greg Yang. Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs V: Tuning large neural networks via shoo hyperparameter transfer. arXiv preprint 244660220200022 2022.
A DETAILS ON THEORY OF BPB-GCN
A.1 SUMMARY OF NOTATIONS
Hyperparameters and Dimensions
P
Number of training nodes
n
Total number of nodes for a graph
No
Input node feature dimension
N
BPB-GCN hidden layer width
Total number of branches
Ow
Le prior regularization strength
T
Temperature
P
ag
Network capacity
LND
Width ratio
N
Network Architecture and Input/Output
À R
Adjacency matrix
A E RAXE
Normalized adjacency matrix by its degree matrix D Input node feature matrix Graph
X E RaX
G (X,
W E NOXN
Hidden layer weight for branch |
" E R
Readout vector for branch Collection of parameters for branch
0 (W),a)
M E
Activation vector for branch I and node μ
J ER
Readout prediction for branch and node μ
ER
Overall readout prediction for node /
y" ER
Overall node label for node μ
H((W) E PXN
Activation feature matrix for branch I.
YER?
Training node labels
FERP
Readout predictions
Statistical Theory
14

---

Published as a conference paper at ICLR 2025
GY)
Energy loss function
7
Partition function
TERPXP
Identity matrix
H(W)
Hamiltonian after integrating out readout aj s
uER'
Order parameters as saddle point solution
u E
Order parameter for branch /
(u)
Hamiltonian function of order parameters
S(u)
Entropy as a function of order parameters
E(u)
Energy function of order parameters
r
Mean squared readout
Tr,
Variance- related readout
(x)
Statistical average of any quantity over the posterior dis- tribution
Kernels
K(W) PxP
Hidden layer weight dependent overall kernel
Ki
Branch / kernel
K PxP
Overall kernel averaged over W's, K
Σ, W.K L
M Px1
Branch 7 kernel column for the P training nodes against test node
W RR
Overall kernel column for the training nodes against test node
lp
Kernel restricted to the P training nodes against P training nodes
((P.v)
Kernel restricted to P training nodes against the test node
((v,))
Kernel restricted to the test node v against test node
Student-teacher Setup
Yi' e R
Teacher network readout prediction for branch I for P nodes
Y" PP
Teacher network overall readout for P nodes
W7 NOXN
Teacher hidden layer weight for branch
of EERV
Teacher readout vector branch
Teacher readout variance
of
Teacher hidden layer weight variance
A.2 KERNEL RENORMALIZATION
Following similar derivations as the first kernel renormalization work Li & Sompolinsky (2021), we will integrate out the weights in the partition function Z f d8 exp -xp(-E(C E()/TT) from the readout layer weights a, to the hidden layer weights Wi's and arrive at an effective Hamiltonian shown in the main text.
15

---

Published as a conference paper at ICLR 2025
First, we linearize the energy in terms of 'ss and transform the integral by introducing the auxiliary variables t", 11 = I.......
- / de
/
II
P
di exp
To
P Σ
itp
N T-1 ΣΣ
(G)
yu
717 t
2 m
VLN
2
1=1 1=0
(22)
Now can integrate out aj's 's as they are linearized and the partition function becomes
2 /
DWc-H(W)
(23)
with effective Hamiltonian
Tw W, ! 2 (K(W)+TT)>> + log (eedddee(e>))e + + TT),
1
H(W)
(24)
20
2
1=0
where
K(W) ト Σ oh No
(ww)
(25)
the P x kernel matrix dependent on the observed P nodes with node features H, = A' XW; and denote p the matrix restricting to the elements generated by the training nodes.
Now we perform the integration on Wi's and introduce the auxiliary variables t"'s again using the Gaussian trick and get
L-1 Z = II dwi / / dtexp exp 1
of (K(W) + 20
ITr(WWWa) + it
(26)
f
T m + G(I)
=
it deep
dt exp
(27)
2
where G(t) is in terms of the kernel averaged over the Gaussian measure
1
x(W)1))
G(t) = log
cxp
(28)
2N
Writing out the integral explicitly, we have
G(t) log /
I-1
II
DWW exp
Σ
th
A'
Wt
Al
2N
NoL
p.ge
1=0
3.1
(29)
where DW; is the Gaussian measure
DW
II
W 20 dWisd
W
(30)
By recognizing that this is just product of the same N integrals on N, dimensional weight vector Wi for fixed j neuron, we get
L-1
02
G(t)) = N log
II 20
dw?
exp
Σ P.P.
the"
et St. + 8r, ) W W.t
(31)
2N
NoL
20
the Ith branch convolution on node Since the network is linear, we are able to perform the Gaussian integration exactly and get
where c/' = - A p. of
N Σ G(t) = 2
a NoL
log(dct(Ir.r +
(32)
16

---

Published as a conference paper at ICLR 2025
where the determinant is over the i, matrix, and thus performing the determinant arrive at
G(I) Σ log(1 + 1- Ki
T. 1)),
(33)
where K n.v of E, SM, is the Gaussian process kernel for branch
Next, we insert L delta functions II
dhz - t" Kit) to enforce the identity G(t) =
Σ: log(1 + hy) and use the Fourier representation of it with auxiliary variables u to get:
Z =
L-1
II
dh exp (itTY
N log(1 + h11) + E
Σ
N
uily 2"
(u TI)s)
(
1
W +
2
20
1=
II t=0
dhaduy exp
Σ
N 2
log(1 + h4) + Σ
N
1 TT
<Σ 1
- + TT) rr)
20
2
(34)
where
oF K2 [A'X' A 1
(35)
No
is the input kernel for branch (. Now as N + oc and a Nv approximation and get the saddle points for 11 as
P
fixed, we can perform the saddle point
o
1 h1
(36)
i
Plugging this back to the equation, we get
z / II dude Mass(0)
(37)
with the effective Hamiltonian
Hess(u) = S(u) + EEu),
(38)
where we call S(u) the entropic term
N log in Σ
N
S(u) = Σ
20 4
(39)
2
and E(u) the energetic term
E(u) = r Σ Euuni uKK + Ti) log det Σ
uiki + T1)
(40)
Therefore, after integrating out Wi, the effective kernel is given by K = Σー LuKi
(41)
and the saddle point equations for ut's 's are determined by
N(1
(K (K + TT) - y
W1
KK
(K
(KKTTT--
Tr(K uiKi L
(42)
L
where we call
r = Y + TI) -1 Ki
(K -TI)-
(43)
I
and
Tr, Tr K
(44)
I
As we will show later, these represent the mean and variance of the readout norm squared respec- tively. In the T = 0 case, the saddle point equation becomes
4KK
K L
N 14 -
K Y
Tr
(45)
L
17

---

Published as a conference paper at ICLR 2025
A.3 PREDICTOR STATISTICS AND GENERALIZATION
We can get the predictor statistics of each branch readout yf (G) on new test node v by considering the generating function:
Z (mI....., ,L) / D0xxp
[
3
((f((G)))-
yyn
2
(46)
Σ im
1
Σ
(G,
T o²
ai
VNL
20 m
Therefore, by taking the derivative with respect to each Ytt we arrive at the statistics for (z) as:
(f) = 0im log Z 7=0
(47)
(G). = 02 log 2 #==
(48)
After integrating out the weights ④ layer by layer, have:
/ Midu exp Σ
N
Z
=
log
20 Mi
2
+ (rY +
Σ moukk
Ki + TT) (i8 + Σ
1!1 I uik
1
2
I
1
log det
Σz uKKi + +TI) - 1
Σ
Ki
(49)
2
2
Here
02 A - 1A'xx" A')((
(50)
No
is the x column kernel matrix for test node v and all training nodes, and
-2
KT [AAX A')
(51)
No
the single matrix element for the test node. Therefore, eventually, we have:
UiIT (K -TT-
(
(52)
L
and
WiKKi
WhiT
-1 WKI.
(8
(K + TT)
(53)
L
L
The predictor statistics of the overall readout f = Σ, fi is given by:
uikt
(F) Σ
(K
(54)
(8(()) Σ u KrK.
Σ
Wkk (K TT)
-uykr Kn
1T[(K TT) kn
(55)
1.P
18

---

Published as a conference paper at ICLR 2025
A.4 STATISTICS OF BRANCH READOUT NORMS
From the partition function Eq 22 variable by
we can relate the mean of readout weights to the auxiliary
(ai)w -i dh
oh of -K+
(56)
VN
VNL
where b is the node feature matrix for the hidden layer nodes. We have
UK u,
(aF) o +T)-
(K -K+T r:o
(57)
L
We can calculate the second-order statistics of a;: the variance is {ōa) δaz) of Tr([
L.B D x = oN Tr(K TI) uKk,
ow -Tr))
(58)
NL
I.
Therefore,
(a (ōuî δ0z) + (ōa) dus) = No + - + o Tr Nu
(59)
Therefore, we have proved the main text claim that the order parameter ut are really the mean squared readout norms of the branches.
B THE NARROW WIDTH LIMIT FOR GENERAL ARCHITECTURES
In this section, we demonstrate that our results on the narth limit can be generalized to any network with branching structures, supported by both theoretical and empirical evidence.
Consider general branching network, with L independent branches each with hidden layer weight Wi and readout weight at
f(z: θ) Σ J f((:; fi(z; (i); VL
0:)= =
1 VNL VNZZZaian (1; Wi),
(60)
1=0
where (x; Wi) represents any ford-forward function on the dataset X, ic. in matrix form the branch / readout is F; क ( XWi). As example, (XWW) NNo A' X'WW, represents / convolutions on the node features in the BPB-GCN architecture. (XW,) VNo A;XW; can also represent an attention head for attention-networks with frozen attention As, well kernel convolution for patch in the CNN architecture. In the residual-ML. example below, Φi represents either the linear or the ReLu branch.
With the general branching network, we can similarly consider the posterior distribution in the Bayesian setup as Eq. 4 with squared error as the likelihood term and T2 regularization as the prior term. Provided the network is linear, the partition function 5 can be integrated in a similar fashion and arrive the saddle point equation
N(1 - ) )= = -YT (K + TI) KKK (K (K Tr[K
K L
(61)
L
with K = = L and K₁ corresponding the NNGP kernel Lee et al. 2018 ). In fact, even with non-linear activations, we conjecture that the kernel renormalizatio still holds with the order parameters u, that satisfies the saddle point equation, and the posterior weight statistics calculated with the renormalized kernel agrees with the experiments as empirically shown in the residual-MLPP section.
B. THE GENERALIZED EQUIPARTITION THEOREM
With the general branching network, we show the narrow width limit result for the student-teacher setup. Specifically we prove the following
19

---

Published as a conference paper at ICLR 2025
The generalized equipartition theorem For the general branching network Eq. 60 where the teacher network width is N₂ with Wi, ~ ~N of and di ~ (0, 5 and the student network with layer width N in the Bayesian regression setup of Eq. 4 at No + and P/N + oo, there exists solution to the saddle point equations 12
temperature, in the limit
03 /ow.
(62)
Proof: The proof is exactly the same as before, except for changing Ki as a general NNGP kernel that corresponds to the feature map di for branch I.
B.2 RESIDUAL-MLP
In the following, we show a simple 2-layer residual network as an example for the general branching
networks, with a linear branch @o (x) _ Wox and ReLU branch Φ1 (x) VNo
- ReLu(W(1))
VNo Therefore, the overall kernel is K (woKo + K1), with the NNGP kernels K₁ given by
Ko Kiincar o² No
Xo-X0
(63)
and the relu kernel (Cho & Saul (2099)
K1 x¹) Krelu (2, 2')
v /Kolz, Kol=', J(),
(64)
2T
where
J(0) = sin(O) + (πr 0) (0)
(65)
and
Ko( x¹)
0 = cos ((
(66)
Ko( x Ko
we observe the same narrow-width limit phenomenon as the BPB-GCN case, where each student branch robustly learns the teacher's corresponding branch but approaches the same GP limit the network width gets larger.
For the student-teacher task (details in Appendix +
Symmetry breanking of branches (0=1.0)
2.50
Reachrr no m tinear branch
225
Teacher narm (ReLu branch) branch
2.00
IIC sempling
HMC sampling (ReLa branch)
1.75
18'0'00
1.50
1.25
1.00
0.5
0.5
10
30
10-
15a 107
Figure 9: Statistical average of student readout norms squared as a function network width from theory and HMC sampling, for the Midiuual-MLD student-teacher task, with fixed aw 1.
As shown in Figure 9 and 10 there a symmetry breaking of branches from the GP-limit to the nirtow-width limit, with the ReLu branch and the linear branch learning their corresponding teacher branch norms as N decreases. Furthermore, narrower network performs better terms of the bias and larger o. corresponds to shorter transition window from the GP-limit to the width limit, similar to the BPB-GCN case. As shown in Figure IT and 12 the trade-off between the bias and variance contribution is more pronounced for small ow s, usually with optimal width for the generalization error. These results support our generalized equipartition theorem and demonstrates that the narrow width limit is general phenomenon for branching networks.
20

---

Published as a conference paper at ICLR 2025
Studant Readout Norms Squared (Linear Branchh) vs
Student Readbut Norms Squared (ReLu Branch)
Youcer rom iiiee boakcT
bracher rem iPtt troncht
Theery o-d
Theary
HBF aud
HMT oud
Theers owl
Therry ond
eMC
HNC
Theery 0.15
Thenry o-15
INC
100
107
10F
159
107
107
107
N
(a)
(b)
Figure 10: [ /N as function of network width N for range of ow values. (a) Linear
branch; (b) Relu branch. The student branch norms with different regularization strengths all con- verge the same teacher readout norm values at narrow width.
Squared Blas vs N
Variance vs
I19eyy expia
0.5
Theary 2ed
0.005
HMC GHO.E
HMC OxB.S
Therry
Thooey 1-3
HMC
0-1.0
HBB owl.2
c=1.2
0.0/5
Theary em1.3
FMC 0-1..
0.1
Theay 2-15
e-1.2
0.055
0.1
0.05
0.0
(a)
(b)
Figure 11: Student residual-MLP network generalization performance a function of network
width N and regularization strength ow. (a) Bias; (b) Variance.
B.3 RESIDUAL-MLP ON CIFARI0
We also perform theory and experiments on the Cifar10 (Krizhevsky al. (2009) dataset for binary classification task (translated to Bayesian regression), with details elaborated in Appendix branches and the branch norms converge at the
Similarly, there exists symmetry breaking narrow width for different o s, as shown in Figure 13 14 Interestingly, the ReLu branch completely dominates at the narrow width, which means the ReLu branch is much more important for the residual-MLP when the training sample size is much bigger than the network width.
Furthermore, we also show in Figure (6 that the dominance of ReLu branch at narrow width drives the bias to decrease, albeit overall genralization is dominated by the variance. Using the average the label predictors over different experimental samples, we show that the test accuracy indeed decreases with larger network width N in Figure 16
21

---

Published as a conference paper at ICLR 2025
Goneralization Error
samplig
Theory cm2.0
Cenerellcation Erter
IIC templing o-10
64
Theory cwll2
IMC temping o=12
Theory c-15
somplig gm1.
033
0.2
10F
167
107
100
Figure 12: Overall generalization as a sum of bias and variance Figure
Citar10 Symmetry brearidng branches
Theary
Theary branch)
IIM cempling neer brench)
HMC sampling (Retu brench)
10F 301
Figure 13:
Statistical average of student readout norms squared function of network width from theory and HMC sampling, for the Cifar 10 dataset with residual-MLP, with fixed ow 1.
C EXPERIMENTAL DETAILS
C.1 STUDENT-TEACHER CSBM
For the student-teacher task, we use the contextual stochastic block model introduced by Deshpande al. 2018) to generate the graph G. The adjacency matrix is given by
1 with probability p Cin/n, nn, if i, S n// with probability p cn//n, if i, 2 n/2 with probability q = Cont n, otherwise
Aij
(67)
where
Cin,out = d vdx
(68)
d is the average degree and À the homophily factor.
The feature vector 7 for particular node y is given by
1 od ūt Eny V
z
(69)
n2
where
~~ No0, N, No0 TNo)
(70)
In the experiment, we use No = 950,d = 20.A = and = 4. The teacher network parameters
variance of 1, width N, 1024, branch norms variance 3g 0.4, B 2 for individual
22

---

Published as a conference paper at ICLR 2025
CifarID Student Readout Norms Squared (Uinear Branch)
Citar10 Student Readout Norms Squared (Relu Branch)
Theery owt
theery 0.1.
HMC aed
HMC avd
Theers not MNC 0w2.
32
Theery MMC ow
10
100
107 100
10
107
10"
157
101
(a)
(b)
Figure 14: 11 / /N function of network width N for range of aw values. (a) Linear branch; (b) Relu branch. The student branch norms with different regularization strengths all converge to the same teacher readout norm values narrow width.
Citar10
Citar10 Vartance
Theery e-10
Theery a-11
IMC 1
IMCS 1.0
0.5
sheery net
0.74
PMCUm23
HMCue23
Bas
0.72
0.70
0.3
0.68
0.69
0.1
10F
10
10a
10'
10'
10"
10"
(a)
(b)
Figure 15: Residual-MLP generalization performance on Cifarl0 as function of network width N and regularization strength ow- (a) Bias; (b) Variance.
element of the readout vector a;- We used total number of nodes 2600 with train ratio 0.65 for semi-supervised node regression. Temperature T 0.0005a2 for each on value in a (0.5, 0.8,, 1.2).
C.2 CORA
For the Cora dataset, we use a random split of the data (total nodes n =2708) into 21% as training set and 79% as test set. We group the classes (1, 4) into group and the rest for the other group for binary node regression, with labels as +1's. The Bayesian theory and HMC sampling follows the same design as in the student-teacher setup. We use temperature and sampling as the sampling becomes more difficult for smaller temperature. This explains the discrepancy of the GP limit bias for different aw values in a (1.0, 1.2,
0.01 for both theory
C.3 STUDENT-TEACHER RESIDUAL-MLP
For the student-teacher task the residual MLP, we use the data matrix X generated from i.i.d unit Gaussian distribution, with input dimension No 1024, training sample size P 1280 and
23

---

Published as a conference paper at ICLR 2025
Citar20 Ceneralization Eroor vs
Crar10 Accuracy vs
HMC
1.15
0.790
Ener
1.0
0.785
1.0
0.700
0.90
Theory dm0.0
0.775
HMC sommalliing
1.859
Theory
0.80
HMC samoling 6-15
0.755
10
10P
10
10
10P
(a)
(b)
Figure 16: (a) Overall generalization error for residual-MLP on Cifar10. (b) Test accuracy on Ci- far10, calculated using the class label by averaging the predictions from HMC samples.
test sample size Prest 320. The teacher network weights are also generated with i.i.d Gaussian, with the hidden layer width N, 1024, hidden layer weight variance o 1, and readout branch variance 87 = (2.4, 0.4) for the two branches. The temperature is set as T = 0.001o w for a range of w values in dw (0.6, 0.8. 1.0, 1.2, 1.2, 1.5).
C.4 CIFAR10
For the Cifarl0 dataset, we filter the category of airplane and that of cat for the regression task. We use random split of training sample size 1000 and test sample size Ptest 1000. The input data is only from the red channel with input dimension No = 1024. The temperature is set 0.005o for a range of ow values o = {1.0, 1.2, 1.5]. Within the compute budget, we observe significant sample correlations, which has a large impact on measuring the bias term. This why we observe relatively large deviation of the bias results from theory Figure T5ta). a).
C.5 HAMILTONIAN MONTE CARLO
The sampling experiments in the paper are all done with Hamiltonian Monte Carlo (Betancour) (017); Neal (2012); Chandra et al. 220TT) simulations, popular method for sampling a proba- bility distribution. HMC has faster convergence the posterior distribution compared to Langevin dynamics. We used Numpyrho to set chains and run the simulations on the GPU cluster. All experiments are within GPU hours budget on Nvidia A100-40G. Due memory constraint, we only sampled up to N = 1024 hidden layer width for the student-teacher CSBM experiment and N = 64 for the Cora experiment. Since we mainly aim to demonstrate the narrow width effect in this paper, this suffices the purpose.
24


Reviews:
[{"review_id": "JZpbmYhQSP", "reviewer": "Reviewer_MLor", "summary": "In this work, the authors challenge the common belief that wider network widths improve generalization by studying the narrow-width limit of Bayesian Parallel Branching Graph Neural Networks (BPB-GNN). Unlike traditional models where increasing width leads to Gaussian Process behavior, authors show that BPB-GNNs with narrow widths exhibit stronger learning per branch due to symmetry breaking in kernel renormalization (when train data points are more compared to width). The series of experiments and theoretical analysis show that narrow-width BPB-GNNs achieve comparable or superior performance to wide-width models in bias-limited cases, with branch readout norms that reflect data properties rather than architectural hyperparameters.", "questions": "[1] did a similar study with other architectures. The only difference is that models are trained using SGD. Others [2-4] have focused on generalization and theoretical bounds. The authors should cite such work and comment on key differences between them.\n\nAdditional question to author, if we limit the width of GNN, the expressivity of model will reduce [5]. So, isn’t this the issue of learnability? In other words, even with wider widths, unstable models will struggle and have lower learnability; do the authors hypothesize that their approach leads to better stability and learnability? Similar to stable TM[6]. As it's obvious that expressivity would be reduced if you narrow the width, and even theoretically, the generalization would vary. In simple words current framework or architecture in no way is turing complete, it will be reduced to finite automata even with unbounded steps. Hence it would be ideal to mention how given approach trads off expressivity for better learnability and stability. Thus, it's important to mention these points and show comparison. \n\nThe authors don’t provide details about hyper-parameter optimization or detailed analysis regarding the statistical significance of the result. Thus, it is very difficult to gauge the overall importance of the result. For instance, the experimental setup does not mention any details about training, the Size of N = width of hidden layer, the weight dimension W_(l), the number of layers l, the size of readout neurons a_l, \n and the optimization steps (to update W and a). How are these learned and chosen? There should be information regarding this. p-value should be reported several trials to show results are statistically relevant. \n\nCan authors also comment on whether the results are true for both cases of homophilous and heterophilous graphs? \n\n\n1.\thttps://arxiv.org/pdf/2305.18411\n\n2.\thttps://arxiv.org/pdf/2207.09408\n\n3.\thttps://proceedings.mlr.press/v202/jaiswal23a/jaiswal23a.pdf\n\n4.\thttps://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=9535500 \n\n5.\thttps://arxiv.org/pdf/1907.03199\n\n6.\thttps://www.sciencedirect.com/science/article/abs/pii/S0020025523016201?via%3Dihub", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Good work showing narrow-width BPB-GNN showcases better robustness and improved generalization compared to infinite-width approaches.\n2. Theoretical justification supports empirical findings\n3. The experimental setup is good.\n4. Weakness is highlighted.", "weaknesses": "Mentioned below\n\n======= After Rebuttal=====\n\nAccept", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bI87NdfMx5", "reviewer": "Reviewer_i3Mq", "summary": "Existing works often study and characterize the infinite-width limit of neural networks. This work proposes to do the opposite of the study of the narrow width limit of neural networks, and the authors used Bayesian Parallel Branching Graph Neural Network (BPB-GNN) as an example. The authors show that when the width is small (as compared to the number of training examples), each branch of BPB-GNN exhibits more robust learning and empirically shows similar or better performance compared to the wide width limit when the bias term dominates in learning.", "questions": "Please address my comments under \"Weaknesses\".", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "I’d like to begin my review by stating the caveat that I am not an expert in this area and I am not very familiar with the existing works; while I have a high-level grasp of the key findings of the paper, I will have to defer the assessment of this work with respect the literature to an expert reviewer.\n\n- The results are pretty intriguing and challenge the common wisdom. As mentioned by the authors, previous works mainly focus on the analysis of the infinite-width limit; studying the narrow width limit is both interesting theoretically and can be potentially more practically useful, as one cannot scale neural networks to infinite width, so the results are mainly theoretical. In contrast, narrow-width can be attainable in real life.\n- The paper is largely well-written, with a clear presentation (there are several minor areas of improvement, though — see “Weaknesses”). Experiments are also conducted on real (albeit toy) datasets, which strengthen the rigor and confidence of the theoretical results.", "weaknesses": "1. My biggest concern is the choice of BPB-GNN, which seems to be a very specific construction that should be better motivated. There are several peculiarities in the chosen architecture: for example, there is a distinction between the “number of branches” and “width of network” as a consequence of the non-weight-sharing branches of BPB-GNN. As the authors show in Fig 5, the behavior of the neural network behaves very differently to different values of L and N. My concerns are: 1) Unlike BPB-GNN, the infinite-width limit that the authors have made extensive references to is not limited to a graph learning situation or GNNs. Although not a weakness per se, I am curious why the authors have deliberately chosen a variant of GNN on a semi-supervised task to perform their analysis rather than something more “vanilla” like CNNs or MLPs on simple, fully-supervised setups like image classification. As a result of my previous question, I wonder to what extent the results would generalize to other architectures with different architectures and levels of supervision. 2) As mentioned above and related to my previous point, the distinction of L and N, which will not be present in a “vanilla” architecture without the branching structure, is particular to BPB-GNN. I would appreciate some discussion on to what extent the results will apply in such a case when such a distinction does not exist.\n2. As the authors mentioned themselves in the limitations, it is unclear to what extent the results would be applicable when the variance term dominates,  which seems more likely for an over-parameterized network? Could it be the case that the robust learning phenomenon and the superiority over a wider network are caused by better regularization from a narrower network when the learning complexity is low, and that benefit will disappear for more complicated tasks? While it is good that the authors acknowledged some of these potential limitations themselves, I believe additional discussions will be beneficial.\n3. Presentation: There are some presentation issues like Fig 1. The different lines (especially the red and orange lines) are difficult to read against the histograms of the same color. The legend fonts are also too small.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1TQK9pxRzU", "reviewer": "Reviewer_MWE6", "summary": "This paper studies the narrow width limit of the Bayesian Parallel Branching Graph Neural Network (BPB-GNN), a summation of several different one-hidden-layer linear Graph Neural Networks. The authors consider the regime that $P, N \\rightarrow \\infty$ but the ratio of $P/N$ varies where P is the number of training samples and N is the width of the networks. They find that when $P/N$ is small, the network can have better performance.", "questions": "- $A^l$ is not defined. Are they all equal to $A$?\n- I don't quite understand how the integration is calculated in Sec 3.3. How do you linearize and integrate the parameters in Eq (25-29)?\n- In Eq. (10), $\\alpha$ should be canceled with $N/P$? I don't see why the entropy term dominates when $\\alpha \\rightarrow 0$ and why it is the NNGP kernel.\n\n---\nAfter rebuttal, raised the score from 3 to 6.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- this paper studies a new narrow-width regime for parallel branching networks using statistical physics tools and shows that the performance of a BPB-GNN in this narrow-width limit can perform better than its wide-width counterpart.", "weaknesses": "- The model considered in the paper is a linear model, which is very restricted considering the NTK and NNGP of nonlinear networks have been analyzed extensively.\n- one of the major results is just a conjecture, which is not proved rigorously.\n- The presentation of the paper is quite dense. There are many derivations in the main paper which are hard to follow and I am not sure if the derivations are rigorous. I would suggest the authors organize the main results and give more explanations and intuition for the results.\n- See other questions below", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1kDoDh8c4s", "reviewer": "Reviewer_fRMt", "summary": "This paper challenges the common belief that wider networks generalize better. The authors focus on the narrow-width regime, where the network's width is significantly smaller than the number of training samples, and demonstrate that BPB-GNNs can achieve comparable or even superior performance to wide networks in bias-limited scenarios. This improved performance is attributed to symmetry-breaking effects in kernel renormalization, which enhances the robustness of each branch in the network. The paper provides a theoretical framework for understanding the narrow-width limit and validates it with empirical results, including experiments on the Cora dataset. The findings suggest that narrow-width BPB-GNNs can offer efficient alternatives to wide networks, highlighting their potential for optimizing network design in various applications.", "questions": "- It is hard to see why $L=2$ directly reduces to the exactly to a 2-layer residual GCN\n\n- In the teacher model (Eq 20) why do you apply a Gaussian distribution to the adjacency matrix? Please correct me if I am wrong.\n\n- From your experiments shown in Figure 5, why don't you show the variance results? Additionally, how do you comment on the claim that increasing the depth can improve performance, despite the oversmoothing problem?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- It is interesting to see the authors introduce the Bayesian Parallel Branching GNN (BPB-GNN) architecture, which differs from previous neural tangent kernel research.\n\n- The analysis presented in Sections 3 and 4 is comprehensive.\n\n- A statistical physics view is applied to the graph neural network, which is novel and interesing.", "weaknesses": "- One main concern is why this work focuses on graph neural networks. It seems that the analysis could be applied to other neural networks, such as MLPs and CNNs. It is hard to see the uniqueness of graph neural networks here.\n\n- The rigor of this work is insufficient. For instance, the definition of $<>$ in Eq. 13 was not introduced.\n\n- To strengthen the experimental section, it would be valuable to include additional real-world datasets (e.g., PubMed, Citeseer) (or tasks beyond node classification).\n\n- To our knowledge, this is the first work to provide a tight bound on the generalization error for GNNs with residual-like structures. However, this may be an overclaim, as other works exist. It would be better to compare with the PAC-Bayes bound for Graph Neural Networks [1].\n\n- A missing reference: [2]\n\n[1] Liao, Renjie, Raquel Urtasun, and Richard Zemel. \"A pac-bayesian approach to generalization bounds for graph neural networks.\" arXiv preprint arXiv:2012.07690 (2020).\n\n[2] Huang, W., Li, Y., Du, W., Yin, J., Da Xu, R.Y., Chen, L. and Zhang, M., 2021. Towards deepening graph neural networks: A GNTK-based optimization perspective. arXiv preprint arXiv:2103.03113.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)