Original: Globally Gated Deep Linear Networks
Revised: Globally Gated Deep Linear Networks

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AffiliationAddress email
Revised: AfﬁliationAddress email

Original: Abstract
Revised: Abstract

Original: Recently proposed Gated Linear Networks (GLNs) present a tractable nonlinearnetwork architecture, and exhibit interesting capabilities such as learning withlocal error signals and reduced forgetting in sequential learning.
Revised: Recently proposed Gated Linear Networks (GLNs) present a tractable nonlinearnetwork architecture, and exhibit interesting capabilities such as learning withlocal error signals and reduced forgetting in sequential learning.

Original: In this work,we introduce a novel gating architecture, named Globally Gated Deep Linear
Revised: In this work,we introduce a novel gating architecture, named Globally Gated Deep Linear

Original: Networks (GGDLNs) where gating units are shared among all processing unitsin each layer, thereby decoupling the architectures of the nonlinear but unlearnedgating and the learned linear processing motifs.
Revised: Networks (GGDLNs) where gating units are shared among all processing unitsin each layer, thereby decoupling the architectures of the nonlinear but unlearnedgating and the learned linear processisng motifs.

Original: We derive exact equations for thegeneralization properties of Bayesian Learning in these networks in the finite-widththermodynamic limit, defined by N, P → ∞ while P/N =O (1) where N and
Revised: Using statistical mechanics, wederive exact equations for the generalization properties of Bayesian Learning inthese networks in the ﬁnite-width thermodynamic limit, deﬁned by N, P ! 1while P/N =

Original:  P are the hidden layers’ width and size of training data sets respectfully.
Revised: O (1) where N and P are the hidden layers’ width and size oftraining data sets respectfully.

Original: We findthat the statistics of the network predictor can be expressed in terms of kernelsthat undergo shape renormalization through a data-dependent order parametermatrix compared to the infinite-width Gaussian Process (GP) kernels.
Revised: We ﬁnd that the statistics of the network predictorcan be expressed in terms of kernels that undergo shape renormalization through adata-dependent order-parameter matrix compared to the inﬁnite-width Gaussian Process (GP) kernels.

Original: Our theoryaccurately captures the behavior of finite width GGDLNs trained with gradientdescent (GD) dynamics.
Revised: Our theory accurately captures the behavior of ﬁnite width GGDLNs trained with gradient descent (GD) dynamics.

Original: We show that kernel shape renormalization gives rise torich generalization properties w.r.t.
Revised: We show that kernel shaperenormalization gives rise to rich generalization properties w.r.t.

Original: network width, depth and L 2 regularizationamplitude.
Revised: network width,depth and L 2 regularization amplitude.

Original: Interestingly, networks with a large number of gating units behavesimilarly to standard ReLU architectures.
Revised: Interestingly, networks with a large numberof gating units behave similarly to standard ReLU architectures.

Original: Although gating units in the model donot participate in the supervised learning, we show the utility of unsupervisedlearning of the gating parameters.
Revised: Although gatingunits in the model do not participate in the supervised learning, we show the utilityof unsupervised learning of the gating parameters.

Original: Additionally, our theory allows the evaluationof the network’s ability for learning multiple tasks by incorporating task-relevantinformation into the gating units.
Revised: Additionally, our theory allowsthe evaluation of the network capacity for learning multiple tasks by incorporatingtask-relevant information into the gating units.

Original: In summary, our work is the first exact theoreticalsolution of learning of a family of nonlinear networks with finite width.
Revised: In summary, our work is the ﬁrstexact theoretical solution of learning of a family of nonlinear networks with ﬁnitewidth.

Original: The richand diverse behavior of the GGDLNs suggests that they are helpful analyticallytractable models of learning single and multiple tasks, in finite-width nonlineardeep networks.
Revised: The rich and diverse behavior of the GGDLNs suggests that they are helpfulanalytically tractable models of learning single and multiple tasks, in ﬁnite-widthnonlinear deep networks.

Original:  Introduction
Revised:  Introduction

Original: Despite the recent advances in machine learning, theoretical understanding of how machine learningalgorithms work is very limited.
Revised: Despite the recent advances in machine learning, theoretical understanding of how machine learningalgorithms work is very limited.

Original: Many current theoretical approaches study neural network generalization in infinitely wide networks [1, 2, 3], where the input-output relation is equivalent to a
Revised: Many current theoretical approaches study neural network generalization in inﬁnitely wide networks [8, 16, 21], where the input-output relation is equivalent to a

Original: Gaussian Process (GP) in function space with a covariance matrix defined by a GP kernel.
Revised: Gaussian Process (GP) in function space with a covariance matrix deﬁned by a GP kernel.

Original: However,this GP limit holds when the network width approaches infinity while the size of the training dataremains finite, severely limiting its applicability to realistic conditions.
Revised: However,this GP limit holds when the network width approaches inﬁnity while the size of the training dataremains ﬁnite, severely limiting its applicability to realistic conditions.

Original: Another line of work focuseson finite-width deep linear neural networks (DLNNs)[4, 5, 6], while applicable in a wider regime,the generalization behavior of linear networks are very limited, and the bias contribution alwaysremains constant with network parameters [4], which fails to capture the behavior of generalizationperformance in general nonlinear networks.
Revised: Another line of work focuseson ﬁnite-width deep linear neural networks (DLNNs)[17, 25, 14], while applicable in a wider regime,the generalization behavior of linear networks are very limited, and the bias contribution alwaysremains constant with network parameters [17], which fails to capture the behavior of generalizationperformance in general nonlinear networks.

Original: Therefore, a tractable nonlinear network architecture is inneed for theoretically probing into the diverse generalization behavior of general nonlinear networks.
Revised: Therefore, a tractable nonlinear network architecture is inneed for theoretically probing into the diverse generalization behavior of general nonlinear networks.

Original: Recently proposed Gated Linear Networks (GLNs) present a tractable nonlinear network architecture[7, 8, 9], with capabilities such as learning with local error signals and mitigating catastrophicforgetting in sequential learning.
Revised: Recently proposed Gated Linear Networks (GLNs) present a tractable nonlinear network architecture[7, 26, 29], with capabilities such as learning with local error signals and mitigating catastrophicforgetting in sequential learning.

Original: Inspired by these recent advances in GLNs, we propose Globally
Revised: Inspired by these recent advances in GLNs, we propose Globally

Original: Gated Deep Linear Networks (GGDLNs) as a simplified GLN structure that preserves the nonlinearproperty of general GLNs, the decoupling of fixed nonlinear gating from learned linear processingunits and the ability to separate the processing of multiple tasks using the gating units.
Revised: Gated Deep Linear Networks (GGDLNs) as a simpliﬁed GLN structure that preserves the nonlinearproperty of general GLNs, the decoupling of ﬁxed nonlinear gating from learned linear processingunits and the ability to separate the processing of multiple tasks using the gating units.

Original: Our GGDLNstructure is different from previous GLNs in several ways.
Revised: Our GGDLNstructure is different from previous GLNs in several ways.

Original: First, the gating units are shared acrosshidden layer unitsand different layers while in previous work each unit has its own set of gatings [10,8, 9].
Revised: First, the gatings are shared across allhidden layer units while in previous work each units have its own set of gatings [10, 26, 29].

Original: Second, we define global learning objective instead of local errors [8, 9].
Revised: Second,we deﬁne global learning objective instead of local errors [26, 29].

Original: These simplificationsallow us to obtain direct analytical expressions of memory capacity and exact generalization error ofthese networks for arbitrary training and testing data, providing  quantitative insight into the effectof learning in nonlinear networks, as opposed to studies of generalization bounds [10], expressivityestimates [9, 11], and indirect quantities relevant to generalization such as the implicit bias of thenetwork [12].
Revised: These simpliﬁcations allow usto obtain direct analytical expressions of memory capacity and exact generalization error of thesenetworks for arbitrary training and testing data, providing high resolution and quantitative insightinto the effect of learning in nonlinear networks, as opposed to studies of generalization bounds [10],expressivity estimates [29, 28], and indirect quantities relevant to generalization such as the implicitbias of the network [18].

Original: Furthermore, the kernel expression of the predictor statistics we propose in this workalso allow us to make qualitative explanations of the generalization and how it’s related to datastructure and network representation for single and multiple tasks.
Revised: Furthermore, the kernel expression of the predictor statistics we propose inthis work allow us to make qualitative explanations of the generalization and how it’s related to datastructure and kernel representation.

Original: First, we introduce the architecture of our GGDLNs and analyze its memory capacity.
Revised: First, we introduce the architecture of our GGDLN and analyze its memory capacity.

Original: We then deriveour theory for generalization properties of GGDLNs, and make qualitative connections betweenthe generalization behavior and the relation between the renormalization matrix and task structure.
Revised: We thenuse statistical mechanics to derive generalization properties of GGDLNs, and make qualitativeconnections between the generalization behavior and the relation between the renormalization matrixand task structure.

Original: Second, we apply our theory to  GGDLNs performing multiple tasks, focusing on two scenarioswhere tasks are defined by different input statistics and different output labels on the same inputsrespectively.
Revised: Second, we apply our theory to understand how GGDLNs perform multiple tasks,focusing on two scenarios where tasks are deﬁned by different input statistics and different outputlabels on the same inputs respectively.

Original: While the effect of kernel renormalization is different in the two cases, we find that forfixed gatings, de-correlation between tasks always improves generalization.
Revised: While the effect of kernel renormalization is different in thetwo cases, we ﬁnd that decorrelation between tasks always improves generalization.

Original:  Globally gated deep linear networks
Revised:  Globally gated deep linear networks

Original: [Figure]
Revised: [Figure]

Original: In GGDLNs, the network input-output relation is defined as follows,
Revised: In GGDLNs, the network input-output relation is deﬁned as follows,

Original: [Equation]
Revised: [Equation]

Original: where x 0 = x is the input, N is the hidden layer width, M is the number of gating units in each layer,and N 0 is the input dimension.
Revised: where x 0 = x is the input, N is the hidden layer width, M is the number of gating units in each layer,and N 0 is the input dimension.

Original: Neurons in each layer has M dendrites, each with an input-dependentglobal gating g m ( x ) shared across all  neurons.
Revised: Neurons in each layer has M dendrites, each with an input-dependentglobal gating g m ( x ) shared across all neurons in the corresponding layer, mimicing the dendriticstructure of cortical neurons.

Original: The m -th dendritic branch of neuron i in the L -thhidden layer connects to neurons in the previous layer with a dendrite-specific weight vector W mL,i (orwith readout weight vector a m for the output neuron), as shown in Fig.1 (a).
Revised: The m -th dendritic branch of neuron i in the L -th hidden layer connectsto neurons in the previous layer with a dendrite-speciﬁc weight vector W mL,i (or with readout weightvector a m for the output neuron), as shown in Fig.1 (a).

Original: Note that although thegatings are fixed during learning, changes in the weights affects how these gatings acts on the hiddenlayer activations, and it is interesting to understand how the learned task interacts with these gatingoperations.
Revised: Note that although the gatings are ﬁxedduring learning, changes in the weights affect how these gatings acts on the hidden layer activations,and it is interesting to understand how the learned task interacts with these gating operations.

Original: Since adding gatings at the input layer is equivalent to expanding the input dimensionand replacing x j by x j g m ( x ) , and learning does not affect how the gatings interact with the input,we do not add gatings at the input layer for simplicity.
Revised: Sinceadding gatings at the input layer is equivalent to expanding the input dimension and replacing x j byx j g m ( x ) , and learning does not affect how the gatings interact with the input, we do not add gatingsat the input layer for simplicity.

Original: Memory Capacity: Memory capacity refers to the number of random samples and labelsforwhich there exists a set of parameters such that the network achieves zero training error (herewe consider the mean squared error, MSE) and the capacity bounds of deep nonlinear networkshas been extensively studied in many recent works [10, 14, 15]. To calculate the  capacity of
Revised: Memory Capacity: Memory capacity refers to the number of random samples and labels a networkcan memorize, and the capacity bounds of deep nonlinear networks has been extensively studied inmany recent works [10, 30, 31]. One advantage of the GGDLN architecture is that their capacity canbe computed analytically in a straightforward way.

Original: GGDLNs, note that the input-output relation given by Eq.1 can be alternatively expressed asf ( x ) 
Revised: The input-output relation given by Eq.1 can bealternatively expressed as f ( x ) = [Equation]

Original: = (cid:80) m 1 , ··· ,m L ,j W eff m 1 , ··· ,m L ,j x eff m 1 , ··· ,m L ,j , which is a linear combination of the effective inputx eff m 1 , ··· ,m L ,j = g m 1 ( x ) g m 2 ( x ) · · · g m L ( x ) x j
Revised: m [Equation] [Equation] ,j , which is a linear functionof the effective input

Original: ( m l = 1 , · · · , M ; j = 1 , · · · , N 0 ), with some effectiveweights W eff which is a complicated function of a and W . Here m l is the index of the gatings inthe l -th layer.
Revised: [Equation]

Original: As the gating units are shared across layers, the effective inputx eff has N 0 (cid:0) M + L − 1 L (cid:1)independent dimensions.
Revised: As the gating units are shared across layers, the effective input has N [Equation] independent dimensions.

Original: This combinatorial term is the number of possible distinct products of Lobjects chosen froma pool of M different objects.
Revised: Here the combinatorial term comes from the number of possible combinations of choosing Lgatings from M total number of gatings with repeat.

Original: Assuming N ≫ M L such that the effective weight
Revised: Assuming N  M L , the capacity is

Original: W eff m 1 , ··· ,m L ,j can take any desired value in the N 0 M L dimensional whole space, so the problem offinding W eff with zero training error is equivalent to a linear regression problem with input x effand the target outputs.
Revised: [Equation]

Original: Therefore, the capacity isequivalent to the number of independent inputdimensions, given by P ≤ N 0 (cid:0)
Revised: where P is the number of training examples.

Original: M + L − 1  L (cid:1) . The above capacity is verified by Fig.1(c-e), showing that the behavior of generalization performancechanges drastically at network capacity, where generalization error and its bias contribution diverge,and the variance contribution shrinks to 0 (see detailed calculation in the next paragraph and Appendix A.3).
Revised: While the total number of weights increase linearlywith L , the capacity increases as O ( L M  1 ) for large L , which is much faster.

Original: This double descent property of the generalization error is similar to previously studied in linearand nonlinear networks. Furthermore, although the output of the network is a linear function of theeffective input x eff , due to the multiplicative nature of the network weights and the gatings, learningin GGDLNs is highly nonlinear and the space of solution for W and a is highly nontrivial, and thenetwork exhibit properties unique to nonlinear networks, as we will show in the following sections.
Revised: In the more generalcase of different sets of gatings in each layer, a similar argument yields P  N 0 M L .

Original: Posterior distribution of network weights: We consider the posterior distribution of the networkweights induced by learning with a Gaussian prior [16, 17, 18, 19], given by
Revised: Posterior distribution of network weights: We consider the posterior distribution of the networkweights induced by learning with a Gaussian prior in a Bayesian learning framework [3, 1, 5, 9],given by

Original: [Equation]
Revised: [Equation]

Original: where Z is the partition function 
Revised: where Z is the partition function Z =

Original: Z = (cid:82) d Θ P ( Θ ) .
Revised: R d ⇥ P ( ⇥ ) .

Original: The first term in the exponent is the mean-squareddeviation of the network outputs on a set of P training data points x µ from their target outputs Y µ ,and the second term is a Gaussian prior  on the network parameters 
Revised: The ﬁrst term in the exponent is the mean-squareddeviation of the network outputs on a set of P training data points x µ from their target outputs y µ , andthe second term is the Gaussian prior term on the network parameters ⇥ =

Original: Θ = { W , a } with amplitudeσ − 2 .
Revised: { W , a } with amplitude  2 .

Original: In this work we focus on the T → 0{+! 
+} limit where the first data dependent term dominates.
Revised: In this work we focus on the T {+! 
+} limit where the ﬁrst data dependent term dominates.

Original: Belowthe network capacity, the distribution of Θ concentrates onto the space of solution that yields zerotraining error, the Gaussian prior then biases the solution space towards weights with smaller L 2norms.
Revised: Belowthe network capacity, the distribution of ⇥ concentrates onto the space of solution that yields zerotraining error, the Gaussian prior then biases the solution space towards weights with smaller L 2norms.

Original: The fundamental properties of the system can be derived from the partition function.
Revised: The fundamental properties of the system can be derived from the partition function.

Original: As thedistribution is quadratic in the readout weights a m,i , it is straightforward to integrate them out, whichyields
Revised: As thedistribution is quadratic in the readout weights a i,m , it is straightforward to integrate out the readoutweights in the partitionfunction , which yields

Original: [Equation]
Revised: [Equation]

Original: where W denotes all the remaining weights in the network, and K L ( W ) is the W dependent P × Pkernel on the training data, defined as K µνL
Revised: where W denotes all the remaining weights in the network, and K L ( W ) is the W dependent P ⇥ Pkernel on the training data, deﬁned as K µ ⌫ L ( W )

Original: [Equation]
Revised: =

Original:  2 M g ( x µ ) ⊤ g ( x ν ))
Revised: (  2 M g ( x µ ) > g ( x ⌫ ))

Original: ( 1 N x µL ( W ) ⊤ x νL ( W )) .
Revised: ( 1 N x µL ( W ) > x ⌫ L ( W )) , whereg ( x ) is an M dimensional vector.

Original: Generalization in infinitely wide GGDLNs: It is well known that in infinitely wide networks where
Revised: Generalization in inﬁnitely wide GGDLNs: It is well known that in inﬁnitely wide networks where

Original: N → ∞ while P remains finite (also referred to as the GP limit), K L ( W ) is self-averaging and doesnot depend on the specific realization of W . It can therefore be replaced by the GP kernel defined as⟨ K L ( W ) ⟩ W where W ∼ N (0 , σ 2 ) [2].
Revised: N ! 1 while P remains ﬁnite (also referred to as the GP limit), K L ( W ) is self-averaging and doesnot depend on the speciﬁc realization of W , and can therefore be replaced by the GP kernel deﬁnedas h K L ( W ) i W where W ⇠ N (0 ,  2 ) [16].

Original: For GGDLNs, the GP kernel for a pair of arbitrary data xand y is given by K GP ( x , y ) =
Revised: For GGDLNs, the GP kernel for a pair of arbitrary datax and y is given by K GP ( x, y ) =

Original: ( σ 2 M g ( x ) ⊤ g ( y ))
Revised: (  2 M g ( x ) > g ( y ))

Original: L K 0 ( x , y ) , where K 0 ( x , y ) =
Revised: L K 0 ( x, y ) , where K 0 ( x, y ) =

Original: σ 2 N 0 x ⊤ y .
Revised:  2 N 0 x > y .

Original: We denotethe P × P kernel data matrix  as K GP where K µνGP = K GP ( x µ , x ν ) , and the input kernel matrixon training data as K 0 where K µν 0{+µ ⌫ 
+} = K 0 ( x µ , x ν ) .
Revised: We denotethe P ⇥ P kernel data matrix , as K GP where K µ ⌫ GP = K GP ( x µ , x ⌫ ) , and the input kernel matrixon training data as K 0 where K {+µ ⌫ 
+} = K 0 ( x µ , x ⌫ ) .

Original: Generalization error measured by MSE including the bias and the mean, ϵ g 
Revised: Generalization error measured by mean squared error (MSE) including the bias and the variance,✏ g =

Original: = ( ⟨ f ( x ) ⟩ Θ − y ( x )) 2 +⟨ δf{+
 + h  f+} ( x ) 2 ⟩ Θ , which depends on the first and second order statistics  of  the predictor.
Revised: ( h f ( x ) i ⇥  y ( x )) {+
 + h  f+} ( x ) 2 i ⇥ , which depends on the ﬁrst and second order statistics ofthe predictor, averaged across the distribution of parameters ⇥ in the solution space.

Original: In the GPlimit, we have ⟨ f ( x ) ⟩ = k GP ( x ) 
Revised: In the GPlimit, we have h f ( x ) i = k GP ( x ) >

Original: ⊤ K − 1 GP Y , and ⟨ δf ( x ) 2 ⟩ = K GP ( x , x ) − k GP ( x ) ⊤ K − 1 GP Y , wherek µGP ( x ) =
Revised: K  1 GP Y , and h  f ( x ) 2 i = K GP ( x, x )  k GP ( x ) > K  1 GP Y , where 112 k µGP ( x ) =

Original: K GP ( x , x µ ) .
Revised: K GP ( x, x µ ) .

Original: Note that the rank of K GP is the same as the capacity of the network, andthe kernel matrix becomes singular as P approaches its capacity (the interpolation threshold), which results in a diverging bias and vanishing variance (Fig.1 (b-d)). The singularity of the kernel at theinterpolation threshold holds even for finite width networks, and similar diverging bias and vanishingvariance are seen in our finite width theory below (Section 3 ) and are confirmed by simulation ofnetworks trained with GD (see Appendix B.1).
Revised: Note that the rank of the kernel is the same as the capacity of the network,and the kernel matrix becomes singular as P approaches its capacity, also known as the interpolationthreshold, which results in a diverging bias and vanishing variance, as shown in Fig.1 The singularityof the kernel at the interpolation threshold holds even for ﬁnite width networks, and similar divergingbias and vanishing variance are seen in our ﬁnite width theory below (Section 3 ) and are conﬁrmedby simulation of networks trained with GD (see Appendix B). This double descent property of thegeneralization error is similar to previously studied in linear and nonlinear networks [6].

Original:  Kernel shape renormalization theory in finite-width GLNs
Revised:  Kernel shape renormalization theory in ﬁnite-width GLNs

Original: We now address the finite width thermodynamic limit, where P, N → ∞ but P/N ∼ O M, L ∼ O (1) .
Revised: We now address the ﬁnite width thermodynamic limit, where P, N ! 1 but P/N ⇠ O (1) .

Original: In this limit, calculating the statistics of the network predictor requires integration over W in Eq.3 .
Revised: In thislimit, calculating the statistics of the network predictor requires integration over W in Eq.4 .

Original: To do so, we apply the previous method of Back-propagating Kernel Renormalization[4] (see Appendix A) to GGDLNs.
Revised: To do so,we apply the previous method of Backpropagating Kernel Renormalization [17] (see Appendix A) tothe GGDLNs.

Original: The partition function for a single hidden layer network is givenby Z = exp( − H 1 ) , where the Hamiltonian H 1 is given by
Revised: The partition function for a single hidden layer network is given by Z = exp (  H 1 ) ,where the Hamiltonian H 1 is given by

Original: [Equation]
Revised: [Equation]

Original: Comparing the matrix ˜ K 1 to K GP , we note that the GP kernel is renormalized by an an M × Mmatrix order parameter U 1 .
Revised: Comparing the matrix ˜ K 1 to K GP , we note that the GP kernel is renormalized by an an M ⇥ Mmatrix order parameter ˆ H 1 .

Original: This order paramter satisfies the self-consistent equation
Revised: This order parameter satisﬁes the self-consistent equation

Original: [Equation]
Revised: [Equation]

Original: (6) where ◦ denotes element-wise multiplication.
Revised: where  denotes elementwise multiplication, and g denotes a P ⇥ M dimensional matrix.

Original: In the linear case (which corresponds to M = 1) , the GP kernel is renormalized by a scalar factor.
Revised: In thelinear case (which corresponds to M = 1) , the GP kernel is renormalized by a scalar factor.

Revised: In the

Original: In the M > 1 case, the effect of renormalization is moredrastic as it changes that not only the amplitude but also the shape of the kernel.
Revised: M > 1 case, the effect of renormalization is more drastic as it changes that not only the amplitudebut also the shape of the kernel.

Original: The renormalizationmatrix has an interesting physical interpretation that relates it to the readout weights a of GGDLNs,
Revised: The renormalization matrix has the interesting physical interpretationthat relates it to the readout weights a of GGDLNs,

Original: [Equation]
Revised: [Equation]

Original: The calculation can be extended to multiple layers with a new order parameter introduced for eachlayer(see Appendix A).
Revised: The calculation can be extended to multiple layers with a new order parameter introduced for eachlayer(see Appendix A).

Revised: However, the extension is only applicable to networks of moderate depthin practice due to the exponential growth of the size of the renormalization matrix.

Original: The predictor statistics for a input x can be expressed in terms of the renormalized kernels, for an L = 1 network
Revised: The predictorstatistics for a input x can be expressed in terms of the renormalized kernels, for an L = 1 network

Original: [Equation]
Revised: [Equation]

Original: where ˜ K 1 ( x , y )
Revised: where ˜ K 1 ( x, y )

Original: = ( 1 M g ( x ) ⊤ U 1 g ( y ))
Revised: = ( 1 M g ( x ) > ˆ H 1 g ( y ))

Original: K 0 ( x , y ) denotes the renormalized kernel for two arbitraryinputs x and y , ˜ K 1 denotes the P × P renormalized kernel on the training data, and ˜ k 1 ( x ) is the Px 1 vector, ˜ k
Revised: K 0 ( x, y ) denotes the renormalized kernel for two arbitraryinputs x and y , ˜ K 1 denotes the P ⇥ P renormalized kernel on the training data, and ˜ k 1 ( x ) is the P ⇥ 1vector, ˜ k

Original: 1 µ ( x ) =
Revised: 1 µ ( x ) =

Original: K ( x , x µ ) .
Revised: K ( x, x µ ) .

Original: The kernel renormalization in GGDLNs changes  the shape of thekernel through the data dependent U 1 , reflecting the nonlinear property of the network, and resultingin more complex behavior of predictor statistics relative to the linear networks, as shown in Section 4.
Revised: The kernel renormalization in GGDLNs changes both the magnitudeand the shape of the kernel through the data dependent ˆ H , reﬂecting the nonlinear property of thenetwork, and resulting in more complex behavior of predictor statistics relative to the linear networksas shown in Section 4.

Original: Our theory describes the properties of the posterior distribution of the network weights induced atequilibrium by  Langevin dynamics with the MSE cost function and the Gaussian prior [4, 20, 21,22]. Simulating this dynamics agrees remarkably well with the simulation (see Appendix B.2).
Revised: Our theory describes the properties of the posterior distribution of the networkweights induced at equilibrium by gradient-based learning under Langevin dynamics with the MSEcost function and the Gaussian prior [17, 22, 32, 24], and simulating this dynamics agrees remarkablywell with the simulation (see Appendix B).

Original: Wewill compare our theoretical results with numerics of the gradient descent (GD) dynamics widelyused in machine learning. We consider the case where the network is initialized with Gaussian i.i.d.weights with variance σ 2 , and the mean and variance of the predictors are evaluated across multipleinitializations (see Appendix C.4 for details).
Revised: We will compare our theoretical results with numericsof a simple gradient descent (GD) dynamics, where the network is initialized with Gaussian i.i.d.weights with variance  2 , and the mean and variance of the predictors are evaluated across multipleinitializations (see Appendix C for details).

Original: As we will show, our theory makes accurate qualitativepredictions for GD dynamics in all examples throughout this paper.
Revised: While we do not expect our theory to agree with GDdynamics quantitatively, as we will show, our theory makes accurate qualitative predictions for GDdynamics in all examples throughout this paper.

Original:  Generalization
Revised:  Generalization

Original: For linear networks the generalization error depends on N, σ 2 and L through the variance only, whilethe mean predictor always assumes the same value asthe mean in the GP limit [4].
Revised: For linear networks the generalization error depends on N,  2 and L through the variance only, whilethe mean predictor is independent on these parameters and assumes the same value as in the GP limit[17].

Original: This is because thescalar kernel renormalization of ˜ k 1 ( x ) is cancelled out in the mean predictor by the renormalizationof the inverse kernel ˜ K − 1 1 .
Revised: This is because the scalar kernel renormalization of ˜ k 1 ( x ) is cancelled out in the mean predictorby the renormalization of the inverse kernel ˜ K  1 1 .

Original: In contrast, for GGDLNs the mean predictor and hence the error bias alsochange with these network parameters due to the matrix nature of  the kernel renormalization(Eq.8) .
Revised: In contrast, for GGDLNs the mean predictorand hence the error bias also change with these network parameters due to the matrix nature of therenormalization of the kernel (Eq.8) .

Original: Below we investigate in detail how matrix renormalization of the kernel affects the generalizationbehavior (especially the bias term) of the network.
Revised: Below we investigate in detail how matrix renormalization ofthe kernel affects the generalization behavior (especially the bias term) of the network.

Original: [Figure]
Revised: 4.1 Networks with single hidden layer

Original: Dependence on network width N : Unlike in DLNs, the bias term in GGDLNs depends on N ,exhibiting different dependence in different parameter regimes.
Revised: Dependence on network width N : Unlike in DLNs, the bias term in GGDLNs depends on N ,exhibiting different dependence in different parameter regimes.

Original: This dependence also varies withthe choice of the gating functions.
Revised: This dependence also varies with the [Figure] choice of the gating functions.

Original: In Fig.2we considera student-teacher learning task [23, 24,25, 26]. We present results of learning a ReLU teacher task in GGDLNs with gatings that havelocalized receptive fields (Fig.2(a) bottom), where the student GGDLN is required to learn the inputoutput relation of a given ReLU teacher.
Revised: In Fig.2we consider learning a ReLU teacher task in GGDLNs withgatings that have localized receptive ﬁelds .

Original: For a ReLU teacher with a single fully connected hiddenlayer (Fig.2(a) top left), gatings with different receptive fields are of equal importance, hence therenormalization does not play a beneficial functional role, and the infinitely wide network performsbetter than finite N .
Revised: For a ReLU teacher with a single fully connected hiddenlayer, gatings with different receptive ﬁelds are of equal importance, hence the renormalization doesnot play a beneﬁcial functional role, hence the inﬁnitely wide network performs better than ﬁnite

Original: As shown in Fig.2(b), bias, variance and generalization error all decrease with

Original: N .
Revised: N .

Revised: As shown in Fig.2(b), bias, variance and generalization error all decrease when N increases.

Original: For a ’local’ ReLU teacher with larger input weights for one subset of input components (thepreferred inputs, Fig.2(a) top right), renormalization improves task performance by the selectiveincrease of the elements of U 1 that correspond to gating units whose receptive fields overlap theteacher preferred inputs (Fig.2(d&e)).
Revised: Fora ’local’ ReLU teacher with larger input weights for one subset of input components (the preferredinputs, as illustrated in Fig.2(a)), renormalization improves task performance by the selective increaseof the elements of ˆ H that correspond to gating units whose receptive ﬁelds overlap the teacherpreferred inputs (as shown in Fig.2(d)).

Original: Hence, narrower networks (with a stronger renormalization)generalize better, and both the bias and the generalization error increase with N (Fig.2(c)).
Revised: Hence, narrower networks (with a stronger renormalization)generalize better, and both the bias and the generalization error increase with N (Fig.2(c)).

Original: Ourtheory not only captures the more complex behavior of generalization (especially bias) as a functionof network width, but also provides qualitative explanation of how generalization is affected by thestructure of the renormalization matrix in different tasks.
Revised: Ourtheory not only captures the more complex behavior of generalization (especially bias) as a functionof network width, but also provides qualitative explanation of how generalization is affected by thestructure of the renormalization matrix in different tasks.

Original: Dependence on weight regularization strength σ 2 : Similar to the dependence on N , generalizationalso exhibits different behavior as a function of the regularization parameter σ in different parameter regimes, with contributions from both the bias and the variance.
Revised: Dependence on weight regularization strength  2 : Similar to the dependence on N , generalizationalso exhibit different behavior as a function of  in different parameter regimes, with contributionsfrom both the bias and the variance.

Original: The dependence of error bias on σalso arises due to the matrix nature of the renormalization.

Original: In Fig.3 , we show parameter regimeswhere the bias can increase (Fig.3 (a-c)) or decrease(Fig.3 (d-f)) with σ on MNIST [13] (Appendix C.2 ). Although the dependence on σ is complicated and diverse, andwe lack a general rule forwhen the qualitative behavior changes, we found that our theory accurately captures the qualitativebehavior of results obtained from GD (Appendix B.3 Fig.3).
Revised: In Fig.3 , we show parameter regimes where the bias canincrease or decrease with  , and results obtained from GD agree qualitatively well with our theory.(Classiﬁcation accuracy is estimated by approximating the statsitics of the predictor as Gaussian withmean and variance given by Eq.8)

Original: In both regimes the variance increaseswith σ as the solution space expands for a weaker regularization.
Revised: In both cases the variance increases with  as the solution spaceexpands for a weaker regularization (large  ).

Original: Specifically in Fig.3 (d-f), due tothe increasing variance (e) and decreasing bias (d), there is a minimum error rate[-((f), Appendix A.
 Eq.49 for how error rate is calculated from the mean and variance of the predictor)-] at intermediate σ ,indicating an optimal level of regularization strength as opposed to linear networks [4], where strongregularization ( σ = 0 ) always results in optimal generalization [Figure]
Revised: Speciﬁcally in Fig.3 (b), due to the increasing varianceand decreasing bias, there is a minimum error rate[-((f), Appendix A.
 Eq.49 for how error rate is calculated from the mean and variance of the predictor)-] at intermediate  , indicating an optimal levelof regularization strength as opposed to linear networks [17], where strong regularization (  = 0 )always results in optimal generalization. a b 0.

Original: of σ .
Revised: b i a s

Original: ( d - f )
Revised: GDtheory 0. 0.4  v a r i an c e 10GDtheory 0. GDtheory e rr o r r a t e b i a s v a r i an c e e rr o r r

Original: Bias decreases as a function of σ while variance increases, leading to an optimal σ with minimum error rate.
Revised: a t e 0.2 0.4 0.6 0.8 0. GDtheory 0.2 0.4 0.6 0.80. GDtheory 0.2 0.4 0.6 0.8 0. GDtheory

Original: Comparable performance to ReLU networks:
Revised: Comparable performance to ReLU networks:

Original: The nonlinear operation of the gatings enables thenetwork to learn nonlinear tasks.
Revised: The nonlinear operation of the gatings enables thenetwork to learn nonlinear tasks.

Original: In Fig.4, we show that although the gatings are fixed during training,the network achieves comparable performance as a fully trained nonlinear (ReLU) network with thesame hidden layer width for classifying even and odd digits in MNIST data when M is sufficientlylarge.
Revised: In Fig.4, we show that although the gatings are ﬁxed during training,the network achieves comparable performance as a fully trained nonlinear (ReLU) network with thesame hidden layer width for classifying even and odd digits in MNIST data when the number ofgatings are sufﬁciently large.

Original: Furthermore, although the gatings are fixed during the supervised training of the GGDLN, they can becleverly chosen to improve generalization performance.
Revised: Furthermore, although the gatings are ﬁxed during the supervised training of the GGDLN, theycan be cleverly chosen to improve generalization performance.

Original: To demonstrate this strategy, we comparedtwo different choices of gatings.
Revised: To demonstrate this strategy, wecompared two different choices of gatings.

Original: Random gatings take the form gm ( x ) = Θ( 1 √ N 0 V ⊤ m x − b ) , where Vm is a 1 × N 0 random vector with standard Gaussian i.i.d.
Revised: Random gatings take the form g ( x ) = ⇥ ( 1 p N 0 V > x  b ) ,where V is a 1 ⇥ N 0 random vector with standard Gaussian i.i.d.

Original: elements, b is a scalar threshold, andΘ( x ) is the heaviside step function.
Revised: elements, b is a scalar threshold,and ⇥ ( x ) is the heaviside step function.

Original: The pretrained gatings are trained on the unlabelled trainingdataset with unsupervised soft k-means clustering, such that the m -th gating
Revised: The pretrained gatings are trained on the unlabelledtraining dataset with unsupervised soft k-means clustering, such that the m -th gating

Original: g m ( x ) outputs theprobability of assigning data x to the m -th cluster (Appendix C.2).
Revised: g m ( x ) outputsthe probability of assigning data x to the m -th cluster (see Appendix C).

Original: As shown in Fig.4, for pretrainedgatings, generalization performance improves with M much faster compared to random gatings, andapproaches the performance of ReLU network at a smaller M .
Revised: As shown in Fig.4, forpretrained gatings, generalization performance improves with M much faster compared to randomgatings, and approaches the performance of ReLU network at a smaller M .

Original: Our theory (Fig.4) and numericalresults of GD dynamics (Appendix B.3 Fig.4) agree qualitatively well.
Revised: Our theory (Fig.4 bottom)and numerical results of GD dynamics (Fig.4 bottom) agree qualitatively well.

Original: The result shows that GGDLNscan still achieve competitive performance on nonlinear tasks while remaining theoretically amenable.
Revised: The result shows thatalthough the nonlinearity of the gated linear network is constrained by a ﬁxed set of gatings, it canstill achieve competitive performance on nonlinear tasks while remaining theoretically amenable.

Original: [Figure]
Revised: [Figure]

Original: 4.2 Kernel shape renormalization in deeper networks
Revised: 4.2 Kernel shape renormalization in deeper networks

Revised: So far, we have focused on L = 1 .

Original: We now consider the effect of the matrix renormalization on GGDLN with more layers.
Revised: We now consider the effect of the matrix renormalization on GGDLN with more layers.

Original: We beginby analyzing the effect of U 1 on the shape of the kernel in deep architectures.
Revised: We begin by analayzing the effect of ˆ H on the shape of the kernel in deeparchitectures.

Original: It is well known thatthe GP kernel of many nonlinear networks flattens as network depth increases [2], ultimately losinginformation about the input and degrading generalization performance.
Revised: It is well known that the GP kernel of many nonlinear networks ﬂattens as networkdepth increases [16], ultimately losing information about the input and degrading generalizationperformance.

Original: Here we show that kernelshape renormalization slows down flattening of kernels by incorporating data relevant informationinto the learned weights.
Revised: Here we show that kernel shape renormalization slows down ﬂattening of kernels byincorporating data relevant information into the learned weights.

Original: To study the shape of the kernel independent of kernel magnitude, we define the normalized kernel
Revised: To study the shape of the kernel independent of kernel magnitude, we deﬁne the normalized kernel

Original: K Lnorm ( x , y ) =
Revised: K Lnorm ( x, y ) =

Original: [Equation]
Revised: ˜ K L ( x,y )

Original: [Equation] [Equation] 2 , where ˜ K L ( x , y ) denotes the renormalized kernel for GGDLNwith L hidden layers.
Revised: ˜ K L ( x,x ) 1 / 2 ˜ K L ( y,y ) 1 / 2 , where ˜ K L ( x, y ) denotes the renormalized kernel for GGDLNwith L hidden layers.

Original: This normalized kernel measures the cosine of the vectors x and y withgeneralized inner product defined by the kernel ˜ K L ( x , y ) , and therefore K Lnorm ( x , y ) ∈ [ − 1 , 1] .
Revised: This normalized kernel measures the cosine of the vectors x and y withgeneralized inner product deﬁned by the kernel ˜ K L ( x, y ) , and therefore K Lnorm ( x, y ) 2 [  1 , 1] .

Original: Forthe GP kernel of GGDLNs, we have K Lnorm ( x , y )
Revised: For the GP kernel of GGDLNs, we have K Lnorm ( x, y )

Original: =
Revised: =

Original: cos( g ( x ) , g ( y ))
Revised: cos( g ( x ) , g ( y ))

Original: L cos( x , y ) .
Revised: L cos( x, y ) .

Original: While K Lnormdepends on the specific choice of gatings in general, in the special case of random gatings withzero threshold gm ( x ) = Θ( 1 √ N 0 V ⊤ m x ) and the number of gatings M → ∞ , we can write K Lnormanalytically as a function of the angle θ between input vectors x and y , given by K Lnorm ( θ ) =(
Revised: While K Lnormdepends on the speciﬁc choice of gatings in general, in the special case of random gatings with zerothreshold g ( x ) = ⇥ ( 1 p N 0 V > x ) and the number of gatings M ! 1 , we can write K Lnorm analytically as a function of the angle ✓ between input vectors x and y , given by K Lnorm ( ✓ ) =

Original: π − θπ ) L cos( θ ) .
Revised: ( ⇡  ✓⇡ ) L cos( ✓ ) .

Original: Thus, K Lnorm ( θ ) converges to a δ -function as L goes to infinity, shrinking to zeroexcept for θ = 0 .
Revised: Thus, K Lnorm ( ✓ ) converges to a  -function as L goes to inﬁnity, shrinking to zero except for ✓ = 0 .

Original: This ’flattening’ effect also empirically holds true for networks with finite M (see Appendix B.4).
Revised: This ’ﬂattening’ effect also empirically holds true for networks with ﬁnite M (see Appendix B).

Original: In Fig.5 , we study the effect of kernel renormalization on the ’flattening’ effect for deep networkstrained to separate even and odd digits in MNIST dataset.
Revised: In Fig.5 , we study the effect of kernel renormalization on the ’ﬂattening’ effect for deep networkstrained to separate even and odd digits in MNIST dataset.

Original: As shown in Fig.5 (a)-(c), the elementsof the renormalized kernel shrink to zero at a much slower rate compared to the GP kernel.
Revised: As shown in Fig.5 (a)-(c), the elements ofthe renormalized kernel shrink to zero at a much slower rate compared to the GP kernel.

Original: (Notethat unlike the variance, that bias is affected only by shape changes, but not by changes intheamplitude of the  kernel, in Fig.5(d) we plot only the bias contribution to the generalization.)
Revised: (Note thatunlike the variance, that bias is affected only by shape changes, but not by changes in the amplitude ofthe kernel, in Fig.5(d) we plot only the bias contribution to the generalization.)

Original: Whilemitigating the flattening of the GP kernel is a general feature of our renormalized kernel for differentparameters, its effect on the generalization performance (especially the bias) may be different fordifferent network parameters.
Revised: While mitigating theﬂattening of the GP kernel is a general feature of our renormalized kernel for different parameters (see Appendix B), its effect on the generalization performance (especially the bias) may be different fordifferent network parameters (Appendix B).

Original: In the specific example in Fig.
Revised: In the speciﬁc example in Fig.

Original: B.4, finite width networks with a less’flattened’ renormalized kernel achieve better performance than the GP limit.
Revised: B, ﬁnite width networkswith a less ’ﬂattened’ renormalized kernel achieve better performance than the GP limit.

Original: Both the GP limit andthe finite width networks have optimal performance at L = 2 in this example.
Revised: Both the GPlimit and the ﬁnite width networks have optimal performance at L = 2 in this example.

Original: [Figure]
Revised: [Figure]

Revised: 5 L 0.

Revised: b i a s a b c d

Original:  GGDLNs for multiple tasks
Revised:  GGDLNs for multiple tasks

Original: In this section, we apply our theory to investigate the ability of GGDLNs to perform multipletasks.
Revised: In this section, we apply our theory to investigate the ability of GGDLNs to perform multipletasks.

Original: We consider two different scenarios below.
Revised: We consider two different scenarios below.

Original: First, different tasks require the network to learninput-output mappings on input data with different statistics.
Revised: First, different tasks require the network to learninput-output mappings on input data with different statistics.

Original: This scenario corresponds to real lifesituations where the training data distribution is non-stationary.
Revised: This scenario corresponds to real lifesituations where the training data distribution is non-stationary.

Original: The tasks can be separated basedon difference in the inputs alone, without any additional top-down information.
Revised: The tasks can be separated basedon difference in the inputs alone, without any additional top-down information.

Original: In this case, thegatings are bottom-up, and are functions of the input data only.
Revised: In this case, thegatings are bottom-up, and are functions of the input data only.

Original: In the second case, different tasksgive conflicting labels for the same inputs, corresponding to the situation where performing the twotasks require additional top-down contextual information, and the information can be incorporatedinto the gating units in GGDLNs.
Revised: In the second case, different tasksgive conﬂicting labels for the same inputs, corresponding to the situation where performing the twotasks require additional top-down contextual information, and the information can be incorporatedinto the gating units in GGDLNs.

Original: In both scenarios,when the gatings are fixed and we modulate thede-correlation by changing network width and thus the strength of the kernel renormalization, wedemonstrate that de-correlation between tasks always leads to better generalization performance.
Revised: In both scenarios, we demonstrate that decorrelation between tasksleads to better generalization performance.

Original: 5.1 Bottom-up gating units
Revised: 5.1 Bottom-up gating units

Original: First we consider learning different tasks defined by vastly different input statistics with bottom-upgatings, using permuted MNIST as an example.[-Previous works have shown that GLNs mitigate [Figure] N= 
-]
Revised: First we consider learning different tasks deﬁned by vastly different input statistics with bottom-upgatings, using permuted MNIST as an example.[-Previous works have shown that GLNs mitigate [Figure] N= 
-]

Original: [Figure]

Original: [Figure]

Original: r a t

Original: i o a

Original: [Figure]

Original: c d

Original: [Figure] r r [Figure] catastrophic forgetting when sequentially trained on permuted MNIST [8].
Revised: Previous works have shown that GLNs mitigatecatastrophic forgetting when sequentially trained on permuted MNIST [26]].

Original: While our theory does notaddress directly the dynamics of sequential learning, we aim to shed light on this question indirectlyby asking how the two tasks interfere with each other when they are learned simultaneously.
Revised: While our theorydoes not address directly the dynamics of sequential learning, we aim to shed light on this questionindirectly by asking how the two tasks interfere with each other when they are learned simultaneously.

Original: We introduce a measure of inter-task interference by noting that after learning the mean predictoron a new data x , Eq.8, is a linear combination of the output labels Y µ of all the training data,and the coefficient of this linear combination, is given by the µ -th coefficient of ˜ k
Revised: We introduce a measure of inter-task interference by noting that after learning the mean predictoron a new data x , Eq.8, is a linear combination of the output labels y µ of all the training data, andthe coefﬁcient of this linear combination, is given by the µ -th coefﬁcient of ˜ k

Original: ( x ) ⊤ ˜ K − 1 .
Revised: ( x ) > ˜ K  1 .

Original: Thus,we define a task-task correlation matrix, via C 
Revised: Thus, wedeﬁne a task-task correlation matrix, via C pq = P

Original: [Equation]
Revised: [Equation]

Original: = (cid:80)

Original: [Equation]
Revised: [Equation]

Original: ),where we assume there are P training examples and P t test data for each task, with a total of ntasks.
Revised: ( p, q = 1 , · · · , n ), wherewe assume there are P training examples and P t test data for each task, with a total of n tasks.

Original: The amplitude of each element C pq measures how much training data of task q contributeto the prediction on the test data of task p .
Revised: The amplitude of each element C pq measures how much training data of task q contribute to the [Figure] r a t i o - [Figure]

Original: Stronger diagonal elements indicates that the networkseparates the processing of data of different tasks ((Fig.6(a)-(b))).
Revised: [Figure] [Figure]

Original: As we show in Fig.6, we cantune the relative strength of the diagonal elements of C smoothly by changing the network width(Fig.6(a)-(d)) or by changing the threshold of the gating (Fig.6 (e)-(f)).
Revised: r a t i o [Figure] rr o r r a t e a [Figure]

Original: In the case where the gatingsare fixed and the network width is changed, an increase in the strength of the diagonal elements(Fig.6(d)) results in better generalization (Fig.6(c)), indicating that the network generalizes betterby processing  data of different tasks more separately through the gating units. However, in the casewhere we change the activation of the gatings by adjusting the threshold, although different tasksare more de-correlated when the thresholdis large due to a set of less overlapping gatings activatedfor each task, generalization error first decreases and then increases again.
Revised: c d e f prediction on the test data of task p . Stronger diagonal elements indicates that the network separatesthe processing of data of different tasks ((Fig.6(a)-(b))). As we show in Fig.6, we can tune the relativestrength of the diagonal elements of C smoothly by changing the network width (Fig.6(a)-(d)) or bychanging the threshold of the gating (Fig.6 (e)-(f)).

Original: This is because for largethreshold the sparsity of the gatings activated for each task limits the nonlinearity of  the network, andtherefore the generalization performance on this nonlinear task. 5.2 Combined top-down and bottom-up gating units
Revised: In both cases an increase in the strength of thediagonal elements (Fig.6(d)&(f)) results in better generalization ((Fig.6(c)&(e))), indicating that thenetwork generalizes better by processing data of different tasks more separately through the gatingunits.

Original: [Figure]
Revised: [Figure]

Original: We now consider learning two tasks that provide conflicting labels on the same input data.
Revised: We now consider learning two tasks that provide conﬂicting labels on the same input data.

Original: The gatingunits combine both top-down task specific signal which informs the system of which task to perform for a given input, and bottom-up signals which, as before, depend on the input.
Revised: The gatingunits combine both top-down task speciﬁc signal which informs the system of which task to performfor a given input, and bottom-up signals which, as before, depend on the input.

Original: In different tasks,different sets of gatings are permitted or forbidden depending on the top-down signal, then the statesof the permitted gatings are further determined as a function of the input x , while the forbiddengatings are set to 0 , meaning that the corresponding dendritic branches do not connect to the previouslayer neurons (Fig.7(a)) in this task.
Revised: In different tasks,different sets of gatings are permitted or forbidden depending on the top-down signal, then the statesof the permitted gatings are further determined as a function of the input x , while the forbidden gatings are set to 0 , meaning that the corresponding dendritic branches do not connect to the previouslayer neurons (Fig.7(a)) in this task.

Original: For a single hidden layer network, with a similar argument asin Section 2, it is straightforward to show that the number of different tasks that can be memorizedis given by n ≤ M and the number of training examples for each task needs to satisfy P ≤ N 0 M p ,where M p is the number of permitted gating units in each task.
Revised: For a single hidden layer network, with a similar argument as in Section 2, it is straightforward to show that the number of different tasks that can be memorized isgiven by n  M and the number of training examples for each task needs to satisfy P  N 0 M perm ,where M perm is the number of permitted gating units in each task.

Original: In the limiting case where a set ofnon-overlapping gating units are permitted in each of the n tasks, the network is equivalent to nsub-networks, each independently performing one task.
Revised: In the limiting case where a setof non-overlapping gating units are permitted in each of the n tasks, the network is equivalent ton sub-networks, each independently performing one task.

Original: In this case M p is limited by M/n , which inturn limits the capacity and the effective input-output nonlinearity for each independent task.
Revised: In this case M perm is limited by M/n ,which in turn limits the capacity and the effective input-output nonlinearity for each independent task.

Original: Weconsider the case where the permitted gatings are chosen randomly for each task and are therefore ingeneral overlapping across tasks.
Revised: We consider the case where the permitted gatings are chosen randomly for each task and are thereforein general overlapping across tasks.

Original: We then investigate, through the data-dependent renormalizationmatrix, how learning modifies the correlation induced by the overlapping gatings.
Revised: We then investigate, through the data-dependent renormalizationmatrix, how learning modiﬁes the correlation induced by the overlapping gatings.

Original: As an example weconsider training on permuted and un-permuted MNIST digits of 0 and 1 ’s. One task is to performbinary classification of the two digits in both permuted and un-permuted data, and the second task isto separate the permuted digits (both 0 and 1) from the un-permuted digits.
Revised: As an example weconsider training on permuted and unpermuted MNIST digits of 0 and 1 ’s. One task is to performbinary classiﬁcation of the two digits in both permuted and unpermuted data, and the second taskis to separate the permuted digits (both 0 and 1) from the unpermuted digits.

Original: The labels of the two tasksare entirely uncorrelated, while the permitted gatings of the two tasks are partially overlapping.
Revised: The labels of the twotasks are entirely uncorrelated, while the permitted gatings of the two tasks are partiallyoverlapping.

Original: Inthis case the renormalized kernel ˜ K 1 can be written as ˜ K pµ,qν =
Revised: In this case the renormalized kernel ˜ K 1 can be written as ˜ K pµ,q ⌫ =

Original: ( 1 M g p ( xµ ) ⊤ U 1 g  q ( xν )) σ 2 N 0 x µ ⊤ x ν .
Revised: ( 1 M g µp ( x ) > ˆ H 1 g ⌫ q ( x ))  2 N 0 x µ > x ⌫ .

Original: Here p, q ∈ { 1 , 2 } is the task index, and µ , ν = 1 , · · · , P are the input index.
Revised: Here p, q 2 { 1 , 2 } is the task index, and µ , ⌫ = 1 , · · · , P are the input index.

Original: The kernel is therefore2 P × 2 P as shown in Fig.7(b); the diagonal blocks (green) correspond to kernels of task 1 and task 2,while the off diagonal blocks (red) correspond to the cross kernels.
Revised: The kernel is therefore 281 2 P ⇥ 2 P as shown in Fig.7(b) ( P = 600 ); the diagonal blocks (green) correspond to kernels of task1 and task 2, while the off diagonal blocks (red) correspond to the cross kernels.

Original: In Fig.7(b) bottom, we show the renormalized kernel with the renormalization matrix U 1 calculatedby solving Eq.6.
Revised: In Fig.7(b) top and middle panels, we show the renormalized kernel with the renormalization matrixˆ H 1 calculated by solving Eq.6 (middle).

Original: Similar results are achieved by by numerically estimating the readout weights fromresults obtained from GD dynamics by using Eq.
Revised: Similar results are achieved by by numerically estimatingthe readout weights from results obtained from GD dynamics by using Eq.

Revised: 7 (top) (for N = 100and N = 500 ).

Original: The results demonstrate that stronger kernelrenormalization achieved in narrower networks suppresses more strongly the correlation betweentasks, reflected by the weaker off-diagonal blocks in Fig.7(b) 
Revised: The results demonstrate that stronger kernel renormalization achieved in narrowernetworks suppresses more strongly the correlation between tasks, reﬂected by the weaker off-diagonalblocks in Fig.7(b) top and middle panels.

Original: A decreasing ratio between the averageamplitudes of the diagonal and off-diagonal blocks show that the de-correlation effect diminishes forlarge N , leading to increasing generalization error with N (Fig.7(c & d))
Revised: A decreasing ratio between the average amplitudes of thediagonal and off-diagonal blocks show that the decorrelation effect diminishes for large N (Fig.7(b)bottom left), leading to increasing generalization error with N .

Original:  Discussion
Revised:  Discussion

Original: [Figure] In this work, we proposed a novel gating network architecture, the GGDLN, amenable to rigoroustheoretical analysis of the network expressivity and generalization performance.
Revised: In this work, we proposed a novel gating network architecture, the GGDLN, amenable to rigoroustheoretical analysis of the network expressivity and generalization performance.

Original: The predictor statistics of GGDLNs can be expressed in terms of kernels that undergo shape renormalization, resultingdiverse behavior of the bias as a function of various network parameters.
Revised: The predictor statistics of GGDLNs can be expressed in terms of kernels that undergo shape renormalization, resultingdiverse behavior of the bias as a function of various network parameters.

Original: This renormalization slowsdown the flattening of the GP kernel in deep networks, suggesting that the loss of input informationas L increases may be prevented in finite-width nonlinear networks.
Revised: This renormalization slowsdown the ﬂattening of the GP kernel in deep networks, suggesting that the loss of input informationas L increases may be prevented in ﬁnite-width nonlinear networks.

Original: Wealso investigate the capabilityof GGDLNs to perform multiple tasks.
Revised: We investigate the capabilityof GGDLNs to perform multiple tasks, and demonstrated that decorrelation between tasks helpsimprove generalization performance in two scenarios.

Original: While our theory is an exact description of the posterior ofweight distribution induced by Langevin dynamics in Bayesian learning, it provides surprisingly wellqualitative agreement with results obtained with GD dynamics for not only the generalization but alsothe kernel representation with matrix renormalization (Fig.7 (b)), largely extending its applicability.
Revised: While our theory is an exact description ofthe posterior of weight distribution induced by Langevin dynamics in Bayesian learning, it providessurprisingly well qualitative agreement with results obtained with GD dynamics for not only thegeneralization but also the kernel representation with matrix renormalization (Fig.7 (b)), largelyextending its applicability.

Original: While our theory currently addresses learning in GGDLNs using a global cost function, exploringthe probability of extending the formalization of the equilibrium distribution to characterize locallearning dynamics by using the properly normalized conditional probability of W l on W k ( k < l )is an ongoing work.
Revised: In our theory, the size of the renormalization matrix increases exponentially with L . Therefore, developing methods to efﬁciently solve for the self-consistent equations in deep networks is ongoing work.

Original: Recent works have shown that multilayer perceptrons (MLPs) with learnedgatings that implements spatial attention have surprisingly good performance on Natural Language Processing (NLP) and computer vision [27].
Revised: Recent works have shown that multilayer perceptrons (MLPs) with learned gatings that implementsspatial attention have surprisingly good performance on Natural Language Processing (NLP) andcomputer vision [19].

Original: Extension of our theory to learnable gatings thatimplements attention mechanisms remains to be explored.
Revised: Extension of our theory to learnable gatings that implements attention mechanisms remains to be explored.

Original: Furthermore, our theory applies to fullyconnected linear networks with gating units, while other types of neural network architectures suchas convolutional neural networks are commonly used for image processing [28, 29, 30, 31, 32].
Revised: Furthermore, our theory applies to fully connected linear networkswith gating units, while other types of neural network architectures such as convolutional neuralnetworks are commonly used for image processing [27, 23, 12, 11, 2].

Original: Incorporating convolutional architecture into our GGDLNs and using the gating units to encodecontext-dependent modification of different feature maps is an interesting direction related to thefast-developing research topic of visual question-answering (VQA) [33, 34, 35] , where answeringdifferent questions about the same image is similar to performing multiple tasks in different contextswith different labels on the same dataset, as we discussed in Section 5.2.
Revised: Incorporating convolutionalarchitecture into our GGDLNs and using the gating units to encode context-dependent modiﬁcationof different feature maps is an interesting direction related to the fast-developing research topic ofvisual question-answering (VQA) [4, 13, 20] , where answering different questions about the sameimage is similar to performing multiple tasks in different contexts with different labels on the samedataset, as we discussed in Section 5.2.

Original: We leave these excitingresearch directions for future work.
Revised: We leave these exciting research directions for future work.

Original: Checklist
Revised: Checklist

Original: The checklist follows the references.

Original: Please read the checklist guidelines carefully for information onhow to answer these questions.

Original: For each question, change the default [TODO] to [Yes] , [No] or[N/A] .

Original: You are strongly encouraged to include a justification to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.

Original: For example:

Original: • Did you include the license to the code and datasets?

Original: [Yes] See Section ?? .

Original: • Did you include the license to the code and datasets?

Original: [No] The code and the data are

Original: proprietary.

Original: • Did you include the license to the code and datasets?

Original: [N/A]

Original: Please do not modify the questions and only use the provided macros for

Original: For all authors...
Revised: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
Revised: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’scontributions and scope?

Original: [Yes] (b) Did you describe the limitations of your work?
Revised: [Yes](b) Did you describe the limitations of your work

Original: [Yes] See Section 3 and 
Revised: [Yes] See Section 3 and 6(

Original: (c) Did you discuss any potential negative societal impacts of your work?
Revised: c ) Did you discuss any potential negative societal impacts of your work

Original: [N/A]
Revised: N/A](d)

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
Revised: Have you read the ethics review guidelines and ensured that your paper conforms tothem

Original: [Yes]
Revised: [Yes

Original: If you are including theoretical results...
Revised: If you are including theoretical results a Did you state the

Revised: full set

Original: (a) Did you state the full set of assumptions of all theoretical results?
Revised: of assumptions of all theoretical results?

Original: [Yes] See Section
Revised: [Yes] See Section

Revised: b)

Revised: Did

Original: (b) Did you include complete proofs of all theoretical results?
Revised: you include complete proofs of all theoretical results

Original: [Yes] See Appendix A
Revised: [Yes] See Appendix A

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi-
Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experi-

Original: mental results (either in the supplemental material or as a URL)?
Revised: mental results (either in the supplemental material or as a URL)?

Original: [Yes] In supplementary
Revised: [Yes] In supplementary

Original: material
Revised: material

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters,
Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters,

Original: how they
Revised: how they

Original: were chosen)?
Revised: were chosen)?

Original: [Yes] See Appendix C
Revised: [Yes] See Appendix C(c)

Original: (c) Did you report error bars (e g., with respect to the random seed after running experiments multiple times)?
Revised: Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original: [Yes] See Fig.
Revised: [Yes] See Fig.

Original: (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
Revised: (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?

Original: [Yes] See Appendix C4.
Revised: [Yes] See Appendix C4.

Original: If you are using existing assets (e g., code, data, models) or curating/releasing new assets...
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?
Revised: (a) If your work uses existing assets, did you cite the creators?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you mention the license of the assets?
Revised: (b) Did you mention the license of the assets?

Original: [Yes] Appendix C
Revised: [Yes] Appendix C

Original: (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [No]
Revised: [No]

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating?

Original: [N/A]
Revised: [N/A](e)

Original: (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
Revised: Did you discuss whether the data you are using/curating contains personally identiﬁableinformation or offensive content?

Original: [N/A]
Revised: [N/A]

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: If you used crowdsourcing or conducted research with human subjects...

Original: (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
Revised: (a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?

Original: [N/A]
Revised: [N/A]

Original: (b) Did you describe any potential participant risks, with links to Institutional Review
Revised: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board (IRB) approvals, if applicable?
Revised: Board (IRB) approvals, if applicable?

Original: [N/A]
Revised: [N/A]

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
Revised: (c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?

Original: [N/A]
Revised: [N/A]