Geometric dynamics of signal propagation predict trainability of transformers

We investigate forward signal propagation and gradient back propagation in deep, randomly initialized transformers, and we use our analysis to propose simple and geometrically meaningful criteria for hyperparameter initialization that ensures trainability of deep transformers. Our approach leverages a dynamical perspective in which we regard the token vectors as particles evolving according to the layerwise map of the transformer. We derive simple update equations for the ensemble-averaged geometry of this many-particle configuration starting from a permutation-symmetric simplex. Our analysis reveals an order-chaos phase transition in the forward direction, which corresponds to converging/diverging token angles, and another transition in the backwards direction between exploding/vanishing gradients. We derive the phase boundaries as a function of the initialization hyperparameters, such as the variance of the weights and the strength of attentional and MLP residual branches. We analytically derive two Lyapunov exponents: a token angle exponent that governs the departure from the edge of chaos and a gradient exponent that governs the rate of exponential growth or decay of backpropagated gradients. We show through experiments that the final test loss at the end of training is correlated with these two exponents at the beginning of training, suggesting that simultaneous vanishing of these two exponents constitute simple criteria for achieving minimal test loss.