Dynamical Isometry and a Mean Field Theory of LSTMs and GRUs

1. Introduction

• Standard hyperparameter searches can miss trainable configurations when those configurations lie in a low-dimensional subspace.Random, grid, and Bayesian search are not guaranteed to succeed under this constraint.Source: §1 · p. 1
• Critical initialization remains trainable on longer sequences where standard initialization begins to fail.The paper illustrates this behavior with peephole LSTMs classifying sequences of MNIST digits.Source: Figure 1 · p. 1, §1 · p. 2
• Initialization settings for long-sequence LSTMs and GRUs occupy a theoretically characterizable subspace whose conditions support signal propagation and stable gradients.Trainability requires relevant input signals to reach the loss while gradients avoid exponential explosion or vanishing.Source: §1 · p. 2
• The analysis combines mean field signal-propagation theory with dynamical isometry to study recurrent architectures including GRUs, LSTMs, and peephole LSTMs.Experiments corroborate the calculations and motivate initialization schemes for long-sequence tasks.Source: §1 · p. 2

2.1. Mean field analysis of neural networks

• Initialization schemes controlling signal propagation have improved benchmark performance in feed-forward and simple gated recurrent networks.Related feed-forward results enabled very deep architectures without batch normalization or skip connections.Source: §2.1 · p. 2
• Mean field and random matrix theory analyze forward signal correlations and Jacobian singular values, with the latter targeted to improve gradient conditioning.Forward propagation is necessary but not sufficient for trainability, while backward analysis controls the first two moments of squared singular values.Source: §2.1 · p. 2
• Modified initialization shows empirical evidence of beneficial rather than adverse effects on generalization error despite neural-network nonconvexity.The concern is that initialization could change the converged parameter-space location and worsen generalization.Source: §2.1 · p. 2

2.2. The exploding/vanishing gradient problem and signal propagation in recurrent networks

• Exploding or vanishing gradients hinder long-sequence training, motivating approaches based on orthogonal weights, gradient modification or clipping, auxiliary losses, and altered recurrent architectures.These methods typically modify the training algorithm, loss, or architecture and are treated as complementary to initialization.Source: §2.2 · p. 2

2.3. Dynamics beyond initialization

• The analysis is formulated at initialization, although large-width results suggest covariance quantities may remain nearly constant during training.The paper notes that finite-width networks can nevertheless gain trainability from satisfying the derived initialization conditions.Source: §2.3 · p. 3

3. Notation

• Matrices use bold uppercase Latin characters, while vectors use bold lowercase Latin characters.Source: §3 · p. 3
• The notation includes normalized traces, Hadamard products, element-wise sigmoid and tanh functions, and diagonal matrices.A diagonal matrix is denoted by D_a when a is placed on the diagonal.Source: §3 · p. 3

4.1. Model description and important assumptions

• The general recurrent formulation covers GRUs, LSTMs, peephole LSTMs, and related architectures under an untied-weights mean-field framework.At infinite width, pre-activations are approximated as jointly Gaussian, yielding deterministic dynamics for moments and correlations.Source: §4.1 · p. 4, §4.2 · p. 4

4.2. Forward signal propagation

• Forward signal propagation is characterized by convergence toward a fixed point, whose rate defines an effective timescale ξ for information retention.Because the fixed point erases distinctions between input sequences, delaying convergence supports longer-range propagation; ξ diverges as χC∗_s approaches 1 from below.Source: §4.2 · pp. 4–5

4.3. Backwards signal propagation - the state-to-state Jacobian

• The state-to-state Jacobian controls gradient stability because length-T gradients depend on matrix polynomials of order T in J.Choosing hyperparameters with mean squared singular value near 1 and small variance keeps powers of J from diverging rapidly.Source: §4.3 · p. 5

4.4. Dynamical Isometry

• Dynamical-isometry conditions require stable forward propagation and controlled Jacobian spectra, although they may need to be satisfied approximately.The conditions are motivated by preventing exploding or vanishing gradients, and no architecture is guaranteed to satisfy all of them exactly.Source: §4.4 · p. 5
• The predicted propagation timescale identifies transitions between high and low training accuracy across architectures and hyperparameter directions.Empirical evidence also indicates that timescales much larger than a task requires can degrade performance, while the Σz = 1 approximation remains predictive across tasks.Source: Figure 2 · p. 6, §4.4 · p. 6
• For LSTMs, calculating signal timescales and Jacobian moments requires integrating over the stationary cell-state distribution.The paper uses specialized calculations because the additional cell state complicates the analysis.Source: §4.4 · p. 6

5.1. Padded MNIST Classification

• For padded MNIST, predicted propagation timescales agree with the transition between successful and unsuccessful training as sequence length and hyperparameters vary.The comparison uses untied-weight GRUs and LSTMs and overlays training accuracy with multiples of ξ.Source: §5.1 · p. 6

5.2. Squared Jacobian spectrum histograms

• The predicted and empirical squared-singular-value moments agree for the state-to-state Jacobian, supporting the calculation under an ergodicity assumption.The spectrum changes dramatically between standard initialization and initialization near the dynamical-isometry conditions.Source: §5.2 · p. 7

5.3. Unrolled MNIST and CIFAR-10

• Critical initializations are evaluated on unrolled MNIST and CIFAR-10 sequence tasks, including a CIFAR-10 case where the best initialization has ξ much smaller than the sequence length.This suggests that successful classification can use information from a subset of the sequence.Source: §5.3 · p. 7, Table 2 · p. 7

5.4. Repeated pixel MNIST and multiple digit MNIST

• As sequences lengthen, critical initialization becomes faster by orders of magnitude than standard initialization, while standard initialization eventually fails and critical initialization still succeeds.For shorter sequences, the two initializations perform equivalently; the result holds for repeated-pixel and multiple-digit MNIST tasks.Source: Figure 4 · p. 7, §5.4 · p. 8

• The derived initialization schemes dramatically improve performance on long sequence tasks by stabilizing input-to-loss signals and loss-to-weight gradients.Source: §6 · p. 8
• The dynamical-isometry-satisfying initialization subspace is multidimensional, so the paper does not provide a principled rule for selecting one preferred initialization.Source: §6 · p. 8
• The analysis indicates that the forget gate is critical, while some other gates or weight matrices may be removable without violating the conditions.Source: §6 · p. 8
• Important open problems include choosing the forward-propagation time scale, quantifying violations of untied weights, and relaxing the constant Σz assumption.Source: §6 · p. 8
• The persistence of dynamical-isometry conditions during training remains unresolved for both feed-forward and recurrent architectures.Source: §6 · p. 8

A.2. Moments of the state-to-state Jacobian and general form of the dynamical isometry conditions

• The appendix characterizes dynamical isometry through normalized traces of powers of JJT and shows that controlling only the first moment is insufficient because singular-value variance can remain large.Source: §A.2 · pp. 10–11
• Under the stated assumptions, the mean squared singular value of the state-to-state Jacobian is linked to χCs, the forward signal-propagation time scale.Source: §A.2 · p. 10
• The general dynamical-isometry conditions can be satisfied when Σz = 1, a0 = 1, and ak≠0 = 0, motivating zero weight variances and large positive forget-gate bias limits.Source: §A.2 · p. 11

A.3. Dynamical isometry conditions for selected architectures

• The resulting conditions are worked out separately for the LSTM, peephole LSTM, and GRU architectures.Source: §A.2 · p. 11, §A.3 · p. 11

A.3.3. LSTM

• For the LSTM analysis, stability and approximation accuracy improve when the output gate is not concentrated near zero.Source: §A.3.2 · p. 11, §A.3.3 · p. 12

A.4. RNN Table supplement

• Table 3 lists additional recurrent architectures in the paper’s general recurrent-network form, with pt denoting the network output at each time step.Source: Table 3 · p. 12

• At large width, isotropy makes the relevant gate variables independently distributed and independent of the state, enabling single-coordinate analysis of the recurrent dynamics.Source: B. Auxiliary Lemmas and Proofs · p. 12
• Lemma 2 establishes an inequality for odd functions under correlated Gaussian inputs, supporting later bounds in the analysis.Source: B. Auxiliary Lemmas and Proofs · pp. 13–14
• The relevant map is convex for the peephole LSTM, which implies that it has at most one stable fixed point.Source: B. Auxiliary Lemmas and Proofs · pp. 14–15

• Because the standard LSTM hidden state depends on a cell state with non-trivial dynamics, its moment evolution cannot generally be written using only a finite Gaussian-variable representation.Source: C. The LSTM cell state distribution · p. 15
• The asymptotic cell-state distribution is a perpetuity, whose stationary distribution can have heavy tails and highly multimodal bulk shapes.Source: C. The LSTM cell state distribution · p. 15
• The stationary distribution’s form depends on the initialization hyperparameters Θ, which determine the distributions of the multiplicative and additive update variables.Source: C. The LSTM cell state distribution · p. 15
• Sampling from the stationary cell-state distribution enables numerical fixed-point and correlation calculations despite the absence of a closed-form density or likelihood.Source: C. The LSTM cell state distribution · pp. 15–16

D.1. Dynamical system

• The large-N calculation agrees well with simulations for the GRU dynamical system when weights are untied.Source: §D.1 · p. 16

D.2. Heatmaps

• Figure 6 shows training accuracy across sequence lengths and hyperparameters for tied-weight networks, with deviations from the untied prediction.Performance generally improves when the predicted forward propagation time scale ξ is high, although tied weights are outside the analysis scope.Source: §D.2 · p. 16, Figure 6 · p. 17

D.3. Sampling the LSTM cell state distribution

• Algorithm 1 efficiently approximates the stationary LSTM cell-state distribution across different initialization hyperparameters.The approximation agrees reasonably with simulations of width-N = 200 networks driven by i.i.d. Gaussian inputs.Source: §D.3 · p. 17, Figure 7 · p. 17

D.4. Critical initializations

• For the unrolled CIFAR-10 task, µf was selected by grid search because solving the task does not require propagation across the entire sequence.The analytical results still constrain the hyperparameter space of potentially useful initializations.Source: §D.4 · p. 18

D.5. Standard initialization

• The standard-initialization comparisons include Glorot uniform, orthogonal, and i.i.d. Gaussian initializations, with an eight-point learning-rate scan from 10^-2 to 10^-5.Validation sets contain 10000 images for MNIST and CIFAR-10.Source: §D.5 · p. 18, §D.5.1 · p. 18