Seeing the forest and the tree: Building representations of both individual and collective dynamics with transformers
Complex time-varying systems are often studied by abstracting away from the
dynamics of individual components to build a model of the population-level
dynamics from the start. However, when building a population-level description,
it can be easy to lose sight of each individual and how each contributes to the
larger picture. In this paper, we present a novel transformer architecture for
learning from time-varying data that builds descriptions of both the individual
as well as the collective population dynamics. Rather than combining all of our
data into our model at the onset, we develop a separable architecture that
operates on individual time-series first before passing them forward; this
induces a permutation-invariance property and can be used to transfer across
systems of different size and order. After demonstrating that our model can be
applied to successfully recover complex interactions and dynamics in many-body
systems, we apply our approach to populations of neurons in the nervous system.
On neural activity datasets, we show that our multi-scale transformer not only
yields robust decoding performance, but also provides impressive performance in
transfer. Our results show that it is possible to learn from neurons in one
animal's brain and transfer the model on neurons in a different animal's brain,
with interpretable neuron correspondence across sets and animals. This finding
opens up a new path to decode from and represent large collections of neurons.