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

Original: Anonymous Author(s)

Original: 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 nonlinear network architecture, and exhibit interesting capabilities such as learning with local 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 units in each layer, thereby decoupling the architectures of the nonlinear but unlearned gating and the learned linear processing 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: We derive exact equations for the generalization properties of Bayesian Learning in these networks in the ﬁnite-width thermodynamic limit, deﬁned by N, P ! 1 while P/N = O (1) where N and

Original: P are the hidden layers’ width and size of training data sets respectfully.
Revised: P are the hidden layers’ width and size of training 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 predictor can be expressed in terms of kernels that undergo shape renormalization through a data-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 shape renormalization 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 number of gating units behave similarly to standard ReLU architectures.

Original: Although gating units in the model donot participate inthe supervised learning, we show the utility of unsupervisedlearning of the gating parameters.
Revised: Although gating units in the model do not participate in supervised learning, we show the utility of 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 allows the evaluation of the network’s ability for learning multiple tasks by incorporating task-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 ﬁrst exact theoretical solution of learning in a family of nonlinear networks with ﬁnite width.

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 helpful analytically tractable models of learning single and multiple tasks, in ﬁnite-width nonlinear 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 learning algorithms 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 Gaussian Process (GP) in function space with a covariance matrix defined by a GP kernel.
Revised: Many current theoretical approaches study inﬁnitely wide networks [1, 2, 3], where the input-output relation is equivalent to a 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 data remains ﬁ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 focuses on ﬁnite-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 always remains constant with network parameters [4], which fails to capture the behavior of generalization performance 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 in need for theoretically probing into the diverse generalization behavior of general nonlinear networks. 36th

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, 8, 9], with capabilities such as learning with local error signals and mitigating catastrophic forgetting 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 nonlinear property of general GLNs, the decoupling of ﬁxed nonlinear gating from learned linear processing units, 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 GGDLN structure is different from previous GLNs in several ways.

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

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 [8, 9].

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 us to obtain direct analytical expressions of memory capacity and exact generalization error of these networks for arbitrary training and testing data, providing quantitative insight into the effect of learning in nonlinear networks, as opposed to studies of generalization bounds [10], expressivity estimates [9, 11], and indirect quantities relevant to generalization such as the implicit bias of the network [12].

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 in this work also allow us to make qualitative explanations of the generalization and how it’s related to data structure and network representation for single and multiple tasks.

Original: First, we introduce the architecture of our GGDLNs and analyze its memory capacity.
Revised: First, we introduce the architecture of our GGDLNs 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 then derive our theory for generalization properties of GGDLNs, and make qualitative connections between the generalization behavior and the relation between the renormalization matrix and 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 GGDLNs performing multiple tasks, focusing on two scenarios where tasks are either deﬁned by different input statistics or different output labels on the same inputs.

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 the two cases, we ﬁnd that for ﬁxed gating functions, 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 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: Each neuron in every layer has M dendrites, each with an input-dependent global gating g m ( x ) shared across all 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 connects to neurons in the previous layer with a dendrite-speciﬁc weight vector W mL,i (or with readout weight vector 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 ﬁxed during learning, changes in the weights affect how these gatings act 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: Since adding gatings at the input layer is equivalent to expanding the input dimension and 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 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 [10, 14, 15].
Revised: Memory Capacity: Memory capacity refers to the maximum number of random (or generic) inputoutput examples for which there exists a set of parameters such that the network achieves zero training error (here we consider the mean squared error, MSE). By deﬁnition, it is irrespective of the learning algorithm. The capacity bounds of deep nonlinear networks has been extensively studied in many recent works [10, 13, 14].

Original: To calculate the capacity of GGDLNs, note that the input-output relation given by Eq.1 can be alternatively expressed asf ( x ) 
Revised: To calculate the capacity of GGDLNs, note that the input-output relation given by Eq.1 can be alternatively expressed as f ( x ) =

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: [Equation] which is a linear combination of the effective input x e ↵ m 1 , · ·· ,m L ,j = g m 1 ( x ) g m 2 ( x ) · [Equation]

Original:  ( m l = 1 , · · · , M ; j = 1 , · · · , N 0 ), with some effectiveweights W eff which is a complicated function of a and W .
Revised: g m L ( x ) x j ( m l = 1 , · · · , M ; j = 1 , · · · , N 0 ), with some effective weights W e ↵ which is a complicated function of a and W .

Original: Here m l is the index of the gatings inthe l -th layer.
Revised: Here m l is the index of the gatings in the l -th layer.

Original: 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.
Revised: As the gating units are shared across layers, the effective input x e ↵ 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: This combinatorial term represents the number of possible combinations of L gatings selected from M total number of gatings.

Original: Assuming N ≫ M L such that the effective weight 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: Assuming N  M L such that the effective weight W e ↵ m 1 , ··· ,m L ,j can take any desired value in the N 0 M L dimensional whole space, the problem of ﬁnding W e ↵ with zero training error is equivalent to a linear regression problem with input x e ↵ and the target outputs.

Original: Therefore, the capacity is equivalent to the number of independent inputdimensions, given by P ≤ N 0 (cid:0)
Revised: Therefore, the capacity is equivalent to the number of independent input dimensions, given by P  N M [Equation]

Original: M + L − 1 L (cid:1) .
Revised: [Equation] L  .

Original: 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: The above capacity is veriﬁed by Fig.1(b), where the training error becomes nonzero above the memory capacity. The generalization behavior also changes drastically at network capacity (Fig.1(ce)), 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: This double descent property of the generalization error is similar to previously studied in linearand nonlinear networks.
Revised: This double descent property of the generalization error is similar to previously studied in linear and nonlinear networks.

Original: 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: Furthermore, although the output of the network is a linear function of the effective input x e ↵ , due to the multiplicative nature of the network weights and the gatings, learning in GGDLNs is highly nonlinear and the space of solution for W and a is highly nontrivial, and the network exhibit properties unique to nonlinear networks, as we will show in the following sections.

Revised: Posterior distribution of network weights: We consider a Bayesian network setup, where the network weights are random variables whose statistics are determined by the training data and network parameters, instead of deterministic variables.

Revised: This probabilistic approach enables us to study the properties of the entire solution space instead of a single solution which may be heavily initialization dependent.

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: We consider the posterior distribution of the network weights induced by learning with a Gaussian prior [15, 16, 17, 18], 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 MSE 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 → 0{+! 
+} limit where the first data dependent term dominates.
Revised: In this work we focus on the T {+! 
+} limit where the ﬁrst term dominates.

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

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 the distribution is quadratic in the readout weights a m,i , it is straightforward to integrate them out, 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 ⇥ P kernel on the training data, deﬁned as K µ ⌫ L ( W ) =

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

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 .
Revised: N ! 1 while P remains ﬁnite (also referred to as the GP limit), K L ( W ) is self-averaging and does not depend on the speciﬁc realization of W .

Original: It can therefore be replaced by the GP kernel defined as⟨ K L ( W ) ⟩ W where W ∼ N (0 , σ 2 )[2].
Revised: It can therefore be replaced by the GP kernel deﬁned as h K L ( W ) i W where W ⇠ N (0 ,  2 )

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 data 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 µν 0 = K 0 ( x µ , x ν ) .
Revised: We denote the P ⇥ P kernel data matrix as K GP where K µ ⌫ GP = K GP ( x µ , x ⌫ ) , and the input kernel matrix on training data as K 0 where K µ ⌫ 0 = K 0 ( x µ , x ⌫ ) .

Original: Generalization error  measured by MSE including the bias and the mean, ϵ g 
Revised: Generalization error is measured by MSE including the bias and the variance contributions, ✏ g = ( h f ( x ) i ⇥  y ( x ))| {z } bias +

Original: = ( ⟨ f ( x ) ⟩ Θ − y ( x )) 2 +⟨ δf ( x ) 2 ⟩ Θ , which depends on the first and second order statistics of the predictor.
Revised: [Equation] f ( [Equation] | {z } variance , which depends on the ﬁrst and second order statistics of the predictor.

Original: In the GPlimit, we have ⟨ f ( x ) ⟩ = k GP ( x ) 
Revised: In the GP limit, we have [Equation] where 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 ) =

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)).
Revised: Note that the rank of K GP is the same as the capacity of the network, and the kernel matrix becomes singular as P approaches its capacity (the interpolation threshold), which results in nonzero training error, diverging bias and vanishing variance contribution to the generalization error (Fig.1 (b-e)).

Original: 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: The singularity of the kernel at the interpolation threshold holds also for ﬁnite width networks, and similar diverging bias and vanishing variance are seen in our ﬁnite width theory below (Section 3 ) and are conﬁrmed by simulation of networks trained with GD (see Appendix B.1,[19]).

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) , M, L ⇠ O (1) .

Original: In this limit, calculating the statistics of the network predictor requires integration over W in Eq.3 .
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 Back-propagating Kernel Renormalization[4] (see Appendix A) to GGDLNs.
Revised: To do so, we apply the previous method of Back-propagating Kernel Renormalization [4] (see Appendix A) to 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 ⇥ M matrix order parameter U 1 .

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

Original: [Equation]
Revised: [Equation]

Original: (6) where ◦ denotes element-wise multiplication.
Revised: where  denotes element-wise multiplication.

Original: In the linear case (which corresponds to M = 1) , the
Revised: In the linear case (which corresponds to M = 1) , the

Original: GP kernel is renormalized by a scalar factor.
Revised: GP kernel is renormalized by a scalar factor.

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: In the M > 1 case, the effect of renormalization is more drastic as it changes that not only the amplitude but 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 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 each layer (see Appendix A).

Original: The predictor statistics 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 a network with L =

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 ) > 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 Px 1 vector, ˜ k1 µ  ( x ) =
Revised: K 0 ( x , y ) denotes the renormalized kernel function for two arbitrary inputs x and y , ˜ K 1 denotes the P ⇥ P renormalized kernel matrix on the training data, and ˜ k 1 ( x ) is a P -dimensional vector, ˜ k µ 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 the shape of the kernel through the data dependent U 1 , reﬂecting the nonlinear property of the network, and resulting 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 network weights induced atequilibrium by Langevin dynamics with the MSE cost function and the Gaussian prior [4, 20, 21,22].
Revised: Our theory describes the properties of the posterior distribution of the network weights induced at equilibrium by Langevin dynamics with the MSE cost function and the Gaussian prior [4, 20, 21, 22].

Original: Simulating this dynamics agrees remarkably well with the simulation (see Appendix B.2).
Revised: Simulating this dynamics agrees remarkably well with the simulation (see Appendix B.2).

Revised: Although our theoretical results do not directly describe the solutions obtained by running gradient descent (GD) dynamics on the training error, it is interesting to ask to what extent the predicted behaviors of our theory are also exhibited by GD dynamics of the same network architectures, as GD-based learning is more widely used.

Original: Wewill compare our theoretical results with numerics of the gradient descent (GD) dynamics widelyused in machine learning.
Revised: We will compare our theoretical results with numerics of GD dynamics throughout the paper.

Original: 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 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 multiple initializations (see Appendix C.5 for details).

Original: As we will show, our theory makes accurate qualitativepredictions for GD dynamics in all examples throughout this paper.
Revised: As we will show, our theory makes accurate qualitative predictions for GD dynamics in all examples in this paper, in the sense that while the exact values may not match, the general trend of how generalization or representation varies with different parameters in different regimes are very similar.

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, while the mean predictor always assumes the same value as in the GP limit [4].

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 predictor by 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 predictor and hence the error bias also change with these network parameters due to the matrix nature of the kernel renormalization (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 of the kernel affects the generalization behavior (especially the bias term) of the network. 4.1 Networks with single hidden layer

Original: [Figure]
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: Feature selection in ﬁnite-width networks: 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 choice of the gating functions.

Original: In Fig.2we consider a student-teacher learning task [23, 24,25, 26].
Revised: In Fig.2 we consider a student-teacher learning task, commonly used for evaluating and understanding neural network performance[23, 24, 25, 26].

Original: 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: We present results of learning a ReLU teacher task in GGDLNs with gatings that have localized receptive ﬁelds (i.e., the activation of each gating unit depends on only a subset of input dimensions, the receptive ﬁeld of all gating units tile the N 0 input dimensions, as shown in Fig.2(a) bottom), where the student GGDLN is required to learn the input-output relation of a given ReLU teacher.

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 hidden layer (Fig.2(a) top left), gatings with different receptive ﬁelds are of equal importance, hence the renormalization does not play a beneﬁcial functional role, and the inﬁnitely wide network performs better than ﬁnite N .

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

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: For a ’local’ ReLU teacher 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 elements in U 1 that correspond to gating units whose receptive ﬁelds overlap the teacher’s 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 both the bias and the generalization error increase with N (Fig.2(c)).

Revised: More generally, the input can represent a set of ﬁxed features of the data, and the ’local’ teacher generate labels depending on a subset of the features.

Revised: Therefore, networks with ﬁnite width are able to select the relevant set of features by adjusting the amplitude in the renormalization matrix U 1 to assign the gating units with different importance for the task, while in the GP limit the network always assigns equal importance to all the gating units.

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: To summarize, our theory not only captures the more complex behavior of generalization (especially bias) as a function of network width, but also provides qualitative explanation of how generalization is affected by the structure 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: Effect of regularization strengths on generalization performance: Similar to the dependence on N , generalization also exhibits different behavior as a function of the regularization parameter  in different parameter regimes, with contributions from both the bias and the variance.

Original: The dependence of error bias on σalso arises due to the matrix nature of the renormalization.
Revised: 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 ).
Revised: In Fig.3 , we show parameter regimes where the bias can increase (Fig.3 (a-c)) or decrease (Fig.3 (d-f)) with  on MNIST dataset [19] (Appendix C.3 ).

Original: 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).
Revised: Although the dependence on  is complicated and diverse, and there lacks a general rule for when the qualitative behavior changes, we found that our theory accurately captures the qualitative behavior of results obtained from GD (Appendix B.3 Fig.3).

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

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
Revised: Speciﬁcally in Fig.3 (d-f), due to the increasing variance (e) and decreasing bias (d), there is a minimum error rate ((f), Appendix [-A.
 Eq.49-]A.3 Eq.50 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 strong regularization (  = 0 ) always results in optimal generalization

Original: [Figure]
Revised: [Figure]

Original: of σ . ( d - f ) Bias decreases as a function of σ while variance increases, leading to an optimal σ with minimum error rate. Comparable performance to ReLU networks:
Revised: GGDLNs with different choices of gatings achieve 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 the network 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 the same hidden layer width for classifying even and odd digits in MNIST data when M is sufﬁciently large (over-parameterization does not lead to over-ﬁtting here, as shown also in other nonlinear networks [27, 28], possibly due to the explicit L 2 prior).

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, they can be cleverly chosen to improve generalization performance.

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

Original: Random gatings take the form g m ( x ) = Θ( 1 √ N 0 V ⊤ m x − b ) , where
Revised: Random gatings take the form g m ( x ) = ⇥ ( 1 p N 0 V > m x  b ) , where

Original: V m is a1 × N 0  random vector with standard Gaussian i.i.d.
Revised: V m is a N 0 -dimensional 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 g m ( x ) outputs theprobability of assigning data x to the m -th cluster (Appendix C.2).
Revised: The pretrained gatings are trained on the unlabelled training dataset with unsupervised soft k-means clustering, such that the m -th gating g m ( x ) outputs the probability of assigning data x to the m -th cluster (Appendix C.3).

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, for pretrained gatings, generalization performance improves with M much faster compared to random gatings, 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) and numerical results of GD dynamics (Appendix B.3 Fig.4) agree qualitatively well.

Original: The result shows that GGDLNscan still achieve competitive performance on nonlinear tasks while remaining theoretically amenable.
Revised: The result shows that GGDLNs can 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: 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 GGDLNs with more layers.

Original: We beginby analyzing the  effectof U 1 on the shape of the kernel in deep architectures.
Revised: We begin by analyzing the renormalization effect on the shape of the kernel in deep architectures.

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 (the kernel function goes to a constant) as network depth increases [2], ultimately losing information about the input and degrading generalization performance.

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 by incorporating 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 ) = [Equation]
Revised: K L ( x , y ) = ˜ 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 GGDLN with 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 with generalized inner product deﬁned by the kernel ˜ K L ( x , y ) , and therefore K L ( x , y ) {+
+} [  1 , 1] .

Original: Forthe GP kernel of GGDLNs, we have K Lnorm ( x , y )
Revised: For the GP kernel of GGDLNs, we have K L ( 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 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 ( θ ) =(
Revised: While K L depends on the speciﬁc choice of gatings in general, in the special case of random gatings with zero threshold g m ( x ) = ⇥ ( 1 p N 0 V > m x ) and the number of gatings M ! 1 , we can write K L analytically as a function of the angle ✓ between input vectors x and y , given by K L ( ✓ ) =

Original: π − θπ ) L cos( θ )  .
Revised: ( ⇡  ✓⇡ ) L cos( ✓ ) , ✓ 2 [  ⇡ , ⇡ ] .

Original: Thus,  K Lnorm ( θ ) converges to a δ -function as L goes to infinity, shrinking to zeroexcept for θ = 0 . 
Revised: Thus, as L ! 1 , K L ( ✓ ) shrinks to zero except for ✓ = 0 . This ’ﬂattening’ effect reﬂects the loss of information in deep networks, as pairs of inputs with different similarities now all have hidden representations that are orthogonal.

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

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 of deep GGDLNs.

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 of the 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 in theamplitude of the kernel, in Fig.5(d) we plot only the bias contribution to the generalization.)
Revised: (Note that unlike the variance, the bias is affected only by shape changes, but not by changes in the amplitude of the 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, its effect on the generalization performance (especially the bias) may be different for different network parameters.

Original: In the specific example in Fig. B.4, finite width networks with a less’flattened’ renormalized kernel achieve better performance than the GP limit.
Revised: In the speciﬁc example in Fig.5, ﬁnite width networks with 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 GP limit and the ﬁnite width networks have optimal performance at L = 2 in this example. 1 K(x,y) 0. p r obab ili t y L=0L=1L=

Original: [Figure]
Revised: L=3L=  0. 1 K(x,y) 0. p r obab ili t y L=0L=1L=

Revised: L=3L=    6 L0.

Revised: r a t i

Revised: o

Revised: o f e l e m en t s <=.

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 multiple tasks.

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 learn input-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 life situations where the training data distribution is non-stationary.

Original: The tasks can be separatedbasedon difference in the inputs alone, without any additional top-down information.
Revised: The tasks can be separated 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, the gatings 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 tasks give conﬂicting labels for the same inputs, corresponding to the situation where performing the two tasks require additional top-down contextual information, and the information can be incorporated into 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 tasksalways leads to better generalization performance.
Revised: In both scenarios, when the gatings are ﬁxed and we modulate the de-correlation by changing network width and thus the strength of the kernel renormalization, we ﬁnd that de-correlation between tasks 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 defined by vastly different input statistics with bottom-upgatings, using permuted MNIST as an example.

Original: Previous works have shown that GLNs mitigate

Original: [Figure]
Revised: [Figure]

Original: [-N= 
 [Figure] [Figure] r a t-]
Revised: [-N= 
 [Figure] [Figure] r a t-]First we consider learning different tasks deﬁned by vastly different input statistics with bottom-up gatings, using permuted MNIST as an example.

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 mitigate catastrophic forgetting when sequentially trained on permuted MNIST [8].

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 theory does not address directly the dynamics of sequential learning, we aim to shed light on this question 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 predictor on a new data x , Eq.8, is a linear combination of the output labels Y µ of all the training data, and the 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, we deﬁne a task-task correlation matrix, via C pq =

Original: [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: where we 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 prediction on the test data of task p .

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

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: As we show in Fig.6, we can tune 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 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 tasksmore separately through the gating units.
Revised: In the case where the gatings are ﬁxed 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 better by processing data of different tasks separately through the gating units.

Original: 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.
Revised: However, in the case where we change the activation of the gatings by adjusting the threshold, although different tasks are more de-correlated when the threshold is large due to a set of less overlapping gatings activated for each task, generalization error ﬁrst decreases and then increases again.

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.
Revised: This is because for large threshold the sparsity of the gatings activated for each task limits the nonlinearity of the network, and therefore the generalization performance on this nonlinear task.

Original: 5.2 Combined top-down and bottom-up gating units
Revised: 5.2 Combined top-down and bottom-up gating units

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 taskspecific 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 gating units combine both top-down task signal which informs the system of which task to perform for a given input [Figure] 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 to0 , 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 states of the permitted gatings are further determined as a function of the input x [Figure] while the forbidden gatings are set to

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: [Figure] and the corresponding dendritic branches do not connect to the previous layer neurons (Fig.7(a)) in this task. 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 is 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 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 set of non-overlapping gating units are permitted in each of the n tasks, the network is equivalent to n 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 p 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 therefore in 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 how learning modiﬁes the correlation induced by the overlapping gatings through the renormalization matrix.

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 we consider training on permuted and un-permuted MNIST digits of 0 and 1 ’s. One task is to classify the two digits in both permuted and un-permuted data, and the second task is to separate the permuted digits (bothand 1) from the un-permuted 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 two tasks are uncorrelated, while the permitted gatings of the two tasks are partially overlapping.

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 ν ))
Revised: ( 1 M g p ( x µ ) > U 1 g q ( x ⌫ ))

Original: σ 2 N 0 x µ ⊤ x ν .
Revised:  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 } are the task indices, and µ , ⌫ = 1 , · · · , P are the input indices.

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 2 P ⇥ 2 P as shown in Fig.7(b) ( P = 600 ); the diagonal blocks (white dashed lines) correspond to kernels of task 1 and task 2, while the off diagonal blocks 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) bottom, we show the renormalized kernel with the renormalization matrix U 1 calculated by solving Eq.6.

Original: Similar results are achieved by by numerically estimating the readout weightsfromresults obtained from GD dynamics by using Eq.
Revised: Similar results are achieved by by numerically estimating Eq.7 with readout weights obtained from GD dynamics (Fig.7(b) top).

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 narrower networks suppresses more strongly the correlation between tasks, reﬂected by the weaker off-diagonal blocks in Fig.7(b).

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 the diagonal and off-diagonal blocks shows that the de-correlation effect diminishes for large N , leading to increasing generalization error with N (Fig.7(c & d)).

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 theoretical 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, resulting diverse 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 slows down the ﬂattening of the GP kernel in deep networks, suggesting that the loss of input information as L increases may be prevented in ﬁnite-width nonlinear networks.

Original: We also investigate the capabilityof GGDLNs to perform multiple tasks.
Revised: We also investigate the capability of GGDLNs to perform multiple tasks.

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 of the posterior of weight distribution induced by Langevin dynamics in Bayesian learning, it provides surprisingly well qualitative agreement with results obtained with GD dynamics for not only the generalization but also the kernel representation with matrix renormalization, largely extending its applicability. There are several limitations of our work.

Revised: Our mean-ﬁeld analysis is accurate in the ‘ﬁnite-width’ thermodynamic limit where both P and N go to inﬁnity, but M and L remain ﬁnite.

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 practice, the size of the renormalization matrix increases as M L , hence for some moderate M , as L increases, any large but ﬁnite N might eventually get the network outside the above thermodynamic regime. The theory also focuses on the equilibrium distribution induced by learning and does not address important questions related to the learning dynamics. Finally, although we have shown qualitative correspondence of the GGDLN properties and standard DNNs with local nonlinearity, as ReLU, a full theory of the thermodynamic limit of DNNs with local nonlinearity is still an open challenge. While our theory currently addresses learning in GGDLNs using a global cost function, exploring the possibility of extending the formalization of the equilibrium distribution to characterize local learning dynamics is an ongoing work.

Original: Recent works have shown that multilayer perceptrons (MLPs) with learnedgatings that implements spatial attention have surprisingly good performance onNatural Language
Revised: Recent works have shown that multilayer perceptrons (MLPs) with learned gatings that implements spatial attention have surprisingly good performance on

Original:  Processing (NLP) and computer vision [27].
Revised: Natural Language Processing (NLP) and computer vision [29].

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].

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: Furthermore, incorporating convolutional architecture [30, 31, 32, 33, 34] into our GGDLNs and using the gating units to encode context-dependent modiﬁcation of different feature maps is an interesting direction related to the fast-developing research topic of visual question-answering (VQA) [35, 36, 37] , where answering different questions about the same image is similar to performing multiple tasks in different contexts with different labels on the same dataset, 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.

Revised: Acknowledgement

Revised: We thank the anonymous reviewers for their helpful comments.

Original: Checklist
Revised: This research is supported by the Swartz Foundation, the NIH grant from the NINDS (No. 1U19NS104653), and the Gatsby Charitable Foundation. We acknowledge the support of a generous gift from Amazon.

Revised: This paper is dedicated to the memory of Mrs. Lily Safra, a great supporter of brain research.

Original: The checklist follows the references.
Revised: The checklist follows the references.

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

Original: For each question, change the default [TODO] to [Yes] , [No] or[N/A] .
Revised: 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.
Revised: You are strongly encouraged to include a justiﬁcation to your answer , either by referencing the appropriate section of your paper or providing a brief inline description.

Original: For example:
Revised: For example:

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

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

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

Original: [No] The code and the data are proprietary.
Revised: [No] The code and the data are proprietary.

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

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

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

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

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

Original: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
Revised: For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions 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(c)

Original: (c) 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) Have you read the ethics review guidelines and ensured that your paper conforms to them?
Revised: [N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Original: [Yes] 
Revised: [Yes] 2.

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

Original: [Yes] See Section
Revised: [Yes] See Section(b)

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

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

Original: If you ran experiments... (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: If you ran experiments... (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 (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they 
Revised: [Yes] In supplementary material (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?

Original: were chosen)? [Yes] See Appendix C (c) Did you report error bars (e g., with respect to the random seed after running experiments multiple times)?
Revised: [Yes] See Appendix C (c) 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., type of GPUs, internal cluster, or cloud provider)?

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

Original: If you are using existing assets (e g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?

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

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

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

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

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

Original: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
Revised: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?

Original: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
Revised: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

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