Understanding and Optimizing Gradients of Neural Networks Using Dynamical Systems Theory

Discovering and responding to associations between temporally distant cues is crucial for an animal's survival. Solving this temporal credit assignment problem allows organisms to bridge the gap between when a stimulus cue arrives and when its effect unfolds. Gradient-based training of recurrent neural circuit models for temporal tasks with long time horizons presents challenges, potentially leading to vanishing or exploding gradients. We have connected this issue to the Lyapunov exponents of the forward dynamics, which describe how perturbations grow or shrink in tangent space. Here, we propose to address gradient instability in recurrent spiking and firing rate networks by controlling the Lyapunov exponents of forward dynamics throughout learning.

We regularize Lyapunov exponents towards zero, ensuring that the corresponding directions in tangent space grow or shrink only slowly, for more robust propagation of learning signals over long time horizons. Our method enhances the success rate of RNNs in typical neuroscience tasks that involve bridging task events across many time steps. Furthermore, we demonstrate that applying gradient flossing during training enhances trainability for challenging temporal credit assignment tasks. We establish a connection between the Lyapunov exponents and the dimensionality of the gradient signal in backpropagation. Additionally, we demonstrate the effectiveness of our approach both on spiking and firing rate networks. Our results suggest that dynamically controlling Lyapunov exponents can significantly improve the stability and effectiveness of RNN training. Although Lyapunov exponents are strictly defined for dynamical systems, like maps or flows that act as endomorphisms, the end-to-end Jacobian of deep feedforward networks can be analyzed and optimized in a comparable way. We speculate on the optimization of neural dynamics in animals over evolutionary timescales to bridge long time horizons.