Trainability, Expressivity and Interpretability in Gated Neural ODEs
Timothy Doyeon Kim Tankut Can Kamesh Krishnamurthy
Understanding how the dynamics in biological and artificial neural networks
implement the computations required for a task is a salient open question in
machine learning and neuroscience. In particular, computations requiring
complex memory storage and retrieval pose a significant challenge for these
networks to implement or learn. Recently, a family of models described by
neural ordinary differential equations (nODEs) has emerged as powerful
dynamical neural network models capable of capturing complex dynamics. Here, we
extend nODEs by endowing them with adaptive timescales using gating
interactions. We refer to these as gated neural ODEs (gnODEs). Using a task
that requires memory of continuous quantities, we demonstrate the inductive
bias of the gnODEs to learn (approximate) continuous attractors. We further
show how reduced-dimensional gnODEs retain their modeling power while greatly
improving interpretability, even allowing explicit visualization of the
structure of learned attractors. We introduce a novel measure of expressivity
which probes the capacity of a neural network to generate complex trajectories.
Using this measure, we explore how the phase-space dimension of the nODEs and
the complexity of the function modeling the flow field contribute to
expressivity. We see that a more complex function for modeling the flow field
allows a lower-dimensional nODE to capture a given target dynamics. Finally, we
demonstrate the benefit of gating in nODEs on several real-world tasks.