marginparsep has been altered.
topmargin has been altered.
marginparpush has been altered.
The page layout violates the ICML style. Please do not change the page layout,
or include packages like geometry, savetrees, or fullpage, which change it for
you. We’re not able to reliably undo arbitrary changes to the style. Please
remove the offending package(s), or layout-changing commands and try again.

The Pathway Race Reduction: Dynamics of Abstraction in Gated Networks

Anonymous Authors1

††footnotetext:  1Anonymous Institution, Anonymous City, Anonymous Region,
Anonymous Country. Correspondence to: Anonymous Author
<EMAIL_ADDRESS>
Preliminary work. Under review by the International Conference on Machine
Learning (ICML). Do not distribute.

###### Abstract

Our theoretical understanding of deep learning has not kept pace with its
empirical success. While model architecture and weight initialization are
known to be critical, we do not yet understand how they affect the nature of
representations learned by deep networks. In this work, we aim to address this
gap by introducing the Gated Deep Linear Network framework that schematizes
how pathways of information flow impact learning dynamics within an
architecture. Crucially, because of the gating, these networks can compute
nonlinear functions of their input. We derive an exact reduction and, for
certain cases, exact solutions to the dynamics of learning. Our analysis
demonstrates that the dynamics of learning in structured networks can be
conceptualized as a neural race with an implicit bias towards shared
representations, which then govern the model’s ability to systematically
generalize, multi-task, and transfer. We validate our key insights on
naturalistic datasets and with relaxed assumptions. Taken together, our work
gives rise to general hypotheses relating neural architecture to learning, and
provides a mathematical approach toward understanding the design of more
complex architectures and the role of modularity and compositionality in
solving real-world problems. The code and results are available at
https://www.saxelab.org/gated-dln.

##  1 Introduction

While neural networks have led to numerous impressive breakthroughs resnet;
transformer; amodei2016deep; baevski2020wav2vec; nature_dqn; alpha_zero, our
theoretical understanding of these models has not advanced at the same pace.
Although we have gained some insight into general principles of network
optimization patel2015probabilistic; Carleo2019; bahri_statistical_2020;
arora2020theory; roberts2021principles, we have a limited understanding of how
the specific choice of architecture—that is, the mesoscale pattern of
connectivity between hidden layers (Fig. 1)—affects a network’s behavior
zagoruyko2016wide; raghu2017expressive; Chizat2018; saxe_mathematical_2019;
tian2019luck. For example, when training networks on the ImageNet dataset,
wide networks perform slightly better on classes reflecting scenes, whereas
deep networks are slightly more accurate on classes related to consumer goods
nguyen2020wide. Understanding the reasons behind these behaviors may lead to
more systematic techniques for designing neural networks.

Figure 1: A multi-modal network composed of simple modules shared across
modalities and tasks. How do shared modules and pathways impact representation
learning and generalization?

Here, we address this gap by introducing and analyzing the Gated Deep Linear
Network (GDLN) framework, which illuminates the interplay between data
statistics, architecture, and learning dynamics. We ground the framework in
the Deep Linear Network (DLN) setting as it is amenable to mathematical
analysis, and previous works (Saxe2014; lampinen2018analytic;
arora2018optimization; saxe_mathematical_2019) have observed that DLNs exhibit
several nonlinear phenomena that are observed in deep neural networks. Because
of the gating mechanism, however, GDLNs can compute nonlinear functions of
their input, making them more expressive than standard deep linear networks.

Our main contributions are: (i) We introduce the GDLN framework (Section 2),
which schematizes how pathways of information flow impact learning dynamics
within an architecture. (ii) We derive an exact reduction and, for certain
cases, exact solutions to the dynamics of learning (Section 3). (iii) Our
analyses reveal the dynamics of learning in structured networks can be
conceptualized as a neural race with an implicit bias towards shared
representations (Section 4.2). (iv) We validate our findings on naturalistic
datasets, with some relaxed assumptions (Section 5).

##  2 Gated Deep Linear Network Framework

A fundamental principle of neural network design is that powerful networks can
be composed out of simple modules goodfellow2016deep; lecun2015deep, and a key
intuition is that a network’s compositional architecture should resonate in
some way with the task to be performed. For instance, a multimodal network
might process each modality independently, before merging these streams for
further processing ngiam2011multimodal; girdhar2022omnivore, and a multi-
tasking NLP model might process its input through a shared encoder before it
splits into task-specific pathways collobert2011natural; liu2019multi;
standley2020tasks. Frequently, a network’s compositional structure can be
conceptualized as an “architecture graph” $\Gamma$, decorated by network
modules, additional interactions, and learning mechanisms.

Here we introduce a class of networks, Gated Deep Linear Networks (GDLNs),
depicted in Fig. 2a, for which we can analytically study the effect of the
architecture graph on learning and generalization. A GDLN is defined as
follows. Let $\Gamma$ denote a directed graph with nodes $\mathbf{V}$ and
edges $\mathbf{E}$, with the structure of $\Gamma$ encoded by functions
$s,t:\mathbf{E}\to\mathbf{V}$ mapping an edge to its source and target node,
respectively. For each $v\in\mathbf{V}$, let $|v|\in\mathbb{N}$ denote the
number of neurons assign to that node, and let $h_{v}\in\mathbb{R}^{|v|}$
denote neural activations for the corresponding network layer. For each edge
$e\in\mathbf{E}$, let $W_{e}$ denote a $|t(e)|\times|s(e)|$ weight matrix
assigned to $e$. An input node of $\Gamma$ is a node with only outgoing edges,
and an output node is a node with only incoming edges. Let
$\textrm{In}(\Gamma),\textrm{Out}(\Gamma)\subset V$ denote the sets of input
nodes and output nodes of $\Gamma$, respectively.

The GDLN associated with $\Gamma$ computes a function as follows: An input
example specifies values $x_{v}\in\mathbb{R}^{|v|}$ for all input nodes
$v\in\textrm{In}(\Gamma)$, and the input nodes are fixed to their values
$h_{v}=x_{v}$ for $v\in\textrm{In}(\Gamma)$. Then, activation propagates to
subsequent layers according to
$h_{v}=g_{v}\sum_{q\in\mathbf{E}:t(q)=v}g_{q}W_{q}h_{s(q)}$ where $g_{v}$ is
the node gate and $g_{q}$ is the edge gate. That is, activity propagates
through the network as in standard neural networks, but modulated by gating
variables that can act at the node and edge level.

In essence, these gating variables enable nonlinear computation from input to
output, and can be interpreted in several ways: The node level gating can be
viewed as an approximate reduction of ReLU dynamics as a ReLU neuron’s
activity can be written as $\max(0,wx)=\textrm{step}(wx)wx$. The gating
variables can also be viewed as context-dependent control signals that
modulate processing in the network. We discuss the extensive connections
between the Gated Deep Linear Network and other approaches in Appendix 6.

In order to keep the analysis mathematically tractable, we assume that the
gating variables are simply specified directly for each input that will be
processed and we consider them to be a part of the dataset.

Figure 2: Formalism and notation. (a) The Gated Deep Linear Network applies
gating variables to nodes and edges in an otherwise linear network. (b) The
gradient for an edge (here using the orange edge in Panel (a) as an example)
can be written in terms of paths through that edge (colored lines). Each path
is broken into the component antecedent $\bar{s}(p,e)$ and subsequent
$\bar{t}(p,e)$ to the edge.

##  3 Gradient Flow Dynamics and Reduction

Having described the design of the network, we now exploit its special
properties to understand learning dynamics and their relationship to network
structure. In particular, we consider training all weights in the network to
minimize the $L_{2}$ loss averaged over a dataset,

$\mathcal{L}(\\{W\\})=\left\langle\frac{1}{2}\sum_{v\in\textrm{Out}(\Gamma)}||y_{v}-h_{v}||_{2}^{2}\right\rangle_{x,y,g}$
(1)

where $y_{v}\in\mathbb{R}^{|v|}$ for $v\in\textrm{Out}(\Gamma)$ are the target
outputs for the output layers in the network, and
$\langle\cdot\rangle_{x,y,g}$ denotes the average over the training dataset
and gating structures. The weights in the network can be updated using
gradient descent. A key virtue of the GDLN formalism is that the gradient flow
equations can be compactly expressed in terms of the paths through the
network.

We first lay out our notation, as illustrated in Fig. 2b for an example
network. A path $p$ is a sequence of edges that joins a sequence of nodes in
the network graph $\Gamma$. Let $\mathcal{P}(e)$ be the set of all paths from
any input node to any output node that pass through edge $e$. Let
$\mathcal{T}(v)$ be the set of all paths terminating at node $v$. We denote
the source and the target node of the path $p$ as $s(p)$ and $t(p)$,
respectively. We denote the component of path $p$ that is subsequent to edge
$e$ (i.e., the path whose source node is the target of $e$, and that otherwise
follows $p$) as $\bar{t}(p,e)$ (for the ‘target’ path of $e$). Similarly, we
denote the component of path $p$ that precedes edge $e$ (i.e., the path whose
target node is the source of $e$, and that otherwise follows $p$) as
$\bar{s}(p,e)$ (for the ‘source’ path of $e$). Overloading the notation, we
will write $W_{p}$ where $p$ is a path to indicate the ordered product of all
weights along the path $p$, with the target of $p$ on the left and the source
of $p$ on the right. Similarly, we write $g_{p}$ where $p$ is a path to denote
the product of the (node and edge) gating variables along the path.

With this notation, the gradient flow equations can be shown to be (full
derivation in Appendix B),

$\displaystyle\tau\frac{d}{dt}W_{e}$ $\displaystyle=$
$\displaystyle-\frac{\partial\mathcal{L}(\\{W\\})}{\partial W_{e}}\quad\forall
e\in E,$ (2) $\displaystyle=$
$\displaystyle\sum_{p\in\mathcal{P}(e)}W_{\bar{t}(p,e)}^{T}\mathcal{E}(p)W^{T}_{\bar{s}(p,e)}$
(3)

where the error term for path $p$ is

$\displaystyle\mathcal{E}(p)=\Sigma^{yx}(p)-\sum_{j\in\mathcal{T}(t(p))}W_{j}\Sigma^{x}(j,p).$
(4)

Here the dataset statistics which drive learning are collected in the
correlation matrices

$\displaystyle\Sigma^{yx}(p)$ $\displaystyle=$ $\displaystyle\left\langle
g_{p}y_{t(p)}x_{s(p)}^{T}\right\rangle_{y,x,g}$ (5)
$\displaystyle\Sigma^{x}(j,p)$ $\displaystyle=$ $\displaystyle\left\langle
g_{j}x_{s(j)}x_{s(p)}^{T}g_{p}\right\rangle_{y,x,g}$ (6)

where $j$ and $p$ index two paths. Hence if there are $N$ paths through the
graph from input nodes to output nodes, there are potentially $N$ distinct
input-output correlation matrices and $N^{2}$ distinct input correlation
matrices that are relevant to the dynamics. Remarkably, no other statistics of
the dataset are considered by the gradient descent dynamics.

Notably, these correlation matrices depend not just on the dataset statistics
($x$ and $y$), but also on the gating structure $g$. The possible gating
structures are limited by the architecture. In this way, the architecture of
the network influences its learning dynamics.

In essence, the core simplification enabled by the GDLN formalism is that the
gating variables $g$ appear only in these data correlation matrices. They do
not appear elsewhere in Eqns. (3)-(4), which otherwise resemble the gradient
flow for a deep linear network Saxe2014; saxe_mathematical_2019. The effect of
the nonlinear gating can thus be viewed as constructing pathway-dependent
dataset statistics that are fed to deep linear subnetworks (pathways).

Figure 3: XoR solution dynamics. (a) The XoR task with positive (red) and
negative (blue) examples. Input-to-hidden weights from ReLU simulations
(magenta) reveal four functional cell types. (b) GDLN with four paths, each
active on one example. (c) Simulations of ReLU dynamics from small weights
(red, 10 repetitions) closely track analytical solutions in the GDLN.
Parameters: $N_{h}=128,\tau=5/2,\sigma_{0}=.0002.$

As a simple example of the power of this framework relative to deep linear
networks, consider the XoR task (Fig. 3a), a canonical nonlinear task that
cannot be solved by linear networks. By choosing the gating structures to
activate a different pathway on each example (Fig. 3b), the gated deep linear
network can solve this task (Fig. 3c blue). Crucially, its dynamics
(analytically obtained in Appendix A based on the reduction in the following
sections) closely approximates the dynamics of a standard ReLU network trained
with backprop (Fig. 3c red). This result demonstrates that the gated networks
are more expressive than their non-gated counterpart, and that gated networks
can provide insight into ReLU dynamics in certain settings. We note that so
far, our analysis does not provide a mechanism to select the gating structure.
We will return to this point in Section 4.2, which provides a perspective on
the gating structures likely to emerge in large networks.

###  3.1 Exact reduction from decoupled initial conditions

Our fundamental goal is to understand the dynamics of learning as a function
of architecture and dataset statistics. In this section, we exploit the
simplified form of the gradient flow equations to obtain an exact reduction of
the dynamics. Our reduction builds on prior work in deep linear networks, and
intuitively, shows that the dynamics of gating networks can be expressed
succinctly in terms of effective independent 1D networks that govern the
singular value dynamics of each weight matrix in the network. The reduced
dynamics can be substantially more compact, as for instance, a weight matrix
of size $N\times M$ has $NM$ entries but only $\textrm{min}(M,N)$ singular
values.

Figure 4: Pathway network solution dynamics. (a) The network contains $M$
different input domains (each consisting of a bank of neurons), $M$ different
output domains, and two hidden layers. The task is to learn a mapping from
each input domain to each output domain. The gating structure gates on one
input and one output pathway. The hidden pathway is always on. (b) Gated
network formalism. There are $M^{2}$ pathways through the network from input
to output. All $M^{2}$ flow through the hidden weight matrix, while only $M$
flow through each input or output weight matrix. This fact causes the hidden
layer to learn faster. (c) Small example dataset with hierarchical structure.
The task of the network is to produce the 7-dimensional output vector for each
of four items. Inputs are random orthogonal vectors for each item. (d) Each
input domain is trained with $K$ output domains (here $K=4$), such that some
input-output routes are never seen in training. (e) Training loss dynamics for
simulated networks from small random weights (red, 10 repetitions), simulated
networks from decoupled initial conditions (green), and theoretical prediction
from Eqn. 15 (blue). The theory matches the decoupled simulations exactly, and
is a good approximation for small random weights. (f) The singular values of
the hidden weight matrix (blue) are larger than those in input or output
matrices by a factor $\sqrt{M}$. Theoretical predictions match simulations
well, particularly for larger singular values. (g) Representational similarity
(or kernel) matrix at the first hidden layer. Inputs from different domains
are mapped to similar internal representations, revealing a shared
representation even for input domains that are never trained with a common
output. (h) Predicted output at the end of training. The network generalizes
perfectly to input-output routes that were never seen during training.
Parameters: $M=7,K=4,\lambda=.02,\sigma_{0}=.2,N_{h}=64.$

To accomplish this, we introduce a change of variables based on the singular
value decomposition of the relevant dataset statistics. Suppose that the
dataset correlation matrices are mutually diagonalizable, such that their
singular value decompositions have the form

$\displaystyle\Sigma^{yx}(p)$ $\displaystyle=$ $\displaystyle
U_{t(p)}S(p)V_{s(p)}^{T}$ (7) $\displaystyle\Sigma^{x}(j,p)$ $\displaystyle=$
$\displaystyle V_{s(j)}D(j,p)V_{s(p)}^{T}$ (8)

where the set of $U$ and $V$ matrices are orthogonal, and the set of $S$ and
$D$ matrices are diagonal. That is, there is a distinct orthogonal matrix
$U_{l}$ for each output layer, a distinct orthogonal matrix $V_{l}$ for each
input layer, and diagonal matrices $S(p),D(p)$ for each path through the
network.

Then, following analyses in deep linear networks Saxe2014, we consider the
following change of variables. We rewrite the weight matrix on each edge as

$\displaystyle W_{e}(t)=R_{t(e)}B_{e}(t)R_{s(e)}^{T}\quad\forall e,$ (9)

where the matrices $B_{e}(t)$ are the new dynamical variables, and the matrix
$R_{v}$ associated to each node $v$ in the graph satisfies $R_{v}^{T}R_{v}=I$.
Further, for output nodes $v$, we require $R_{v}=U_{v}$, the output singular
vectors in the diagonalizability assumption. Similarly, for input nodes, we
require $R_{v}=V_{v}$.

Inserting (7)-(9) into (3)-(4) shows that the dynamics for $B_{e}$ decouple:
if all $B_{e}(0)$ are initially diagonal, they will remain so under the
dynamics (full derivation in Appendix B). For this decoupled initialization,
the dynamics are

$\tau\frac{d}{dt}B_{e}=\sum_{p\in\mathcal{P}(e)}B_{p\setminus
e}\left[S(p)-\sum_{j\in\mathcal{T}(t(p))}B_{j}D(j,p)\right]$ (10)

where $B_{p\setminus e}=B_{\bar{t}(p,e)}B_{\bar{s}(p,e)}$ is the product of
all $B$ matrices on path $p$ after removing edge $e$ (see Appendix B).

In essence, this reduction removes competitive interactions between singular
value modes, such that the dynamics of the overall network can be described by
summing together several “1D networks,” one for each singular value.
Intuitively, this reduction shows that learning dynamics depend on several
factors.

Input-output correlations

Other things being equal, a pathway learning from a dataset with larger input-
output singular values will learn faster. This fact is well known from prior
work on deep linear networks Saxe2014.

Pathway counting

Other things being equal, a weight matrix corresponding to an edge that
participates in many paths (such that the sum contains many terms) will learn
faster. This fact is less obvious, as it becomes relevant only if one moves
beyond simple feed-forward chains to study complex architectures and gating.

We now turn to examples that verify and illustrate the rich behavior and
consequences of these dynamics.

##  4 Applications and consequences

To fix a specific scenario with rich opportunities for generalization, we
consider a “routing” setting, as depicted in Fig. 4a. In this setting, a
network receives inputs from $M$ different input domains and produces outputs
across $M$ different output domains. The goal is to learn to map inputs from a
specific input domain to a specific output domain, with no negative-
interference from other input-output domain pairs. There are thus $M^{2}$
possible tasks which can be performed, each corresponding to mapping one of
the $M$ input domains to one of the $M$ output domains.

We assume that the target input-output mapping from the active input domain to
the active output domain is the same for all pathways, and defined by a
dataset with input correlations $\langle xx^{T}\rangle=VDV^{T}$ and input-
output correlations $\langle yx^{T}\rangle=USV^{T}$. For the simulations in
this section, we take the dataset to contain four examples, and the target
output to be a 7-dimensional feature vector with hierarchical structure (Fig.
4c), but note that the theory is more general.

To investigate the possibility of structured generalization, we consider a
setting where only a subset of input-output pathways are trained. That is,
each input domain is trained with only $K\leq M$ output domains, as depicted
in Fig. 4d, such that some input-output pathways are never observed during
training.

We consider solving this task with a two-hidden layer gated deep linear
network depicted in Fig. 4a. We emphasize that this task is fundamentally
nonlinear, because inputs on irrelevant input domains must be ignored. We take
the gating structure to gate off all first layer pathways except the relevant
input domain, and to gate off all third layer pathways but the one to the
relevant output domain. As shown in Fig. 4b, this scheme results in $M^{2}$
pathways through this network that must be considered in the reduction. The
resulting pathway correlations are simply given by the original dataset,
scaled be the probability that each path is active (see Appendix C)

$\displaystyle\Sigma^{yx}(p)$ $\displaystyle=$
$\displaystyle\frac{1}{KM}USV^{T}$ (11) $\displaystyle\Sigma^{x}(j,p)$
$\displaystyle=$
$\displaystyle\begin{cases}\frac{1}{KM}VDV^{T}\quad\textrm{if}~{}j=p\\\
0\quad\textrm{otherwise}\end{cases}.$ (12)

Crucially, we have a simple “pathway counting” logic behind the reduction: the
first and third layer weights are active in $K$ paths (all tasks originating
from a given input domain or terminating at a given output domain,
respectively), while second layer weights are active in $KM$ trained paths.
This fact causes the second layer weights to learn more rapidly.

Assuming that weights start out roughly balanced in each first layer and third
layer weigh matrix (a reasonable assumption when starting from small random
weights), this yields the reduced dynamics (Appendix C)

$\displaystyle\tau\frac{d}{dt}B_{1}$ $\displaystyle=$
$\displaystyle\frac{1}{M}B_{2}B_{1}\left[S-B_{2}B_{1}^{2}D\right]$ (13)
$\displaystyle\tau\frac{d}{dt}B_{2}$ $\displaystyle=$ $\displaystyle
B_{1}^{2}\left[S-B_{2}B_{1}^{2}D\right]$ (14)

where $B_{1}$ describes the input and output pathway weights singular values,
and $B_{2}$ describes the hidden layer weight singular values.

We note that the quantity $MB_{1}^{2}-B_{2}^{2}$ is conserved under the
dynamics. Defining the constant $C=MB_{1}(0)^{2}-B_{2}(0)^{2}$, we can
therefore write the dynamics as

$\tau\frac{d}{dt}B_{2}=\frac{1}{M}(B_{2}^{2}+C)\left[S-\frac{1}{M}B_{2}(B_{2}^{2}+C)D\right].$
(15)

Remarkably, this equation reveals that the dynamics of this potentially large,
gated, multilayer network with arbitrary numbers of hidden neurons can be
reduced to a single scalar for each singular value in the dataset. Each
diagonal element of this equation provides a separable differential equation
that may be integrated to give an exact formal solution.

Figure 4e compares the training error dynamics predicted by Eqn. 15 to full
simulations starting from small random weights (i.e. scaling Xavier
initialiation weights by 0.2), verifying that our reduction is a good
description of dynamics starting from small random weights. Furthermore, as
shown in Appendix C, the hidden layer weight singular values change by a
factor of $\sqrt{M}$ more than the input or output weights, as verified in Fig
4f.

Figure 5: Pathway race dynamics. (a) The same routing task can be solved using
a variety of gating schemes that differ in their use of shared
representations. Every input-output combination can be given a dedicated
pathway such that $P=1$ tasks flow through each (left), groups of two input
domains and two output domains can share a pathway such that $P=4$ tasks flow
through each (middle), or all $P=M^{2}$ tasks can run through a shared
representation (right). (b) Singular value dynamics as a function of the
number of pathways $P$ flowing through the hidden layers. Networks that share
more structure learn faster. Consequently, in a single network where subparts
share structure to different degrees, the maximally shared dynamics dominate
the race between pathways. (c) Error on trained (blue) and untrained (red)
pathways as a function of the fraction of output domains $K$ trained with each
input domain $M$. When few outputs are trained per input domain, the race
dynamics do not strongly favor shared structure and so error on untrained
domains is large. When $\sim 40\%$ of output domains are trained with each
input, the shared structure solutions are sufficiently faster to reliably
dominate the race and yield generalization to untrained domains. Parameters:
$M=10,\lambda=.05/K,\sigma_{0}=.2,N_{h}=64.$

###  4.1 Shared representations and generalization

With this description of the training dynamics of the network, we can then ask
what representations emerge in the network over training. One way of
interrogating the nature of representations in the network is to compute the
representational similarity between different input examples from the same
input domain, and across different input domains. Specifically, we compute the
dot product between the neural activity in the first hidden layer in response
to different inputs. As shown in Fig. 4g, the pathway network learns a shared
representation, in which each individual example maps to the same
representation regardless of what input domain it arrives on. That is, the
gating and learning dynamics enable the network to learn a representation that
is invariant to input domain, and which is abstract in the sense that the
representation contains no information about what input domain produced it.
This shared representation supports zero-shot generalization to untrained
input-output pathways, as shown in Fig. 4h.

The intuition behind obtaining zero-shot generalization is as follows: Say
that we are evaluating the network on a new input-output pair of domains. As
long as the network has been trained on examples from the current input domain
(in conjunction with any output domain), the network will map it to the shared
representation. Similarly, as long as the network has been trained on examples
from the current output domain, it will be able to map this shared
representation to the output. In this way, training on a subset of $M^{2}$
tasks is enough to obtain strong generalization to all $M^{2}$ tasks.

In essence, this solution accomplishes a factorization of the problem into two
interacting but distinct components: the gating variables represent what input
domain links to what output domain, providing information about “where”
signals should go; while the neural activity represents “what” task-relevant
input was presented, regardless of where it came from or where it should be
routed to. This factorization can permit generalization to untrained pathways
provided the gating structure is configured appropriately.

Figure 6: Pathway race vs. NTK regime: Initialization, shared representation
formation, and zero-shot generalization. (a) Training loss for pathway
networks with different initialization scales $\sigma_{0}$. Small $\sigma_{0}$
yields pathway race dynamics with stage-like drops through training. Large
$\sigma_{0}$ yields NTK-like dynamics with rapid, exponential learning curves.
(b) Error for trained (blue) and untrained (red) input-output domain
combinations as a function of initialization scale. While performance on
trained domains is excellent for all scales, zero-shot generalization only
emerges in the pathway race regime. (c) Representational similarity between
inputs presented to different domains, for small (left column), medium (middle
column) and large (right column) initialization scales. At small
initialization scales, internal representations in the first hidden layer are
similar even across domains, indicating one common shared representation.
Large initialization scales place the network in the NTK regime where random
initial connectivity persists throughout learning, yielding distinct high-
dimensional random representations for each domain. (d) Network output for all
input-output combinations for three initialization scales (labeled in panel
c). Because networks in the NTK regime do not learn a shared representation
for different input domains, they do not generalize to untrained pathways.

###  4.2 Pathway race dynamics: Implicit bias toward shared representations

The reductions so far have assumed that the gating structure is specified a
priori, and furthermore, that different domains connect to the same singular
value modes in the hidden layer. That is, the gating structure provides the
opportunity for learning a shared representation, but this is not obligatory:
different parts of the hidden representation could learn distinct pathways,
despite all being gated on. Remarkably, the dynamics from small random weights
track the trajectory predicted for maximally shared representation, suggesting
that the full solution dynamics rapidly converge to the submanifold of
decoupled weights. Why are shared representations favored under the dynamics?

To investigate this, we note that the same task and architecture typically
permit several gating schemes and singular value mode connectivity patterns
that each would obtain zero training error. As shown in Fig. 5a (left), for
instance, the routing task could be solved with an alternative gating
structure in which each input-output route receives a dedicated pathway that
is gated on only when the task is to connect that specific input-output route.
This gating scheme would still obtain zero training error, but does not yield
any representation sharing. Other partial sharing schemes are possible; for
instance, representations could be shared across groups of two input and two
output domains (Fig. 5a, middle). What is the impact of these choices on
learning dynamics and generalization?

Taking the case where all routes are trained ($K=M$) for simplicity, with no
sharing, each pathway participates in just one of the $M^{2}$ total trained
pathways, compared to the fully shared solution where the input and output
layers participate in $M$ pathways and the hidden layer participates in
$M^{2}$. From this, we can see that greater sharing leads to faster learning.
In a combined network that produces its output using both shared and non-
shared representations, the dynamics of each pathway will race each other to
solve the task; and hence the most-shared structure will dominate.

To see this quantitatively, we repeat a similar derivation to the preceding
section for networks with varying degrees of pathway overlap. In particular,
we parameterize the degree of pathway overlap with the parameter $P$ that
counts the number of pathways flowing through a given hidden layer. The
resulting reduction is

$\displaystyle\tau\frac{d}{dt}B_{1}$ $\displaystyle=$
$\displaystyle\frac{\sqrt{P}}{M^{2}}B_{2}B_{1}\left[S-B_{2}B_{1}^{2}D\right]$
(16) $\displaystyle\tau\frac{d}{dt}B_{2}$ $\displaystyle=$
$\displaystyle\frac{P}{M^{2}}B_{1}^{2}\left[S-B_{2}B_{1}^{2}D\right],$ (17)

which shows that the learning rate in all layers increases as $P$ increases.
Solution dynamics for a range of degrees of sharing $P$ are plotted in Fig.
5b, which show that greater degrees of sharing reliably leads to faster
singular value dynamics.

Hence, dynamics in GDLNs take the form of a pathway race: when many gating
schemes coexist in the same network, the ones that share the most
structure–and hence learn the fastest–will come to dominate the solution.
Therefore gradient flow dynamics in complex network architectures has an
implicit bias toward extracting shared representations where possible. The
strength of this bias increases as each input domain is trained with more
output domains. As shown in Fig. 5c, shared representations begin to dominate
reliably when roughly 40% of input-output routes are trained, enabling
generalization to unseen input-output routes.

###  4.3 Impact of initialization

The training and generalization dynamics of deep networks are known to depend
on the weight initialization. Here we show that initial weight variance exerts
a pronounced effect on the emergence of shared representations, and hence
generalization abilities. As observed in a number of theoretical and empirical
works, neural networks can operate in two different initialization regimes
Chizat2018; bahri_statistical_2020; flesch_orthogonal_2022. Sufficiently wide
networks initialized with large variance initializations enter the Neural
Tangent Kernel regime, where training dynamics follow a simple linear
dynamical system and error trajectories exhibit exponential approach to their
asymptote jacot2018neural; Lee2019; arora_exact_2019. Intuitively, in this
regime, the initial strong random connectivity in the network provides
sufficiently rich features to learn the task without substantially changing
internal representations. In this setting, deep networks behave like kernel
machines with a fixed kernel (the neural tangent kernel). By contrast,
networks initialized with sufficiently small variance initializations learn
rich task-specific representations, and their dynamics as we have seen can be
more complex mei_mean_2018; rotskoff_parameters_2018; sirignano_mean_2020;
saxe_mathematical_2019.

To show the effect of this transition in our setting, we train pathway
networks starting from different random matrices with singular value
$\sigma_{0}$. For a range of initialization scales, all networks converge to
zero training error (Fig. 6a). As expected, large initialization scales lead
to NTK-like exponential dynamics, while small initialization scales lead to
progressive stage-like drops in the error consistent with prior analyses of
deep linear networks in the rich regime. Critically, initialization has a
dramatic impact on generalization (Fig. 6b), and only small initializations
are capable of zero-shot generalization to untrained routes. To understand
why, we visualize the representational similarity structure for several
networks in Fig. 6c, which shows a transition from shared to independent
representations as networks move from the rich to the lazy regime. Finally,
Fig. 6d shows the breakdown in generalization in the NTK regime. Hence our
pathway race reduction can describe learning in the rich feature learning
regime, with non-trivial generalization behavior.

Figure 7: Experimental results. (a) Each input domain receives inputs that
have been subjected to one of $M$ input transformations. The target output for
each output domain is also transformed by one of $M$ output transformations.
Here the visualization uses rotations in the input and permutations in the
output. Only a subset of input-output transformation pairs are seen during
training. (b) Error for trained (blue) and untrained (orange) input-output
domain pairs as a function of the percentage of trained pathways ($K/M$) on
the CIFAR dataset with $M^{2}=1600$ total tasks. (c) Error on MNIST with
$M^{2}=10^{4}$ total tasks. Training accuracy is always high while zero-shot
transfer to untrained pathways becomes as good as the training performance
when $\approx$40% of pathways are trained. (d) Error as a function of
initialization scale. While performance on trained domains is good for all
scales, zero-shot generalization only emerges at small inits.

##  5 Experiments

So far, we have described mathematical principles relating network
architecture to learning dynamics and the nature of learned representations in
GDLNs. We demonstrated how the network architecture affects generalization in
pathway networks when applied to a simple toy dataset. In this section, we
qualitatively validate our key findings on naturalistic datasets.
Specifically, we test if GDLNs can exhibit strong zero-shot generalization
performance on untrained input-output domain pairs (Section 4.1) and if the
link between initialization scale and zero-shot generalization (Section 4.3)
holds when training on naturalistic datasets.

Several datasets and benchmarks have been proposed to evaluate systematic
generalization in neural networks johnson2017clevr; bahdanau2018systematic;
Sinha2019CLUTRRAD; lake2019compositional; ruis2020benchmark;
Sinha2020EvaluatingLG. However, these naturalistic datasets generally require
the network to learn multiple capacities (like spatial reasoning, logical
induction, etc.) and come with a fixed (and often unclear) extent of training
in different domains. To develop a setting that remains close to the theory,
we create new datasets by composing transformations on popular vision
datasets. Having fine-grained control over the dataset generation mechanism
enables us to understand the effect of parameters like the number of
input/output domains.

Briefly, as depicted in Figure 7a, starting from a base dataset with $n$
inputs and outputs $\\{(x^{\mu},y^{\mu})\\},\mu=1,\cdots,n$, we generate new
tasks by applying one of $M$ input transformations
$f^{\textrm{in}}_{i},i=1,\cdots,M$ and one of $M$ output transformations
$f^{\textrm{out}}_{j},j=1,\cdots,M$ (full details deferred to Appendix D due
to space constraints). The task from input domain $i$ to output domain $j$
thus has samples
$\\{(f^{\textrm{in}}_{i}(x^{\mu}),f^{\textrm{out}}_{j}(y^{\mu}))\\}$. We use
the MNIST (deng2012mnist) and CIFAR-10 datasets (krizhevsky2009learning) as
base datasets, and rotations and permutations as transformations. Relative to
our previous experiments, these datasets add real data correlations and
distinct transformations on each domain that make finding a shared
representation challenging.

###  5.1 Results

In Figure 7, we evaluate the zero-shot generalization performance of gated
deep linear networks on untrained pathways and study the effect of
initialization on their performance. Full model and training details are given
in Appendix D. Fig. 7(b,c) shows mean accuracy, over trained and untrained
pathways, as a function of the fraction of datasets that the model used for
training, for CIFAR ($M=40$) and MNIST ($M=100$) respectively. Training
accuracy is always high while the zero-shot transfer to untrained pathways
becomes near-perfect when $\approx$40% of pathways are trained. In Fig. 7d, we
report the error for trained (blue) and untrained (orange) input-output domain
combinations as a function of initialization scale (gain ratio). While
performance on trained domains is good for all scales, zero-shot
generalization only emerges for smaller scales, as in the pathway race regime.
These observations validate our findings (from Section 4) on naturalistic
datasets.

##  6 Related Work

Our work is closely related to several areas in machine learning: Deep Linear
Networks, the study of dynamics of learning in neural networks, and modular
neural networks.

Deep Linear Networks: Baldi1989; Fukumizu1998; Saxe2014;
arora2018optimization; lampinen2018analytic; saxe_mathematical_2019 showed
that deep linear networks exhibit several nonlinear phenomena that are
observed in deep neural networks and proposed studying the dynamics of deep
linear networks as a surrogate for understanding the dynamics in the deep
neural networks. (Baldi1989) described the loss landscape, while (Saxe2014)
developed the theory of gradient descent learning in deep linear neural
networks and provided exact solutions to the nonlinear dynamics. Motivated by
their observation about the similarity in dynamics of linear and non-linear
networks and the feasibility of analyzing the gradients in the linear network,
we ground the proposed framework in the deep linear network setting.

Many works have studied the dynamics of deep networks under specific
assumptions like linear separability of data des2018learning, deep ReLU
networks tian2019luck; straat2019line, Tensor Switching Networks
NIPS2016_b1563a78 (generalization of ReLU to tensor-valued hidden units), the
Neural Tanget Kernel limit jacot2018neural; fort2020deep, the Mean Field limit
mei_mean_2018; rotskoff_parameters_2018; sirignano_mean_2020 and Ensemble
Networks fort2019deep to name a few (see Carleo2019; bahri_statistical_2020;
arora2020theory for reviews). Similar to these works, we also focus on a
specific subset of deep networks, gated deep linear networks, that captures a
nonlinear relationship between a network’s input-output maps and its
parameters, while being amenable to theoretical analysis.

Within the setup of deep linear networks, several works have focused on the
analysis of convergence rate Saxe2014; arora2018convergence;
arora2018optimization; du2019width, on understanding inductive biases like
implicit regularization laurent2018deep; ji2018gradient;
gunasekar2018implicit; saxe_mathematical_2019; arora2019implicit and
understanding generalization dynamics lampinen2018analytic; poggio2018theory;
pmlr-v119-huh20a. Veness2021GatedLN; Budden2020GaussianGL propose and study
the Gated Linear Networks (GLNs) as a class of backpropagation-free neural
architectures using geometric mixing. While GLNs appear similar to GDLNs,
there are several differences. GDLN is a framework that schematizes how
pathways of information flow impact learning dynamics within architecture and
studies networks trained using back-propagation. Additionally, GLNs are good
at online learning and continual learning, while in this work, we use the GDLN
framework for understanding zero-shot generalization capabilities. In the GLN
model, a neuron is defined as a gated geometric mixer of the output of linear
networks in the previous layer, while in the GDLN model, the neurons are
linear networks where the input is the output of the linear network along the
previous path. In the GLN setup, multiple input-to-output connections (in
successive layers) can be active for the same input, while in the GDLN setup,
only one input-output connection (in successive layers) is active for one
input.

In this work, we propose modeling the model architecture as a graph and study
how pathways of information flow impact learning dynamics within an
architecture. Previous works have also proposed analyzing neural networks as
directed graphs using the Complex Network Theory boccaletti2006complex.
scabini2021structure analyzed the structure and performance of fully connected
neural networks, zambra2020emergence focused on the emergence of motifs in
fully connected networks, testolin2020deep study deep belief networks using
techniques from Complex Network Theory literature and la2021characterizing
focused on convolution and fully connected networks, with ReLU non-linearity.

Our Gated Deep Linear Network framework is closely related to areas like
modular networks (happel1994design; sharkey1997modularity; auda1999modular;
johnson2017clevr; santoro2017simple), routing networks (rosenbaum2019routing)
and mixture of experts (jacobs1991adaptive; jordan1994hierarchical;
chen1999improved; yuksel2012twenty). In these works, the common theme is to
learn a set of modules (or experts) that can be composed (or selected) using a
controller (or a router). The modules are generally instantiated as neural
networks, while the controller can either be a neural network or a hand-
designed policy. These approaches have been prominently used in natural
language processing (shazeer2017outrageously; lepikhin2020gshard;
fedus2021switch; lewis2021base), computer vision (ahmed2016network;
gross2017hard; yang2019condconv; wang2020deep) and reinforcement learning
(multi-task-rl-with-soft-modularization; sodhani2021multi; andreas2017modular;
goyal2019reinforcement; he2016opponent; Goyal2021RecurrentIM).

Our work is also related to previous works in systematic generalization
bahdanau2018systematic; bahdanau2019closure; lake2019compositional;
ruis2020benchmark; gontier2020measuring and multi-task learning
caruana1997multitask_learning;
zhang2014facial_landmark_detection_by_deep_multitask_learning;
kokkinos2017ubernet;
radford2019language_models_are_unsupervised_multitask_learners;
ruder2017overview; end_to_end_multi_task_learning_with_attention;
mott2019towards; electronics9091363. Specifically, we explore the role of
model architecture and weight initialization on models’ ability to exhibit
systematic generalization and multi-task learning.

##  7 Conclusion

A core intuition behind deep learning holds that the structure of naturalistic
problems is compositional, and this structure should be reflected in
architectural connectivity. Here we have introduced the Gated Deep Linear
Network framework, which reveals how architecture and nonlinearity–in the form
of simple gating–control pathways of information flow that govern learning
dynamics, representation learning, and ultimately generalization abilities.
Our exact reductions and solutions show that learning dynamics take the form
of a race, with greater representational reuse causing faster learning,
imparting a bias toward shared representations. We validate our key insights
on naturalistic datasets and with relaxed assumptions. An interesting future
research direction will be to explore mechanisms for inferring the optimal
architecture and gating for a given setup.

## Acknowledgements

We thank Hannah Sheahan, Timo Flesch, Devon Jarvis, and Olivier Delalleau for
their feedback and comments. This work was supported by a Sir Henry Dale
Fellowship from the Wellcome Trust and Royal Society (grant number
216386/Z/19/Z) to A.S., and the Sainsbury Wellcome Centre Core Grant from
Wellcome (219627/Z/19/Z) and the Gatsby Charitable Foundation (GAT3755). A.S.
is a CIFAR Azrieli Global Scholar in the Learning in Machines & Brains
program.

##  Appendix A Simple Nonlinear Classification

###  A.1 The XoR task

We consider the XoR task with $P=4$ data points that lie at $[\pm 1~{}\pm
1]^{T}$ as depicted in Fig. 3a. The task is exclusive-or on the input bits,
with a target output of $y=1$ if true and $y=-1$ otherwise. We solve this with
a GDLN containing four pathways (Fig. 3b), with each pathway active on exactly
one of the four examples. By symmetry, all pathways will have the same loss
dynamics and so we need only solve one. Consider the pathway active for the
example $x=[1~{}1]^{T},y=1$, that is, whose gating variable $g_{1}=1$ when
this example is presented and $g_{1}=0$ on the remaining three examples. The
relevant dataset correlations are

$\displaystyle\Sigma^{yx}$ $\displaystyle=$ $\displaystyle\langle
g_{1}yx^{T}\rangle=1/P\begin{bmatrix}1&1\end{bmatrix}$ (18)
$\displaystyle\Sigma^{x}$ $\displaystyle=$ $\displaystyle\langle
g_{1}xx^{T}\rangle=1/P\begin{bmatrix}1&1\\\ 1&1\end{bmatrix}$ (19)
$\displaystyle\Sigma^{y}$ $\displaystyle=$ $\displaystyle\langle
g_{1}yy^{T}\rangle=1/P$ (20)

The singular value decomposition of the input-output correlations yields the
nonlinear singular value $s=\sqrt{2}/P$ and input singular vector
$v=[1/\sqrt{2}~{}1/\sqrt{2}]^{T}$. Applying this singular vector to
diagonalize the input correlations yields the associated input variance
$d=v^{T}\Sigma^{x}v=2/P$. The effective singular value dynamics of this
pathway is given by the deep linear network dynamics with these correlations
(see Saxe2014; saxe_mathematical_2019), yielding

$a(t)=\frac{s/d}{1-(1-\frac{s}{da_{0}})e^{-2st/\tau}}$ (21)

where $a(t)$ is the singular value in the product of both weight matrices in
the pathway, and $a_{0}$ is the initial effective singular value, related to
the initialization variance. Finally the loss for this pathway is the loss
trajectory $l(t)$ of the associated deep linear network. The total loss, by
symmetry, is the loss from all four pathways $\mathcal{L}(t)=Pl(t)$,

$\displaystyle\mathcal{L}(t)$ $\displaystyle=$
$\displaystyle\frac{1}{2}-Psa(t)+\frac{P}{2}da(t)^{2}$ (22) $\displaystyle=$
$\displaystyle\frac{1}{2}-\sqrt{2}a(t)+a(t)^{2}.$ (23)

This analytical expression is exact for GDLNs initialized in the decoupled
regime, and it agrees closely with the dynamics of standard ReLU networks
trained end-to-end on the task starting from small random weights (Fig. 3c).
Hence GDLNs can learn nonlinear input-output tasks, and in certain settings,
describe the dynamics of standard ReLU networks when the gating structure is
chosen appropriately.

###  A.2 Nonlinear Contextual Classification

As another simple example, consider a nonlinear contextual classification
problem that cannot be solved using deep linear networks but can be solved
using the gated deep linear network, again highlighting that the gated
networks are more expressive than their non-gated counterpart.

Consider receiving two-dimensional inputs $x\in\mathbb{R}^{2}$ where each
component $x_{i},i=1,2$ is drawn from a uniform distribution between -1 and 1.
The task of the network is to classify stimuli based either on the first or
second input component. That is, the target output is $y=x_{c}$ in context
$c\in{1,2}$, and each context appears with probability 1/2. In this simple
scenario (a variant of the XoR task), the same input must be treated in two
different ways depending on context, and nonlinearity is required for solving
it correctly.

Now we must choose a gating structure. If we choose a single pathway that is
always active ($g=1$ for all samples), then we recover a deep linear network.
The resulting correlation matrices are

$\displaystyle\Sigma^{yx}$ $\displaystyle=$ $\displaystyle\langle
yx^{T}\rangle=[1/6~{}1/6]$ (24) $\displaystyle\Sigma^{x}$ $\displaystyle=$
$\displaystyle\langle xx^{T}\rangle=1/3I$ (25)

where $I$ is the $2\times 2$ identity matrix. Under the resulting dynamics,
the total weights converge to the linear least squares solution
$W^{tot}=\Sigma^{yx}(\Sigma^{x})^{-1}=[1/2~{}1/2]$, the best solution
attainable by the linear network.

Alternatively, we can set the gating variables such that a different pathway
is active in each context. We then have the collection of correlation matrices

$\displaystyle\Sigma^{yx}(1)$ $\displaystyle=$ $\displaystyle[1/6~{}0]$ (26)
$\displaystyle\Sigma^{yx}(2)$ $\displaystyle=$ $\displaystyle[0~{}1/6]$ (27)
$\displaystyle\Sigma^{x}(1,1)$ $\displaystyle=$
$\displaystyle\Sigma^{x}(2,2)=1/6I$ (28) $\displaystyle\Sigma^{x}(1,2)$
$\displaystyle=$ $\displaystyle\Sigma^{x}(2,1)=0.$ (29)

We thus see that each pathway faces a subproblem defined by just one context.
For this simple case, the pathways converge to their respective linear least
squares solutions. In particular, $W^{tot}(1)=[1~{}0]$ and
$W^{tot}(2)=[0~{}1]$, such that each pathway picks out the correct input
coordinate. In combination with the gating scheme, these weights exactly solve
this nonlinear task, showing that gated linear networks are more expressive
than linear networks. Interestingly, neuroimaging and electrophysiological
recordings from this paradigm suggest that this type of solution is observed
in the human and primate brain, as well as in standard ReLU networks trained
in the “rich” feature learning regime flesch_orthogonal_2022.

##  Appendix B Gradient flow dynamics

The gradient flow equations are

$\displaystyle\tau\frac{d}{dt}W_{e}$ $\displaystyle=$
$\displaystyle-\frac{\partial\mathcal{L}(\\{W\\})}{\partial W_{e}}\quad\forall
e\in E$ (30) $\displaystyle=$
$\displaystyle\left\langle\sum_{p\in\mathcal{P}(e)}g_{p}W_{\bar{t}(p,e)}^{T}\left[y_{t(p)}x^{T}_{s(p)}-h_{t(p)}x_{s(p)}^{T}\right]W^{T}_{\bar{s}(p,e)}\right\rangle_{y,x,g}$
(31) $\displaystyle=$
$\displaystyle\left\langle\sum_{p\in\mathcal{P}(e)}g_{p}W_{\bar{t}(p,e)}^{T}\left[y_{t(p)}x^{T}_{s(p)}-\sum_{j\in\mathcal{T}(t(p))}g_{j}W_{j}x_{s(j)}x_{s(p)}^{T}\right]W^{T}_{\bar{s}(p,e)}\right\rangle_{y,x,g}$
(32) $\displaystyle=$
$\displaystyle\sum_{p\in\mathcal{P}(e)}W_{\bar{t}(p,e)}^{T}\left[\left\langle
g_{p}y_{t(p)}x^{T}_{s(p)}\right\rangle_{y,x,g}-\sum_{j\in\mathcal{T}(t(p))}W_{j}\left\langle
g_{j}x_{s(j)}x_{s(p)}^{T}g_{p}\right\rangle_{y,x,g}\right]W^{T}_{\bar{s}(p,e)}$
(33) $\displaystyle=$
$\displaystyle\sum_{p\in\mathcal{P}(e)}W_{\bar{t}(p,e)}^{T}\left[\Sigma^{yx}(p)-\sum_{j\in\mathcal{T}(t(p))}W_{j}\Sigma^{x}(j,p)\right]W^{T}_{\bar{s}(p,e)},$
(34)

where we have simply rearranged terms and used the linearity of expectation.

The dynamics reduction can then be obtained by applying the change of
variables,

$\displaystyle\tau\frac{d}{dt}W_{e}$ $\displaystyle=$
$\displaystyle\sum_{p\in\mathcal{P}(e)}W_{\bar{t}(p,e)}^{T}\left[\Sigma^{yx}(p)-\sum_{j\in\mathcal{T}(t(p))}W_{j}\Sigma^{x}(j,p)\right]W^{T}_{\bar{s}(p,e)}$
(35) $\displaystyle\tau\frac{d}{dt}\left(R_{t(e)}B_{e}R_{s(e)}^{T}\right)$
$\displaystyle=$
$\displaystyle\sum_{p\in\mathcal{P}(e)}\left(U_{t(p)}B_{\bar{t}(p,e)}R_{t(e)}^{T}\right)^{T}\left[U_{t(p)}S(p)V_{s(p)}^{T}-\right.$
(37)
$\displaystyle\left.\sum_{j\in\mathcal{T}(t(p))}U_{t(j)}B_{j}V_{s(j)}^{T}V_{s(j)}D(j,p)V_{s(p)}^{T}\right]\left(R_{s(e)}B_{\bar{s}(p,e)}V_{s(p)}^{T}\right)^{T}$
$\displaystyle\tau\frac{d}{dt}B_{e}$ $\displaystyle=$
$\displaystyle\sum_{p\in\mathcal{P}(e)}B_{\bar{t}(p,e)}\left[S(p)-\sum_{j\in\mathcal{T}(t(p))}B_{j}D(j,p)\right]B_{\bar{s}(p,e)}$
(38)

where we have used the fact that $R_{v}^{T}R_{v}=I$ for all nodes, and the
fact that $U_{t(j)}=U_{t(p)}$ by definition of the set $\mathcal{T}(t(p))$.
From this we see that if the $B$ variables are initially diagonal they remain
so under the dynamics. In this case, the dynamics decouple and each element
along the diagonal of the $B$ matrices evolves independently of the rest.

##  Appendix C Routing task and network reduction

To understand the dynamics in the pathway network, we first collect the
relevant input statistics. We have

$\displaystyle\Sigma^{yx}(p)$ $\displaystyle=$ $\displaystyle\left\langle
g_{p}y_{t(p)}x_{s(p)}^{T}\right\rangle$ (39) $\displaystyle=$
$\displaystyle\textrm{Pr}(g_{p}=1)USV^{T}$ (40) $\displaystyle=$
$\displaystyle\frac{1}{KM}USV^{T}$ (41) $\displaystyle\Sigma^{x}(j,p)$
$\displaystyle=$ $\displaystyle\left\langle
g_{j}x_{s(j)}x_{s(p)}^{T}g_{p}\right\rangle$ (42) $\displaystyle=$
$\displaystyle\begin{cases}\frac{1}{KM}VDV^{T}\quad\textrm{if}~{}j=p\\\
0\quad\textrm{otherwise}\end{cases}$ (43)

because there are $KM$ total trained paths from input to output and all
pathways are gated off except for the active pathway.

Inserting these data statistics into (10), and assuming that initial singular
values are equal for all input domains and all output domains (a reasonable
approximation when starting from small random weights), we can track only the
variables $B_{1},B_{2},$ and $B_{3}$ encoding the singular values in the
input, hidden, and output weights respectively.

Next, we note that the first and third layer weights are active in $K$ tasks
(all tasks originating from a given input or terminating at a given output,
respectively), while second layer weights are active in all $KM$ tasks. This
yields the reduced dynamics

$\displaystyle\tau\frac{d}{dt}B_{1}$ $\displaystyle=$
$\displaystyle\frac{1}{M}B_{3}B_{2}\left[S-B_{3}B_{2}B_{1}D\right]$ (44)
$\displaystyle\tau\frac{d}{dt}B_{2}$ $\displaystyle=$ $\displaystyle
B_{3}B_{1}\left[S-B_{3}B_{2}B_{1}D\right]$ (45)
$\displaystyle\tau\frac{d}{dt}B_{3}$ $\displaystyle=$
$\displaystyle\frac{1}{M}B_{2}B_{1}\left[S-B_{3}B_{2}B_{1}D\right].$ (46)

If we consider ‘balanced’ initial conditions where $B_{1}(0)=B_{3}(0)$, we
have

$\displaystyle\tau\frac{d}{dt}B_{1}$ $\displaystyle=$
$\displaystyle\frac{1}{M}B_{2}B_{1}\left[S-B_{2}B_{1}^{2}D\right]$ (47)
$\displaystyle\tau\frac{d}{dt}B_{2}$ $\displaystyle=$ $\displaystyle
B_{1}^{2}\left[S-B_{2}B_{1}^{2}D\right],$ (48)

recovering Eqns. (13)-(14) of the main text.

To estimate the ratio of singular values in the first layer to that in the
second, we consider its time derivative and calculate the steady state. We
have

$\displaystyle\frac{d}{dt}B_{2}/B_{1}$ $\displaystyle=$ $\displaystyle
B_{1}\left[S-B_{2}B_{1}^{2}D\right]-\frac{1}{M}B_{2}^{2}/B_{1}\left[S-B_{2}B_{1}^{2}D\right]$
(49) $\displaystyle 0$ $\displaystyle=$ $\displaystyle
B_{1}-\frac{1}{M}B_{2}^{2}/B_{1}$ (50) $\displaystyle B_{2}$ $\displaystyle=$
$\displaystyle\sqrt{M}B_{1}.$ (51)

Hence if training continues for long times (such that the error term does not
become zero), the shared portion of the pathway changes more by a factor
$\sqrt{M}$ (and this ratio does not depend on $K$).

We can extend this analysis to the case where all input-output routes are
trained but the gating structure sends only $P$ paths through each hidden
weight matrix, as considered in Section 4.2. With this gating scheme, the
reduction is

$\displaystyle\tau\frac{d}{dt}B_{1}$ $\displaystyle=$
$\displaystyle\frac{\sqrt{P}}{M^{2}}B_{2}B_{1}\left[S-B_{2}B_{1}^{2}D\right]$
(52) $\displaystyle\tau\frac{d}{dt}B_{2}$ $\displaystyle=$
$\displaystyle\frac{P}{M^{2}}B_{1}^{2}\left[S-B_{2}B_{1}^{2}D\right],$ (53)

and the singular value ratio is

$\displaystyle\frac{d}{dt}B_{2}/B_{1}$ $\displaystyle=$
$\displaystyle\frac{P}{M^{2}}B_{1}\left[S-B_{2}B_{1}^{2}D\right]-\frac{\sqrt{P}}{M^{2}}B_{2}^{2}/B_{1}\left[S-B_{2}B_{1}^{2}D\right]$
(54) $\displaystyle 0$ $\displaystyle=$ $\displaystyle
PB_{1}-\sqrt{P}B_{2}^{2}/B_{1}$ (55) $\displaystyle B_{2}$ $\displaystyle=$
$\displaystyle P^{\frac{1}{4}}B_{1}.$ (56)

This ratio scales from 1 to $\sqrt{M}$ as the number of shared paths goes from
$P=1$ (no sharing) to $P=M^{2}$ (full sharing). Hence for this architecture,
greater sharing causes larger weight changes in the hidden pathway.

##  Appendix D Experimental details and further results

This section contains details and hyperparameter settings for the simulations
reported in Section 5, as well as additional visualization of results in
Figures 9 and E.1.

We start by explaining the general procedure for transforming the existing
datasets and then describe the new dataset instances that we create.

Consider a dataset $D=(X,Y)$, defined as a tuple of inputs $X$ and targets
$Y$. The dataset $D$ has $n$ datapoints, that is, $X\in R^{n\times n_{dim}}$,
where $n_{dim}$ is the dimensionality of each input111While images are multi-
dimensional arrays, they can be represented as flattened 1-d arrays.. The
dataset has $n_{class}$ unique classes, referred by their indices
$\\{0,1,\cdots,n_{class}-1\\}$. We want to create new datasets by transforming
the given dataset $D$. We assume that we have a list of $M$ input
transformations $f_{i}^{input}:R^{n_{dim}}\to R^{n_{dim}}\forall
i\in\\{0,\cdots,M-1\\}$ and $M$ output transformations
$f_{j}^{output}:R^{n_{class}}\to R^{n_{class}}\forall j\in\\{0,\cdots,M-1\\}$.
Now, we can define a new dataset, $D_{i,j}$, as $(X_{i},Y_{j})$, where
$X_{i}=f_{i}^{input}(X)$, $Y_{j}=f_{j}^{output}(Y)$ i.e any transformation of
the given dataset is a new dataset. We can apply $M$ transformations on the
input and $M$ transformations on the output to obtain $M^{2}$ datasets.

We consider the following two operations for input transformations: rotation
of the input image and permutation of pixels. For the $i^{th}$ rotation
transformation, we rotate the input images by an angle $\theta=180i/M$
degrees. For the $i^{th}$ permutation transformation, we apply a random
permutation matrix to the flattened input. Each of these transformations
provides input to one input domain. We use the permutation operation as the
output transformation, implying that each new output transformation
corresponds to a $n_{class}$-way classification task.

We use the MNIST (deng2012mnist) and CIFAR-10 datasets
(krizhevsky2009learning) to create three datasets:

MNIST-Permuted-Input-Permuted-Output-40:

$40$ transformations on both the input and the output, leading to a total of
$1600$ datasets. The input transformation is permutation of pixels and the
output transformation is permutation of the targets.

MNIST-Permuted-Input-Permuted-Output-100:

$100$ permutation transformations on both the input and the output, leading to
a total of $10^{4}$ datasets.

MNIST-Rotated-Input-Permuted-Output-40

40 rotation transformations on the input and 40 permutation transformations on
the output.

CIFAR-Rotated-Input-Permuted-Output-40

40 rotation transformations on the input and 40 permutation transformations on
the output.

In the case of CIFAR-Rotated-Input-Permuted-Output-40 dataset, use a pre-
trained ResNet18 resnet model222We use the following code for pre-training the
models: https://github.com/akamaster/pytorch_resnet_cifar10 to map the images
into $512$ dimensional vectors. We pretrain the ResNet18 model on full
CIFAR-10 dataset, freeze the pre-trained model and use the first two residual
blocks to encode the images from the transformed datasets. The output of the
(frozen) ResNet encoder is used as input to the gated network.

###  D.1 Model and training

The model consists of $M$ encoders, denoted as ($\\{\phi_{i}\forall
i\in\\{1,\cdots,M\\}\\}$) and $M$ decoders, denoted as ($\\{\psi_{i}\forall
i\in\\{1,\cdots,M\\}\\}$) and a shared hidden layer $\theta$. In the GDLN
framework, the encoders correspond to the input nodes, the decoders correspond
to the output nodes and the connection from an encoder, to the hidden layer,
to the decoder corresponds to a path. The encoders, decoders and the shared
hidden layer are all instantiated as linear networks. Given a dataset,
$D_{i,j}$, or $(x_{i},y_{j})$, we compute the prediction using the following
function: $\psi_{j}(\theta(\phi_{i}(x_{i})))$ 333Note that we overload the
notation to represent the components and the computation using the same
symbol.

###  D.2 Training and Evaluation Setup

Following the setup in Section 4, we train a subset of input-output domains
such that each input domain is trained with only $K\leq M$ output domains,
resulting in $M\times K$ trained pathways and $M\times(M-K)$ untrained
pathways. During evaluation, we report the performance on both the trained
pathways and the untrained pathways.

Figure 8: Experimental results. Error for trained (blue) and untrained
(orange) input-output domain pairs as a function of the percentage of trained
pathways ($K/M$) on: (i) CIFAR dataset, where input is rotated and output is
permuted, with $M^{2}=1600$ total tasks, (ii) MNIST dataset, where input and
output, both are permuted, with $M^{2}=10^{4}$ total tasks, (iii) MNIST
dataset, where input and output, both are permuted, with $M^{2}=1600$ total
tasks, and (iv) MNIST dataset, where input is rotated and output is permuted,
with $M^{2}=1600$ total tasks. (in the order of left to right). Training
accuracy is always high while zero-shot transfer to untrained pathways becomes
as good as the training performance when $\approx$40% of pathways are trained.

Figure 9: Experimental results. Error for trained (blue) and untrained
(orange) input-output domain pairs as a function of the percentage of trained
pathways ($K/M$) on: (i) MNIST dataset, where input and output, both are
permuted, with $M^{2}=1600$ total tasks, and (ii) MNIST dataset, where input
is rotated and output is permuted, with $M^{2}=1600$ total tasks (in the order
of left to right). These models are trained with the leaky-ReLU non-linearity
with the negative slope parameter set to be $0.01$, thus making the model non-
linear. The training accuracy is always high while zero-shot transfer to
untrained pathways becomes as good as the training performance when
$\approx$25% of pathways are trained.

Figure 10: Experimental results. Error as a function of initialization scale
for trained (blue) and untrained (orange) input-output domain pairs when
training on $K/M$ fraction of pathways, where $K=10$ on (i) MNIST dataset,
where input and output, both are permuted, with $M^{2}=10^{4}$ total tasks,
(i) MNIST dataset, where input and output, both are permuted, with
$M^{2}=1600$ total tasks, and (iii) MNIST dataset, where input is rotated and
output is permuted, with $M^{2}=1600$ total tasks. (in the order of left to
right). While performance on trained domains is reasonably good for all
scales, zero-shot generalization only emerges at small inits.

##  Appendix E Additional Implementation Details

###  E.1 Libraries

We use the following open-source libraries:

  1. 1.

PyTorch 2019pytorch444https://pytorch.org/

  2. 2.

Hydra Yadan2019Hydra555https://github.com/facebookresearch/hydra

  3. 3.

Numpy harris2020array666https://numpy.org/

  4. 4.

Pandas reback2020pandas777https://pandas.pydata.org/

  5. 5.

ResNet Implementation for CIFAR10/CIFAR100 in PyTorch
Idelbayev18a888https://github.com/akamaster/pytorch_resnet_cifar10

  6. 6.

Xplogger
Sodhani_xplogger_Logging_utility_2022999https://github.com/shagunsodhani/xplogger

Table 1: Hyperparameter values common across all the task distributions Hyperparameter |  Hyperparameter values
---|---
Batch size (per task) |  8
Number of epochs |  $1000$
Number of classes |  $10$
Dimensionality of the hidden layer |  $128$
Gain Ratio for weights |  $10.0,1.0,0.1,0.01,0.001,0.0001$
Initial value for bias |  $0.0$
Optimizer |  SGD
Learning Rate |  $0.0001$
Momentum |  $0.9$
Table 2: Hyperparameter values for MNIST-Permuted-Input-Permuted-Output-40 (different from the values described in Table 1 Hyperparameter |  Hyperparameter values
---|---
Number of input transformations |  $40$
Number of output transformations |  $40$
Number of encoders |  $40$
Number of decoders |  $40$
Table 3: Hyperparameter values for MNIST-Permuted-Input-Permuted-Output-100 (different from the values described in Table 1 Hyperparameter |  Hyperparameter values
---|---
Number of input transformations |  $100$
Number of output transformations |  $100$
Number of encoders |  $100$
Number of decoders |  $100$
Table 4: Hyperparameter values for MNIST-Rotated-Input-Permuted-Output-100 (different from the values described in Table 1 Hyperparameter |  Hyperparameter values
---|---
Number of input transformations |  $100$
Number of output transformations |  $100$
Number of encoders |  $100$
Number of decoders |  $100$
Table 5: Hyperparameter values for CIFAR-Rotated-Input-Permuted-Output-40 (different from the values described in Table 1 Hyperparameter |  Hyperparameter values
---|---
Number of input transformations |  $40$
Number of output transformations |  $40$
Number of encoders |  $40$
Number of decoders |  $40$
Dimensionality of the hidden layer |  $1024$