Original: Seeing the forest and the tree: Building representations of both individual and collective dynamics with transformers
Revised: Seeing the forest and the tree: Building representations of both individual and collective dynamics with transformers

Original: Abstract
Revised: Abstract

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: On neural activity datasets, we show that our multi-scale transformer not only yields robust decoding performance, but also provide 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.
Revised: 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.

Original:  Introduction
Revised:  Introduction

Original: Complex systems (such as the brain) contain multiple individual elements (e g. neurons) that interact dynamically to generate their outputs.
Revised: Complex systems (such as the brain) contain multiple individual elements (e g. neurons) that interact dynamically to generate their outputs.

Original: The process by which these local interactions give rise to large-scale behaviors is important in many domains of science and engineering, from ecology [1, 2] and social networks [3–5] to microbial interactions [6] and brain dynamics [7].
Revised: The process by which these local interactions give rise to large-scale behaviors is important in many domains of science and engineering, from ecology [1, 2] and social networks [3–5] to microbial interactions [6] and brain dynamics [7].

Original: A natural way to model the activity of a system is to build a collective or population-level view, where we consider individuals (or channels) jointly to determine the dynamics of the population (Figure 1(A)).
Revised: A natural way to model the activity of a system is to build a collective or population-level view, where we consider individuals (or channels) jointly to determine the dynamics of the population (Figure 1(A)).

Original: In many cases, studying systems from this population-level perspective has provided important insights into collective computations and emergent behaviors [7]; however, it is also possible to lose sight of the contributions of different individuals’ dynamics to the final prediction or inference.
Revised: In many cases, studying systems from this population-level perspective has provided important insights into collective computations and emergent behaviors [7]; however, it is also possible to lose sight of the contributions of different individuals’ dynamics to the final prediction or inference.

Original: This is important in many settings where the dynamics of  individuals are significant, either due to their different functional roles in the system [8, 9, 6], or due to shift in their dynamics because of sensor displacement or corruption [10–12].
Revised: 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], or due to shift in their dynamics because of sensor displacement or corruption [11–13].

Original: 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.
Revised: 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.

Original: Here, we present a new framework for modeling time-varying observations of a system which uses dynamic embeddings of individual channels to construct a population-level view (Figure 1(B-C)).
Revised: Here, we present a new framework for modeling time-varying observations of a system which uses dynamic embeddings of individual channels to construct a population-level view (Figure 1(B-C)).

Original: Our model, which we dub Embedded Interaction Transformer or EIT , decomposes multi-channel time-series by first learning rich features from individual time-series before incorporating information and learned interactions across different individuals in the population.
Revised: 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.

Original: One critical benefit of our model is spatial/individual separability : it builds a population-level representation from embeddings of individual channels, which naturally leads to channel-level permutation invariance.
Revised: One critical benefit of our model is spatial/individual separability : it builds a population-level representation from embeddings of individual channels, which naturally leads to channel-level permutation invariance.

Original: In domain generalization tasks, this means a trained model can be tested with permuted channels or entirely different numbers of channels.
Revised: In domain generalization tasks, this means a trained model can be tested with permuted channels or entirely different numbers of channels.

Original: To first understand how the model captures different types of interactions and dynamics, we experiment with synthetic many-body systems.
Revised: To first understand how the model captures different types of interactions and dynamics, we experiment with synthetic many-body systems.

Original: Under different types of interactions, we show that our model can be applied to successfully recover the dynamics of known systems.
Revised: Under different types of interactions, we show that our model can be applied to successfully recover the dynamics of known systems.

Original: Next, we turn our attention to recordings from the nervous system [13], 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.
Revised: 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.

Original: The stability of the neural representations and collective dynamics found in data from these animals [13–15] makes them an ideal testbed to examine the generalization of our approach.
Revised: The stability of the neural representations and collective dynamics found in data from these animals [14–16] makes them an ideal testbed to examine the generalization of our approach.

Original: 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.
Revised: 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.

Original: The performance of the linear decoder based on pre-trained weights in the across-animal transfer outperforms many models that are trained and tested on the same animal.
Revised: The performance of the linear decoder based on pre-trained weights in the across-animal transfer outperforms many models that are trained and tested on the same animal.

Original: This result provides an exciting path forward in neural decoding  and highlights the utility of our architecture.
Revised: 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.

Original: The main contributions of this work are:
Revised: The main contributions of this work are as follows:

Original: • We proposed a novel framework that considers individual dynamics before collective dynamics, and realized it with a multi-stage transformer: Embedded Interaction Transformer ( EIT ). EIT learns both population representation and individual representation through decoupling the dynamics of individuals and their interactions in multi-variate time-series.
Revised: • In Section 3, 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.

Original: • We introduced methods for generalization across datasets of different input size (number of channels) and ordering.
Revised: • In Section 3.3, we propose methods for generalization across datasets of different input size (number of channels) and ordering.

Original: We further proposed to measure the alignments of time-series in different datasets through the Wasserstein divergence.
Revised: 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.

Original: • We applied EIT to both many-body systems and neural systems recordings.
Revised: In Section 4, we apply EIT to both many-body systems and neural activity recordings.

Original: After demonstrating our model’s robust decoding performance, we validated its ability to transfer individual dynamics by performing domain generalization across different populations of neurons and finding neuron alignments.
Revised: After demonstrating our model’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.

Original:  Background and Related Work 2.1 Transformers
Revised:  Background and Related Work 2.1 Transformers

Original: Self-attention.
Revised: Self-attention.

Original: Transformers have revolutionized natural language processing with the mechanism of self-attention [16], which helps the model to learn long-range interactions across different elements (represented by tokens) in a sequence.
Revised: Transformers have revolutionized natural language processing (NLP) through the mechanism of self-attention [17], which helps the model to learn long-range interactions across different elements (represented by tokens) in a sequence.

Original: Consider a sequence X
Revised: Consider a sequence X

Original: =
Revised: =

Original: [ x 1 x 2 . . .
Revised: [ x 1 x 2 . . .

Original: x N ] T ∈ R N × d consisting of N tokens of d -dimensions.
Revised: x N ] T ∈ R N × d consisting of N tokens of d -dimensions.

Original: The self-attention mechanism is built on the notion of queries Q , keys K , and values V , which are linear projections of the token embeddings.
Revised: The self-attention mechanism is built on the notion of queries Q , keys K , and values V , which are linear projections of the token embeddings.

Original: Let
Revised: Let

Original: Q = XW Q , K = XW K , V = XW V , where Q ∈ R N × d q , K ∈ R N × d k , V ∈ R N × d v .
Revised: Q = XW Q , K = XW K , V = XW V , where Q ∈ R N × d q , K ∈ R N × d k , V ∈ R N × d v .

Original: The attention operation can be written as:
Revised: The attention operation can be written as:

Original: [Equation]
Revised: [Equation]

Original: where softmax( · ) denotes a row-wise softmax normalization function.
Revised: where softmax( · ) denotes a row-wise softmax normalization function.

Original: Multi-head self-attention (MSA).
Revised: Multi-head self-attention (MSA).

Original: 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.
Revised: 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.

Original: Each of the h heads provides a d v dimensional output, and the outputs from all heads are concatenated into a d v · h vector, which is further projected to produce the final representation.
Revised: Each of the h heads provides a d v dimensional output, and the outputs from all heads are concatenated into a d v · h vector, which is further projected to produce the final representation.

Original: The full operation of a transformer of L layers can be written as below:
Revised: The full operation of a transformer of L layers can be written as below:

Original: [Equation]
Revised: [Equation]

Original: where Z 0 is the summation of the individual embedding of each data token and the positional embedding E pos 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.
Revised: where Z 0 is the summation of the individual embedding of each data token and the positional embedding E pos 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.

Original: Transformers for multi-channel time-series.
Revised: Transformers for multi-channel time-series.

Original: Multi-channel time-series, as a data format that consists of sequences of data, are a natural candidate for modeling with the transformer architecture.
Revised: Multi-channel time-series are a natural candidate for modeling with the transformer architecture.

Original: One common approach of modeling many time-varying data streams with transformers is to first aggregate features across all channels into a combined representation at the beginning of the model.
Revised: One common approach of modeling many time-varying data streams with transformers is to first aggregate features across all channels into a combined representation at the beginning of the model.

Original: The population dynamics is thus learnt with the resulting embedding via a temporal transformer for the purpose of inference.
Revised: The population dynamics is thus learnt with the resulting embedding via a temporal transformer for the purpose of inference.

Original: There are many complex ways to create this embedding:[17] and [18]  extract embeddings from multivariate physical systems with Koopman operators before feeding the resulting representation into a temporal transformer; [19] use a graph neural network to embed interconnected-structures to perform skeleton-based action recognition; [20] use a convolutional architecture to extract image embeddings before feeding them into a transformer.
Revised: There are many complex ways to create this embedding: [18] and [19] extract embeddings from multivariate physical systems with Koopman operators before feeding the resulting representation into a temporal transformer; [20] use a graph neural network to embed interconnected-structures to perform skeleton-based action recognition; [21] use a convolutional architecture to extract image embeddings before feeding them into a transformer.

Original: Another approach is to re-design the attention block such that the attention operation is computed both along the temporal and spatial dimension.
Revised: Another approach is to re-design the attention block such that the attention operation is computed both along the temporal and spatial dimension.

Original: Many variants of this ‘spatial-temporal’ attention block have been shown effective: [21] propose a non-autoregressive module to generate queries for time series forecasting; [22] use stacked spatial-temporal modules for traffic flow forecasting; GroupFormer [23] embed spatial and temporal context in parallel to create a clustered attention block for group activity recognition.
Revised: Many variants of this ‘spatial-temporal’ attention block have been shown effective: [22] propose a 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 in parallel to create a clustered attention block for group activity recognition.

Original: 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-related attention map and the temporally-related attention map, which yields a representation of the individual which is completely free of spatial interactions.
Revised: 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-related attention map and the temporally-related attention map, which yields a representation of the individual which is completely free of spatial interactions.

Original: Spacetime attention in video transformers.
Revised: Spacetime attention in video transformers.

Original: With the advances of the Vision Transformer [24] as a new way to extract image embeddings, many ‘spatial-temporal transformer’ architectures have been developed in the video domain [25–27].
Revised: With the advances of the Vision Transformer [25] as a new way to extract image embeddings, many ‘spatial-temporal transformer’ architectures have been developed in the video domain [26–28].

Original: Such works explore and propose interesting solutions for how to organize spatial attention and temporal attention with either coupled (series) [27] and factorized (parallel) attention blocks [25], as well as how to create better tokens for videos by creating three-dimensional spatio-temporal ‘tubes’ as the tubelet tokenizations [25].
Revised: 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 as how to create better tokens for videos by creating three-dimensional spatio-temporal ‘tubes’ as the tubelet tokenizations [26].

Original: 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.
Revised: 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.

Original: Object-centric representation learning.
Revised: Object-centric representation learning.

Original: There are many representation learning methods in video [28–31] or physics-guided complex systems [32] that also aim to learn interactions between discrete objects observed in the data.
Revised: There are many representation learning methods in video [29–32], physics-guided complex systems [33], and behavior modeling [34, 35] that also aim to learn interactions between discrete objects observed in the data.

Original: In these approaches, interactions between different objects are typically learned in a joint manner by combining the inferred representation of the objects into a common representation.
Revised: In these approaches, interactions between different objects are typically learned in a joint manner by combining the inferred representation of the objects into a common representation.

Original: In contrast, our frameworkoperates on pre-separated channels and aims to decompose dynamics into two parts, forming a representation of the dynamics of each individual source or object in addition to a representation of the interactions across many sources.
Revised: In contrast, our framework aims to decompose dynamics into two parts, forming a representation of the dynamics of each individual source or object in addition to a representation of the interactions across many sources.

Original: Thus, one could imagine using our approach to build enhanced representations of object dynamics in other vision and behavioral neuroscience applications [33, 34].
Revised: Thus, one could imagine using our approach to build enhanced representations of object dynamics in other vision and behavioral neuroscience applications [30, 34].

Original: It would be interesting to see if an enhanced or individual representation could be used to further improve the performance of upstream object localization and inference approaches.
Revised: It would be interesting to see if an enhanced or individual representation could be used to further improve the performance of upstream object localization and inference approaches.

Original: 2.2 Neural decoding and the challenge of generalization
Revised: 2.2 Neural decoding and the challenge of generalization

Original: The brain is an incredibly complex system.
Revised: The brain is an incredibly complex system.

Original: Neural activity is guided both by a tight interplay of membrane dynamics and ionic currents within a single neuron [35–38] as well as interactions with other neurons over larger distributed circuits [39–41].
Revised: Neural activity is guided both by a tight interplay of membrane dynamics and ionic currents within a single neuron [36–39] as well as interactions with other neurons over larger distributed circuits [40–42].

Original: Understanding how populations of neurons work together to represent their inputs is an important objective in modern neuroscience.
Revised: Understanding how populations of neurons work together to represent their inputs is an important objective in modern neuroscience.

Original: Accordingly, constructing representations that can explain neural population activity is an active area of research with many existing approaches.
Revised: Accordingly, constructing representations that can explain neural population activity is an active area of research with many existing approaches.

Original: Discriminative models produce a representation that optimizes the capability of predicting external covariates (e.g. the subject’s perception or behavior) or  temporally-adjacent  activity, and include work such as MYOW [42] and the Neural Data Transformer (NDT) [43].
Revised: 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].

Original: Generative approaches instead attempt to model neural activity as driven by latent factors [44], with success using switching linear dynamical systems [45] and recurrent neural networks [46].
Revised: 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 [47].

Original: Recently, SwapVAE [47] emerged as a hybrid of these two approaches, combining a latent factor model of neural activity with self-supervised learning techniques to predict across data augmentations.
Revised: Recently, SwapVAE [48] emerged as a hybrid of these two approaches, combining a latent factor model of neural activity with self-supervised learning techniques to predict across data augmentations.

Original: However, all of these methods produce representations exclusively at the population level.
Revised: However, all of these methods produce representations exclusively at the population level.

Original: The challenges of transfer across recordings of neural population activity.
Revised: The challenges of transfer across recordings of neural population activity.

Original: More often than not, neural recordings are collected from different sets of neurons within a brain region of interest.
Revised: More often than not, neural recordings are collected from different sets of neurons within a brain region of interest.

Original: 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 [48, 49].
Revised: 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 [49, 50].

Original: Many existing approaches attempt to address this challenge by learning population-level descriptions of neural data and finding ways to align their latent representations [49, 13].
Revised: Many existing approaches attempt to address this challenge by learning population-level descriptions of neural data and finding ways to align their latent representations [50, 14].

Original: For example, [50] proposed a manifold alignment method to continuously update a neural decoder over time as neurons shift in and out of the recording; [51] proposed an adversarial domain adaptation for stable brainmachine interfaces; [52] proposed a robust alignment of cross-session recordings of neural population activity through unsupervised domain adaptation; [46] fit each set of neurons with its own encoder and used dynamics to inform alignment in the latent space during training.
Revised: For example, [51] proposed a 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 brainmachine interfaces; [53] proposed a robust alignment of cross-session recordings of neural population activity through unsupervised domain adaptation; [47] fit each set of neurons with its own encoder and used dynamics to inform alignment in the latent space during training.

Original: 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 [15], 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: In addition,
Revised: In addition,

Original: EIT can potentially be used to find neuron-level correspondence (or functional correspondence) through its decomposed representation of the population-level activity.
Revised: EIT can potentially be used to find neuron-level correspondence (or functional correspondence) through its decomposed representation of the population-level activity.

Original:  Methods 3.1 Formulation and Setup
Revised:  Methods 3.1 Formulation and Setup

Original: In this work, we consider time-varying prediction from multi-variate time-series observations.
Revised: In this work, we consider time-varying prediction from multi-variate time-series observations.

Original: Consider X ∈ R N × T as an input datum that measures the dynamics of N individual channels over T points in time.
Revised: Consider X ∈ R N × T as an input datum that measures the dynamics of N individual channels over T points in time.

Original: Let x ij denote the value of channel i at the j -th time point, where 1 ≤ i ≤ N and 1 ≤ j ≤ T , and let X ( i ) denote the time series of the i -th channel.
Revised: Let x ij denote the value of channel i at the j -th time point, where 1 ≤ i ≤ N and 1 ≤ j ≤ T , and let X ( i ) denote the time series of the i -th channel.

Original: We consider time-varying labels y
Revised: We consider time-varying labels y

Original: =
Revised: =

Original: { y 1 , . . .
Revised: { y 1 , . . .

Original: , y T } where y j ∈ R d is the label at the j -th time point.
Revised: , y T } where y j ∈ R d is the label at the j -th time point.

Original: Our aim is to find a function which takes input X to predict target variables y .
Revised: Our aim is to find a function which takes input X to predict target variables y .

Original: Instead of jointly considering individual time-series immediately, we instead decompose the model into two functions f : R T → R T and g : R N → R N , so that the optimization problem becomes:
Revised: Instead of jointly considering individual time-series immediately, we instead decompose the model into two functions f : R T → R T and g : R N → R N , so that the optimization problem becomes:

Original: [Equation]
Revised: [Equation]

Original: where Proj projects latent representations to predict the labels, and ℓ ( · , · ) is a loss measuring the discrepancy between the labels and the time-varying output from the decoder.
Revised: where Proj projects latent representations to predict the labels, and ℓ ( · , · ) is a loss measuring the discrepancy between the labels and the time-varying output from the decoder.

Original: In words, f is applied to each time-series independently, while g combines individual embeddings across the population at a single time step.
Revised: In words, f is applied to each time-series independently, while g combines individual embeddings across the population at a single time step.

Original: Crucially, this decomposition provides f invariance to the number of channels fed into the model.
Revised: Crucially, this decomposition provides f invariance to the number of channels fed into the model.

Original: In practice, we map each x ij to a higher dimension to increase capacity for inference.
Revised: In practice, we map each x ij to a higher dimension to increase capacity for inference.

Original: 3.2 Approach
Revised: 3.2 Approach

Original: As shown in Figure 1(B), EIT decomposes collections of temporal data through two transformers to build the representations of the dynamics of individuals and  their population-level interactions.
Revised: As shown in Figure 1(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.

Original: A temporal-spatial module that disentangles individual  and  interactions.
Revised: A temporal-spatial module that disentangles individual dynamics and their interactions.

Original: To separate interactions at the population level from the individual dynamics
Revised: To separate interactions at the population level from the individual dynamics

Original: , we start by building representations from individual time-series.
Revised: , we start by building representations from individual time-series.

Original: For the i -th channel, we obtain an initial embedding of T temporal tokens as Z ( i ) 

Revised: For the i -th channel, we obtain an initial embedding of T temporal tokens as Z ( i )

Original: =
Revised: =

Original: [ Embed ( x i 1 ) , Embed ( x i 2 ) , ..., Embed ( x iT )] , where Embed : R → R M and Z ( j )∈ R T × M .
Revised: [ Embed ( x i 1 ) , Embed ( x i 2 ) , ..., Embed ( x iT )] , where Embed : R → R M and Z ( j ) 0 ∈ R T × M .

Original: We apply a transformer with multi-head attention to the resulting sequence using Equation (2).
Revised: We apply a transformer with multi-head attention to the resulting sequence using Equation (2).

Original: Let (cid:98) Z
Revised: Let (cid:98) Z

Original: ( i ) ∈ R T ×
Revised: ( i ) ∈ R T ×

Original: M denote the final embeddings with elements ˆ z ij obtained from the first transformer in our model and (cid:98) Z ∈ R N × T ×
Revised: M denote the final embeddings with elements ˆ z ij obtained from the first transformer in our model and (cid:98) Z ∈ R N × T ×

Original: M be the combined embedded output.
Revised: M be the combined embedded output.

Original: After building representations of dynamics at the scale of individual time-series, we then pass the representations into a population-level transformer to capture interactions between all of the individual time-series.
Revised: After building representations of dynamics at the scale of individual time-series, we then pass the representations into a population-level transformer to capture interactions between all of the individual time-series.

Original: To do this, we take the output of the first module (cid:98) Z and slice it along the temporal axis to yield an embedding of each time point in terms of the different channels.
Revised: To do this, we take the output of the first module (cid:98) Z and slice it along the temporal axis to yield an embedding of each time point in terms of the different channels.

Original: Denote this new sequential reslicing of the data at the j -th timepoint as V ( j )=
Revised: Denote this new sequential reslicing of the data at the j -th timepoint as V ( j )=

Original: [ˆ z 1 j , ˆ z 2 j , ..., ˆ z Nj ] + E pos , where E pos ∈ R N × M is the fixed positional embedding that is dependent on neuron i ’s identity.
Revised: [ˆ z 1 j , ˆ z 2 j , ..., ˆ z Nj ] + E pos , where E pos ∈ R N × M is the fixed positional embedding that is dependent on neuron i ’s identity.

Original: After applying the attention block operations in Equation (2) again, we then arrive at our final embedding of our time-series as (cid:98) V ∈ R N × T × M .
Revised: After applying the attention block operations in Equation (2) again, we then arrive at our final embedding of our time-series as (cid:98) V ∈ R N × T × M .

Original: Let ˆ v ij be the final representation of the observation x ij .
Revised: Let ˆ v ij be the final representation of the observation x ij .

Original: A projection module that preserves individual identity.
Revised: A projection module that preserves individual identity.

Original: To form a population representation with individual representation, we design a projection module that preserves individual identities.
Revised: To form a population representation from the embeddings of individual channels, we design a projection module that preserves individual identities.

Original: The projection module Proj( · ) consists of two parts: 1) For representations of individuals z ij ∈ R M for the i individual and j -th time-point, Proj first bottleneck the individual latents with a function bottleneck : R M → R to reduce the dimensions.
Revised: The projection module Proj( · ) consists of two parts: 1) For representations of individuals z ij ∈ R M for the i individual and j -th time-point, Proj first bottleneck the individual latents with a function bottleneck : R M → R to reduce the dimensions.

Original: To form a population representation, Proj then concatenates the representation across individuals at the j -th time point, and learns another projection project : R N → R d to infer the label y j .
Revised: To form a population representation, Proj then concatenates the representation across individuals at the j -th time point, and learns another projection project : R N → R d to infer the label y j .

Original: 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.
Revised: 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.

Original: Learning representations through a multi-stage loss.
Revised: Learning representations through a multi-stage loss.

Original: To build an unified representation that can provide both a description at the individual-level and the population-level, we compose a loss function as follows: 
Revised: To build an unified representation that can provide both a description at the individual-level and the population-level, we compose a loss function as follows: L

Original: [Equation]
Revised: [Equation]

Original: where Proj-Int( · ) and Proj-Ind( · ) denote two projection modules after the population and individual transformers, respectively; ℓ can be set to either classification loss ( CrossEntropy ) or regression loss ( MSE ), depending on the type of prediction target we consider; α is a scaling factor that determines how much emphasis is placed on prediction from individual representations in the first half of the network.
Revised: where Proj-Int( · ) and Proj-Ind( · ) denote two projection modules after the population and individual transformers, respectively; ℓ can be set to either classification loss ( CrossEntropy ) or regression loss ( MSE ), depending on the type of prediction target we consider; α is a scaling factor that determines how much emphasis is placed on prediction from individual representations in the first half of the network.

Original: Remark:
Revised: Remark:

Original: By setting α > 0 , we train a projection module on the intermediate individual-level embeddings that provides a purely individual-based representation of the population.
Revised: By setting α > 0 , we train a projection module on the intermediate individual-level embeddings that provides a purely individual-based representation of the population.

Original: Empirically, this more flexible architecture also seems to improve the stability of training when compared with setting α = 0 .
Revised: Empirically, this more flexible architecture also seems to improve the stability of training when compared with setting α = 0 .

Original: Note that the first part of Equation 4 is a special case of Equation 3 when both the individual module and interaction module are transformer encoders.
Revised: Note that the first part of Equation 4 is a special case of Equation 3 when both the individual module and interaction module are transformer encoders.

Original: 3.3 Generalization across domains through linear decoding and functional alignment
Revised: 3.3 Generalization across domains through linear decoding and functional alignment

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Decoding from a new population of different dimension and ordering.
Revised: Decoding from a new population of different dimension and ordering.

Original: In many domains, it is difficult to reliably preserve channels across datasets.
Revised: In many domains, it is difficult to reliably preserve channels across datasets.

Original: Our model instead provides a flexible framework for domain generalization as the first transformer ( f ) operates only on individual timeseries without considering the entire population.
Revised: Our model instead provides a flexible framework for domain generalization as the first transformer ( f ) operates only on individual timeseries without considering the entire population.

Original: Specifically, to transfer to a new dataset, we can learn a new projection h ( · ) that acts on the concatenated pre-trained embeddings (outputs of f ) to build predictions about downstream target variables y :
Revised: Specifically, to transfer to a new dataset, we can learn a new projection h ( · ) that acts on the concatenated pre-trained embeddings (outputs of f ) to build predictions about downstream target variables y :

Original: [Equation]
Revised: [Equation]

Original: (5) where h is a decoder mapping representations to predicted labels, (cid:98) z ij are the representations formed for the i individual at the j -th time point, and y j is the label at j -th time point.
Revised: (5) where h is a decoder mapping representations to predicted labels, (cid:98) z ij are the representations formed for the i individual at the j -th time point, and y j is the label at j -th time point.

Original: When we restrict h to be linear, this approach is similar to the evaluation methods in self-supervised learning [53, 54], where a linear decoder is used to evaluate the quality of the frozen representation based on various different downstream tasks.
Revised: When we restrict h to be linear, this approach is similar to the evaluation methods in self-supervised learning [54, 55], where a linear decoder is used to evaluate the quality of the frozen representation based on various different downstream tasks.

Original: Functional alignment through  representational similarity analysis.
Revised: Functional alignment through Wasserstein-based representational similarity analysis.

Original: Our model also provides a flexible framework to understand correspondences that might exist between channels in different recordings.
Revised: Our model also provides a flexible framework for identifying correspondences that might exist between channels in different recordings.

Revised: For instance, in the case of neural recordings, we can use EIT to find correspondence across neurons in different recording sessions.

Original: Recall that for each individual time-series X ( i ) ∈ R T , our model builds a representation space of size
Revised: Recall that for each individual timeseries X ( i ) ∈ R T , our model builds a representation space of size

Original: (cid:98) Z ( i ) ∈ R T × M .
Revised: (cid:98) Z ( i ) ∈ R T × M .

Original: To characterize the similarity across different individual time-series in our datasets, we pass many samples through the  encoder to estimate the distribution of (cid:98) Z ( i ) ∈ R T × M .
Revised: To characterize the similarity across different channels, we pass many samples through the individual encoder to estimate the distribution of (cid:98) Z ( i ) ∈ R T × M .

Original: We denote this distribution of individual i as R i .
Revised: We denote this distribution of individual i as R i .

Original: We used the Wasserstein divergence ( W ), a measure of distributional distance motivated by optimal transport (OT) [55, 56] to obtain robust measures of similarity (More details in Appendix 1.1).
Revised: We then use the Wasserstein divergence ( W ), a measure of distributional distance motivated by optimal transport (OT) [56, 57] to obtain robust measures of similarity (More details in Appendix 1.1).

Original: For one set of distributions { R (1)1 , R (1)2 , . . .
Revised: For one set of distributions { R (1)1 , R (1)2 , . . .

Original: , R (1) N 1 } from domain (1) and another set of distributions { R (2)1 , R (2)2 , . . .
Revised: , R (1) N 1 } from domain (1) and another set of distributions { R (2)1 , R (2)2 , . . .

Original: , R (2) N 2 } from domain (2), we compute the divergence between all pairs of individuals, which yields a matrix D ∈ R N 1 × N 2 where D [ i ][ j ] =
Revised: , R (2) N 2 } from domain (2), we compute the divergence between all pairs of individuals, which yields a matrix D ∈ R N 1 × N 2 where D [ i ][ j ] =

Original: W ( R (1) i , R (2) j ) .
Revised: W ( R (1) i , R (2) j ) .

Original: When correspondence between channels is known, we can compare against the distance matrix to quantify how well our encoder preserves functional properties.

Original:  Experiments 4.1 Synthetic experiment: observing superposed many-body systems
Revised:  Experiments 4.1 Synthetic experiment: observing superposed many-body systems

Original: We first tested our model on synthetic many-body systems [57] (see [58, 59] for other possible synthetic datasets with additional complexity).
Revised: We first tested our model on synthetic many-body systems [58] (see [59, 60] for other possible synthetic datasets with additional complexity).

Original: As shown in Figure.2(A), for a 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.
Revised: As shown in Figure 2(A), for a 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.

Original: 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 [60, 61].
Revised: 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 [61, 62].

Original: To test whether our model is able to capture, disentangle, and decode the body dynamics, we create an experiment of ‘superposed’ many-body systems (as shown in Figure.2(B)), where two separate many-body systems are overlaid with each other as to create a system where some connections exists (bodies within the same system), while some do not (bodies from different systems).
Revised: To test whether our model is able to capture, disentangle, and decode the body dynamics, we create an experiment of ‘superposed’ many-body systems (as shown in Figure 2(B)), where two separate many-body systems are overlaid with each other as to create a system where some connections exists (bodies within the same system), while some do not (bodies from different systems).

Original: Generating non-chaotic and chaotic systems.
Revised: Generating non-chaotic and chaotic systems.

Original: We generatedthe non-chaotic two-body systems and chaotic three-body systems in a relatively stable near-circular way [62], where the trajectories are obtained with the Explicit Runge-Kutta method [63].
Revised: We generated non-chaotic two-body systems and chaotic three-body systems in a relatively stable near-circular way [63], where the trajectories are obtained with the Explicit Runge-Kutta method [64].

Original: We refer the readers to Appendix 2.1 for the details of the set of second-order differential equations that guide the movement.
Revised: We refer the readers to Appendix 2.1 for the details of the set of second-order differential equations that guide the movement.

Original: As shown in Figure.2(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 or interaction information are observed .
Revised: As shown in Figure 2(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 or interaction information are observed .

Original: 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.
Revised: 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.

Original: Note that we use a 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 ).
Revised: Note that we use a 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 ).

Original: The model is trained to recover these two factors.
Revised: The model is trained to recover these two factors.

Original: For the two-body system and the three-body system, we sample 10 consecutive observations with a gap of 0 .
Revised: For the two-body system and the three-body system, we sample 10 consecutive observations with a gap of 0 .

Original: 2 and 50 consecutive observations with a gap of 0 .
Revised: 2 and 50 consecutive observations with a gap of 0 .

Original: within time span [0 , 10] , correspondingly.
Revised: within time span [0 , 10] , respectively.

Original: For each system that consists of different body types and interactions, we generate 500 trials to provide large variability of the movements.
Revised: For each system that consists of different body types and interactions, we generate 500 trials to provide large variability of the movements.

Original: We performed a 80/20 % train/test split throughout, where the testing set contains unseen combinations of body weights that determine the trajectory.
Revised: We performed a 80/20 % train/test split throughout, where the testing set contains unseen combinations of body weights that determine the trajectory.

Original: [Table]
Revised: [Table]

Original: Model performance.
Revised: Model performance.

Original: We evaluate the performance of EIT by investigating its ability to decode the body types (‘Type’), body interaction (‘Int’), as well as all possible combinations (‘All’) of them.
Revised: We evaluate the performance of EIT by investigating its ability to decode the body types (‘Type’), body interaction (‘Int’), as well as all possible combinations (‘All’) of them.

Original: The experiments are performed under the domain generalization scheme where the body mass distribution of the testing set (T) is different from that of the training set (S).
Revised: The experiments are performed under the domain generalization scheme where the body mass distribution of the testing set (T) is different from that of the training set (S).

Original: We benchmark EIT with a transformer  with the same depth and amount of heads that considers the population (create the population representation) right from the start of the model.
Revised: We benchmark EIT with a transformer [44] with the same depth and amount of heads that considers the population (create the population representation) right from the start of the model.

Original: As shown inTable 1,
Revised: As shown in

Original:  EIT consistently outperforms the baseline model (BM) in all cases, where our model showed significant advantage over the baseline when predicting ≈ 100 classes (the ‘Type’ and ‘All’ in three-body setting are 90 and 360 classes, respectively).
Revised: Table 1, EIT consistently outperforms the baseline model (BM) in all cases, where our model showed significant advantage over the baseline when predicting ≈ 100 classes (the ‘Type’ and ‘All’ in threebody setting are 90 and 360 classes, respectively).

Original: The advantages of the architecture are especially apparent when it comes to the three-body chaotic systems: as the system trajectories provide much higher variability and are highly sensitive to the initial conditions (as body types), our model provide a significantly betterinference performance by analyzing the individual dynamics separately to provide a stable representation space for each individuals.
Revised: The advantages of the architecture are especially 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.

Original: We visualized both the individual and population latent space in Figure 2(C) for the two-body systems.
Revised: We visualized both the individual and population latent space in Figure 2(C) for the two-body systems.

Original: Our individual-module successfully distinguishes different types of bodies, while the populationmodule further separates different systems by learning the interaction of bodies.
Revised: Our individual-module successfully distinguishes different types of bodies, while the populationmodule further separates different systems by learning the interaction of bodies.

Original: We refer readers to
Revised: We refer readers to

Revised: Appendix 2.1 for visualizations of the three-body systems.

Revised: On the top-right corner, we show the OT

Original: [Table] [Table]
Revised: [Table] [Table]

Revised: distance matrix that is produced by our individual-based evaluation method on the test set.

Original:  can clearly quantify the similarity between bodies, tells types 
Revised: The matrix can clearly quantify the similarity between bodies, and reveals the different types of individuals.

Original: 4.2 Decoding behavior from neural populations
Revised: 4.2 Decoding behavior from neural populations

Original: Datasets.
Revised: Datasets.

Original: The Mihi-Chewie reaching dataset [13] consists of stable behavior-based neural responses across different neuron population groups and animals, and thus is used in previous methods for across-animal decoding [52, 51, 13] and neural representation learning [42, 47].
Revised: The Mihi-Chewie reaching dataset [14] consists of stable behavior-based neural responses across different neuron populations and animals, and thus is used in previous methods for acrossanimal decoding [53, 52, 14] and neural representation learning [43, 48].

Original: Mihi-Chewie is a spike sorted dataset where two rhesus macaques, Chewie (‘C’) and Mihi (‘M’), are trained to perform a simplified reaching task to one of eight targets, while their neural activities in the primary motor cortex were simultaneously recorded.
Revised: Mihi-Chewie is a spike sorted dataset where two rhesus macaques, Chewie (‘C’) and Mihi (‘M’), are trained to perform a simplified reaching task to one of eight targets, while their neural activities in the primary motor cortex were simultaneously recorded.

Original: The dataset contains two recordings of different sets of neurons for each of the subjects, for a total of four sub-datasets.
Revised: The dataset contains two recordings of different sets of neurons for each of the subjects, for a total of four sub-datasets.

Original: The Jenkins’ Maze dataset [64] contains  the primary motor and dorsal premotor cortices activities of a rhesus macaque named Jenkins (‘J’).
Revised: 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 a rhesus macaque named Jenkins (‘J’).

Original: In this dataset, J reaches towards targets while avoiding the boundaries of maze that appears on the screen.
Revised: In this dataset, J reaches towards targets while avoiding the boundaries of maze that appears on the screen.

Original: Since this dataset provides more complex behaviour trajectories with a complete collection of continuous labels (movement velocity and hand position), we used this datasetas to examine if our model creates an individual representation space that is rich enough to decode continuous targets.
Revised: Since this dataset provides more complex behaviour trajectories with a 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.

Original: All neural datasets are sorted and binned into 100 ms bins by counting the number of spikes each neuron emits during that time frame to generate a time-series for each neuron.
Revised: 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.

Original: Experimental setup.
Revised: Experimental setup.

Original: The temporaltransformer andthe spatial transformerboth have a depth ofand 6 heads.
Revised: Both the temporal and spatial transformer have a depth of 2 and 6 heads.

Original: We set the dimension of single neuron representation to be 16 through the transformer training, and bottleneck it to be 1d when evaluating the activity representation, which makes the activity representation equal to the total number of neurons.
Revised: We set the dimension of single neuron representation to be 16 through the transformer training, and bottleneck it to be 1d when evaluating the activity representation, which makes the activity representation equal to the total number of neurons.

Original: We train EIT end-to-end with an Adam optimizer of learning rate 0 .
Revised: We train EIT end-to-end with an Adam optimizer of learning rate 0 .

Original: 0001 for 400 epochs.
Revised: 0001 for 400 epochs.

Original: We benchmark our models’ performance with a MLP, bi-directional GRU [65], NDT [43], and a supervised variant of the NDT  that we implemented for sequential data decoding.
Revised: We benchmark our models’ performance with a MLP, bi-directional GRU [66], NDT [44], and a supervised variant of the NDT (NDT-Sup) that we implemented for sequential data decoding.

Original: For domain generalization tasks, we re-trained the first-layer of NDT, and compared our model against numbers obtained when training on the same individuals.
Revised: For domain generalization tasks, we re-trained the first-layer of NDT, and compared our model against numbers obtained when training on the same individuals.

Original: We also tested EIT ’s performance when considering only non-sequential data and compared it withthe recent self-supervised methods [42, 47] in Appendix 3.2.
Revised: 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.

Original: Decoding performance.
Revised: Decoding performance.

Original: As shown in Table 2, our model provides strong decoding performance on various benchmarks in terms ofboth its classification of reach (Mihi-Chewie) and continuous decoding of position and velocity in the regression task studied (Jenkins).
Revised: As shown in Table 2, our model provides strong decoding performance on various benchmarks, both in terms of its classification of reach (Mihi-Chewie) and continuous decoding of position and velocity in the regression task studied (Jenkins).

Original: For the classification task on Mihi-Chewie, we followed [42, 47] and tested the model’s robustness under both shorter sequence setting ( T = 2 ) and longer sequence setting ( T = 6 ), while for J’s Maze we evaluated our model on the full sequence with a regression loss.
Revised: For the classification task on Mihi-Chewie, we followed [43, 48] and tested the model’s robustness under both shorter sequence setting ( T = 2 ) and longer sequence setting ( T = 6 ), while for J’s Maze we evaluated our model on the full sequence with a regression loss.

Original: Both our model and NDT [43] provide good performance, and in some cases, the EIT (T) temporal transformer does quite well on these tasks.
Revised: Both our model and NDT [44] provide good performance, and in some cases, the EIT (T) temporal transformer does quite well on these tasks.

Original: Our results suggests that our model is competitive on these diverse decoding tasks from different neural populations in multiple animals.
Revised: Our results suggests that our model is competitive on these diverse decoding tasks from different neural populations in multiple animals.

Original: Generalization task # 1: across recording session and animal.
Revised: Generalization task # 1: across recording session and animal.

Original: After validating the model’s decoding performance, we investigate its generalization across different sessions and animals.
Revised: After validating the model’s decoding performance, we investigate its generalization across different sessions and animals.

Original: As the number of neurons in different recordings varies, we cannot use the model in an end-to-end manner.
Revised: As 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 a different size.

Original: Instead we train our model on one recording, freeze the weights, and apply the channel-invariant portion of our model to the new data as described in Section 3.3 (Figure 3(A)).
Revised: Instead, we train our model on one recording, freeze the weights, and apply the channel-invariant portion of our model to the new data as described in Section 3.3 (Figure 3(A)).

Original: In our experiments, we re-train only a final linear classification layer to predict the labels on the new dataset.
Revised: In our experiments, we re-train only a final linear classification layer to predict the labels on the new dataset.

Original: Even when tested on new populations of neurons, our model provides impressive overall performance in decoding across sessions and animals (Table 2, bottom).
Revised: Even when tested on new populations of neurons, our model provides impressive overall performance in decoding across sessions and animals (Table 2, bottom).

Original: Here, we compare with a re-trained NDT (with a new input layer) and with models trained and tested on the same animal (within domain, top).
Revised: Here, we compare with a re-trained NDT (with a new input layer) and with models trained and tested on the same animal (within domain, top).

Original: 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.
Revised: 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.

Original: These results suggests that our temporal transformer has learned information about the firing patterns of neurons that can be transferred across populations in the same and different animals.
Revised: These results suggests that our temporal transformer has learned information about the firing patterns of neurons that can be transferred across populations in the same and different animals.

Original: To further investigate the underlying factors that contribute to our ability to generalize across populations, we visualized the latent space of individual neurons that have either a closer distribution (in terms of their OT distance) or distant distribution in Figure 3(B).
Revised: To further investigate the underlying factors that contribute to our ability to generalize across populations, we visualized the latent space of individual neurons that have either a closer distribution (in terms of their OT distance) or distant distribution in Figure 3(B).

Original: In multiple cases, we found that the embedding of neurons in the training set (M1) had close neighbors in our test set (M2) with similar tuning  (bottom right of each latent embedding).
Revised: In multiple cases, we found that the embedding of neurons in the training set (M1) had close neighbors in our test set (M2, same animal at a different time point) with similar tuning profiles (bottom right of each latent embedding).

Original: At the same time, neurons that were more distant had orthogonal (or distinct) functional tuning.
Revised: At the same time, neurons that were more distant had orthogonal (or distinct) functional tuning.

Original: To characterize this functional alignment at a population level, we measured the divergence between all pairs of neurons in the train and test set (Figure 3(C)).
Revised: To characterize this functional alignment at a population level, we measured the divergence between all pairs of neurons in the train and test set (Figure 3(C)).

Original: When sorting neurons in both conditions by their firing rate (bottom), we found that neurons already have a good overall global correspondence in terms of their latent embeddings. This provides some evidence that the learnt latent representation of individuals can accurately reflect how active the neurons are in terms of their overall firing rate and that neurons in both conditions can potentially be aligned.
Revised: When sorting neurons in both conditions by their firing rate (bottom), we found that the learned neuron representations have a 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.

Original: These experiments open up a lot of possibilities for finding functional groups in the individual embeddings and  that EIT could be used to find aligned neurons across datasets.
Revised: These experiments open up a 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.

Original: [Table]
Revised: [Table]

Original: Generalization task # 2: across behaviours.
Revised: Generalization task # 2: across behaviours.

Original: To test the models ability to transfer when the total class distribution shifts between the train and the test, we tested the model when we fix the neurons but the targets are different between the train ( classes) and the test (all 8 classes).
Revised: To test the model’s ability to 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 8 classes.

Original: Again, we can test generalization in this condition by training another linear decoder to decode a different set of labels.
Revised: 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).

Original: Our results in Table 3 suggest that EIT can work well when tested in new behavioral conditions and provides significant gaps over other approaches, even when the training data is limited.
Revised: Our results in Table 3 suggest that EIT can work well when tested in new behavioral conditions and provides significant gaps over other approaches, even when the training data is limited.

Original:  Conclusion
Revised:  Conclusion

Original: In this work, we proposed EIT as a model to learn representations of  individual and collective dynamics in multi-variate time-series data.
Revised: In this work, we introduced EIT as a model for learning representations of both individual and collective dynamics from multi-variate time-series data.

Original: 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 permutationinvariant design.
Revised: In our experiments on both synthetic manybody 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-invariant design.

Original: While our model provides a novel strategy for building representations of collective and individual dynamics, we note that there are also limitations and areas for improvement:
Revised: While our model provides a novel strategy for decomposing complex dynamics, we note that there are also limitations and areas for future work:

Original: • Treating individual systems separately not only requires prior knowledge about their separation [29], but also might restrict the model’s representational capacity (also discussed in [66]) due to the limited size of the final representation space.
Revised: • 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.

Original: One possible solution to both challenges is to stack multiple individual-interaction modules [22], which would sacrifice the spatial/individual separability of the individual module.
Revised: One possible solution to both challenges is to stack multiple individual-interaction modules [23], which would sacrifice the spatial/individual separability of the individual module.

Original: Instead, one could combine the proposed framework with individual/object disentanglement methods to both decompose mixtures of sources and to learn richer representation for individuals.
Revised: Instead, one could combine the proposed framework with individual/object disentanglement methods to both decompose mixtures of sources and to learn richer representations for individuals.

Original: •  EIT utilizedversatile transformer encoders for both the individual and interaction modules.
Revised: • Replacing transformers with task-specific architectures: EIT utilized transformer encoders for both the individual and interaction modules.

Original: It is possible to replace theversatile transformer encoders with a task-specific architecture to process specific types of interactions, such as graph-based encoders [67, 22] or recurrent encoders [68], which might improve performance on certain tasks.
Revised: However, it is possible to replace the transformer encoders with a task-specific architecture to process specific types of interactions, such as graph-based encoders [68, 23] or recurrent encoders [69], which might improve performance on certain tasks.

Original: Additionally, the individual module of EIT models individual dynamics in a deterministic way, but in certain cases (e
Revised: Additionally, the individual module of EIT models individual dynamics in a deterministic way, but in certain cases (e

Original: g. multi-armed bandit tasks for neural activities), it might be more appropriate to model dynamics and interactions in a probabilistic manner (such as in [69]).
Revised: g. multi-armed bandit tasks for neural activities), it might be more appropriate to model dynamics and interactions in a probabilistic manner (such as in [70]).

Original: • The generalizability of EIT can be further improved by: (1) Training the whole network in a self-supervised way, perhaps using recent advances in transformers for vision [70, 71] and  contrastive approaches for neural activity [42, 47], which would eliminate the dependency on labels. (2) Exploring ways of improving domain transfer and individual alignment by using structured transport methods [72, 73].
Revised: • Reducing the need for labels through self-supervised training: While the individual representations of EIT generalize reasonably well across animals in the neural activity experiments, the model currently relies on 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].

Original: Impacts in neuroscience.
Revised: 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.

Original: Recently, foundation models have emerged as a key tool in natural language processing [74, 75] and vision [76]. These models are trained using large amounts of data from diverse sources, which have demonstrated an impressive ability to generalize to downstream  tasks.
Revised: In many domains, the ability to integrate large amounts of data from different sources into a single pre-trained model has enabled impressive advances on many downstream tasks [73–75]. With such a 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, 77] to decode behaviors in diverse contexts and complex tasks.

Original: However, building foundation models for brain decoding is challenging because recordings are collected from different population of neurons, which has so far prevented effective pre-training and transfer. Here we show that by decoupling  learningfrom a neural population into two stages, EIT can learn generalizable representations across different neuronal populations, which is a significant first step towards  a foundation model for neuroscience.
Revised: 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 a significant first step towards building a foundation model for neural decoding that can learn from diverse sets of neurons in different brains and contexts.

Original: Acknowledgements
Revised: Acknowledgements

Original: This project was supported by NIH award 1R01EB029852-01, NSF award IIS-2039741, the NSF
Revised: This project was supported by NIH award 1R01EB029852-01, NSF award IIS-2039741, the NSF

Original: Graduate Research Fellowship Program (GRFP) for MD, as well as generous gifts from the Alfred
Revised: Graduate Research Fellowship Program (GRFP) for MD, as well as generous gifts from the Alfred

Original: Sloan Foundation, the McKnight Foundation and the CIFAR Azrieli Global Scholars Program.
Revised: Sloan Foundation, the McKnight Foundation, and the CIFAR Azrieli Global Scholars Program.