Seeing the forest and the tree: Building representations of both individual and collective dynamics with transformers
Authors: Ran Liu
Mehdi Azabou
Max Dabagia
Jingyun Xiao
Eva L Dyer
Venue: NeurIPS 2022 Accept (2022)
Keywords: Transformers
population dynamics
multi-variate time-series
neural activity
many-body systems

Abstract: 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 they contribute 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 model not only yields robust decoding performance, but also provides impressive performance in transfer across recordings of different animals without any neuron-level correspondence. By enabling flexible pre-training that can be transferred to neural recordings of different size and order, our work provides a first step towards creating a foundation model for neural decoding.

Seeing the forest and the tree: Building representations of both individual and collective dynamics with transformers
Ran Liu, Mehdi Azabou, Max Dabagia, Jingyun Xiao, Eva L. Dyer* Georgia Institute of Technology
Abstract
Complex time-varying systems are often studied by abstracting away from the dynamics of individual components build model of the population-level dy- namies from the start. However, when building a population-level description, it can be easy to lose sight of each individual and how they contribute the larger picture. In this paper, we present novel transformer architecture for learning from time-varying data that builds descriptions of both the individual as well 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-invant invvrriacce 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 pop- ulations of neurons in the nervous system. On neural activity datasets, we show that our model not only yields robust decoding performance, but also provides impressive performance in transfer across recordings of different animals without any neuron-level correspondence. By enabling flexible pre-training that can be transferred to neural recordings of different size and order, work provides a first step towards creating foundation model for neural decoding.
Introduction
Complex systems (such as the brain) contain multiple individual elements (e.g. neurons) that interact dynamically to generate their outputs. The process by which these local interactions give rise to large-scale behaviors is important in many domains of science and engineering, from ecology 11
and social networks 5
to microbial interactions [6] and brain dynamics [7].
A natural way to model the activity of system is to build collective or population view, where we consider individuals (or channels) jointly to determine the dynamics of the population (Figure A))) In many cases, studying systems from this population-level perspective has provided important insights into collective computations and emergent behaviors 77; however, is also possible to lose sight of the contributions of different individuals' dynamics to the final prediction or inference. This is important in many settings where the dynamics of different individuals may be of interest, either due to their different functional roles in the system [8] 9 6 10], due to shift in their dynamics because of sensor displacement or corruption [TT Moving forward, we need methods that can build good population-level representations while also providing an interpretable view of the data at the individual level.
Contact: rliu361 @ @gatech.edu, evadyer@ gatech.edu. Project page and code: https://ncrdslaab htt.rrdddllagiihuuu.///////
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Here, we present a new framework for modeling vimy-ng observations of a system which uses dynamic embeddings of individual chamnels construct a population-Ievel view (Figure[]B.C),, Our model, which we dub Embedded Interaction Transformer or EIT, decomposes population dynamics by first learning rich features from individual time-series before incorporating information and learned interactions across different individuals in the population. One critical benefit of our model is spatial/indivddaal separability: it builds a population-leve representation from embeddings of individual channels, which naturally leads to channel-level permutation invariance. In domain generalization tasks, this means a trained model can be tested with permuted channels or entirely different numbers of channels.
To first understand how the model captures different types of interactions and dynamics, we exper- iment with synthetic many-body systems. Under different types of interactions, we show that our model can be applied to successfully recover the dynamies of known systems.
We then turn our attention to recordings from the nervous system [14), where we have access to readouts from populations of neurons in the primary motor cortex of two rhesus macaques that are performing the same underlying center-out reaching motor task. The stability of the neural representations and collective dynamics found in data from these animals [4 makes them an ideal testbed to examine the generalization of our approach. We show that, remarkably, generalization not only occurs across recording sessions (with different sets of neurons) from the same animal, but that we can also transfer across animals through only a simple linear readout. The performance of the linear decoder based on pre-trained weights the across- across-animal transfer outperforms many models that are trained and tested on the same animal. This result provides an exciting path forward in neural decoding across large and diverse datasets from different animals and highlights the utility of our architecture.
The main contributions of this work are follows:
we introduce Embedded Interaction Transformer (EIT): a novel framework for learning from multi-variate time-series data that decomposes the dynamics of individual channels and their interactions through a two-stage transformer architecture.
In Section 3
In Section we propose methods for generalization across datasets of different input size (number of channels) and ordering. We show that by decomposing individuals and interactions, our architecture can be used to find functional correspondence between channels in different datasets by measuring the similarity in their embeddings with the Wasserstein divergence.
In Section we apply EIT to both many-body systems and neural activity recordings. After demonstrating our model's 's robust decoding performance, we validate its ability to transfer individual dynamics by performing domain generalization across different neural recordings and investigate the alignments of neurons across different populations.
2 Background and Related Work
2.1 Transformers
Transformers have revolutionized natural language processing (NLP) through the mechanism of self-attention [17], which helps the model Iearn long-range interactions across dif- ferent elements (represented by tokens) in a sequence. Consider a sequence X = [x x² x xNI" e N consisting of N tokens of d-dimensions. The aelf-attention mechanism is built on the notion of queries Q, keys K, and values V, which linear projections of the token embeddings. Let Q = XW K = XWR,V V = XWy, where Q E RNNd, K E IR Xxd V E NN... The attention operation can be written as:
Self-attention.
Attention (Q, V) = softmax (xT)))
(1)
where softmax(-) denotes row-wise softmax normalization function.
Multi-head self-attention (MSA). Rather than only learning one set of keys, queries, and values for our tokens, MSA allows each head to find different patterns in the data that are useful for inference. Each of the h heads provides a d, dimensional output, and the outputs from all heads are concatenated into a d f vector, which is further projected to produce the final representation.

---

A state space view populations
Time
Charne nzurons)
Ind
Population Encoder
e
Ind
III III
(8)
Individual Encoder
Interaction Encader
NM
A state space view individuals
Proj Module
Bottleneck, Concatenate, Proyect
Lobel space (V)
individual Transformer
interaction
tronsformee
(C) EIT Model Training
Figure 1: Embedded Interaction Transformer (EIT) J. (A) A traditional state-space view would treat the collective dynamics as population right from the beginning. and use population encoder to learn how the dynamics progress along time. (B) EIT learns dynamic embeddings of each channel with individual encoder at the beginning. After embedding each channel's dynamics, feed them into interaction encoder build population representation. The two encoders work together to build representation space that is richer than that of the traditional method. (C) The detailed architecture: EIT consists of an individual transformer that processes data for each individual, an interaction transformer that processes embeddings at each timepoint, and projection modules the end of both transformers.
The full operation of transformer of L layers can be written below:
Zo = [Embed(x' Embed(x") Embed x + Epos
(2)
Ze+1 Ze MSA(Z.) FF (Z)). { {0. (O..... 13
where Zo the summation of the individual embedding of each data token and the positional embedding Epus that helps to retain the positional information, and each transformer layer is the combination of the MSA operation (MSA(-)), the projection (FF())) and residual connections.
Transformers for multi-channel time-series. Multi-channel time-series are a natural candidate for modeling with the transformer architecture. One common approach of modeling many time-varying data streams with transformers to first aggregate features across all channels into combined representation at the beginning of the model. The population dynamics is thus learnt with the resulting embedding via temporal transformer for the purpose of inference. There are many complex ways to
create this embedding: [18] and [99] extract embeddings from multivariate physical systems with Koopman operators before feeding the resulting representation into temporal transformer; [20] use graph neural network to embed interconnected-strrccii
to perform Akeleton-based action recognition; [2T] use a convolutional architecture to extract image embeddings before feeding them into a transformer.
Another approach re-design the attention block such that the attention operation is computed both along the temporal and spatial dimension. Many variants of this have been shown effective:
attention block [2 propose non-autoregressive module to generate queries for time series forecasting; (23] use stacked spatial-temporal modules for traffic flow forecasting: GroupFormer 24 embed spatial and temporal context parallel create clustered attention block for group activity recognition. While our approach also makes use of the high-level idea of separating spatial (individual) and temporal information, crucially, we restrained the direct computation between the spatially-an leeeaae attention map and the temporally-relan of the individual which is completely free of spatial interactions.
attention map, which yields a representation
Spacetime attention in video transformers. With the advances of the Vision Transformer as new way to extract image embeddings, many apatial-tempont transformer' architectures have been developed the video domain 28). Such works explore and propose interesting solutions for how to organize spatial attention and temporal attention with either coupled (series) [28] and factorized (parallel) attention blocks (26), as well how to create better tokens for videos by creating three-dimensional apatio-temport 'tubes' as the tubelet tokenizations (26). However, these methods

---

leverage inductive biases that are specific to images and videos, and do not process potentially separatable channels that we are interested in.
Object-centric representation learning. There are many representation learning methods in video 32) physics-guided complex systems [33], and behavior modeling [34] 33] that also aim to learn interactions between discrete objects observed in the data. In these approaches, interactions between different objects are typically learned in joint manner by combining the inferred representation of the objects into a common representation. In contrast, our framework aims to decompose dynamics into two parts, forming representation of the dynamics of each individual source or object in addition to a representation of the interactions across many sources. Thus, one could imagine using our approach to build enhanced representations of object dynamics in other vision and behavioral neuroscience applications [30 34]. would interesting to see if an enhanced or individual representation could be used to further improve the performance of upstream object localization and inference approaches.
2.2 Neural decoding and the challenge of generalization
The brain is an incredibly complex system. Neural activity is guided both by tight interplay of membrane dynamics and ionic currents within a single neuron [36 as well as interactions with other neurons over larger distributed circuits [0,42] Understanding how populations of neurons work together to represent their inputs an important objective in modern neuroscience.
Accordingly. constructing representations that can explain neural population activity is an active area of research with many existing approaches. Discriminative models build representations that align or predict temporally-adjacent and masked neural activity, and include work such as MYOW [43] and the
Neural Data Transformer (NDT)
44 Generative approaches instead attempt to model neural activity
as driven by latent factors 45 with success using switching linear dynamical systems [46] and
recurrent neural networks 477. Recently, SwapVAE [48] emerged as a hybrid of these two approaches, combining latent factor model of neural activity with self-supervised learning techniques to predict across data augmentations. However, all of these methods produce representations exclusively at the population level.
The challenges of transfer across recordings of neural population activity. More often than not, neural recordings are collected from different sets of neurons within a brain region of interest. As a result, it is very difficult to generalize across neural recordings and the integration of many neural datasets remains a critical and yet open problem 44 50).
Many existing approaches attempt to address this challenge by learning population-level descriptions of neural data and finding ways align their latent representations [50] 14]. For example, (51) proposed manifold alignment method to continuously update a neural decoder over time as neurons shift in and out of the recording;
[52) proposed an adversarial domain adaptation for stable brain- machine interfaces; (53) proposed a robust alignment of cross-session recordings of neural population 47 fit each set of neurons with its own encoder and used dynamics to inform alignment in the latent space during training. While such approaches have provided useful insights into the stability of neural representations across subsets of distinct neurons in the same and different brains solving the same task [16). the representations used for alignment are all formed jointly across the entire neural recording and thus each new dataset needs a different encoder or unique re-mapping, making it inherently unscalable when applied to many new datasets. EIT provides a novel strategy for decomposing neural dynamics that can be applied to arbitrary numbers of new recordings. making it possible to scale to large numbers of datasets. In addition, EIT can potentially be used to find neuron-level correspondence (or functional correspondence) through its decomposed representation of the population-level activity.
activity through unsupervised domain adaptation;
Methods
3.1 Formulation and Setup
In this work, we consider time-varying prediction from multi-variate time-series observations. Con- sider X E NxT as an input datum that measures the dynamics of N individual channels over T points in time. Let ziy denote the value of channel at the j-th time point, where 1 < i S N and

---

1 i T, and let X denote the time series of the i-th channel. We consider time-varying labels y = (VI.....Tr] where yj E R the label the j-th time point.
Our aim to find a function which takes input X to predict target variables y- Instead of jointly considering individual time-series immediately, we instead decompose the model into two functions : R and g IN + R2, so that the optimization problem becomes:
T Σ
Proj [g (((("")),..,,,,(,))))))) yi
min Proi
Ex
(3)
where Proj projects latent representations to predict the labels, and E(:,) is a loss measuring the discrepancy between the labels and the time-varying output from the decoder. In words, f is applied to each time-series independently, while combines individual embeddings across the population at single time step. Crucially, this decomposition provides f invariance to the number of channels fed into the model. In practice, we map each Zij to a higher dimension to increase capacity for inference.
3.2 Approach
As shown in Figure B), EIT decomposes collections of temporal data through two transformers. The first learns representations of the dynamics of individual channels and the second learns their population-level interactions.
temporal-spatial module that disentangles individual dynamics and their interactions. To separate interactions at the population level from the individual dynamics,
start by building
representations from individual time-series. For the i-th channel, we obtain an initial embedding of T
temporal tokens as Z [Embed Embed .... Embed (F }] where Embed R RM
and z E ER TRTx We apply a transformer with multi-head attention to the resulting sequence using Equation ( Let 20) E TRM denote the final embeddings with elements 2, obtained from the first transformer in model and 7 R xTxM be the combined embedded output.
After building representations of dynamics at the scale individual time-series, we then pass the representations into population-leved transformer to capture interactions between all of the individual time-series. To do this, we take the output of the first module Z and slice it along the temporal axis yield an embedding of each time point in terms of the different channels. Denote this new sequential
reslicing of the data at the j-th timepoint v-(3) [81j. Žgj, ÊN₃] Epus- where Epus E R x the fixed positional embedding that dependent on neuron i's identity. After applying the attention block operations in Equation (2) again, we then arrive at our final embedding of our time-series as V e N
Let diy be the final representation of the observation Tiy-
A projection module that preserves individual identity. To form a population representation from the embeddings individual channels, we design a projection module that preserves individual identities. The projection module Proj(-) consists of two parts: 1) For representations of individuals 2i E RM for the i individual and j-th time-point, Proj first bottleneck the individual latents with function bottleneck R.M + I reduce the dimensions. 2) To form a population representation, Proj then concatenates the representation across individuals at the j-th time point, and learns another projection project: N -> R infer the label y,- Through this operation, we could obtain both the population representation that could be used for inference and the individual representation that is affected in the minimal level.
Learning representations through a multi-stage loss. can provide both a description at the individual-level and the population-level. we compose a loss function as follows:
To build unified representation that
Liotal (nrooj GN u;) + a ((Pr.)- EN,)),
(4)
where (roj-Int(-) and Proj-Ind(-) denote two projection modules after the population and individual transformers, respectively; can set to either classification loss (CrossFEnt or regression loss (MSE), depending the type of prediction target consider; is scaling factor that determines
5

---

how much emphasis is placed on prediction from individual representations in the first half of the network.
By setting a > O, we train a projection module on the intermediate individual-level embeddings that provides purely individual-based representation of the population. Empirically, this more flexible architecture also seems to improve the stability of training when compared with setting a O. Note that the first part of Equation 4 is a special case of Equation 3 when both the individual module and interaction module are translormer encoders.
Remark:
3.3 Generalization across domains through linear decoding and functional alignment
A key element of our framework is that, after training, we can use the first part of the network (f) that we have learned for individual-level dynamics, on a new dataset of arbitrary size (number of neurons) and ordering. Here, we describe: (i) the linear probing approach for transfer across different sized or shuffled populations, and (ii) ways to characterize alignment of functional properties of different individual channels that are studied.
Decoding from a new population of different dimension and ordering. In many domains, it difficult to reliably preserve channels across datasets. Our model instead provides a flexible framework for domain generalization as the first transformer (f) operates only on individual time- series without considering the entire population. Specifically, to transfer to new dataset, we can learn new projection h(-)) that acts on the concatenated pre-trained embeddings (outputs of f) to build predictions about downstream target variables y:
h arg min ((((())j...,,NN))))) v));
(5)
where h is a decoder mapping representations to predicted labels, 2 are the representations formed for the individual at the j-th time point, and yj is the label at j-th time point. When we restrict h to be linear, this approach is similar to the evaluation methods self-supervised learning [54] 55), where linear decoder is used to evaluate the quality of the frozen representation based various different downstream tasks.
Functional alignment through Wasserstein- representational similarity analysis. Our model also provides a flexible framework for identifying correspondences that might exist between channels in different recordings. For instance, in the case of neural recordings, we can use EIT find correspondence across neurons in different recording sessions. Recall that for each individual time- series ( E R² our model builds a representation space of size 2 E RTxM To characterize the similarity across different channels, we pass many samples through the individual encoder to estimate the distribution of (() E T We denote this distribution of individual as R.
We then use the Wasserstein divergence (W), measure of distributional distance motivated by optimal transport (OT) [56] 57 to obtain robust measures similarity (More details in Appendix [.T]]. For one set of distributions (R("), ( .................... ..... ) from domain (1) and another set of distributions
{1(22) (2) ..... N2 from domain (2), we compute the divergence between all pairs of individuals,
which yields a matrix D E gN,,N,, where D|i = W(R ,R)")).
4
Experiments
4.1
Synthetic experiment: observing superposed many-body systems
We first tested our model on synthetic many-body systems [58] (see 59 60 for other possible synthetic datasets with additional complexity). As shown in Figure AA) for many-body system, a body's observed trajectory is decided by its own properties (its mass and starting point), and its own intended dynamics are further affected by its interaction with other bodies. Furthermore, for certain many-body systems (e.g. k-body systems for k > 3) the trajectories are chaotic in nature, which is aligned with the stochasticity of neural activity [6] 62). To test whether our model able to capture, disentangle, and decode the body dynamics, we create an experiment of 'superposed" many-body systems (as shown in Figure B)). where two separate many-body systems are overlaid with each other to create system where some connections exists (bodies within the same system), while some do not (bodies from different systems).

---

(A)
(1) body moss
(B)
'Superposed' 2-body System
(C)
69)
Observe
ob)
oby
object abject 2 start
(2) body type
Somple
Individusi lotents
evomple
(3) Interoction
Predict
Population latents
2 body System
The color rbody type)? interaction?
Figure 2: Synthetic Experiments. (A) We show an example of a two-body system, where the system is roughly controlled by three factors: the body mass, the body type, and the body interaction. (B) We create a "superposed" two-hody system by sampling initial conditions from the distribution, observing the produced trajectories over shorter time intervals, and predicting the body type and their interaction from the trajectories. both the individual latent space and the population latent space. The upper right corner shows how our OT evaluation metric can align and distinguish bodies of the same type (red denotes lower discrepancy).
(C) We visualize
Generating non-chaotic and chaotic systems. We generated non-chaotie two-body systems and chaotic thrce-body systems in a relatively stable near-circular way [63], where the trajectories are obtained with the Explicit Runge- Kutta method [64]. We refer the readers to Appendix 2. for the details of the set of second-order differential equations that guide the movement.
As shown in Figure B)) we create the 'superposition' of many-body systems in the following pipeline: two independent many-body systems are sampled, and then their trajectories without additional type interaction information are observed. When training the model, we aim to generate a latent representation for each body at each time that predicts the underlying factors that guide the body movements. Note that we use different distribution of body mass to separate the training set (Source) and testing set (Target), which makes the generated body trajectories purely controlled by two factors: 1) where the body initially starts (as the body types); and 2) which body is this body interacting with (as the interaction). The model trained to recover these two factors.
For the two-body system and the three-body system, we sample 10 consecutive observations with gap of 0.2 and 50 consecutive observations with a gap of 0.05 within time span [0. 10), respectively. For each system that consists of different body types and interactions, we generate 500 trials to provide large variability of the movements. We performed a 80/20% train/test split throughout, where the testing set contains unseen combinations of body weights that determine the trajectory.
Model performance. We evaluate the performance of EIT by investigating its ability to decode the body types ("Type"), body interaction ( ('Int"), well as all possible combinations ('All') of them. The experiments performed under the domain generalization scheme where the body mass distribution of the testing set (T) different from that of the training set (S).
Table 1: Decoding performance of EIT on many-body systems when compared to a base- line model (BM).
2-body
3-body
BM
EIT
BM
EIT
'In'
93.97
97.28
19.99
91.70
'Type" 'AIl'
97.40
99.93
4.03
37.76
We benchmark EIT with a transformer same depth and amount of heads that considers the population (create the population representation) right from the start of the model. As shown in T EIT consistently outperforms the baseline model (BM) in all cases, where our model showed significant advantage over the bascline when predicting se 100 classes (the 'Type' and 'All' in three- body setting are 90 and 360 classes, respectively). The advantages of the architecture apparent when it comes to the three-body chaotic systems: as the system trajectories provide much more variability and are highly sensitive to the initial conditions (as body types), our model provides significantly better performance by analyzing the individual dynamics separately to provide a stable representation space for each individual.
with the
87.90
94.72
2.03
41.99
especially
We visualized both the individual and population latent space in Figure for the two-body systems. Our individualmmodun successfully distinguishes different types of bodies, while the population- module further separates different systems by learning the interaction of bodies. We refer readers to Appendix 2.1 for visualizations of the bhrye-body systems. On the top-right corner, show the OT
7

---

Table 2: Performance behavioral decoding from populations of neurons in the motor cortex.
1. Decoding performance neural dutasets
Mihi-Chewie (T
Mihi- Chii--Chewie (T = 6)
Maze (R2x100)
C-1
C-2
M-1
M-2
C-1
C-2
M-1
M-2
Vel
Pos
MLP
74.22
74.54
78.17
74.42
78,91
90.74
87.90
84.11
60. 60
78.59
GRU
75.78
75.46
79.96
74.22
84.72
90.12
85.98
78.81
79.97
94.01
NDT
59.90
58.56
73.02
70.74
64,93
63.73
80.56
71.96
76.01
90.07
NDT-Sup
80.47
80.56
83,93
80.79
87.33
94.29
96.83
91.47
82.02
94.88
EIT (T)
76.04
81.25
81.15
71.71
83.33
88.27
93.65
86.82
71.18
87.09
EIT 79.69
82.41
86.51
81.61
88.36
92.59
95.24
91.57
82.15
94.77
H. Generalization performance trained on one population. texted on another (more Appendix
3.2
NDT (et2
75,52
77.78
73.06
85,24
92,46
86.18
X
x
NDT IDeeraai (M-2)
74.48
70.37
76.78
86.28
89.20
91.99
x
EIT (C-2)
79.17
82.94
75.24
81.42
92.33
91.34
EIT (M-2)
78.13
81.02
84.13
84.72
91.06
93.25
X
X
distance matrix that is produced by our individual-basse evaluation method on the test The matrix can clearly quantify the similarity between bodies, and reveals the different types of individuals.
4.2 Decoding behavior from neural populations
Datasets. The Mihi-Chewin reaching dataset [14] consists of stable behavior-basee neural responses across different neuron populations and animals, and thus used in previous methods for across- animal decoding (53 52 14 and neural representation learning 43 48). Mihi-Chewie is a spike
sorted dataset where two rhesus macaques, Chewie ('C') and Mihi ('M'), ) are trained perform simplified reaching task to one of eight targets, while their neural activities in the primary motor cortex were simultaneously recorded. The dataset contains two recordings of different sets of neurons for each of the subjects, for total of four sub-datasets. The Jenkins' Maze dataset [65] is a dataset in the Neural Latents Benchmark (45) that contains activity from the primary motor and dorsal premotor cortex of rhesus macaque named Jenkins ('J"). In this dataset, reaches towards targets while avoiding the boundaries of maze that appears on the screen. Since this dataset provides more complex behaviour trajectories with complete collection of continuous labels (movement velocity and hand position), we used this dataset to examine if our model creates an individual representation space that is rich enough to decode continuous targets. All neural datasets are sorted and binned into 100 ms intervals by counting the number of spikes each neuron emits during that time frame to generate a time-series for each neuron.
Experimental setup. Both the temporal and spatial transformer have depth of 2 and 6 heads. We set the dimension of single neuron representation to be 16 through the transformer training, and bottleneck it to be Id when evaluating the activity representation, which makes the activity representation equal to the total number of neurons. We train EIT end-to-end with an Adam optimizer of learning rate 0.0001 for 400 epochs. We benchmark our models' performance with a MIP, bi-directional GRU (66). NDT 44, and supervised variant of the NDT (NDT-Sup) that we implemented for sequential data decoding. For domain generalization tasks, re-trained the first-layer NDT, and compared our model against numbers obtained when training on the same individuals. We also tested EIT 's performance when considering only non-sequential data and compared it with recent self-supervised methods
43 48 in Appendix 3.2
Decoding performance. As shown in Table 2 our model provides strong decoding performance on various benchmarks, both in terms of its classification of reach (Mihi-Chewil and continuous decoding of position and velocity the regression task studied (Jenkins). For the classification task on
we followed 43 48 and tested the model's 's robustness under both shorter sequence setting (7=2) and longer sequence setting (T=6), while for I's Maze we evaluated our model on the full sequence with regression loss. Both our model and NDT [4] provide good performance, and in some cases, the EIT (T) temporal transformer does quite well on these tasks. Our results suggests that our model is competitive on these diverse decoding tasks from different neural populations in multiple animals.
8

---

(A)
MIIITOGS Train
M22111 31 90-2 (.d.ddd55
Dirtte
MIJAEZZO Trais
M224600 2012 deo0n
M22985) (1 231
100
Figure 3: Experiments on multi-neuron recordings from motor cortex. (A) A schematic of two transfer conditions considered for neural recordings. (B) Along each row (from left to right), we display query neuron, its first and second nearest neighbor. and high firing neuron that has large divergence (the divergence d is also reported for all examples). The average firing rate per reach direction (tuning curve) is shown in the bottom right of each embedding. (C) The couron-to-neuron Wasserstein divergences for M1-M2 encodes the similarity between the neuron-level embeddings (red denotes similar and blue is dissimilar). Below, we show the firing rates for the population along the rows (top) and columns (below) in descending order.
Generalization task # I: across recording session and animal. After validating the model's decoding performance, the number of neurons in different recordings varies, we cannot use the model in an end-to-end manner when testing on a new neural population of different size. Instead, we train our model on one recording. freeze the weights, and apply the channel-invariany portion of our model to the new data as described in Section 3.3 (Figure (()) In our experiments, we re-train only final linear classification layer to predict the labels on the new dataset.
investigate its generalization across different sessions and animals. As
Even when tested in decoding across sessions and animals (Table bottom). Here, we compare with a re-trained NDT (with a new input layer) and with models trained and tested the same animal (within domain, top). In these cases, we find that through a linear readout from EIT we can actually outperform many models that are trained from scratch within domain. These results suggests that our temporal transformer has learned information about the firing patterns of neurons that can be transferred across populations the same and different animals.
populations of neurons, model provides impressive overall performance
To further investigate the underlying factors that contribute to our ability to generalize across popula- tions, visualized the latent space of individual neurons that have either closer distribution (in terms of their OT distance) or distant distribution in Figure 3
BB).. B). multiple cases, we found that the embedding of neurons in the training set (M1) had close neighbors in test (M2, same animal at different time point) with similar tuning profiles (bottom right of each latent embedding). At the same time, neurons that were more distant had orthogonal (or distinct) functional tuning.
To characterize this functional alignment a population level, we measured the divergence between all pairs of neurons in the train and test set (Figure C)). When sorting neurons in both conditions by their firing rate (bottom), we found that the learned neuron representations have good overall global correspondence in terms of their latent embeddings, which suggests that the learned latent representation of individuals contains sufficient information about neurons' overall firing rate. These experiments open up lot of possibilities for finding functional groups in the individual embeddings and suggest that EIT could be used to find correspondence between neurons in different datasets.
Table Performance on Mihi- Chewie reuch decod- ing task when trained on two targets and tested on all eight targets.
Generalization task #2: across behaviours. To test the model's ability transfer when the overall class distribution shifts between the train and the test, we trained on a limited amount of data with a limited set of targets (2 classes) and tested on all classes. Again, we can test generalization in this condition by training a linear layer to decode from the individual transformer (trained on 2 classes) on the new test condition (full 8 classes), Our results
C-1
C-2
M-1
M-2
MLP GRU NDT-Sup EIT
74.28
74.00
83.33
75.81
74.64
70.67
78.49
79.30
75.72
71.00
84,41
82.80
82.25
83.00
91.94
85.22

---

in T suggest that EIT can work well when tested in new behavioral conditions and provides significant gaps over other approaches, even when the training data limited.
Conclusion
In this work, we introduced EIT model for learning representations of both individual and collective dynamics from multi-variate time-series data. In our experiments on both synthetic many- body systems and real-world neural systems, we demonstrated that our model not only provides state-of-the-art decoding performance, but also leads to impressive domain generalization thanks to its permutation-invariint design.
While our model provides a novel strategy for decomposing complex dynamics, we note that there are also limitations and areas for future work:
The downsides of building individual embeddings: Training an individual module at the beginning not only requires prior knowledge about the separation of the system (30), but also might restrict the model's representational capacity (also discussed in [67] due to the limited size of the final representation space. One possible solution to both challenges is to stack multiple sedividual-interaatte separability of the individual module. Instead, one could combine the proposed framework with individual/object disentanglement methods to both decompose mixtures sources and to learn richer representations for individuals.
modules [23]. which would sacrifice the spatial/individual
Replacing transformers with task-specific architectures: EIT utilized transformer encoders for both the individual and interaction modules. However, it is possible replace the transformer encoders with a task-specific architecture to process specific types of interactions, such as graph-basee encoders [68] [23] or recurrent encoders [69], which might improve performance on certain tasks. Additionally, the individual module of EIT models individual dynamics in a deterministic way, but in certain cases (e.g. multi-armed bandit tasks for neural activities), it might be more appropriate to model dynamics and interactions in probabilistic manner (such as in 77)).
Reducing the need for labels through self-supervised training: While the individual represen- tations of EIT generalize reasonably well across animals in the neural activity experiments, the model currently relies labels to guide this functional alignment. Moving forward, EIT could be further extended through training the whole network in a self-supervised way, perhaps using masking and completion tasks [71 72) or contrastive approaches for neural activity 43 48). Self-supervised training not only would eliminate the dependency on labels, but also might improve generalization.
How EIT provides a new paradigm that can help advance foundation models for brain decoding. In many domains, the ability to integrate large amounts of data from different sources into single pre-trained model has enabled impressive advances on many downstream tasks (73)75). With such model for neural data analysis, we could similarly build powerful decoders that could leverage the large amounts of open neural data currently being generated [76] 777 to decode behaviors in diverse contexts and complex tasks. Here we show that by decoupling our learning into two stages and building an encoder that processes single neurons first, we can apply the front-end of our model to new collections of neurons without any modifications or re-training regardless of the mismatch of the inputs. We see this as significant first step towards building a foundation model for neural decoding that can learn from diverse sets of neurons in different brains and contexts.
Acknowledgements
This project was supported by NIH award IRO1EB029852-01, NSF award IIS223994,,, the NSF Graduate Research Fellowship Program (GRFP) for MD, as well as generous gifts from the Alfred Sloan Foundation, the McKnight Foundation, and the CIFAR Azrieli Global Scholars Program.
References
[1] A. R. Ives, B. Dennis, K. Cottingham, and S. R. Carpenter, "Estimating community stability and ecological interactions from time-series data," Ecological Monographs, vol. 73, no. 2.
10

---

pp. 301-330, 2003.
[2] D. M. Gordon, "The ecology of collective behavior," PLOS Biology, vol. 12, no. p. e1001805, 2014.
[3] C. Cattuto, A. Barrat, A. Baldassarri, G. Schehr, and V. Loreto, "Collective dynamics of social annotation," Proceedings of the National Academy of Sciences, vol. 106, no. 26, pp. 10511- 10515, 2009.
[4] M. Moussaid, J. E. Kämmer, P. P. Analytis, and H. Neth, "Social influence and the collective dynamics of opinion formation," PloS one, vol. 8, no. 11, p. e78433, 2013.
[5] L.. Pinheiro, F. C. Santos, and J. M. Pacheco, "Linking individual and collective behavior in adaptive social networks," Physical Review Letters, vol. 116, no. 12, p. 128702, 2016.
[6] S. van Vliet, C. Hauert, K. Fridberg, M. Ackermann, and Dal Co. "Global dynamics of microbial communities emerge from local interaction rules," PLOS Computational Biology, vol. 18, no. 3, e e1009877, 2022.
[7] S. Vyas, M. D. Golub, D. Sussillo, and K. V. Shenoy, "Computation through neural population dynamics," " Annual Review of Neuroscience, vol. 43, pp. 249-275, 2020.
[8] N. Li, T.-W. Chen, Z. V. Guo, C. R. Gerfen, and K. Svoboda, "A motor cortex circuit for motor planning and movement," Nature, vol. 519, no. 7541, pp. 51-56, 2015.
[9] Y. Li and M. Meister, "Functional cell types in the mouse superior colliculus," bioRxiv, 2022.
[10] A. Schneider, M. Azabou, L. McDougall-Vige
D. B. Parks, S. Ensley, K. Bhaskaran T. J. Nowakowski, E. L.. Dyer, and K. B. Hengen, "Transcriptomic cell type structures in vivo neuronal activity across multiple time scales.." bioRxiv, 2022.
[11] V. S. Polikov, P. A. Tresco, and W. M. Reichert, "Response of brain tissue chronically implanted neural electrodes," Journal of Neuroscience Methods, vol. 148, no. 1, pp. 1-18, 2005.
[12] W. M. Grill, S. E. Norman, and R. V. Bellamkonda, "Implanted neural interfaces: biochallenges and engineered solutions," Annual Review of Biomedical Engineering, vol. 11, pp. 1-24, 2009.
[13] D. McCreery, Pikov, and R. Troyk, "Neuronal loss due to prolonged controlled-current stimulation with chronically implanted microelectrodes in the cat cerebral cortex," Journal of Neural Engineering, vol. 7, no. 3, p. 036005, 2010.
[14] E. L. Dyer, M. G. Azar, M. G. Perich, H. L. Fernandes, S. Naufel, L. E. Miller, and K. P. Körding, "A cryptography-base -based approach for movement decoding," Nature Biomedical Engineering, vol. 1, no. 12, pp. 967-976, 2017.
[15] A. Gallego, M. G. Perich, S. N. Naufel, C. Ethier, S. A. Solla, and L. E. Miller, "Cortical population activity within preserved neural manifold underlies multiple motor behaviors." Nature Communications, vol. 9, no. 1, pp. 1 1-13, 2018.
[16] J. A. Gallego. M. G. Perich, R. H. Chowdhury, A. Solla, and L. E. Miller, "Long-term stability of cortical population dynamics underlying consistent behavior," Nature Neuroscience, vol. 23, 2, pp. 260-270, 2020.
[17] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin, "Attention all you need," Advances in neural information processing systems, pp. 5998-6008, 2017.
[18] N. Geneva and N. Zabaras, "Transformers for modeling physical systems," Neural Networks, vol. 146, pp. 272-289, 2022.
[19] Bai and Tahmasebi, "Characterization of groundwater contamination: transformer-hased deep learning model," Advances in Water Resources, p. 104217, 2022.
11

---

[20] C. Plizzari, M. Cannici, and M. Matteucci, "SSatiil temporal transformer network for skeleton- based action recognition," in International Conference on Pattern Recognition, pp. 694-701, Springer, 2021.
[21] R. Girdhar, J. Carreira, C. Doersch, and A. Zisserman, "Video action transformer network," Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 244-253, 2019.
[22] K. Chen, G. Chen, D. Xu, L. Zhang, Y. Huang, and A. Knoll, "Nast: non-autoregressive apatial-tempore transformer for time series forecasting," arXiv preprint arXiv:2102.05624, 2021.
[23] M. Xu, W. Dai, C. Liu, X. Gao, W. Lin, G.-J. Qi, and H. Xiong, "Spatial-temporal transformer networks for traffic flow forecasting," arXiv preprint arXiv:2001.02908 2020.
[24] S. Li, Q. Cao, L. Liu, K. Yang, S. Liu, J. Hou, and S. Yi, "Groupformer: Group activity recognition with clustered spatial-temporal transformer." in Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 13668-13677, 2021.
[25] A. Dosovitskiy, L. Beyer, A. Kolesmikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby, "An image worth 16x16 words: Transformers for image recognition at scale," in International Conference on Learning Representations, 2021.
[26] A. Arnab, M. Dehghani, G. Heigold, C. Sun, M. Lučié, and C. Schmid, "Vivit: A video vision transformer," in 2021 IEEE/CVF International Conference on Computer Vision (ICCV),
pp. 6816-6826, 2021.
[27] Z. Liu, J. Ning, Y. Cao, Y. Wei, Z. Zhang, S. Lin, and H. Hu, "Video swin transformer," Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022.
[28] G. Bertasius, H. Wang, and L. Torresani, "Is space-time attention all you need for video understanding?," in International Conference on Machine Learning, PMLR, 2021.
[29] J.-T. Hsieh, B. Liu, D.-A. Huang, L. F. Fei-Fei, and J. C. Niebles, "Learning to decompose and disentangle representations for video prediction," " Advances in neural information processing systems, vol. 31, 2018.
[30] A. Kosiorek, H. Kim, Y. W. Tch, and I. Posner, "Sequential attend, infer, repeat: Generative modelling of moving objects," Advances in Neural Information Processing Systems, vol. 31, 2018.
[31] Y.-F Wu, J. Yoon, and S. Ahn, "Generative video transformer: Can objects be the words?," in International Conference on Machine Learning, pp. 11307-11388, PMLR, 2021.
[32] G. Singh, Y.-F Wu, and S. Ahn, "Simple unsupervised object-centric learning for complex and naturalistic videos." arXiv preprint ar 22tv:22205.10665 2022.
[33] P. Battaglia, R. Pascanu, M. Lai, D. Jimenez Rezende, al., "Interaction networks for learning about objects, relations and physics," Advances in neural information processing systems, vol. 29, 2016.
[34] A. Wu, E. K. Buchanan, M. Whiteway, M. Schartner, G. Meijer, J.-P. Noel, E. Rodriguez, C. Everett, A. Norovich, E. Schaffer, et al., "Deep graph pose: a semi-suppervii deep graphical model for improved animal pose tracking," Advances in Neural Information Processing Systems, vol. 33, pp. 6040-6052, 2020.
[35] M. Azabou, M. Mendelson, M. Sorokin, S. Thakoor, N. Ahad, C. Urzay, and E. L. Dyer, "I "Learning behavior representations through multi-timescale bootstrapping." arXiv preprint arX arXiv:2206.07041 2022.
[36] A. L. Hodgkin and A. F. Huxley, "A quantitative description of membrane current and its application to conduction and excitation nerve," The Journal of Physiology, vol. 117, no. p. 500, 1952.
12

---

[37] W. M. Kistler, W. Gerstner, and J. L. v. Hemmen, "Reduction of the hodgkin-huxley equations single-variable threshold model," Neural Computation, vol. 9, no. 5, pp. 1015-1045, 1997.
[38] R. D. Traub, R. K. Wong, R. Miles, and H. Michelson, "A model of ca3 hippocampal pyramidal neuron incorporating voltage-clamp data on intrinsic conductances," Journal of Neurophysiology, vol. 66, no. 2, pp. 635-650, 1991.
[39] M. R. Panahi, G. Abrevaya, J.-C. Gagnon-Audet, V. Voleti, I. Rish, and G. Dumas, "Generative models of brain dynamics-a review," arXiv preprint arXiv:21 /2147, 2021.
[40] D. Semedo, E. Gokcen, C. K. Machens, A. Kohn, and M. Byron, "Statistical methods for dissecting interactions between brain areas," Current Opinion in Neurobiology, vol. 65, pp. 59-69, 2020.
[41] M. G. Perich and K. Rajan, "Rethinking brain-wide interactions through multi-region 'network of networks' models," Current Opinion in Neurobiology, vol. 65, pp. 146-151, 2020.
[42] S. L. Keeley, D. M. Zoltowski, M. C. Aoi, and J. W. Pillow, "Modeling statistical dependencies in multi-region spike train data," " Current Opinion in Neurobiology, vol. 65, pp. 194-202, 2020.
[43] M. Azabou, M. G. Azar, R. Liu, C.-H. Lin, E. C. Johnson, K. Bhaskaran-Narr M. Dabagia, K. B. Hengen, W. Gray-Roncal, M. Valko, and E. Dyer, "Mine your own view: Self-supervised learning through across-sample prediction," arXiv preprint 2rX1v:2122.10166, 2021.
[44] J. Ye and C. Pandarinath, "Representation learning for neural population activity with Neural Data Transformers," Neurons, Behavior, Data analysis, and Theory, Aug. 2021.
[45] F. Pei, J. Ye, D. M. Zoltowski, A. Wu, R. H. Chowdhury, H. Sohn, J. E. O'Doherty, K. V. Shenoy, M. T. Kaufman, M. Churchland, M. Jazayeri, L. E. Miller, J. Pillow, I. M. Park, E. L. Dyer, and C. Pandarinath, "Neural latents benchmark 21: Evaluating latent variable models of neural population activity," Advances in Neural Information Processing Systems (NeurIPS), Track on Datasets and Benchmarks, 2021.
[46] J. Nassar, S. W. Linderman, M. Bugallo, and I. M. Park, "Tree-structured recurrent switching linear dynamical systems for multi-scale modeling," in International Conference on Learning Representations, 2019.
[47] C. Pandarinath, D. J. O'eaa J. Collins, R. Jozefowicz, S. D. Stavisky, J. C. Kao, E. M. Trautmann, M. T. Kaufman, S. Ryu, L. R. Hochberg, al., "Inferring single-trial neural population dynamics using sequential auto-encoders," Nature Methods, vol. 15, no. 10, pp. 805- 815, 2018.
[48] R. Liu, M. Azabou, M. Dabagia, C.-H. Lin, M. Gheshlaghi Azar, K. Hengen, M. Valko, and E. Dyer, "Drop. swap, and generate: A self-supervised approach for generating neural activity." Advances in Neural Information Processing Systems, vol. 34, 2021.
[49] A. E. Urai, Doiron, A. M. Leifer, and A. K. Churchland, "Large-scale neural recordings call for new insights to link brain and behavior," Nature Neuroscience, pp. 1-9, 2022.
[50] M. Dabagia, K. P. Kording, and L. Dyer, "Comparing high-dimensionaal neural recordings by aligning their -iw-dimensional latent representations," arXiv preprint arXiv:2205.08.4 08413,, 2022.
[51] A. D. Degenhart, W. E. Bishop. E. R. Oby, C. Tyler-Kabara, S. M. Chase, A. P. Batista, and B. M. Yu, "Stabilization of a brain-computer interface via the alignment of low-dimensional spaces of neural activity," Nature Biomedical Engineering, vol. 4, no. 7, pp. 672-685, 2020.
[52] A. Farshchian, J. A. Gallego, J. P. Cohen, Y. Bengio, L.. F. Miller, and S. A. Solla, "Adversarial domain adaptation for stable brain-machine interfaces." 2019.
[53] J. Jude, M. G. Perich, L. F. Miller, and M. H. Hennig, "Robust alignment of cross-session recordings of neural population activity by behaviour via unsupervised domain adaptation." arXiv preprint wwXiv:2202.06/333 2022.
13

---

[54] T. Chen, S. Kornblith, M. Norouzi. and G. Hinton, "A simple framework for contrastive learning of visual representations." in International Conference on Machine Learning, Pp. 1597-1607, PMLR, 2020.
[55] J.-B. Grill, F. Strub, F. Altché, C. Tallec, P. H. Richemond, E. Buchatskaya, C. Doersch, B. A. Pires, Z. D. Guo, M. G. Azar, el al., "Bootstrap your own latent: A new approach to self-supervised learning," arXiv preprint 2xXii:2006.07733, 2020.
[56] M. Cuturi, "Sinkhorn distances: Lightspeed computation of optimal transport." Advances in neural information processing systems, vol. 26, 2013.
[57] G. Peyré, M. Cuturi, el al., "Computational optimal transport: With applications to data science," Foundations and Trends@ in Machine Learning, vol. 11, no. 5-6, pp. 355-607, 2019.
[58] R. Jastrow, "Many-body problem with strong forces," Physical Review, vol. 98, no. 5, p. 1479, 1955.
[59] T. Kipf, E. Fetaya, K.-C. Wang. M. Welling, and R. Zemel, "Neural relational inference for interacting systems," in International Conference on Machine Learning, pp. 2pp22688-66 PMLR, 2018.
[60] C. Graber and A. Schwing, "Dynamic neural relational inference for forecasting trajectories," in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops, pp.
2020.
[61] L. Duncker, G. Bohner, J. Boussard, and M. Sahani, "Learning interpretable continuous-ting models of latent stochastic dynamical systems," in International Conference on Machine Learning, pp. 1734,, PMLR, 2019.
[62] T. D. Kim, T. Z. Luo, W. Pillow, and C. Brody, "Inferring latent dynamics underlying neural population activity via neural differential equations," in International Conference on Machine Learning, pp. 5551-5561, PMLR, 2021.
[63] S. Greydanus, M. Dzamba, and J. Yosinski, "Hamiltonian neural networks," Advances in Neural Information Processing Systems, vol. 32, 2019.
[64] R. Dormand and P. J. Prince, family of embedded runge-kutta formulae," Journal of Computational and Applied Mathematics, vol. 6, no. 1, pp. 19-26, 1980.
[65] M. M. Churchland, J. P. Cunningham, M. T. Kaufmann S. Ryu, and K. V. Shenoy, "Cortical preparatory activity: representation of movement or first cog in dynamical machine?," Neuron, vol. 68, no. 3, pp. 387-400, 2010.
[66] C. Xiong, S. Merity, and R. Socher, "Dynamic memory networks for visual and textual question answering." in International Conference on Machine Learning, pp. 2397-2406, PMLR, 2016.
[67] S. Dutta, T. Gautam, S. Chakrabarti, and T. Chakraborty, "Redesigning the transformer archi- tecture with insights from multi-particle dynamical systems," Advances in Neural Information Processing Systems, vol. 34, 2021.
[68] S. Guo, Y. Lin, N. Feng, C. Song, and H. Wan, "Attention based spatial-temporal graph convolutional networks for traffic flow forecasting," in Proceedings of the AAAI conference on artificial intelligence, vol. 33, pp. 922-929, 2019.
[69] T. Guo, T. Lin, and N. Antulov-Fantulla "Exploring interpretable 1stm neural networks over vulti-variable data," in International conference on machine learning, pp. 2494-2504, PMLR, 2019.
[70] A. Shalova and I. Oseledets, "Tensorized transformer for dynamical systems modeling." arXiv preprint arXiv:2006.088: 03445, 2020.
[71] K. He, X. Chen, Xie, Li, P. Dollár, and R. Girshick, "Masked autoencoders are scalable vision learners," arXiv preprint arXiv:2111.06:: 2021.
14

---

[72] H. Bao, L. Dong, and F. Wei, "Beit: Bert pre-training of image transformers," arXiv preprint 2.2::22006.08244, 2021.
[73] Devlin, M.-W. Chang, K. Lee, and K. Toutanova, "Bert: Pre-training of deep bidirectional transformers for language understanding," arXiv preprint arXiv: 04805, 2018.
[74] T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, Shyam, G. Sastry, A. Askell, et al., "Language models are few-shot learners," arXiv preprint
1405, 2020.
[75] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, et al., "Learning transferable visual models from natural language supervi- sion," " in International Conference on Machine Learning, pp. 8748-8763, PMLR, 2021.
[76] H. Siegle, Jia, S. Durand, Gale, Bennett, Graddis, G. Heller, T. K. Ramirez, H. Choi, J. A. Luviano, el al., "Survey of spiking in the mouse visual system reveals functional hierarchy," Nature, vol. 592, no. 7852, pp. 86-92, 2021.
[77] S. E. de Vries, J. A. Lecoq, M. A. Buice, P. A. Groblewski, G. K. Ocker, M. Oliver, D. Feng, N. Cain, P. Ledochowitsch, D. Millman, et al., "A large-scale standardized physiological survey reveals functional organization the mouse visual cortex," Nature Neuroscience, vol. 23, no. 1,
pp. 138-151, 2020.
[78] M.-H. Guo, J.-X. Cai, Z.-N. Liu, T.J. Mu, R. R. Martin, and S.-M. Hu, "Pet: Point cloud transformer," Computational Visual Media, vol. 7, no. 2, pp. 187-199, 2021.
[79] I. Higgins, L. Matthey, A. Pal, C. Burgess, X. Glorot, M. Botvinick, S. Mohamed, and A. Ler- chner, "beta-vae: Learning basic visual concepts with a constrained variational framework," International Conference on Learning Representations, 2016.
15


Reviews:
[{"review_id": "ys6D7SKwAbW", "reviewer": "Reviewer_6wnv", "summary": "This paper presents a transformer architecture-based time-varying data learning method. It consists of a temporal-spatial module and a projection module to learn the descriptions of both the individuals and collective population and then leverages the interactions to predict the target behavior. The experimental evaluation is performed on synthetic many-body systems and real neural response datasets.\n", "questions": "Apart from the presentation issues mentioned in the above weak point summary, comparing the Transformer and conventional sequence models on univariate time series in the proposed framework would be useful for assessing the significance of the proposed method.  ", "limitations": "No potential negative societal impact is found.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "--- Strong points:\n\nThis paper proposes a disentangled representation learning which models individual channel time series and their interactions. \n\n\n--- Weak points:\n\nThe proposed method is built on top of existing ideas and lacks novelty and insights. Meanwhile, some design choices are not well justified theatrically or experimentally. \n\n(a) The problem formulation is confusing. As is mentioned in Sec. 3.1, the input observation and target sequence are both in the time range from 1 to T. Depending on the conditioned observations on which the prediction of the target at a certain timestamp is, it can lead to different types of problem setups, e.g., forecasting, filtering or smoothing. These problems have different directions of solutions and the author is expected to explicitly specify the problem type. \n\n(b) The optimization problem (3) is ill-defined in that the input into g() is of the shape N*T, while the function g() is defined as a mapping from R^N to R^N. Meanwhile, this definition is different from the loss function (4), thereby giving the impression that the optimization problem seems not clearly defined in the first place and the subsequently proposed method can hardly be solid enough. \n\n(c) The temporal-spatial attention and then slicing along different dimensions of the output representations have been extensively studied in the architectures like RNN and CNN [1, 3]. It is also not novel to explore disentangled individual representations of multi-channel time series [2, 3]. The Transformer is applied in this paper as an off-the-shelf encoder. Overall, the paper seems to bear marginal technical contributions regarding model architectures. \n\n[1] Qin, Yao, et al. \"A dual-stage attention-based recurrent neural network for time series prediction.\" Proceedings of the 26th International Joint Conference on Artificial Intelligence. 2017.\n\n[2] Guo, Tian, Tao Lin, and Nino Antulov-Fantulin. \"Exploring interpretable LSTM neural networks over multi-variable data.\" International conference on machine learning. PMLR, 2019.\n\n[3] Assaf, Roy, et al. \"Mtex-CNN: Multivariate time series explanations for predictions with convolutional neural networks.\" 2019 IEEE International Conference on Data Mining (ICDM). IEEE, 2019.\n\n(d) The motivation for using the Transformer on individual time series is weak. The Transformer is intended to capture pair-wise relations between tokens/steps in a multi-dimensional sequence. As far as I know, recent works mostly applied the Transformer or variants to multi-dimension time series [4, 5]. However, applying Transformer on univariate time series riks overcomplicating the model and overfitting as well. It would be more convincing to see the comparison between conventional sequence models and Transformer in the proposed framework.\n\n[4] Tang, Binh, and David S. Matteson. \"Probabilistic transformer for time series analysis.\" Advances in Neural Information Processing Systems 34 (2021): 23592-23608.\n\n[5] Zhou, Tian, et al. \"FEDformer: Frequency enhanced decomposed transformer for long-term series forecasting.\" arXiv preprint arXiv:2201.12740 (2022).\n\n(d) A projection module and two different project operations are described in Sec. 3.2. It is necessary to clarify their difference and connection. Meanwhile, the projection module consists of some plain operations, e.g., dense transformation and concatenation. Technically, this module bears marginal novelty or new insights. \n\n(e) In the synthetic data experiment, there is no reference for the baseline method and thus it is difficult to evaluate the significance of the proposed method.   \n\n(f) The hyper-parameter search process is missing for baselines. It is important to describe the process, since it significantly affects the performance and the validity of the comparison. \n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "p6ewxQerXc", "reviewer": "Reviewer_xZtG", "summary": "The paper presents a neural architecture that models a dynamical system by building representations for both individual and collective population dynamics separately. The first transformer operates only on individual time series and the second transformer acts as an interaction network to come up with the collective representation of the models. The models are trained to classify or regress the target values. Authors then test this model on a toy datasets and neural activity modelling.", "questions": "- The method doesn't seem to infer or decompose entities from raw observational data. It requires the knowledge of separate channels and feed them separately to individual transformer which a simpler problem than a cocktail party problem. Doesn't it restrict its applicability for e.g. in many cases we don't know the number of interacting agents.\n- How do you infer the components/objects/neurons whose dynamics needs to be modelled individually?\n- In the works of [1], [2], [3], [4] the models first try to infer individual entities from the observed frames in videos and then try to model dynamics of each entity individually. Moreover, an interaction module is used to explicitly model objects interaction and combined influence on their dynamics. [3] even uses transformer for modelling dynamics. It would be interesting to position this paper w.r.t them.\n- Can interaction transformer introduce some kind of bias towards the order of channels?\n- Since it’s not a generative model, it’s hard to quantitatively see how individual channels are independent. Could there be a test to measure how independent individual dynamics are? Can authors try permuting the channels and see if the results on the test set remain the same?\n- [Related to above question]: Did you try dropping one or a few channel to see how models behave?\n- How does interaction network $g$ capture complex interaction without any specific inductive bias for capturing such interactions for e.g. as shown in [5]. - Could there be any specific interaction module to account for the interactions between individual agents? Would it work for physics based interactions between agents e.g. collision between objects or if agents exhibits some magnetic push/pull?\n\n[1] https://arxiv.org/abs/1806.04166\n[2] https://arxiv.org/abs/1806.01794\n[3] https://arxiv.org/abs/2107.09240\n[4] https://arxiv.org/abs/2205.14065\n[5] https://arxiv.org/abs/1612.00222\n", "limitations": "The model applicability seems to be limited to the problems where the exact number of agents contributing to overall population dynamics need to be known beforehand.", "rating": 5, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n- The paper is well written and easy to follow.\n- The paper demonstrates the clear benefits of decomposing dynamics modelling into  individual dynamics and collective dynamic for better transfer learning.\n- The application to neural activity datasets is interesting.\n\n**Weaknesses:**\n- The method is limited in terms of novelty as the individual modules for capturing individual dynamics and collective dynamics are commonly used transformer architectures.\n- The method is only applicable to the data where the existence and number of separate channels in the input data is known. The limits it applicability to models where the number of interacting agents in the systems are not known.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "bMiaXpzs1jI", "reviewer": "Reviewer_7m6L", "summary": "Authors introduce a method to decompose the representations of a population of neurons of the individual dynamics from the population dynamics which allows decoding generalization to unseen populations of neurons.\n\nThe results are presented on a behavioral set of benchmark datasets where the author's method outperforms (by a meaningful margin) on 6/10 tasks.\n\nOverall, the work is impactful in that generalization to unseen populations of neurons is a big unblock for the neuroscience community.", "questions": "- there wasn't a lot on related work on the comp neuro side. are there related papers that attempt something similar?", "limitations": "Have not addressed society impacts directly... however, this is HUGELY (positively) impactful for applications like prosthesis where the prosthesis design needs to generalize to different neural populations for each individual.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- motivation is very strong\n- the approach is well explained by breaking up the dynamic into f, g which enable the decomposition desired and a general approach where anything can be used for f or g\n- formulation is simple which is a huge plus.\n\nWeaknesses:\n- would have liked to see more benchmarks, but understand that there is limited availability of this.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "OPh0NVPmnEr", "reviewer": "Reviewer_Cc4d", "summary": "The authors present a novel approach for inference over collective systems, through a separable transformer-based neural network architecture. The novelty lies in learning both the individual components of the system (the trees) as well as the population-level dynamics (the forest). The authors correctly identify that separating individual from system level dynamics has previously been investigated (e.g., in the GroupFormer archetecture [Li, Shuaicheng, et al]) but set themselves apart from previous work by noting that they learn representations at both scales in a quasi-independent manner or, in the authors words, they 'restrain the direct computation' between the two learned maps. To move between the two maps, population and individual, a projection module is learned simultaneously which provides access to both representations. \n\nCrucially, by disentangling learning over the collective into two independent representations, the authors demonstrate that they can use embeddings of individual-scale dynamics for transfer learning on new datasets. This can be applied to new datasets of arbitrary size and/or ordering due to the permutation-invariance property (a consequence of separating individual and population learning tasks). \n\nThe authors test their approach on two datasets: a synthetic many-body simulation (of two/three masses in orbit) and neuronal behavioural data. They find strong performance both on the synthetic dataset and on the neuronal data. Importantly, the authors find that they can use a model trained on one dataset and apply the invariant portion of this model to a new unseen dataset. This retrained model outperforms many others which learn from scratch. \n\nLi, Shuaicheng, et al. \"Groupformer: Group activity recognition with clustered spatial-temporal transformer.\" Proceedings of the IEEE/CVF International Conference on Computer Vision. 2021.\n", "questions": "The authors mention in their conclusion that by treating individual systems separately, the individuals have limited dimensionality. I would have also expected that approaches that decouple learned representations of individual from the population have poorer performance on population-scale prediction tasks. My reasoning is that population-based learning can be improved by leveraging/constraining learned representations on accurate relational/individual representations. Do the authors have any thoughts on this? \n\nI believe that decoupling the learning of individual from population allows more efficient scaling to large datasets. Is this correct? ", "limitations": "The authors highlight two limitations of their work: reliance on deterministic trajectories at the individual-scale and lost information due to treating systems independently. I have asked for clarification on whether population-based predictions also suffer due to decoupling of learning on individuals from population, and whether this is an additional limitation of the approach.\n\n\n", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This paper is well structured and written, provides a novel approach to an interesting problem and is considered in its discussion of its own strengths and weaknesses. I think the work will be seen as significant, especially given its framework with respect to challenges in neuroscience. \n\nSpecifically, by focussing on learning representations of individual dynamics distinct from population dynamics, the authors rightly point to the added benefit of transfer learning by levering permutation-invariance of the separated system. I see this form of transfer learning as the central result of the approach and the authors did well to emphasis this. \n\nMy main concern is that the synthetic dataset is a relatively simple example when compared against the application of neuronal data. The two-body system will not have particularly complex dynamics and I did not find it particularly compelling. Whereas I appreciate the three-body system will have chaotic dynamics under some parameter values, it wasn't clear to me what proportion of the orbits were expected to be chaotic. Additionally, predicting the trajectories and underlying parameters for a three-body system seems an easy problem when compared to the neural dataset. I suggest the authors introduce a third example of intermediate difficulty, either with significantly more individuals/channels, time-varying interactions, or a typical dataset from previous work on neural relational inference (see, for example, [Kipf, Thomas, et al.] and [Graber, Colin, and Alexander Schwing]). \n\nThere are some minor languages points. I would consider rewording sentence on lines 86-87. Additionally, the second challenge in section \"Learning across recordings or individuals\" is not especially clear. I think the authors are referring here to heterogeneity in labels, both within (e.g. a brain), across a single species/system, and across different species/systems. I think the reference to number of neurons (on line 98) is referencing number of channels/individuals. \n\nKipf, Thomas, et al. \"Neural relational inference for interacting systems.\" International Conference on Machine Learning. PMLR, 2018.\n\nGraber, Colin, and Alexander Schwing. \"Dynamic neural relational inference for forecasting trajectories.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops. 2020.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept