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

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

Original: AfﬁliationAddress email
Revised: AffiliationAddress 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 processisng 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 processing motifs.

Original: 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 = 
Revised: 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

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

Original: 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.
Revised: 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.

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

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

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

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

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

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

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

Original: 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.
Revised: 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.

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 inﬁnitely wide networks [8, 16, 21], where the input-output relation is equivalent to a
Revised: Many current theoretical approaches study neural network generalization in infinitely wide networks (3; 35; 34), where the input-output relation is equivalent to a

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

Original: 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.
Revised: 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.

Original: 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.
Revised: Another line of work focuseson finite-width deep linear neural networks (DLNNs)(29; 32; 31), while applicable in a wider regime,the generalization behavior of linear networks are very limited, and the bias contribution alwaysremains constant with network parameters (29), 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, 26, 29], 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(25; 28; 26), 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 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.
Revised: 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.

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 gatings are shared across allhidden layer units  while in previous work each units have its own set of gatings [10, 26, 29].
Revised: First, the gating units are shared acrosshidden layer units and different layers while in previous work each unit has its own set of gatings(13; 28; 26).

Original: Second,we deﬁne global learning objective instead of local errors [26, 29].
Revised: Second, we define global learning objective instead of local errors (28; 26).

Original: These simpliﬁcations allow usto obtain direct analytical expressions of memory capacity and exact generalization error of thesenetworks for arbitrary training and testing data, providinghigh 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].
Revised: Thesesimplifications allow us to obtain direct analytical expressions of memory capacity and exact generalization error of these networks for arbitrary training and testing data, providing quantitative insightinto the effect of learning in nonlinear networks, as opposed to studies of generalization bounds (13),expressivity estimates (26; 12), and indirect quantities relevant to generalization such as the implicitbias of the network (24).

Original: 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.
Revised: Furthermore, the kernel expression of the predictor statistics we propose inthis work also allow us to make qualitative explanations of the generalization and how it’s related todata structure and network representation for single and multiple tasks.

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

Original: 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.
Revised: 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.

Original: Second, we apply our theory tounderstand 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.
Revised: 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.

Original: While the effect of kernel renormalization is different in thetwo cases, we ﬁnd that decorrelation between tasks always improves generalization.
Revised: While the effect of kernel renormalization is different in the two cases, we find that forfixed gatings, de-correlation 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 deﬁned as follows,
Revised: In GGDLNs, the network input-output relation is defined 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 allneurons in the corresponding layer, mimicing the dendriticstructure of cortical neurons.
Revised: Neurons in each layer has M dendrites, each with an input-dependentglobal gating g m ( x ) shared across all neurons.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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 theGGDLN architecture is that their capacity canbe computed analytically in a straightforward way.
Revised: Memory Capacity: Memory capacity refers to the number of random samples and labels forwhich 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 (13; 27; 23). To calculate the capacity of

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

Original:  m [Equation] [Equation] ,j , which is a linear functionof the effective input
Revised: = (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

Original: [Equation]
Revised: ( 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.

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

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

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

Original: [Equation]
Revised: 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.

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

Original: While the total number of weights increase linearlywith L , the capacity increases as O ( L M  1) for large L , which is much faster.
Revised: 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).

Original: In the more generalcase of different sets of gatings in each layer, a similar argument yields P  N 0 M L .
Revised: 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.

Original: 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
Revised: Posterior distribution of network weights: We consider the posterior distribution of the networkweights induced by learning with a Gaussian prior (11; 10; 30; 8), given by

Original: [Equation]
Revised: [Equation]

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

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

Original: 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 priorterm on the network parameters⇥ =
Revised: 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

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

Original: In this work we focus on the T [-! 
-] limit where the ﬁrst data dependent term dominates.
Revised: In this work we focus on the T [-! 
-]→ 0 limit where the first 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 i,m , it is straightforward to integrate out the readoutweights in the partitionfunction , which yields
Revised: As thedistribution is quadratic in the readout weights a m,i , it is straightforward to integrate them out, whichyields

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, deﬁned as K µ ⌫ L ( W )
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, defined as K µνL

Original: =
Revised: [Equation]

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

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

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

Original: 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].
Revised: 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 definedas ⟨ K L ( W ) ⟩ W where W ∼ N (0 , σ 2 ) (35).

Original: For GGDLNs, the GP kernel for a pair of arbitrary datax and y is given by K GP ( x, y ) =
Revised: For GGDLNs, the GP kernel for a pair of arbitrarydata x 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 [-µ ⌫ 
-] = K 0 ( x µ , x ⌫ ) .
Revised: Wedenote the P × P kernel data matrix as K GP where K µνGP = K GP ( x µ , x ν ) , and the input kernelmatrix on training data as K 0 where K [-µ ⌫ 
-]µν 0 = K 0 ( x µ , x ν ) .

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

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

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

Original:  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 ) =
Revised: ⊤ K − 1 GP Y , and ⟨ δf ( x ) 2 ⟩ = K GP ( x , x ) − k GP ( x ) ⊤ K − 1 GP Y , wherek µGP ( x ) =

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

Original: 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].
Revised: 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).

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

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

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

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

Original: The partition function for a single hidden layer network is given by Z = exp (  H 1 ) ,where the Hamiltonian H 1 is given by
Revised: The partition function for a single hidden layer network is givenby 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 ˆ H 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 U 1 .

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

Original: [Equation]
Revised: [Equation]

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

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

Original: In the

Original:  M > 1 case, the effect of renormalization is more drastic as it changes that not only the amplitudebut also the shape of the kernel.
Revised: 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.

Original: The renormalization matrix has the interesting physical interpretationthat relates it to the readout weights a of GGDLNs,
Revised: The renormalizationmatrix has an interesting physical interpretation that 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).

Original: 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 predictorstatistics for a input x can be expressed in terms of the renormalized kernels, for an L = 1 network
Revised: The predictor statistics 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 ) > ˆ H 1 g ( y ))
Revised: = ( 1 M g ( x ) ⊤ U 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 P ⇥ 1vector, ˜ 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 Px 1 vector, ˜ k

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

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

Original: The kernel renormalization in GGDLNs changesboth 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.
Revised: The kernel renormalization in GGDLNs changes the shape ofthe kernel through the data dependent U 1 , reflecting the nonlinear property of the network, andresulting in more complex behavior of predictor statistics relative to the linear networks, as shown in Section 4.

Original: Our theory describes the properties of the posterior distribution of the networkweights induced at equilibrium bygradient-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).
Revised: Our theory describes the properties of the posterior distribution of the network weightsinduced at equilibrium by Langevin dynamics with the MSE cost function and the Gaussian prior(29; 7; 6; 5). Simulating this dynamics agrees remarkably well with the simulation (see Appendix B.2).

Original: 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).
Revised: We will compare our theoretical results with numerics of the gradient descent (GD) dynamicswidely used in machine learning. We consider the case where the network is initialized with Gaussiani.i.d. weights with variance σ 2 , and the mean and variance of the predictors are evaluated acrossmultiple initializations (see Appendix C.4 for details).

Original: 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.
Revised: As we will show, our theory makes accuratequalitative predictions for GD dynamics 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 is independent on these parameters and assumes the same value as  in the GP limit[17].
Revised: For linear networks the generalization error depends on N, σ 2 and L through the variance only,while the mean predictor always assumes the same value as the mean in the GP limit (29).

Original: 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 .
Revised: Thisis because the scalar kernel renormalization of ˜ k 1 ( x ) is cancelled out in the mean predictor by therenormalization of the inverse kernel ˜ K − 1 1 .

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

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

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

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 with the [Figure] choice of the gating functions.
Revised: This dependence also varies with thechoice of the gating functions.

Original: In Fig.2we consider  learning a ReLU teacher task in GGDLNs withgatings that have localized receptive ﬁelds .
Revised: In Fig.2we consider a student-teacher learning task (3; 2; 2; 9). Wepresent results of learning a ReLU teacher task in GGDLNs with gatings that have localized receptivefields (Fig.2(a) bottom), where the student GGDLN is required to learn the input-output relation of agiven ReLU teacher.

Original: 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
Revised: For a ReLU teacher with a single fully connected hidden layer (Fig.2(a) topleft), gatings with different receptive fields are of equal importance, hence the renormalization doesnot play a beneficial functional role, and the infinitely wide network performs better than finite N .

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

Original: 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)).
Revised: For a ’local’ ReLUteacher with larger input weights for one subset of input components (the preferred inputs, Fig.2(a)top right), renormalization improves task performance by the selective increase of the elementsof U 1 that correspond to gating units whose receptive fields overlap the teacher preferred inputs(Fig.2(d&e)).

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 boththe 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: Our theory not only captures the morecomplex behavior of generalization (especially bias) as a function of network width, but also providesqualitative explanation of how generalization is affected by the structure of the renormalization matrixin different tasks.

Original: 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.
Revised: 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: The dependence of error bias on σalso arises due to the matrix nature of the renormalization.

Original: 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)
Revised: In Fig.3 , we show parameter regimeswhere the bias can increase (Fig.3 (a-c)) or decrease (Fig.3 (d-f)) with σ on MNIST (4) (Appendix C.2 ). Although the dependence on σ is complicated and diverse, and we 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).

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

Original: 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.
Revised: 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 (29), wherestrong regularization ( σ = 0 ) always results in optimal generalization [Figure]

Original: b i a s
Revised: of σ .

Original: 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
Revised: ( d - f )

Original:  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
Revised: Bias decreases as a function of σ while variance increases, leading to an optimal σ with minimum error rate.

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 ﬁ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.
Revised: 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.

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

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

Original: 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.
Revised: Random gatings take the form g m ( x ) = Θ( 1 √ N 0 V ⊤ m x − b ) , where V m 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 unlabelledtraining dataset with unsupervised soft k-means clustering, such that the m -th gating
Revised: The pretrained gatings are trained on the unlabelled trainingdataset with unsupervised soft k-means clustering, such that the m -th gating

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

Original: 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 .
Revised: 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 .

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

Original: 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.
Revised: The result shows that GGDLNscan still 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

Original: 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 begin by analayzing the effect of ˆ H on the shape of the kernel in deeparchitectures.
Revised: We begin byanalyzing the effect of U 1 on the shape of the kernel in deep architectures.

Original: 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.
Revised: It is well known that the GP kernel of many nonlinear networks flattens as network depth increases (35), ultimately losinginformation about the input and degrading generalization performance.

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

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

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

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

Original: ˜ 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.
Revised: [Equation] [Equation] 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 deﬁned by the kernel ˜ K L ( x, y ) , and therefore K Lnorm ( x, y ) 2 [  1 , 1] .
Revised: 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] .

Original: For the GP kernel of GGDLNs, we have K Lnorm ( x, y )
Revised: Forthe 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 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 ( ✓ ) =
Revised: While K Lnormdepends on the specific choice of gatings in general, in the special case of random gatings withzero threshold g m ( 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 ( θ ) =(

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

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

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

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

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

Original: 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).
Revised: 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.

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

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

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

Original: [Figure]
Revised: [Figure]

Original: 5 L 0.

Original: 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 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.
Revised: 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.

Original: In both scenarios,  we demonstrate that decorrelation between tasksleads to better generalization performance.
Revised: 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.

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

Original: First we consider learning different tasks deﬁned by vastly different input statistics with bottom-upgatings, using permuted MNIST as an example.
Revised: First we consider learning different tasks defined by vastly different input statistics with bottom-upgatings, using permuted MNIST as an example.

Original: Previous works have shown that GLNs mitigatecatastrophic forgetting when sequentially trained on permuted MNIST [26]].{+mitigate [Figure] N= 
 [Figure]+}
Revised: Previous works have shown that GLNs {+mitigate [Figure] N= 
 [Figure]+}

Revised: [Figure]

Revised: r

Revised: [Figure]

Revised: t

Revised: i

Revised: [Figure]

Revised: a

Revised: [Figure]

Revised: c d

Revised: [Figure]

Revised: r r

Revised: [Figure]

Revised: catastrophic forgetting when sequentially trained on permuted MNIST (28).

Original: 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.
Revised: While our theory doesnot 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, andthe coefﬁcient of this linear combination, is given by the µ -th coefﬁcient 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,and the coefficient of this linear combination, is given by the µ -th coefficient of ˜ k

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

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

Original: [Equation]
Revised: [Equation]

Revised: = (cid:80)

Original: [Equation]
Revised: [Equation]

Original: ( 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.
Revised: ),where we assume there are P training examples and P t test data for each task, with a total of ntasks.

Original: 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]
Revised: 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 .

Original: [Figure]

Original: [Figure]

Original: r a t i o

Original: [Figure]

Original: rr o r r a t e a [Figure] 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))).
Revised: Stronger diagonal elements indicates that the networkseparates the processing of data of different tasks ((Fig.6(a)-(b))).

Original: 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)).
Revised: 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)).

Original: 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.
Revised: 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 threshold is large due to a set of less overlapping gatings activatedfor each task, generalization error first decreases and then increases again. 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

Original: [Figure]
Revised: [Figure]

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

Original: 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.
Revised: 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.

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 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.
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 forbiddengatings 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 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

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

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

Original: 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.
Revised: 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.

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

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

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

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

Original: 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.
Revised: 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.

Original: 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).
Revised: In Fig.7(b) bottom, we show the renormalized kernel with the renormalization matrix U 1 calculatedby solving Eq.6.

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

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

Original: 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.
Revised: 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)

Original: 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 .
Revised: 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))

Original:  Discussion
Revised:  Discussion

Original:  In this work, we proposed a novel gating network architecture, the GGDLN, amenable to rigoroustheoretical analysis of the network expressivity and generalization performance.
Revised: [Figure] 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 ﬂ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.
Revised: 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.

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

Original: While our theory is an exact descriptionofthe posterior of weight distribution induced by Langevin dynamics in Bayesian learning, it providessurprisingly well qualitative agreement with results obtained with GD dynamics for not onlythegeneralization but also the  kernel representation with matrix renormalization (Fig.7 (b)), largelyextending its applicability.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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].
Revised: 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 (22).

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

Original: 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].
Revised: 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 (19; 17; 21; 18; 20).

Original: 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.
Revised: 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) (16; 15; 14) , 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.

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

Original: Checklist
Revised: Checklist

Revised: The checklist follows the references.

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

Original: For all authors...
Revised: For each question, change the default [TODO] to [Yes] , [No] , or[N/A] . 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. For example: • Did you include the license to the code and datasets?

Revised: [Yes] See Section ?? .•

Revised: Did you include the license to the code and datasets?

Revised: [No] The code and the data areproprietary.•

Revised: Did you include the license to the code and datasets?

Revised: [N/A]

Revised: Please do not modify the questions and only use the provided macros for your answers.

Revised: Note that the

Revised: Checklist section does not count towards the page limit.

Revised: In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.

Revised: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’scontributions and scope?
Revised: (a) Do the main claims made in the abstract and introduction accurately reflect 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 6( c
Revised: [Yes] See Section 3 and 6(c)

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

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

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

Original: [Yes
Revised: [Yes]2.

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

Original: full set

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

Original: [Yes] See Section b) Did you include complete proofs 
Revised: [Yes] See Section 3(b) Did you include complete proofs of all theoretical results?

Original: of all

Original: theoretical

Original: results

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

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- mental results (either in the supplemental material or as a URL)?
Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

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

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?

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

Original:  Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
Revised: (c) Did you report error bars (e.g., with respect to the random seed after running experi- ments 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 C

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’reusing/curating?
Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?

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

Original:  Did you discuss whether the data you are using/curating contains personally identiﬁableinformation or offensive content?
Revised: (e) Did you discuss whether the data you are using/curating contains personally identifiable information 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, ifapplicable?
Revised: (a) Did you include the full text of instructions given to participants and screenshots, if applicable?

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 amountspent on participant compensation?
Revised: (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

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