summaries: This paper tackles the problem of finding a sparse network architecture before training, so as to facilitate training on resource-constrained platforms. To this end, the authors propose to replace dense layers with series of sparsely-connected linear layers. They then study how to initialize such sparse layers to avoid gradient vanishing. Furthermore, they propose an approach to finding the best topology by measuring how well the sparse layers can approximate random weights of their dense counterparts.

summaries: This paper proposes to replace dense layers with multiple sparse linear layers.  The idea is that the product ABC (for A, B, C all sparse matrices)  can accurately approximate a dense matrix D, but A(B(Cx))) requires much less work than Dx.  The paper then continues with the assumption that topology of the sparse matrices should be fixed before training, and that given this assumption we would like to find the "best" fixed topology to pick.  The paper introduces a new task to determine the "best" topology - that of reconstructing a random dense matrix.  On somewhat of a tangent the paper also introduces a minor modification to the Xavier initialization scheme that works better for deep stacks of sparse layers.

abstractText: Long training times of deep neural networks are a bottleneck in machine learning research. The major impediment to fast training is the quadratic growth of both memory and compute requirements of dense and convolutional layers with respect to their information bandwidth. Recently, training ‘a priori’ sparse networks has been proposed as a method for allowing layers to retain high information bandwidth, while keeping memory and compute low. However, the choice of which sparse topology should be used in these networks is unclear. In this work, we provide a theoretical foundation for the choice of intra-layer topology. First, we derive a new sparse neural network initialization scheme that allows us to explore the space of very deep sparse networks. Next, we evaluate several topologies and show that seemingly similar topologies can often have a large difference in attainable accuracy. To explain these differences, we develop a data-free heuristic that can evaluate a topology independently from the dataset the network will be trained on. We then derive a set of requirements that make a good topology, and arrive at a single topology that satisfies all of them.

authors:
Sajid Anwar
Kyuyeon Hwang
Wonyong Sung

title: Structured pruning of deep convolutional neural networks

authors:
Wenlin Chen
James T. Wilson
Stephen Tyree
Kilian Q. Weinberger
Yixin Chen

title: Compressing neural networks with the hashing

venue: trick. CoRR,

authors:
Caiwen Ding
Siyu Liao
Yanzhi Wang
Zhe Li
Ning Liu
Youwei Zhuo
Chao Wang
Xuehai Qian
Yu Bai
Geng Yuan
Xiaolong Ma
Yipeng Zhang
Jian Tang
Qinru Qiu
Xue Lin
Bo Yuan

title: Circnn: Accelerating and compressing deep neural networks using block-circulant weight matrices

venue: In Proceedings of the 50th Annual IEEE/ACM International Symposium on Microarchitecture,

venue: Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics,

authors:
Song Han
Huizi Mao
William J. Dally

title: Deep compression: Compressing deep neural network with pruning, trained quantization and huffman

venue: coding. CoRR,

authors:
Song Han
Jeff Pool
John Tran
William J. Dally

title: Learning both weights and connections for efficient neural networks

authors:
Babak Hassibi
David G. Stork

title: Second order derivatives for network pruning: Optimal brain surgeon

venue: Advances in Neural Information Processing Systems

authors:
Mikael Henaff
Arthur Szlam
Yann LeCun

title: Orthogonal rnns and long-memory

venue: tasks. CoRR,

authors:
Andrew G. Howard
Menglong Zhu
Bo Chen
Dmitry Kalenichenko
Weijun Wang
Tobias Weyand
Marco Andreetto
Hartwig Adam

title: MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications

authors:
Mihailo Isakov
Alan Ehret
Michel A. Kinsy

title: Closnets: Batchless dnn training with on-chip a priori sparse neural topologies

authors:
Yann LeCun
John S. Denker
Sara A. Solla

title: Optimal brain damage

authors:
Baoyuan Liu
Min Wang
Hassan Foroosh
Marshall Tappen
Marianna Pensky

title: Sparse convolutional neural networks

venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,

authors:
Dominic Masters
Carlo Luschi

title: Revisiting small batch training for deep neural networks

authors:
Sören Pirk
Mohi Khansari
Yunfei Bai
Corey Lynch
Pierre Sermanet

title: Online object representations with contrastive learning

authors:
Ameya Prabhu
Girish Varma
Anoop M. Namboodiri

title: Deep expander networks: Efficient deep networks from graph theory

authors:
Ryan A. Robinett
Jeremy Kepner

title: Radix-net: Structured sparse matrices for deep neural networks

authors:
Andrew Saxe
James Mcclelland
Surya Ganguli

authors:
Duncan J. Watts
Steven H. Strogatz

title: Collective dynamics of ’small-world

venue: networks. Nature,

authors:
Wei Wen
Chunpeng Wu
Yandan Wang
Yiran Chen
Hai Li

title: Learning structured sparsity in deep neural networks

authors:
Ming Yuan
Yi Lin

title: Model selection and estimation in regression with grouped variables

venue: Journal of the Royal Statistical Society Series B, 68:49–67,

heading: 1 INTRODUCTION

text:
Training deep neural networks requires both powerful hardware and a significant amount of time. Long training times are a significant bottleneck to deep learning research, as researchers typically iteratively design and test new architectures for a specific problem. While a lot of research has been dedicated to accelerating inference, we investigate training as (1) accelerating training can speed up research iteration, (2) evolutionary algorithms for DNN architecture exploration are increasingly being used as an alternative to domain expertise (Jaderberg et al., 2017), and network training is moving to edge devices (Pirk et al., 2019). Unfortunatelly, the memory requirements of dense, convolutional and recurrent layers grow quadratically with layer information bandwidth 1. In other words, doubling the size of layer inputs and outputs quadruples the size of the layer. This causes majority of the networks to be memory-bound, making DNN training impractical without batching, a method where training is performed on multiple inputs at a time and updates are aggregated per batch. While batching alleviates the pressure on DRAM bandwidth, it can decrease model accuracy (Masters & Luschi, 2018) especially when scaling training on large clusters (Akiba et al., 2017). Furthermore, larger models in off-chip memory become dominant energy cost (Han et al., 2015a), complicating on-line training on battery-power devices.
Conventional dense and convolutional layers do not offer the user to individually tune layer size and the number of layer inputs and outputs. In this work, we seek a method to decouple the information bandwidth from layer expressivity. Such a method would allow us to (1) speed up training networks by storing them in on-chip memory, (2) remove the memory bottleneck and the need for batching, (3) allow more efficient training on distributed systems, and (4) reduce the energy consumption due to the excessive compute and storage requirements of modern DNNs, potentially allowing us to move training to edge devices. Several works have proposed a priori structured sparsity (Prabhu et al., 2017; Isakov et al., 2018) or weight sharing (Ding et al., 2017) to allow training simpler but ‘wider’ models. A priori sparsity, where the sparse network topology is selected before training has started, is a promising approach that allows the user to finely and transparently tune the ratio of information bandwidth to memory requirements. If the topology is structured, efficient software
1We define a layer’s information bandwidth as the number of independent signals passing through that layer.
or hardware implementations can be built to accelerate processing with dense network performance . However, before custom architectures or low-level kernels can be built, a general theory of why certain topologies perform – or underperform – is needed. To the best of our knowledge, no work yet tackles the question of the existence of a ‘best’ topology for sparse neural network training. This paper provides an answer on how a topology should be selected.
Our contributions are as following:
• We propose a sparse cascade architecture that can replace dense or convolutional layers without affecting the rest of the network architecture.
• We develop a sparse neural network initialization scheme that allows us to train very deep sparse networks without suffering from the vanishing gradient effect.
• We evaluate sevaral topologies on a matrix reconstruction task and show that the choice of topology has a strong effect on attainable network accuracy.
• In order to evaluate topologies independently of a dataset, we develop a data-free heuristic for predicting the expressiveness of a given sparse network.
• From the heuristic, we derive requirements that make a good topology, and settle on a single family of sparse networks.

heading: 2 RELATED WORK

text:
We classify methods that arrive at a sparse network into those that enforce sparsity before, during, or after training. The following is a brief description of each class.
Enforcing sparsity after training: In this class of methods, certain weights are zero-ed out after training has finished. This approach has the benefit of first discovering the baseline model accuracy, allowing the training mechanism to evaluate the accuracy to size trade-off. Since training is performed using the dense model, only inference can benefit from these post-training pruning methods. One of the early pruning methods are Optimal Brain Damage (LeCun et al., 1990) and Optimal Brain Surgeon (Hassibi & Stork, 1993), where authors remove weights based on the second derivative of the loss w.r.t. to each weight. The insight here is that removing a weight causes some perturbation in the network, and by picking weights with the smallest second derivative of the loss, the effect of the perturbation on the network functionality is be minimized. DeepCompression (Han et al., 2015a;b) uses a similar approach, but replaces the Hessian-based metric with weight magnitudes. Authors show that high (>95%) sparsity can be achieved as long as networks are finetuned after pruning to restore performance. Alternatively, in (Liu et al., 2015) authors decompose convolutional layers into a set of per-channel basis kernels, which are applied to input feature maps, and a sparse kernel matrix that mixes outputs of the basis kernels into the output feature maps. However, all of these methods lead to unstructured sparsity that is difficult to take advantage of. Structured sparsity, where some assumptions can be made on the structure of sparse kernels, has been explored as a way to improve the execution efficiency of sparse structures on GPUs and CPUs. In (Anwar et al., 2015), authors use particle filters to prune whole channels or kernels. Similarly, Kadetotad et al. explore structured intra-layer sparsity, where instead of individual weights, small blocks of weights are pruned.
Enforcing sparsity during training: Instead of pruning after training, pruning can also be applied during training. This has the benefit of potentially reducing the computational load of the training phase, however, the device performing training must still be able to store the whole dense model at the beginning of training. L1 regularization or L1 weight decay is known to cause sparsity during training, as unlike in the case of L2 regularization, all weights will equally be incentivized to approach zero. However, L1 weight decay often causes a decrease in accuracy, and the sparsity is unstructured. In (Wen et al., 2016), authors use Group Lasso (Yuan & Lin, 2006) regularization to enforce the sparsity of more coarse-grained structures instead of individual weights.
Enforcing sparsity before training: Model size can also be reduced before training has started. We focus on layer-level methods and not architecture level approaches, as they are orthogonal. Majority of works reducing the size of layers before training have focused either on a priori sparsity or weight reuse. On the weight reuse side, HashedNets (Chen et al., 2015) use a hash function to group multiple weights and have them share and train a single value. CirCNN (Ding et al., 2017) uses block-circulant matrices for storing weights, where elements are shared in a predictable manner and
Fourier transforms are used for inference, reducing the computational complexity of both inference and training.
On the a priori sparsity side, several topologies have been proposed in literature. Deep Expander Networks (X-Nets) (Prabhu et al., 2017) replace dense layers with sparse layers with the expander graph topology. Authors give guarantees of each input neuron being connected to each output neuron within a logarithmic number of layers. Similarly, RadiX-Nets (Robinett & Kepner, 2019) build on X-nets but use the radix topology instead of graph expanders. Alternatively, ClosNets (Isakov et al., 2018) replace a single dense layer with a cascade of three sparse layers with the Clos topology. Clos topology guarantees full connectivity, and has a tunable parameter for the path diversity between all inputs and outputs. While deep expander networks grow in depth with the number of neurons per layer, ClosNets grow in width.
None of the above a priori sparse network works give a definitive answer to which topology maximizes performance per weight. In this work we aim to answer that question.

heading: 3 APPROACH

text:
The number of parameters in a neural network layer is decided by the number of input and output neurons (in case of fully-connected networks), or the number of input and output channels (in the case of convolutional networks). This prevents decoupling the network bandwidth (i.e. number of inputs or outputs of a certain layer) and the parameter count. We propose that sparsifying layers can allow users to train wider networks without the quadratic growth in network size. Our approach is to replace each fully-connected layer with a cascade of sparsely-connected layers, as shown in Figure 1. The cascade topology and depth are selected at design time so that the number of parameters in the cascade is lower than in the original dense layer. Hidden layer neurons in the cascade have a linear activation function, while the output neurons use the activation of the original network. The cascade needs to have certain properties such as connectivity between all cascade input-output pairs, a small parameter count, and hardware efficiency. In Section 7 we explore the topology requirements further.
Similarly, a priori pruning can be applied to convolutional networks. In conventional CNN layers, each input channel is connected to each output channel by a convolution. For a filter of size f × f , c input and k output channels, the convolutional layer has f2ck parameters. Since the number of input and output channels c and k directly control both the information bandwidth and the size of the network, we propose to disentangle the number of input/output features and the number of convolutional filters. We adopt the architecture of MobileNets (Howard et al., 2017) and break up convolutional layers into depthwise and pointwise convolutions, as seen in Figure 2a. In the original MobileNets, the majority of parameters belong to the pointwise convolutions, which are simply dense neural networks applied to each ‘pixel’ individually. We propose to prune only the pointwise convolutions in the same manner we prune dense layers.

heading: 4 INITIALIZING A PRIORI SPARSE NEURAL NETWORKS

text: Initializing deep neural networks highly affects the training accuracy of the models, but is often not given as much attention as network architecture. If a network is suboptimally initialized, it may train slower than a correctly initialized one or may even not train at all. This problem has been extensively researched in the context of deep neural networks (Glorot & Bengio, 2010; He et al., 2015a) and recurrent neural networks (Henaff et al., 2016). Our proposed approach replaces fully-connected layers with cascades of sparse layers. As such, the new network may be several times deeper than the original one, and may not train as efficiently. This problem is further compounded by the fact that our networks are a priori sparse. Our tests show that deep sparse networks initialized with common initialization schemes like Xavier initalization (Glorot & Bengio, 2010) completely fail to learn. By observing the activation and error values of deep sparse networks, we deduce that these networks suffer from the vanishing gradient problem, i.e., with each successive layer, the variance of both the activations and the errors drops exponentially. In this section, we develop a new initialization scheme for a priori sparse networks that alleviates the vanishing gradient problem by taking layer sparsity into account.

heading: 4.1 SPARSE XAVIER INITIALIZATION

text:
In Appendix A, we briefly cover the original derivation of the Xavier initialization. Here we generalize it to apply to sparse networks as well. We construct a sparse layer from a matrix W ∈ Rm×n by multiplying it element-wise with the mask M ∈ {0, 1}m×n with sparsity s ∈ [0, 1]. For a random topology, each element Mij of the mask is set as Mij = Ber(s), where Ber is the Bernoulli distribution. For a layer Wi with nin input and nout output neurons, an output neuron’s activation variance depends on the variance of each input neuron, each weight connected to it, and the number of input neurons (Appendix Equation 21). For sparse networks, each output neuron is on average only connected to nin(1− s) neurons, hence we update Equation 21 as:
σ2(an+1) = nin(1− s)σ2(an)σ2(Wn+1) σ2(δn) = nout(1− s)σ2(δn+1)σ2(Wn+1)
(1)
Updating the Xavier initialization to take sparsity into account, we write our sparse initialization as: W ∼ U [ − √ 6√
(nj + nj+1)(1− s) ,
√ 6√
(nj + nj+1)(1− s)
] (2)
We test the new initialization on the MNIST dataset with networks of different sparsities and depths (Figure 2b). We train randomly-connected networks with 256 neurons in the hidden layers, 1 to 20 hidden layers, and with sparsities between 0 and 255/256. Using the sparse Xavier initialization, we are able to train deep sparse networks. For very sparse networks (sparsity of 63/64 and higher), often there exists no path between certain inputs and outputs, limiting trainability. A better, non-random topology with the same amount of parameters may however be able train.

heading: 5 TOPOLOGY EXPLORATION

text:
The choice of topology has a strong impact on both the accuracy and the parallelizability of a sparse network. In this section we aim to (1) answer how two topologies can be compared, and (2) create a metric for evaluating a topology independently of a task, given that we can assume nothing about training data beforehand. We first devise a task that allows us to experimentally evaluate a topology. We choose a matrix reconstruction problem where an original matrix WO ∈ Rn×n is reconstructed as a product of l sparse matrices Si ∈ Rn×n with adjacency matrices Mi ∈ [0, 1]n×n as:
L(WO,M1, ...,Ml) = min ∥∥∥∥∥WO − l∏
i=1
(Si Mi) ∥∥∥∥∥ (3) The matrix WO must be random, so that the network cannot abuse any regularities within it. The topology we seek should perform well independently of the matrix structure. Though topologies derived for a specific task may perform better, on average across all tasks, the general topology should achieve the best results. A number of loss functions can be used for this evaluation but we restrict ourselves to L2 loss for now. To gain intuition into the problem, we use the reconstruction loss in Equation 3 on a number of common topologies. Figure 3 illustrates the impact of topology
choice on overall network accuracy. Our reasoning is that as topologies can underperform on an unseen task, there exists a ‘best’ topology, one that on average achieves optimal performance. In this and the following section we aim find that topology.
Though we may arrive at an optimal topology purely through evolution, this approach has several issues: (1) evaluating a topology on random matrices is inherently noisy, and we would have to reconstruct many random matrices in order to compare the fitness of two different topologies, (2) the evolved topology is only a point solution, and we would have to rerun the process for a topology with a different number of inputs, outputs, parameters, or layers, and (3) the evolved topologies tell us nothing of the underlying reasons for why a certain topology is underperforming. Therefore, we aim to develop a heuristic that can accurately predict the quality of a topology, so that by analyzing the heuristic we can arrive at the root cause of why certain topologies are underperforming, and produce conclusions on how to construct better performing ones.

heading: 5.1 L0 CONSTRAINT SATISFACTION

text:
We revisit Equation 3 and define the sparse decomposition Wd as:
Wd = l∏ i=1 Si Mi (4)
From here we can write an individual element Wd i,j of matrix Wd from Equation 4 by looking at paths between input i and output j as:
Wd i,j = P (i,j)∑ p p∏ emn∈p wmn (5)
where P (i, j) is the set of all paths from input neuron i to output neuron j. According to Equation 5, in order to satisfy WO i,j = Wd i,j , at least one edge in P (i, j) must be set to a specific value. Given an edge emn with weight wmn that exists on some path between nodes i and j:
wmn = ( WO i,j −
P (i,j)∑ p, emn /∈p p∏ exy wxy )/( P (i,j)∑ p,emn∈p p\emn∏ exy wxy ) (6)
Due to the difficulty of analyzing the quality of topologies using an L2 loss, one approach is to use L0 loss. Here, we task the matrix decomposition Wd with perfectly reconstructing as many individual elements of WO as possible. The networks are still trained using SGD with L1 loss. After training converges, we can count the number of elements in WO −Wd where |WO −Wd| < for some arbitrarily small ∈ R. As wmn cannot take multiple values, it can only satisfy a single constraint. Therefore, if |e| is the number of edges in the network, and |csat| is the number of satisfiable constraints, we can say that 0 ≤ |csat| ≤ |e|. At best, the topology can solve as many constraints as it has edges. As we will show in Section 7, this is not possible to achieve without several modifications to the sparse network.
Equation 6 allows us to reframe the problem of finding the minimal L0 reconstruction loss as a bipartite maximal matching problem. First, we create a set of input-output constraints C = {(i, o) | i ∈ I, o ∈ O}, and a set of weights W = {(wli,j) | Ml i,j = 1}. An edge w ∈ W is connected to all constraints (i, o) ∈ C where w ∈ P (i, o). The problem of finding the largest number of elements in WO that can be perfectly reconstructed by Wd becomes a bipartite maximal matching problem, and can be solved in linear time. However, we find that the maximal number of constraints matched with edges is not a good heuristic for the performance of a topology. Constraint satisfaction counting heuristics fail to account for topology fairness, and equally rate topologies that balance the satisfaction of all input-output pairs, and topologies that greedily satisfy only a subset of input-output-pairs.

heading: 6 GRAPH CONTROLLABILITY

text: In this section, we present a continious method of evaluating the capability of a topology to reconstruct a certain graph. While averaging reconstruction results of many random matrices may provide an estimate of a topology quality, we develop a data-free approach that (1) evaluates a topology in a single pass, and (2) is not sensitive to the randomness of generated matrices.

heading: 6.1 NEURONS NEED ONLY LEARN INPUT RATIOS:

text:
In this work we focus on networks that use ReLU activations between sparse cascades, and linear activations inside them. Both ReLU and linear functions are homogenous, i.e. f(cx) = cf(x) for constants c ∈ R. Take a neuron with activation a and n inputs with activations and weights z,w ∈ Rn. Activation a can be written as:
a = f(wz) (7) where f is a homogenous function. We can extract the magnitude and normalize the vector w:
a = mf(vz), w = mv, ||v||22 = 1 (8) Since the neurons are using homogeneous activation functions, we can shift the job of learning neuron magnitude to the next layer. Now the neuron is only tasked with learning the input ratios.
In the previous section we have used L0 loss to measure the number of constraints a network can solve. Here we see that form×n constraints that exist when reconstructing a matrixW ∈ Rm×n, n of those constraints are magnitude constraints, and (m− 1)n are ratio constraints. In other words, a neuron withm inputs hasm−1 ratio and 1 magnitude constraint. In Appendix B we give a practical way of eliminating magnitude constraints and only measuring the number of ratio constraints.

heading: 6.2 NEURON CONTROL

text:
We define controllability of a certain neuron nw.r.t. to an inputs a and b as the ability of an optimizer to set the ratio of a to b at n. We give an inductive definition on three graph primitives: (1) individual input neurons, (2) a neuron with two inputs and one output, and (3) a neuron with one input and two outputs, and show how controllability propagates through them. We show how any sparse neural network can be decomposed into these primitives and how control can be calculated on neurons with larger numbers of inputs and outputs.
Definition 6.1. For a neuron n connected to a set of inputs I, we define the controllability of input a ∈ I relative to b ∈ I at n as Ca/b(n). If Ca/b(n) = 1, the optimizer can set the ratio a/b at neuron n to any value, without impacting any other ratios already set in the network.
Lemma 6.1. For inputs a, b ∈ I and any neuron n, controllability Ca/b(n) is bounded as:
0 ≤ Ca/b(n) + Cb/a(n) ≤ 1 (9)
This is understandable since the optimizer can only set the ratio of inputs a and b at neuron n to a single value, hence controllability of a to b plus controllability of b to a cannot be greater than 1.
Lemma 6.2. For input neurons a, b, c ∈ I, controllability of ratio a/b at neuron c is:
Ca/b(c) = 0 (10)
This is obvious since the optimizer has no control over network inputs, and we will use this lemma as the base case of our induction.
Lemma 6.3. For an neuron n connected to a set of inputs I, total controllability of n is limited as: I∑ a I∑ b Ca/b(n) ≤ |I| − 1 (11)
This lemma is a direct result of Section 6.1 limit on the number of ratios.
We now analyze two graph primitives that show how control propagates through the graph.
Theorem 6.4. Control aggregation: For a neuron n with two input neurons i and j connected with weights win and wjn, where Ii and Ij are the sets of inputs directly or indirectly connected to neurons i and j, the controllability Ca/b(n), a, b ∈ Ii ∪ Ij is:
Ca/b(n) = min(1, max(0, (Ca/b(i) + Ca/b(j) + ∆Ca/b(n)))) (12)
where I∑ a I∑ b ∆Ca/b(n) ≤ { 0, neither win or wjn are tunable 1, at least one of the weights is tunable
(13)
Intuitively, neuron n inherits the controllabilities of inputs i and j, and if at least one weight is tunable, can additionally control the ratio between the two input neurons. This allows it to set an additional ratio between any of the inputs in Ii ∪ Ij . If the loss function is quadratic, instead of using this additional ratio to solve a single constraint, the optimizer may want to partially satisfy multiple constraints. Hence, we allow added controllability ∆Ca/b to have a value in [0, 1]. Notice that ∆Ca/b(n): (1) abides by Lemma 6.1, and (2) the optimizer can tune all |Ii ∪ Ij |2 individual values in ∆C. In corollary C.0.1 we extend this Lemma for the case where n has multiple inputs.
Theorem 6.5. Control fannout: For a neuron n and two output neurons x and y connected with constant connections such that n = x = y, x and y controllabilities Ca/b(x) and Ca/b(y) abide by:
∀a, b ∈ I, Ca/b(x) + Ca/b(y) = Ca/b(n) (14)
In other words, neuron n’s control is split across outputs. The optimizer chooses how best to split this control, i.e., it does not have to be fair. In Appendix D, we show how graphs can be decomposed so that we can apply Theorems 6.4 and 6.5.
Corollary 6.5.1. For a neuron n with a set of output neurons O = {x1, ..., xk} connected with constant connections, output neuron controllabilities Ca/b(xj) abide by:
∀a, b ∈ I, O∑ xj Ca/b(xj) = Ca/b(n) (15)
We can reframe Equation 15 using trainable ratios:
∀a, b ∈ I, ∀xj ∈ O, Ca/b(xj) = rabj(n)Ca/b(n) (16)
∀a, b ∈ I, |O|∑ j=1 rabj(n) = 1, 0 ≤ rabj(n) ≤ 1 (17)

heading: 6.3 TRAINING CONTROLLABILITY

text:
Finally, we decompose the sparse cascade’s topology so that we can apply Theorems 6.4 and 6.5, in order to write out the equations for the controllability of each output neuron with respect to each input neuron.
Take a cascade with L layers. Each layer has |nl| sparsely connected neurons, with |n0| being the number of cascade inputs, and |nL| being the number of cascade outputs. We define layer l controllability tensor Cl ∈ [0, 1]n0×n0×nl as a tensor where the element Cli,j,k represents neuron nlk’s control over ratio n 0 i /n 0 j . The added controllability tensor ∆C l+1 ∈ [0, 1]n0×n0×nl+1×nl represents the added controllability added by the tunable connections between layer l and l + 1, as per Theorem 6.4. The element ∆Cli,j,k,m represents the additional control over i/j provided by the edge between nlm and n l+1 k . Each tensor ∆Ci,j,:,m abides by Equation 28. The ratio tensor Rl ∈ [0, 1]n0×n0×nl+1×nl represents the control split from Equation 16 with Ri,j,k,m representing the portion of controllability Ci/j(nlm) passed on to Ci/j(n l+1 k ).
We can propagate controllability through the network as: Cl+1 = Cl Rl + ∑ m ∆Cli,j,k,m (18)
where the operation is defined as:
: Rq×q×r × Rq×q×r×s → Rq×q×s, (C R)i,j,k = Ci,j,:Ri,j,k,: (19)
The controllability tensor C, added controllability tensor ∆C and ratio tensor R still have to abide by constraints in Lemmas 6.1, 6.2, and Equations 13, 16.
If nin is the number of cascade inputs, we define the controllability loss as: L(CL) = ∑ k ((nin − 1)− ∑ i ∑ j CLi,j,k) 2 (20)
i.e., if a certain neuron output k has control over nin − 1 ratios, that neuron’s loss is 0. We can now minimize this loss to discover how much control each cascade output has over the cascade inputs.

heading: 7 DERIVING BETTER TOPOLOGIES

text: In this section, we analyze the results of the controllability heuristic and answer (1) how topologies can be improved on the matrix reconstruction task, and (2) if there exists a definitive answer to what topology performs the best.

heading: 7.1 THE NEED FOR SKIP CONNECTIONS

text:
Similar to skip connections used in ResNet networks (He et al., 2015b), we propose that skip connections in can significantly improve the performance of sparse cascades, though for a different reason.
As mentioned in Section 6.1, if a network uses homogenous activation functions, each neuron only needs to learn the ratio of inputs at that neuron, as the job of learning the magnitude can be shifted to the layer above. Since the number of learnable ratios at a neuron is one less than the number of inputs of the same neuron, that neuron does not need to have all of it’s connections trainable. One of the connections can have a constant value of 1, and the network performance will not be impacted. This effect is particularly noticable in butterfly networks, where each neuron only has 2 inputs and 2 outputs, hence 50% of connections are wasted. We replace one connection per neuron with a skip connection valued at 1, and show their performance in Figure 4. See that skip connections significantly improve the performance of butterfly and hypercube networks.

heading: 7.2 THE NEED FOR INPUT-OUTPUT PAIR EQUALITY

text:
While skip connections help topologies achieve similar performance (see hypercube and butterfly in Figure 4), some topologies still outperform / underperform. We turn to our controllability heuristic for an explanation of this behavior. We train our controllability network from Section 6.3 with topologies from Figure 4. The trained network produces CL, the controllability tensor of the last layer of the network. This is a 3D tensor where CLi,j,k specifies the optimizer’s control of input ratio n0i /n 0 j at output neuron n L k . Since we optimize for the number of ratios set, and not the specific
configuration, we sum CL in the second dimension with KLi,k = ∑ CLi,:,k. We plot the resulting K matrices in Figure 5 (a-d).
The total controllability a network achieves is equal to the sum of K, and can be interpreted as the total ‘brightness’ of Figures 5 (a-d). With skip connections and a topology that does not oversaturate certain input-output pairs, we can guarantee that the number of controllable ratios is equal to the number of trainable edges in the graph. Notice that the hypercube and torus figures have significant variance, while Clos and butterfly networks are smooth. This means Clos and butterfly networks do not prioritize any specific input-output pair. Since torii and hypercubes are examples of smallworld networks (Watts & Strogatz, 1998), these networks find it easier to satisfy closer input-output pairs. Even though these networks are trained with L2 loss, where outlier pairs are incentivized
to approach the mean, torus and hypercube networks show signifcant vraiance. This is important since when given a matrix WO which is to be reconstructed, permuting WO columns or rows in a targeted way may improve reconstruction accuracy. Ideally, the position of an input within the topology should not impact the performance of the network. In Figure 5, we give four examples: hypercubes whose controllability matrix has high mean and high variance, Clos with high mean and low variance, butterfly with low mean and low variance, and torus with low mean and high variance.

heading: 7.3 THE NEED FOR SHALLOWNESS

text: Saxe et al. (2013) explore the dynamics of training deep linear networks, and show that deep linear networks have at most a constant time slower convergence compared to shallow linear networks. This effect (1) may still be detrimental to training, and (2) to the best of our knowledge, has not been studied for the case of sparse deep linear networks. Hence, when choosing between two otherwise equivalent topologies (i.e., topologies with the same controllability mean and variance), we should choose the shallower one. Furthermore, observing Figure 3, we see that after a certain depth, butterfly, torus, and hypercube networks lose performance with depth, despite gaining parameters. This is likely due to an issue with initializing very sparse deep networks, as sparse random initializations may be more vulnerable to noise compared to dense networks. On the other hand, constant-depth topologies such as Clos (with depth 3) and low rank (with depth 2) eventually outperform all other variable-depth topologies. Similarly, our ideal topology should have constant depth.

heading: 7.4 THE NEED FOR HIGH INFORMATION BANDWIDTH

text: In Figure 4 we notice that for small parameter counts, the Clos topology is outperformed by both butterflies and hypercubes. By analyzing the controllability matrices of low-parameter Clos networks, we see that this behavior stems from the limited information bandwidth of Clos networks. In Figure 8a, we show an example of a Clos network that underperforms due to limited bandwidth.

heading: 7.5 ONE TOPOLOGY TO RULE THEM ALL

text: We evaluate different topologies with the above criterions, namely: (1) a topology should use skip connections, (2) the controllability matrix of a topology should have no variance, (3) the topology depth should not change with the number of parameters, and (4) the topology should have high information bandwidth, independent of the number of parameters. All of the above topologies can satisfy constraint (1) given skip connections, however, only Clos and butterfly satisfy constraint (2). Since butterfly, hypercube and torus topologies grow in depth as the parameter budget grows, while Clos grows in width, only Clos satisfies constraint (3). However, while butterfly, hypercube, and torus satisfy requirement (4), Clos does not. Hence, we propose a topology we call parallel butterfly, which satisfies all of these constraints. Parallel butterfly consists p butterfly networks of maximum depth d, where d is a metaparameter selected ahead of time. All p networks are connected to the same inputs, and sum their outputs. With a small number of parameters at it’s disposal, the parallel butterfly sets p = 1, and grows in depth up to d layers. Afterwards, it grows in width by increasing the parameter p. In Figure 4 we show that parallel butterfly outperforms all other topologies. In Appendix Figure 8b, we give an example of a parallel butterfly topology where p = 2.

heading: 8 CONCLUSION

text: In this work, we have explored accelerating DNN training by pruning networks ahead of time. We proposed replacing dense and convolutional layers using sparse cascades with topologies selected ahead of time. We presented an a priori sparse neural network initialization scheme that allows us to train very deep networks without the vanishing gradient problem. Since networks are pruned before the model has seen any training data, we investigated topologies that maximize accuracy over any domain. We have developed a data-free heuristic that can evaluate the sparse network’s control of outputs with respect to inputs, allowing us to assess the expressiveness of a given topology. We have extracted several requirements that make for a good topology, such as the need for skip connections, information bandwidth, shallowness, and input-output pair equality. Finally, we have proposed a topology we call parallel butterfly as the ideal topology for training a priori sparse networks, and have experimentally shown that it outperforms other considered topologies.

heading: B MEASURING THE NUMBER OF SOLVABLE RATIO CONSTRAINTS

text: On a practical note, one way to test how many ratios a network can learn is to append a ‘diagonal layer’ to the end of the network (i.e., a new layer with a single neuron attached to each output), as seen in Figure 6. The diagonal layer is a diagonal matrix whose only trainable elements are on the main diagonal, and all other values are 0. When training a network, this diagonal layer can only learn magnitudes, and not ratios between signals, because each neuron only has one input and cannot ‘mix’ any signals. This gives us an easy way of measuring the number of ratios a network can correctly express: we train a network with L1 loss until it converges. We then count the number of constraints k the network has satisfied. These constraints can be ratio constraints or magnitude constraints. If we have n output neurons, we know that the last layer will have satisfied all magnitude constraints. Hence, the number of ratios the network can satisfy is k − n. For example, the network in Figure 6 (right, though true for left too) can satisfy three out of the 4 absolute constraints. 2 of those are magnitude constraints, meaning it can only satisfy one ratio constraint. That ratio is calculated at neuron n, so either neuron x or y can get a correct ratio of inputs, but not both. Of course, with L2 loss, the network will settle for a solution that doesn’t satisfy either, but picks some middle ground.

heading: C CONTROLLABILITY COROLLARIES

text:
Corollary C.0.1. For a neuron n with a set of input neurons {i1, ..., ik} connected with t trainable and k − t constant connections, controllability Ca/b(n) is:
Ca/b(n) = min(1,max(0, k∑ j Ca/b(ij) + ∆Ca/b(n))) (27)
where I∑ a I∑ b ∆Ca/b(n) ≤ min(t, k − 1) (28) Notice that as at least one connection is constant, the network can make full use of all trainable connections.

heading: D DECOMPOSING GRAPHS

text:
We briefly show how graphs can be decomposed so that we can apply Theorems 6.4 and 6.5. In Figure 7 we see a 3-1-2 graph decomposed into a graph where each neuron has at most 2 inputs
or outputs. We can apply Theorem 6.4 to subgraph {a, b,m} and {m, c, n}, and Theorem 6.5 to subgraph {n, x′, y′}. Since neurons x and y only have one input, their controllability is identical to that of x′ and y′, respectively.

heading: E EXAMPLE TOPOLOGIES

title: NEUROFABRIC: IDENTIFYING IDEAL TOPOLOGIES


summaries: This paper touches the signal processing/long term propagation problem in gated recurrent neural networks from the mean field theory. The paper starts from a dynamic system view of the recurrent neural networks and calculates the time scale of converging to the fixed point. In order to avoid the system to converge to the fixed point, the authors utilize some initialization strategy to keep the time scale to infinity. The authors also relate the time scale to state-to-state Jacobians.

summaries: The aim of this paper is to suggest randomized initializations for the various weights of a recurrent neural network (GRUs and various LSTMs are covered), such that training these networks gets to a successful start, when the model is trained on long sequences. Instead of being heuristic, their approach follows first principles of analyzing signal propagation through time, using ideas from statistical thermodynamics (mean field approximations). Some experiments, on toy datasets, validate their approach.

abstractText: Training recurrent neural networks (RNNs) on long sequence tasks is plagued with difficulties arising from the exponential explosion or vanishing of signals as they propagate forward or backward through the network. Many techniques have been proposed to ameliorate these issues, including various algorithmic and architectural modifications. Two of the most successful RNN architectures, the LSTM and the GRU, do exhibit modest improvements over vanilla RNN cells, but they still suffer from instabilities when trained on very long sequences. In this work, we develop a mean field theory of signal propagation in LSTMs and GRUs that enables us to calculate the time scales for signal propagation as well as the spectral properties of the state-to-state Jacobians. By optimizing these quantities in terms of the initialization hyperparameters, we derive a novel initialization scheme that eliminates or reduces training instabilities. We demonstrate the efficacy of our initialization scheme on multiple sequence tasks, on which it enables successful training while a standard initialization either fails completely or is orders of magnitude slower. We also observe a beneficial effect on generalization performance using this new initialization.

authors:
Devansh Arpit
Bhargav Kanuparthi
Giancarlo Kerg
Nan Rosemary Ke
Ioannis Mitliagkas
Yoshua Bengio

title: h-detach: Modifying the lstm gradient towards better optimization

venue: arXiv preprint arXiv:1810.03023,

authors:
Yoshua Bengio
Patrice Simard
Paolo Frasconi

title: Learning long-term dependencies with gradient descent is difficult

venue: IEEE transactions on neural networks,

authors:
Bo Chang
Minmin Chen
Eldad Haber
Ed H. Chi

title: AntisymmetricRNN: A dynamical system view on recurrent neural networks

venue: In International Conference on Learning Representations,

authors:
Minmin Chen
Jeffrey Pennington
Samuel S Schoenholz

title: Dynamical isometry and a mean field theory of rnns: Gating enables signal propagation in recurrent neural networks

venue: arXiv preprint arXiv:1806.05394,

authors:
Kyunghyun Cho
Bart van Merrienboer
Caglar Gulcehre
Dzmitry Bahdanau
Fethi Bougares
Holger Schwenk
Yoshua Bengio

title: Learning phrase representations using rnn encoder–decoder for statistical machine translation

venue: In Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing (EMNLP),

authors:
Felix A Gers
Nicol N Schraudolph
Jürgen Schmidhuber

title: Learning precise timing with lstm recurrent networks

venue: Journal of machine learning research,

authors:
Xavier Glorot
Yoshua Bengio

title: Understanding the difficulty of training deep feedforward neural networks

venue: In Proceedings of the thirteenth international conference on artificial intelligence and statistics,

title: Implicit renewal theory and tails of solutions of random equations

venue: The Annals of Applied Probability,

authors:
Charles M Goldie
Rudolf Grübel

title: Perpetuities with thin tails

venue: Advances in Applied Probability,

title: Almost surely asymptotic freeness for jacobian spectrum of deep network

venue: arXiv preprint arXiv:1908.03901,

authors:
Sepp Hochreiter
Jürgen Schmidhuber

title: Long short-term memory

venue: Neural computation,

authors:
Arthur Jacot
Franck Gabriel
Clément Hongler

title: Neural tangent kernel: Convergence and generalization in neural networks

venue: arXiv preprint arXiv:1806.07572,

authors:
Cijo Jose
Moustpaha Cisse
Francois Fleuret

title: Kronecker recurrent units

venue: arXiv preprint arXiv:1705.10142,

authors:
Jaehoon Lee
Lechao Xiao
Samuel S Schoenholz
Yasaman Bahri
Jascha Sohl-Dickstein
Jeffrey Pennington

title: Wide neural networks of any depth evolve as linear models under gradient descent

authors:
Alexander G de G Matthews
Mark Rowland
Jiri Hron
Richard E Turner
Zoubin Ghahramani

title: Gaussian process behaviour in wide deep neural networks

venue: arXiv preprint arXiv:1804.11271,

authors:
Razvan Pascanu
Tomas Mikolov
Yoshua Bengio

title: On the difficulty of training recurrent neural networks

venue: In International Conference on Machine Learning,

authors:
Jeffrey Pennington
Samuel Schoenholz
Surya Ganguli

title: Resurrecting the sigmoid in deep learning through dynamical isometry: theory and practice

venue: In Advances in neural information processing systems,

authors:
Jeffrey Pennington
Samuel S Schoenholz
Surya Ganguli

title: The emergence of spectral universality in deep networks

venue: arXiv preprint arXiv:1802.09979,

authors:
Ben Poole
Subhaneil Lahiri
Maithra Raghu
Jascha Sohl-Dickstein
Surya Ganguli

title: Exponential expressivity in deep neural networks through transient chaos

authors:
Andrew M Saxe
James L McClelland
Surya Ganguli

title: Exact solutions to the nonlinear dynamics of learning in deep linear neural networks

venue: arXiv preprint arXiv:1312.6120,

authors:
Samuel S Schoenholz
Justin Gilmer
Surya Ganguli
Jascha Sohl-Dickstein

title: Deep information propagation

venue: arXiv preprint arXiv:1611.01232,

authors:
Trieu H Trinh
Andrew M Dai
Thang Luong
Quoc V Le

title: Learning longer-term dependencies in rnns with auxiliary losses

venue: arXiv preprint arXiv:1803.00144,

authors:
Eugene Vorontsov
Chiheb Trabelsi
Samuel Kadoury
Chris Pal

title: On orthogonality and learning recurrent networks with long term dependencies

authors:
Scott Wisdom
Thomas Powers
John Hershey
Jonathan Le Roux
Les Atlas

title: Full-capacity unitary recurrent neural networks

venue: In Advances in Neural Information Processing Systems,

authors:
Lechao Xiao
Yasaman Bahri
Jascha Sohl-Dickstein
Samuel S Schoenholz
Jeffrey Pennington

title: Dynamical isometry and a mean field theory of cnns: How to train 10,000-layer vanilla convolutional neural networks

venue: arXiv preprint arXiv:1806.05393,

authors:
Ge Yang
Samuel Schoenholz

title: Mean field residual networks: On the edge of chaos

title: Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation

authors:
Greg Yang
Sam S. Schoenholz

title: Deep Mean Field Theory: Layerwise Variance and Width Variation as Methods to Control Gradient Explosion

venue: ICLR Workshop,

authors:
Greg Yang
Jeffrey Pennington
Vinay Rao
Jascha Sohl-Dickstein
Samuel S

title: Schoenholz. A Mean Field Theory of Batch Normalization

heading: 1 INTRODUCTION

text:
A common paradigm for research and development in deep learning involves the introduction of novel network architectures followed by experimental validation on a selection of tasks. While this methodology has undoubtedly generated significant advances in the field, it is hampered by the fact that the full capabilities of a candidate model may be obscured by difficulties in the training procedure. It is often possible to overcome such difficulties by carefully selecting the optimizer, batch size, learning rate schedule, initialization scheme, or other hyperparameters. However, the standard strategies for searching for good values of these hyperparameters are not guaranteed to succeed, especially if the trainable configurations are constrained to a low-dimensional subspace of hyperparameter space, which can render random search, grid search, and even Bayesian hyperparameter selection methods unsuccessful.
In this work, we argue that for long sequence tasks, the trainable configurations of initialization hyperparameters for LSTMs and GRUs lie in just such a subspace, which we characterize theoretically. In particular, we identify precise conditions on the hyperparameters governing the initial weight and bias distributions that are necessary to ensure trainability. These conditions derive from the observation that in order for a network to be trainable, (a) signals from the relevant parts of the input sequence must be able to propagate all the way to the loss function and (b) the gradients must be stable (i.e., they must not explode or vanish exponentially). As shown in Figure 1, training of recurrent networks with standard initialization on certain tasks begins to fail as the sequence length increases and signal propagation becomes harder to achieve. However, as we shall show, a suitably-chosen initialization scheme can dramatically improve trainability on such tasks.
We study the effect of the initialization hyperparameters on signal propagation for a very broad class of recurrent architectures, which includes as special cases many state-of-the-art RNN cells, including the GRU (Cho et al. (2014)), the LSTM (Hochreiter & Schmidhuber (1997)), and the peephole LSTM (Gers et al. (2002)). The analysis is based on the mean field theory of signal propagation developed in a line of prior work (Schoenholz et al. (2016); Xiao et al. (2018); Chen et al. (2018); Yang et al. (2019)), as well as the concept of dynamical isometry (Saxe et al. (2013); Pennington
200 400 600 800 Sequence length
0.0
0.2
0.4
0.6
0.8
1.0
Ac cu
ra cy
Critical initialization Standard initialization
Figure 1: Critical initialization improves trainability of recurrent networks. Test accuracy for peephole LSTM trained to classify sequences of MNIST digits after 8000 iterations. As the sequence length increases, the network is no longer trainable with standard initialization, but still trainable using critical initialization.
et al. (2017; 2018)) that is necessary for stable gradient backpropagation and which was shown to be crucial for training simpler RNN architectures (Chen et al. (2018)). We perform a number of experiments to corroborate the results of the calculations and use them to motivate initialization schemes that outperform standard initialization approaches on a number of long sequence tasks.

heading: 2 BACKGROUND AND RELATED WORK

heading: 2.1 MEAN FIELD ANALYSIS OF NEURAL NETWORKS

text:
Signal propagation at initialization can be controlled by varying the hyperparameters of fullyconnected (Schoenholz et al. (2016); Yang & Schoenholz (2017)) and convolutional (Xiao et al. (2018)) feed-forward networks, as well as for simple gated recurrent architectures (Chen et al. (2018)). In all these cases, such control was used to obtain initialization schemes that outperformed standard initializations on benchmark tasks. In the feed-forward case, this enabled the training of very deep architectures without the use of batch normalization or skip connections.
By forward signal propagation, we specifically refer to the propagation of correlations between inputs at the start of an input sequence through the hidden states of the recurrent network at later times, as will be made precise in Section 4.2. This appears to be a necessary but not sufficient condition for trainability on certain tasks. Another is the stability of the gradients, which depend on the state-to-state Jacobian matrix, as discussed in Bengio et al. (1994). In our context, the goal of the backward propagation analysis is to improve the conditioning of the Jacobian by controlling the first two moments of its squared singular values. Forward signal propagation and the spectral properties of the Jacobian at initialization can be studied using mean field theory and random matrix theory (Poole et al. (2016); Schoenholz et al. (2016); Yang & Schoenholz (2017; 2018); Xiao et al. (2018); Pennington et al. (2017; 2018); Chen et al. (2018); Yang et al. (2019)).
As neural network training is a nonconvex problem, using a modified initialization scheme could lead to convergence to different points in parameter space in a way that adversely affects the generalization error. We provide some empirical evidence that this does not occur, and in fact, the use of initialization schemes satisfying these conditions has a beneficial effect on the generalization error.

heading: 2.2 THE EXPLODING/VANISHING GRADIENT PROBLEM AND SIGNAL PROPAGATION IN RECURRENT NETWORKS

text:
The exploding/vanishing gradient problem is a well-known phenomenon that hampers training on long time sequence tasks (Bengio et al. (1994); Pascanu et al. (2013)). Apart from the gating mechanism, there have been numerous proposals to alleviate the vanishing gradient problem by constraining the weight matrices to be exactly or approximately orthogonal (Pascanu et al. (2013); Wisdom et al. (2016); Vorontsov et al. (2017); Jose et al. (2017)), or more recently by modifying some terms in the gradient (Arpit et al. (2018)), while exploding gradients can be handled by clipping (Pascanu et al. (2013)). Another recently proposed approach to ensuring signal propagation in long sequence tasks introduces auxiliary loss functions (Trinh et al. (2018)). This modification of the loss can be seen as a form of regularization. Chang et al. (2019) study the connections between recurrent networks and certain ordinary differential equations and propose the AntisymmetricRNN that can capture long term dependencies in the inputs. While many of these approaches have been quite successful, they
typically require modifying the training algorithm, the loss function, or the architecture, and as such exist as complementary methods to the one we investigate here. We postpone the investigation of a combination of techniques to future work.

heading: 2.3 DYNAMICS BEYOND INITIALIZATION

text: While the analysis below applies to the network at initialization, it was recently shown that as the network width becomes large, parameters move very little from their initial values (Jacot et al. (2018); Lee et al. (2019)). A non-trivial and surprising consequence is that covariance matrices like those studied here remain constant during training. While Jacot et al. (2018); Lee et al. (2019) do not specifically study RNNs, the analysis of Yang (2019) suggests the conclusions are likely to carry over. Moreover, even for finite-width networks, simply satisfying the conditions derived in this work at initialization can lead to a dramatic improvement in trainability as we show below.

heading: 3 NOTATION

text: We denote matrices by bold upper case Latin characters and vectors by bold lower case Latin characters. Dx denotes a standard Gaussian measure. The normalized trace of a random N × N matrix A, 1NEtr(A), is denoted by τ(A). ◦ is the Hadamard product, σ(·) is a sigmoid function and both σ(·), tanh(·) act element-wise. We denote by Da a diagonal matrix with a on the diagonal.

heading: 4 MEAN FIELD ANALYSIS OF SIGNAL PROPAGATION AND DYNAMICAL ISOMETRY

heading: 4.1 MODEL DESCRIPTION AND IMPORTANT ASSUMPTIONS

text:
We begin with a general description of recurrent architectures that can be specialized to the GRU, LSTM and peephole LSTM among others. We denote the state of a recurrent network at time t by st ∈ RN with s0i ∼ D0, and a sequence of inputs to the network by {z1, ..., zT }, zt ∈ RN .
The state evolution of the network is given by
st = f(st−1, {u1k}, ..., {utk}, zt) (1a)
where f is an element-wise, affine function of st−1 and {utk} is a set of pre-activations utk ∈ RN , k ∈ K defined for a set of subscripts K. They are given by
utk = Wks t−1 + Ukz t + bk (1b)
where Wk,Uk ∈ RN×N ,bk ∈ RN . We define additional pre-activations
utk2 = Wk2Dgk(utk)s t−1 + Uk2z t + bk2 (1c)
where gk : RN → RN is an element-wise function and utk is defined as in eqn. (1b) 1. In cases where there is no need to distinguish between variables of the form (1b) and (1c) we will refer to both as
1Variables of the form (1c) will be present only in the GRU (see Table 1).
utk. This state update equation describes all the architectures studied in this paper, as well as many others, as detailed in Tables 1,3. The pre-activations utk, which are typically used to construct gates in recurrent network with the application of appropriate nonlinearities, will be of interest because they will tend in distribution to Gaussians as the network width is taken to infinity. We assume Wk,ij ∼ N (0, σ2k/N), Uk,ij ∼ N (0, ν2k/N), bk,i ∼ N (µk, ρ2k) i.i.d. and denote Θ =⋃ k {σ2k, ν2k , ρ2k, µk}. As in Chen et al. (2018), we make the untied weights assumption that Wk is independent of st. Although the assumption of untied weights may seem counter-intuitive in the context of recurrent networks, recall that we are only interested in characterizing various statistics of the network, not in computing the entire input-output map, for which using tied weights would certainly be crucial. For merely characterizing certain statistics, the untied weights assumption actually has long history of yielding correct predictions, and therefore it is not at all unnatural to adopt: 1) when computing the covariance of gradients in an MLP, the weights during backpropagation can be treated as independently sampled from the weights used during the forward propagation (Schoenholz et al., 2016; Yang & Schoenholz, 2017); 2) when computing the Jacobian singular value distribution of an MLP, the pre-activations h = Wx can be treated as the result of multiplying an iid copy W′ of W (i.e. h = W′x) (Pennington et al., 2017; 2018; Hayase, 2019); 3) a wide, randomly initialized MLP corresponds to the same Gaussian process whether or not its weights across layers are tied (Poole et al., 2016), and similarly, a simple RNN evolves the same whether or not its weights are tied across time, as long as its width is large (Chen et al., 2018). Indeed, we provide empirical evidence that calculations performed under this assumption still have considerable predictive power, even in cases where it is violated.

heading: 4.2 FORWARD SIGNAL PROPAGATION

text:
We now study how correlations between inputs propagate into the hidden states of a recurrent network. This section follows closely the development in Chen et al. (2018). We now consider two sequences of normalized inputs {zta}, {ztb} with zero mean and covariance R = R (
1 Σz Σz 1
) fed into two
copies of a network with identical weights, and resulting in sequences of states {sta}, {stb}. We consider the time evolution of the moments and correlations
µts = E[stia] (2a) Qts ≡ Qts,aa = E[stiastia] (2b)
Σt2s C t s + (µ t s) 2 = E[stiastib] (2c)
where we define Σt2s = Q t s − (µts)2.
Returning to the pre-activations defined in eqn. (1). We make the mean field approximation2 that the pre-activations are jointly Gaussian at the infinite width N →∞ limit:(
utkia utkib
) ∼ N (( µk µk ) , ( Qtk − µ2k )( 1 Ctk Ctk 1 )) (3)
where the second moment Qtk is given by
Qtk = E[utkiautkia] = σ2kE[stiastia] + ν2kE[ztiaztia] + ρ2k + µ2k = σ2kQts + ν2kR+ ρ2k + µ2k (4a)
and defining Σt2k = Q t k − µ2k
Ctk = E[utkiautkib]− µ2k Qtk − µ2k = σ2k ( Σt2s C t s + µ 2 s ) + ν2kRΣz + ρ 2 k Qtk − µ2k . (4b)
The variables utk2ia, u t k2ib given by eqn. (1c) are also asymptotically Gaussian, with their covariance detailed in Appendix A. We will subsequently drop the vector subscript i since all elements are
2For feed-forward networks, it has been shown using the Central Limit Theorem for exchangeable random variables that the pre-activations at all layers indeed converge in distribution to a multivariate Gaussian (Matthews et al. (2018)).
identically distributed and f ,gk act element-wise, and the input sequence index in expressions that involve only one sequence.
For any l ≤ t, the ulk are independent of sl at the large N limit. Combining this with the fact that their distribution is determined completely by µl−1s , Q l−1 s , C l−1 s and that f is affine, one can rewrite eqn. (2a-c) using eqn. (1a-c) as the following deterministic dynamical system
(µts, Q t s, C t s) =M(µt−1s , Qt−1s , Ct−1s , ..., µ1s, Q1s, C1s ) (5)
where the dependence on Θ and the data distribution has been suppressed. In the peephole LSTM and GRU, the form will be greatly simplified toM(µt−1s , Qt−1s , Ct−1s ). In Appendix E we compare the predicted dynamics to simulations, showing good agreement.
One can now study the fixed points of eqn. (5) and the rates of convergence by linearizing the dynamics. The fixed point of eqn. (5) is pathological, in the sense that any information that distinguishes two input sequences is lost. Therefore, delaying the convergence to the fixed point should allow for signals to propagate across longer time horizons. Quantitatively, the rate of convergence of eqn. (5) to a fixed point gives an effective time scale for signal propagation from the inputs to the hidden state at later times.
While the dynamical system is two-dimensional and analysis of convergence rates should be performed by linearizing the full system and studying the smallest eigenvalue of the resulting matrix, in practice as in Chen et al. (2018) this eigenvalue appears to always corresponds to the Cts direction
3 . Hence, if we assume convergence of Qts, µ t s we need only study
Cts =MC(µ∗s, Q∗s, Ct−1s ) (6)
whereMC also depends on expectations of functions of {u1k}, ...{u t−2 k } that do not depend on Ct−1s . While this dependence is in principle on an infinite number of Gaussian variables as the dynamics approach the fixed point, MC can still be reasonably approximated in the case of the LSTM as detailed in Appendix D. We show that in the case of the peephole LSTM this map is convex in Appendix C. This can be shown for the GRU by a similar argument. It follows directly that it has a single stable fixed point in these cases.
The rate of approach to the fixed point, χC∗s , is given by linearizing around it. Setting C t s = C ∗ s + ε t, we have
Ct+1s = C ∗ s + ε t+1 =MC(µ∗s, Q∗s, C∗s + εt) =MC(µ∗s, Q∗s, C∗s ) + χC∗s ε t +O((εt)2). (7)
The time scale of convergence to the fixed point is thus given by
ξC∗s = − 1
logχC∗s , (8)
which diverges as χC∗s approaches 1 from below. Due to the detrimental effect of convergence to the fixed point described above, it stands to reason that a choice of Θ such that χC∗s = 1− δ for some small δ > 0 would enable signals to propagate from the initial inputs to the final hidden state when training on long sequences.

heading: 4.3 BACKWARDS SIGNAL PROPAGATION - THE STATE-TO-STATE JACOBIAN

text:
We now turn to controlling the gradients of the network. A useful object to consider in this case is the asymptotic state-to-state transition Jacobian
J = lim t→∞
∂st+1
∂st .
J and powers of it will appear in the gradient as the covariance dynamics described in Section 4.2 approach their fixed point (specifically, the gradient of a network trained on a sequence of length T will depend on a matrix polynomial of order T in J), hence we desire to control the squared singular
3This observation is explained in the case of fully-connected feed-forward networks in Schoenholz et al. (2016).
value distribution of J. This distribution is described by moments of the form mJJT ,n = τ((JJT )n), with mean and variance of the distribution given by
mJJT ,1 = E ∑
k∈K∪{0}
ak (9)
σ2JJT = mJJT ,2 −m 2 JJT ,1 = E ∑ k,l∈K∪{0} 2akal − a20 − E ∑ k∈K∪{0} ak 2 (10) where the scalars {ak}k∈K∪{0} are architecture and hyperparameter dependent and are given in Appendix A.2.
Forward and backward signal propagation are in fact intimately related, as the following lemma shows: Lemma 1. For a recurrent neural networks defined by (1), the mean squared singular value of the state-to-state Jacobian defined in (16) and χCs that determines the time scale of forward signal propagation (given by (7)) are related by
mJJT ,1 = χC∗s =1,Σz=1 (11)
Proof. See Appendix C.
If one can find a setting of the hyperparameters such that mJJT ,1 is close to 1 and σ2JJT is small, then the spectrum of powers of J should remain well-conditioned ensuring that gradients will not explode or vanish.

heading: 4.4 DYNAMICAL ISOMETRY

text:
Combining the results of Sections 4.2 and 4.3, we conclude that an effective choice of initialization hyperparamters should satisfy
χC∗s = 1 (12a)
mJJT ,1 = 1 (12b)
σ2JJT = 0. (12c)
We refer to these as dynamical isometry conditions. eqn. (12)a ensures stable signal propagation from the inputs to the loss, and eqn. (12)b-c are motivated by the additional requirement of preventing the gradients from exploding/vanishing. Both of these conditions appear to be necessary in order for a network to be trainable. Combining eqns. (9, 10, 11) we find that if Σz = 1, the dynamical isometry conditions are satisfied if a0 = 1, ak 6=0 = 0, which can be achieved by setting ∀k : σ2k = 0 and taking µf →∞. This motivates the general form of the initializations used in the experiments. We demand that these equations are only satisfied approximately since for a given architecture, there may not be a value of Θ that satisfies them all. Additionally, even if such a value exists, the optimal value of χC∗s for a given task may not be 1. There is some empirical evidence that if the characteristic time scale defined by χC∗s is much larger than that required for a certain task, performance is degraded.
The typical values of Σz will depend on the data set, yet satisfying the dynamical isometry conditions is simplified if Σz = 1 due to Lemma 1. It is thus a natural choice, yet we acknowledge that a more comprehensive treatment should consider the case of general Σz . We also find empirically that the results obtained under the Σz = 1 assumption prove to be predictive and enable training on a number of tasks without requiring a detailed analysis of Σz .

heading: 4.5 THE STATIONARY DISTRIBUTION OF THE LSTM CELL STATE

text:
In architectures that were considered in previous works, as well as the peephole LSTM and GRU, all the quantities required to calculate the dynamical isometry conditions through eqns. (9, 10) can be written in terms of integrals over a finite number of Gaussian variables at the large width and time
limit and are thus easy to evaluate. This is not the case in the standard LSTM due to the presence of the cell state. By unrolling the cell state update equation
ct = σ(utf ) ◦ ct−1 + σ(uti) ◦ tanh(utr) (13)
in time, we find that ct depends on the entire sequence of utk from the first time step 4.
From eqn. 13 we find that the asymptotic cell state distribution is that of a perpetuity, which is a random variable X that obeys X d= XY + Z where Y, Z are random variables and d= denotes equality in distribution. The stationary distribution of a perpetuity generally does not have a closed form, and neither does its likelihood (Goldie (1991)). In practice, one can overcome this difficulty by sampling from the stationary cell state distribution, as described in Appendix D.

heading: 5 EXPERIMENTS

heading: 5.1 PADDED MNIST CLASSIFICATION

text:
The calculations presented above predict a characteristic time scale ξ (defined in (8)) for forward signal propagation in a recurrent network. It follows that on a task where success depends on propagation of information from the first time step to the final T -th time step, the network will not be trainable for T ξ. In order to test this prediction, we consider a classification task where the inputs are sequences consisting of a single MNIST digit followed by T − 1 steps of i.i.d Gaussian noise and the targets are the digit labels. By scanning across certain directions in hyperparameter space, the predicted value of ξ changes. We plot training accuracy of a network trained with untied weights after 1000 iterations for the GRU and 2000 for the LSTM, as a function of T and the hyperparameter values, and overlay this with multiples of ξ. As seen in Figure 2, we observe good agreement between the predicted time scale of signal propagation and the success of training. As expected, there are some deviations when training without enforcing untied weights, and we present the corresponding plots in the supplementary materials.
5.2 UNROLLED MNIST AND CIFAR-10 MNIST CIFAR-10standard LSTM 98.6 58.8 critical LSTM 98.9 61.8
Table 2: Test accuracy on unrolled MNIST and CIFAR-10.
The calculation results motivate critical initializations that we test on standard long sequence benchmarks. The details of the initializations are presented in Appendix E. We unroll an MNIST digit into a sequence of length 784 and
train a critically initialized peephole LSTM with 600 hidden units. We also train a critically initialized
4In contrast, the hidden state distribution of the GRU at time t for instance can be expressed as an integral over {ut−1k } alone (which are Gaussian at the large t fixed point) through eqn. 2.

heading: 102 103 104

text: LSTM with hard sigmoid nonlinearities on unrolled CIFAR-10 images feeding in 3 pixels at every time step, resulting in sequences of length 1024. We also apply standard data augmentation for this task. We present accuracy on the test set in Table 2. Interestingly, in the case of CIFAR-10 the best performance is achieved by an initialization with a forward propagation time scale ξ that is much smaller than the sequence length, suggesting that information sufficient for successful classification may be obtained from a subset of the sequence.

heading: 5.3 REPEATED PIXEL MNIST AND MULTIPLE DIGIT MNIST

text: In order to generate longer sequence tasks, we modify the unrolled MNIST task by repeating every pixel a certain number of times and set the input dimension to 7. To create a more challenging task, we also combine this pixel repetition with concatenation of multiple MNIST digits (either 0 or 1), and label such sequences by a product of the original labels. In this case, we set the input dimension to 112 and repeat each pixel 10 times. We train a peephole LSTM with both a critical initialization and a standard initialization on both of these tasks using SGD with momentum. In this former task, the dimension of the label space is constant (and not exponential in the number of digits like in the latter). In both tasks, we observe three distinct phases. If the sequence length is relatively short the critical and standard initialization perform equivalently. As the sequence length is increased, training with a critical initialization is faster by orders of magnitude compared to the standard initialization. As the sequence length is increased further, training with a standard initialization fails, while training with a critical initialization still succeeds. The results are shown in Figure 3.

heading: 6 DISCUSSION

text:
We have derived initialization schemes for recurrent networks motivated by ensuring stable signal propagation from the inputs to the loss and of gradient information from the loss to the weights. These schemes dramatically improve performance on long sequence tasks.
The subspace of initialization hyperparameters Θ that satisfy the dynamical isometry conditions is multidimensional, and there is no clear principled way to choose a preferred initialization within it. It would be of interest to study this subspace and perhaps identify preferred initializations based on additional constraints. One could also use the satisfaction of the dynamical isometry conditions as
a guiding principle in simplifying these architectures. A direct consequence of the analysis is that the forget gate, for instance, is critical, while some of the other gates or weights matrices can be removed while still satisfying the conditions. A related question is the optimal choice of the forward propagation time scale ξ for a given task. As mentioned in Section 5.2, this scale can be much shorter than the sequence length. It would also be valuable understand better the extent to which the untied weights assumption is violated, since it appears that the violation is non-uniform in Θ, and to relax the constant Σz assumption by introducing a time dependence.
Another compelling issue is the persistence of the dynamical isometry conditions during training and their effect on the solution. It has been recently shown that in sufficiently wide networks the gradient flow dynamics in function space are effectively linear (Jacot et al. (2018); Lee et al. (2019)) and under these same conditions the dynamical isometry conditions persist during training. Understanding the finite width and learning rate corrections to such calculations could help extend the analysis of signal propagation at initialization to trained networks in this regime.

heading: Appendices

heading: APPENDIX A DETAILS OF RESULTS

text:
A.1 COVARIANCES OF utk2
The variables utk2ia, u t k2ib given by 1c are asymptotically Gaussian at the N →∞, with
Qtk2 = σ 2 k2 ∫ g2k(ua)DkQts + ν2k2R + ρ 2 k2 + µ2k2 (14a)
Ctk2 =
( σ2k2 ∫ gk(ua)gk(ua)Dk (Σt2s Cts + µ2s)
+ν2k2RΣz + ρ 2 k2 ) Qtk2 − µ 2 k2
(14b)
where Dk is a Gaussian measure on (ua, ub) corresponding to the distribution in eqn. (3).

heading: A.2 MOMENTS OF THE STATE-TO-STATE JACOBIAN

text:
The moments of the squared singular value distribution are given by the normalized traces
mJJT ,n = τ((JJ T )n).
Since U t′ , t′ < t is independent of st, if we index by k, k′ the variables defined by eqn. (1b) and (1c) respectively we obtain
J = ∂f ∂s∗ + ∑ k ∂f ∂u∗k Wk + ∑ k′ ∂f ∂u∗k′2 Wk′2 ( Dgk(u∗k′ ) + Wk ′Dgk′(u∗k′ )Ds ∗ ) . (15)
Under the untied assumption Wk,Wk2 are independent of s t,utk at the large N limit, and are also independent of each other and their elements have mean zero. Using this and the fact that f acts element-wise, we have
mJJT ,1 = τ(JJ T ) = E ∑ k∈K∪{0} ak (16)
where
ak =  D20 k = 0 σ2kD 2 k uk is given by (1b) σ2k2D 2 k ( g2k(u ∗ k)
+σ2ks ∗2g′2k (u ∗ k)
) uk is given by (1c)
(17)
and D0 = ∂f∂s∗ , Dk = ∂f ∂u∗k . The values of D0, Dk for the architectures considered in the paper are detailed in Appendix A.3.
Controlling the first moment of JJT is not sufficient to ensure that the gradients do not explode or vanish, since the variance of the singular values may still be large. This variance is given by
σ2JJT = mJJT ,2 −m 2 JJT ,1.
The second moment mJJT ,2 can be calculated from (15), and is given by mJJT ,2 = E ∑
k,l∈K∪{0}
2akal − a20 (18)
where the ak are defined in eqn. (17). We arrange the scalars a0, {ak} into a vector a. In Appendix B we show that the results of the calculations in this section match the empirical spectrum of the Jacobian.
For all the architectures considered in this paper, we find that D0 = σ(u∗f ) while Dk are finite as ∀k : σ2k → 0. Combining this with (11), (16), (18), we find that if Σz = 1 the dynamical isometry conditions are satisfied if a0 = 1, ak 6=0 = 0, which can be achieved by setting ∀k : σ2k = 0 and taking µf → ∞. This motivates the general form of the initializations used in the experiments 5 although there are many other possible choices of Θ such that the ak 6=0 vanish.
Given the above general form of the dynamical isometry conditions for recurrent networks, we now provide the detailed forms that apply to the LSTM, peephole LSTM and GRU.

heading: A.3 DYNAMICAL ISOMETRY CONDITIONS FOR SELECTED ARCHITECTURES

text:
We specify the form of χC∗s ,Σ and a for the architectures considered in this paper:
A.3.1 GRU
χC∗s ,Σ = E  σ(u ∗ fa)σ(u ∗ fb) + σ 2 f ( tanh(u∗r2a) tanh(u ∗ r2b ) +h∗ah ∗ b ) σ′(u∗fa)σ ′(u∗fb)
+σ2r2(1− σ(u ∗ fa))(1− σ(u∗fb)) tanh ′(u∗r2a) tanh ′(u∗r2b)
( σ(u∗r1a)σ(u ∗ r1b
) +σ2r1h ∗ ah ∗ bσ ′(u∗r1a)σ ′(u∗r1b)
)
a =  σ2(u∗f ) σ2f ( tanh2(u∗r2) +Q ∗ h ) σ′2(u∗f )
σ2r1σ 2 r2 h∗2(1− σ(u∗f ))2 tanh ′2(u∗r2)σ ′2(u∗r1)
σ2r2(1− σ(u ∗ f )) 2 tanh′2(u∗r2)σ 2(u∗r1)

heading: A.3.2 PEEPHOLE LSTM

text:
χC∗s ,Σ = E [
σ(u∗fa)σ(u ∗ fb) + σ 2 i σ ′(u∗ia)σ ′(u∗ib) tanh(u ∗ ra) tanh(u ∗ rb)
σ2fc ∗ ac ∗ bσ ′(u∗fa)σ ′(u∗fb) + σ 2 rσ(u ∗ ia)σ(u ∗ ib) tanh ′(u∗ra) tanh ′(u∗rb)
]
a =  σ2(u∗f ) σ2i σ ′2(u∗i ) tanh 2(u∗r) σ2fc
∗2σ′2(u∗f ) σ2rσ 2(u∗i ) tanh ′2(u∗r)  5In the case of the LSTM, we also want to prevent the output gate from taking very small values, as explained
in Section A.3.
0.0 0.5 1.0 1.5 2.0 2.5
Typical Hyperparameters
0.85 0.90 0.95 1.00 1.05 1.10 1.15
Approximate Dynamical Isometry
Figure 4: Squared singular values of the state-to-state Jacobian in eqn. (15) for two choices of hyperparameter settings Θ. The red lines denote the empirical mean and standard deviations, while the dotted lines denote the theoretical prediction based on the calculation described in Section 4.3. Note the dramatic difference in the spectrum caused by choosing an initialization that approximately satisfies the dynamical isometry conditions.
A.3.3 LSTM
χC∗s ,Σz = E  σ(u∗fa)σ(u ∗ fb) + σ 2 oσ ′(u∗oa)σ ′(u∗ob) tanh(c ∗ a) tanh(c ∗ b)
+σ(u∗oa)σ(u ∗ ob) tanh ′(c∗a) tanh ′(c∗b)  σ2fc∗ac∗bσ′(u∗fa)σ′(u∗fb)+σ2i σ′(u∗ia)σ′(u∗ia) tanh(u∗ra) tanh(u∗rb) +σ2rσ(u ∗ ia)σ(u ∗ ia) tanh ′(u∗ra) tanh ′(u∗rb)

a =  σ2(u∗f ) σ2i σ 2(u∗o) tanh ′2(c∗)σ′2(u∗i ) tanh 2(u∗r) σ2fσ 2(u∗o) tanh
′2(c∗)c∗2σ′2(u∗f ) σ2rσ 2(u∗o) tanh ′2(c∗)σ2(u∗i ) tanh
′2(u∗r) σ2oσ ′2(u∗o) tanh 2(c∗)  When evaluating 7 in this case, we write the cell state as ct = tanh−1 ( st
σ(ot)
) , and
assume σ(o t)
σ(ot−1) ≈ 1 for large t. The stability of the first equation and the accuracy of
the second approximation are improved if ot is not concentrated around 0.

heading: A.4 RNN TABLE SUPPLEMENT

heading: APPENDIX B SQUARED JACOBIAN SPECTRUM HISTOGRAMS

text:
To verify the results of the calculation of the moments of the squared singular value distribution of the state-to-state Jacobian presented in Section 4.3 we run an untied peephole LSTM for 100 iterations with i.i.d. Gaussian inputs. We then compute the state-to-state Jacobian and calculate its spectrum. This can be used to compare the first two moments of the spectrum to the result of the calculation, as well as to observe the difference between a standard initialization and one close to satisfying the dynamical isometry conditions. The results are shown in Figure 4. The validity of this experiment rests on making an ergodicity assumption, since the calculated spectral properties require taking averages over realizations of random matrices, while in the experiment we instead calculate the moments by averaging over the
eigenvalues of a single realization. The good agreement between the prediction and the empirical average suggests that the assumption is valid.

heading: APPENDIX C AUXILIARY LEMMAS AND PROOFS

text:
Proof of Lemma 1. Despite the fact that each ukai as defined in 1 depends in principle upon the entire state vector sa, at the large N limit due to the isotropy of the input distribution we find that these random variables are i.i.d. and independent of the state. Combining this with the fact that f is an element-wise function, it suffices to analyse a single entry of st = f(st−1, {u1k}, ..., {utk}), which at the large t limit gives
MC(µ∗s, Q∗s, Cs) = E [f(sa, U∗a )f(sb, U∗b )]− (µ∗s)2
Q∗s − (µ∗s)2
where U∗a = {u∗ka|k ∈ K, u∗ka ∼ N (µk, Q∗k − µ2k} and U∗b is defined similarly (i.e. we assume the first two moments have converged but the correlations between the sequences have not, and in cases where f depends on a sequence of {u1k}, ..., {utk} we assume the constituent variables have all converged in this way). We represent u∗ka, u ∗ kb via a Cholesky decomposition as
u∗ka = Σkzka + µ ∗ k (19a)
u∗kb = Σk ( C∗kzka + √ 1− (C∗k)2zkb ) + µ∗k (19b)
where zka, zkb ∼ N (0, 1) i.i.d. We thus have ∂u∗kb ∂Ck
=√ Q∗k − (µ∗k)2 ( zka − Ck√
1−C2k zkb ) . Combining this with the fact that∫
Dzg(z)z = ∫ Dzg′(z) for any g(z), integration by parts gives for any
g1, g2
∂ ∂Ck ∫ DzkaDzkbg1(u∗ka)g2(u∗kb)
= ∫ DzkaDzkbg1(u∗ka) ∂g2(u∗kb)
∂u∗kb ∂u∗kb ∂Cc = Σ2k ∫ DzkaDzkb ∂g1(u∗ka)
∂u∗ka
∂g2(u∗kb)
∂u∗kb
(20)
Denoting Σ2s = Q ∗ s − (µ∗s)2 and defining Σk,Σk2 similarly, we have
dCka dCs = σ2kaΣ 2 s Σ2ka (21)
dCkb dCs
= σ2kbΣ 2 s Σ2kb ∗∫ [ gk(u∗ka)gk(u∗kb)
+ σ2k2(Σ 2 sCs+(µ ∗ s) 2) Σ2k2 ∂gk(u ∗ ka)gk(u ∗ kb) ∂Ck dCk dCs
] Dk
=
σ2k2 Σ2s
Σ2k2 ∗ ∫  gk(u∗ka)gk(u∗kb) +σ2k ( Σ2sCs +µ∗2s ) g′k(u ∗ ka)g ′ k(u ∗ kb) Dk where in the last equality we used 20. Using 20 again gives
∂MC(µ∗s, Q∗s, Cs) ∂Ck = Σ2k
∫ DzkaDzkb ∂f(u∗ka)
∂u∗ka
∂f(u∗kb)
∂u∗kb
We now note, using 4, that if Cs = 1,Σz = 1 we obtain Ck = 1 and thus
dCk2 dCs Cs=1,Σz=1 = σ2k2Σ 2 s Σ2k2
( ∫ g2k(u ∗ k)Dk
+σ2kQ ∗ s ∫ (g′k(u ∗ k)) 2Dk ) ∂MC(µ∗s, Q∗s, Cs)
∂Ck Cs=1,Σz=1 = Σ2k
∫ Dzk( ∂f(u∗k)
∂u∗k )2
∂MC(µ∗s, Q∗s, Cs) ∂Cs Cs=1,Σz=1
= E [ ( ∂f(s) ∂s )2 ]
combining the above equations with 21 and comparing the result to 16 completes the proof.
Lemma 2. For any odd function g(x), ( x y ) ∼ N (( µ µ ) ,Σ ( 1 C C 1 )) , 0 ≤
C ≤ 1 we have Eg(x)g(y) ≥ 0.
Proof of Lemma 2. For C = 1 the proof is trivial. We now assume 0 ≤ C < 1. We split up R2 into four orthants and consider a point (a, b) with a, b ≥ 0. We have g(a)g(b) = g(−a)g(−b) = −g(a)g(−b) = −g(−a)g(b) ≥ 0. We will show that p(a, b) + p(−a,−b) > p(a,−b) + p(−a, b) where p is the probability density function of (x, y) and hence the points where the integrand is positive will contribute more to the integral than the ones where it is negative. Plugging these points into p(x, y) gives
p(a, b) + p(−a,−b) p(a,−b) + p(−a, b) = e−α[(a−µ) 2−2C(b−µ)(a−µ)+(b−µ)2] + e−α[(a+µ) 2−2C(b+µ)(a+µ)+(b+µ)2] e−α[(a−µ)2+2C(b+µ)(a−µ)+(b+µ)2] + e−α[(a+µ)2+2C(b−µ)(a+µ)+(b−µ)2]
where α is some positive constant that depends on the determinant of the covariance (since C < 1 the matrix is invertible and the determinant is positive).
= exp (2αCab) exp (−2αCab) cosh (2µα(1− C)(a+ b)) cosh (2µα(1− C)(a− b)) ≥ 1
where the last inequality holds for 0 ≤ C < 1. It follows that the positive contribution to the integral is larger than the negative one, and repeating this argument for every (a, b) in the positive orthant gives the desired claim (if a = 0 or b = 0 the four points in the analysis are not distinct but the inequality still holds and the integrand vanishes in any case).
Lemma 3. The map 5 is convex in the case of the peephole LSTM.
Proof of Lemma 3. We have M(Ctc) = ∫ f taf t b ((Q ∗ c − (µ∗c)2)Ctc + µ∗2c ) + ∫ itai t b ∫ rtar t b + 2 ∫ f t ∫ it ∫ rtµ∗c − µ∗2c Q∗c − µ∗2c .
From the definition of utkb and C t k we have
∂utkb ∂Ctc
= √ Q∗k − (µ∗k)2 ( zka −
Ctk√ 1− Ct2k zkb ) ∂Ctk ∂Ctc = ( zka − Ctk√ 1− Ct2k zkb ) σ2k(Q ∗ c − µ∗2c )√ Q∗k − µ∗2k .
and using ∫ Dxg(x)x = ∫ Dxg′(x) we then obtain for any g(x)
∂ ∂Ctc
∫ DzkaDzkbg(utka)g(utkb) = ∫ DzkaDzkbg(utka)g′(utkb) ∂utkb ∂Ctc
= σ2k(Q ∗ c − µ∗2c ) ∫ DzkaDzkbg′(utka)g′(utkb)
(22)
(Q∗c − µ∗2c ) ∂2M(Cc) ∂C2c = ∂2 ∂C2c
[∫ fafb ( (Q∗c − (µ∗c)2)Cc + µ∗2c ) + ∫ iaib ∫ rarb ] = ∂2 ∫ fafb
∂C2c ((Q∗c − (µ∗c)2)Cc + µ∗2c ) + 2 ∂∂Cc
∫ fafb(Q ∗ c − (µ∗c)2)
+ ∂2
∫ iaib
∂C2c
∫ rarb + ∂ ∫ iaib
∂Cc
∂ ∫ rarb ∂Cc + ∫ iaib ∂2 ∫ rarb ∂C2c
From 22 and non-negativity of some of the integrands
≥ ∂2 ∫ fafb
∂C2c
( (Q∗c − (µ∗c)2)Cc + µ∗2c ) + ∂2 ∫ iaib
∂C2c
∫ rarb + ∫ iaib ∂2 ∫ rarb
∂C2c .
From Lemma 2 we have ∫ rarb ≥ 0 and ∂ 2 ∫ gagb
∂C2c = α
∫ g′′ag ′′ b ≥ 0 for g = f, i, r. We
thus have
∂2M(Cc) ∂C2c ≥ 0
for 0 ≤ Cc ≤ 1. Convexity of this map has a number of consequences. One immediate one is that the map has at most one stable fixed point.

heading: APPENDIX D THE LSTM CELL STATE DISTRIBUTION

text:
Algorithm 1 LSTM hidden state moment fixed point iteration using cell state sampling
function FIXEDPOINTITERATION(µt−1s , Qt−1s ,Θ, ns, niters) Qt−1k ← CALCULATEQK(Qt−1s ,Θ) . Using 4 Initialize c ∈ Rns for i← 1 to niters do
ui,uf ,ur ← SAMPLEUS(Qt−1k ,Θ) . Using 3 c← UPDATEC(c,ui,uf ,ur) . Using 13
end for (µts, Q t s)← CALCULATEM(µt−1s , Qt−1s ,Θ, c) . Using 5
return (µts, Qts) end function
As mentioned in the main text, the cell state differs substantially from other random variables that appear in this analysis since it cannot be expressed as a function of a finite number of variables that are Gaussian at the large N and t limit (see Table 1). Since at this limit the uti are independent, by examining the cell state update eqn. (13) we find that the asymptotic cell state distribution is that of a perpetuity, which obeys
X d = XY + Z
for some random variables Y, Z. Its stationary distributions will have heavy tails Goldie (1991). Due to the bounds on Y, Z in the case of the LSTM, one expects based on results in (Goldie & Grübel (1996)) that the tails of the stationary cell state distribution will be exponential in the case of hard sigmoid nonlinearities or decay like x−x in the case of soft sigmoid nonlinearities, and this tail behavior is indeed observed in simulations. Aside from the tails, the bulk of the distribution can take a variety of different forms and can be highly multimodal, depending on the choice of Θ which in turn determines the distributions of Y, Z.
We overcome this difficulty by sampling from the stationary cell state distribution. For a given value of Qh, the variables utk appearing in (13) can be sampled since their distribution is given by (3) at the large N limit. The update equation (13) can then be iterated and the resulting samples approximate well the stationary cell state distribution for a range of different choices of Θ, which result in a variety of stationary distribution profiles (see Appendix E3). The fixed points of (5) can then be calculated numerically as in the deterministic cases, yet care must be taken since the sampling introduces stochasticity into the process.
An example of the fixed point iteration eqn. (5) implemented using sampling is presented in Algorithm 1. The correlations between the hidden states can be calculated in a similar fashion. In practice, once the number of samples ns and sampling iterations niters is of order 100 reasonably accurate values for the moment evolution and the convergence rates to the fixed point are obtained (see for instance the right panel of Figure 2). The computational cost of the sampling is linear in both ns, niters (as opposed to say simulating a neural network directly in which case the cost is quadratic in ns).

heading: APPENDIX E ADDITIONAL EXPERIMENTS AND DETAILS OF EXPERIMENTS

heading: E.1 DYNAMICAL SYSTEM

text: We simulate the dynamics of Qts in eqn. (5) for a GRU using inputs with Σ t z = 0 for t < 10 and Σtz = 1 for t ≥ 10. Note that this evolution is independent of the evolution of Cts. The results show good agreement in the untied case between the calculation at the large N limit and the simulation, as shown in Figure 5.

heading: E.2 HEATMAPS

text: In Figure 6 we present results of training on the same task shown in Figure 2 with tied weights, showing the deviations resulting from the violation of the untied weights assumption.

heading: E.3 SAMPLING THE LSTM CELL STATE DISTRIBUTION

text: As described in Appendix D, calculating the signal propagation time scale and the moments of the state-to-state Jacobian for the LSTM requires integrating with respect to the stationary cell state distribution. The method for doing this is described in Algorithm 1. As is shown in Figure 7, this distribution can take different forms based on the choice of initialization hyperparameters Θ, but in all cases we have studied the proposed algorithm appears to provide a reasonable approximation to this distribution efficiently. The simulations are obtained by feeding a network of width N = 200 with i.i.d. Gaussian inputs.

heading: E.4 CRITICAL INITIALIZATIONS

text:
Peephole LSTM:
µi, µr, µo, ρ 2 i , ρ 2 f , ρ 2 r, ρ 2 o, ν 2 i , ν 2 f , ν 2 r , ν 2 o = 0
µf = 5 σ2i , σ 2 f , σ 2 r , σ 2 o = 10 −5
LSTM (Unrolled CIFAR-10 task):
µi, µr, µo, ρ 2 i , ρ 2 f , ρ 2 r, ρ 2 o, ν 2 f , ν 2 o = 0
ν2i , ν 2 r , σ 2 o , µf = 1 σ2i , σ 2 f , σ 2 r = 10 −5
The value of µf was found by a grid search, since for this task information necessary to solve it did not require signal propagation across the entire sequence. In other words, classification of an image can be achieved with access only to the last few rows of pixels. The utility of the analytical results in this case, as mentioned in the text, is to greatly constrain the hyperparameter space of potentially useful initializations from theoretical considerations.

heading: E.5 STANDARD INITIALIZATION

text:
LSTM and peephole LSTM:
Kernel matrices (corresponding to the choice of ν2k) : Glorot uniform initialization Glorot & Bengio (2010)
Recurrent matrices (corresponding to the choice of σ2k): Orthogonal initialization (i.i.d. Gaussian initialization with variance 1/N also used giving analogous performance)
µi, µr, µo, ρ 2 i , ρ 2 f , ρ 2 r, ρ 2 o = 0
µf = 1

heading: E.5.1 LONG SEQUENCE TASKS

text: Learning rate scan: 8 equally spaced points between 10−2 and 10−5. Validation set: 10000 images for MNIST and CIFAR-10. In Table 2, the results for the standard LSTM on MNIST were reproduced from Arpit et al. (2018) and the results on CIFAR-10 were reproduced from Trinh et al. (2018).