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  individual components and their dynamics and building a model of the population-leveldynamics 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 collective description, it can beeasy to lose sight of each individual and how different individuals contribute to thelarger 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: Here, we present a novel transformer architecture for learning fromtime-varying data by building descriptions of both the individual as well as thecollective 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 many individuals into ourmodel at the onset, we develop a separable architecture that operates on individualtime-series first before passing them forward; this induces a permutation-invarianceproperty 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 complexinteractions 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 thatour multi-scale transformer not only yields robust decoding performance, but alsoprovide impressive performance in transfer.
Revised: On neural activity datasets, we show that our multi-scale transformer not only yields robust decoding performance, but also provide impressive performance in transfer.

Original: Our results show that it is possibleto learn from neurons in one animal’s brain and transfer the model on neurons ina different animal’s brain, with interpretable neuron correspondence across setsand animals.
Revised: 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.

Original: This finding opens up a new path to decode from and represent largecollections of neurons.
Revised: This finding opens up a new path to decode from and represent large collections of neurons.

Original:  Introduction
Revised:  Introduction

Original: Complex systems (such as the brain) contain multiple individual elements (e g. neurons) that interactdynamically 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 tolarge-scale behaviors is important in many domains of science and engineering, from ecology [1, 2]and social networks [3, 4, 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 providedimportant insights into collective computations and emergent behaviors [7]; however, it is alsopossible to lose sight of the contributions of different individuals’ dynamics to the final prediction orinference.
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, eitherdue to their different functional roles in the system [8, 9, 6], or due to shift in their dynamics becauseof sensor displacement or corruption [10, 11, 12].
Revised: 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].

Original: Moving forward, we need methods that can buildgood population-level representations while also providing an interpretable view of the data at theindividual 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 usesdynamic 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-channeltime-series by first learning rich features from individual time-series before incorporating informationand learned interactions across different individuals in the population.
Revised: 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.

Original: One critical benefit of ourmodel is spatial/individual separability : it builds a population-level representation from embeddingsof 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 domaingeneralization tasks, this means a trained model can be tested with permuted channels or entirelydifferent number 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 ourmodel 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 toreadouts from populations of neurons in the primary motor cortex of two rhesus macaques thatare performing the same underlying center-out reaching motor task.
Revised: 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.

Original: The stability of the neuralrepresentations and collective dynamics found in data from these animals [13, 14, 15] makes them anideal testbed to examine the generalization of our approach.
Revised: 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.

Original: We show that, remarkably, generalizationnot only occurs across recording sessions (with different sets of neurons) from the same animal, butthat 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 thelinear decoder based on pre-trained weights in the across-animal transfer outperforms many modelsthat 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 neuraldecoding and highlights the utility of our architecture.
Revised: This result provides an exciting path forward in neural decoding and highlights the utility of our architecture.

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

Original: • We proposed a novel framework that considers individual dynamics before collective dynamics, and realized it with a multi-stage transformer:
Revised: • We proposed a novel framework that considers individual dynamics before collective dynamics, and realized it with a multi-stage transformer:

Original: Embedded Interaction Transformer ( EIT ).EIT learns both population representation and individual representation through decouplingthe dynamics of individuals and their interactions in multi-variate time-series.
Revised: 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.

Original: • With this new framework, we proposed methods for generalization across datasets ofdifferent input size  and ordering.
Revised: • We introduced methods for generalization across datasets of different input size (number of channels) and ordering.

Original: We further proposed to measure the alignments of unknownindividuals through the Wasserstein divergence.
Revised: We further proposed to measure the alignments of time-series in different datasets through the Wasserstein divergence.

Original: • We applied EIT to both many-body systems and neural systems recordings.
Revised: • We applied EIT to both many-body systems and neural systems recordings.

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

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

Original: 2.1 Background on Transformers

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

Original: Transformers have revolutionized natural language processing with the mechanismof self-attention, which helps the model to learn long-range interactions across different elementsor tokens in a sequence.
Revised: 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.

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

Original: =
Revised: =

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

Original: , x N ]  ∈ R N × d, consisting of Ntokens 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, keys, andvalues.
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 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: Let 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 forour 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 concatenatedinto 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 positionalembedding E pos that helps to retain the positional information, and each transformer layer is thecombination 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 thatconsists of sequences of samples, are a natural choice to model with the transformer architecture.
Revised: Multi-channel time-series, as a data format that consists of sequences of data, 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 combineall channels into a collective 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 dynamicsis 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:[16] and [17] extracted embeddingsfrom multivariate physical systems with Koopman operators before feeding the resulting representation into a temporal transformer; [18] used a graph neural networkbefore transformers to embedinterconnected-structures to perform skeleton-based action recognition; [19] used convolutionalarchitecture to extract image embeddings before feeding them into  transformer.
Revised: 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.

Original: Another approach is to re-design the attention block such that the attention operation is computedboth 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’ attentionblock have been shown effective: [20] designed a non-autoregressive module to generate queriesfor time series forecasting; [21] used stacked spatial-temporal modules for traffic flow forecasting; GroupFormer [22] embed spatial and temporal context in parallel to create a clustered attention blockfor group activity recognition.
Revised: 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.

Original: While our approach also makes use of the high-level idea of separating spatial (individual) and temporal information, crucially, we restrained the direct computationbetween the spatially-related attention map and the temporally-related attention map, which yields arepresentation 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  Vision Transformer [23] asa new way to extract image embeddings, many ‘spatial-temporal transformer’ architectures aredeveloped in the video domain [24, 25, 26].
Revised: 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].

Original: Such works provided interesting explorations on howto organize spatial attention and temporal attention with  coupled  attention blocks [26] or factorizedattention blocks [24], and how to create better tokens for videos with tubelet tokenizations [24].
Revised: 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].

Original: However, as image patches do not possess the channel-specific dynamics that we are interested in,processing videos with spacetime attention presents a very different problem.
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. Object-centric representation learning.

Revised: 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: 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:  2.2 Neural decoding and the challenge of generalization
Revised: In contrast, our framework operates 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. Thus, one could imagine using our approach to build enhanced representations of object dynamics in other vision and behavioral neuroscience applications [33, 34]. 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. 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 ofmembrane dynamics and ionic currents [27, 28, 29, 30] as well as interactions with other neuronsover larger distributed circuits [31, 32, 33].
Revised: 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].

Original: Understanding how  neurons work together to representtheir 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 areaof 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 whichis a good predictor of  external covariates (e.g. the subject’s perception or behavior) or temporallyadjacent activity, and include work such as MYOW [34] and the Neural Data Transformer [35].
Revised: 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].

Original: Generative approaches instead attempt to model neural activity as driven by latent factors (whichmight include external covariates, but could be unknown) [36] and make use of switching lineardynamical systems [37] and recurrent neural networks [38].
Revised: 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].

Original: Recently, SwapVAE [39] emerged as ahybrid of these two approaches, combining a latent factor model of neural activity with self-supervisedlearning techniques to predict across data augmentations.
Revised: 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.

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

Original: Learning across recordings or individuals.
Revised: 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.

Original: Generalizing acrossdifferent neural recordings isa critical but challenging task [40, 41].
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 [48, 49].

Original: The challenge is two-fold: 1) Tracking individual neuronsover a long-term recording is nearly impossible, as the detected subset of neurons, even at the samephysical location, could shift from time to time [10, 11, 12]; 2) We labels to map the recorded neuronsfrom different subsets or different individuals, as the specific neurons might be completely different. Many existing approaches attempt to conquer this challenge through the population point of view:[42] proposed a manifold alignment method to continuously update a neural decoder over time asneurons may shift in and out of the recording; [43] proposed an adversarial domain adaptation forstable brain-machine interfaces; [44] proposed a robust alignment of cross-session recordings ofneural population activity through unsupervised domain adaptation; [38] fit each set of neurons withits own encoder and used dynamics to inform alignment in the latent space during training. 
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 [49, 13]. 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. 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.

Original: Crucially,these works all rely on a subset of common neurons across recordings which can be matched [42].
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: However, such a subset of course does not exist across animals.
Revised: In addition,

Original: Although there are works assumethat the neural activities are intrinsically stochastic [45, 46], we believe it is reasonable to expect thatthere exists individuals that share similar representations when producing similar behaviour.
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 
Revised:  Methods 3.1 Formulation and Setup

Original: 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 Tpoints 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 and1 ≤ 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 labelsy = { y 1 , . . .
Revised: We consider time-varying labels y = { 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 jointlyconsidering individual time-series immediately, we instead decompose the model into two functionsf : 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 thediscrepancy 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 appliedto each time-series independently, while g combines individual embeddings across the population ata 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 fedinto 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 tobuild the representations of the individual time-series and also the population-level interactions.
Revised: 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.

Original: A temporal-spatial module that disentangles individual and interactions.
Revised: A temporal-spatial module that disentangles individual and 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 fromindividual 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
Revised: For the i -th channel, we obtain an initial embedding of T temporal tokens as

Original: Z ( i ) 

Revised: 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 )∈ 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)
Revised: Let (cid:98) Z

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

Original: After building representations of dynamics at the scale of individual time-series, we then pass therepresentations into a population-level transformer to capture interactions between all of the individualtime-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 toyield 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 sequentialreslicing 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 × Mis 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 attentionblock 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 representationwith individual representation, we design a projection module that preserves individual identity.
Revised: To form a population representation with individual representation, we design a projection module that preserves individual identities.

Original: Theprojection module Proj( · ) consists of two parts: 1) For representations of individuals z ij ∈ R Mfor the i individual and j -th time-point, Proj first bottleneck the individual latents with a functionbottleneck : 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 thenconcatenates the representation across individuals at the j -th time point, and learns another projectionproject : R N → R d to inference of 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 thepopulation representation that could be used for inference and the individual representation that isaffected 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 thatcan provide both a description at the individual-level and the population-level, we compose a lossfunction 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:

Original: [Equation]
Revised: [Equation]

Original: where Proj-Int( · ) and Proj-Ind( · ) denote two projection modules after the population and individualtransformers, 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 determineshow much emphasis is placed on prediction from individual representations in the first half of thenetwork.
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-levelembeddings 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  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 .

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 ) thatwe 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 sizedor shuffled populations, and (ii) ways to characterize alignment of functional properties of differentindividual 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, itis 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 canlearn a new projection h ( · ) that acts on the concatenated pre-trained embeddings (outputs of f ) tobuild 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 formedfor 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 hto be linear, this approach is similar to the evaluation methods in self-supervised learning [47, 48],where a linear decoder is used to evaluate the quality of the frozen representation based on variousdifferent downstream tasks.
Revised: 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.

Original: Functional alignment through representational similarity analysis.
Revised: Functional alignment through representational similarity analysis.

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

Original: Recall that for each individual time-series X ( i ) ∈ R T , our model builds a representationspace of size
Revised: Recall that for each individual time-series 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 inour 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 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 .

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 optimaltransport (OT) [49, 50] to obtain robust measures of similarity (More details in Appendix 1.1).
Revised: 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).

Original: Forone 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 betweenchannels is known, we can compare against the distance matrix to quantify how well our encoderpreserves functional properties.
Revised: When correspondence between channels is known, we can compare against the distance matrix to quantify how well our encoder preserves functional properties.

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

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

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

Original: As shown in Figure.2(A), for amany-body system, a body’s observed trajectory is decided by its own properties (its mass andstarting 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 arechaotic in nature, which is aligned with the stochasticity of neural activity [45, 46].
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 [60, 61].

Original: To test whetherour 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 systemsare overlaid with each other as to create a system where some connections exists (bodies within thesame 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 generated the non-chaotic two-body systemsand chaotic three-body systems in a relatively stable near-circular way [52], where the trajectories areobtained with the Explicit Runge-Kutta method [53].
Revised: We generated the 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].

Original: We refer the readers to Appendix 2.1 for thedetails 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 followingpipeline: two independent many-body systems are sampled , and then their trajectories withoutadditional 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 generatea latent representation for each body at each time that predicts the underlying factors that guide thebody 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 bytwo factors: 1) where the body initially starts (as the body types ); and 2) which body is this bodyinteracting 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: 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 agap of 0 .
Revised: 2 and 50 consecutive observations with a gap of 0 .

Original: 05 , correspondingly, for the two-body system and the three-body system, within time span[0 , 10] .
Revised: within time span [0 , 10] , correspondingly.

Original: For each system that consists of different body types and interactions, we generate 500 trialsto 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 performanceof EIT by investigating its ability to decode the bodytypes (‘Type’), body interaction (‘Int’), as well as allpossible 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 experimentsare performed under the domain generalization schemewhere 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 samedepth 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 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 in Table 1,
Revised: As shown in Table 1,

Original: EIT consistently outperforms the baseline model (BM) in all cases, where our model showedsignificant advantage over the baseline when predicting ≈ 100 classes (the ‘Type’ and ‘All’ inthree-body setting are 90 and 360 classes, respectively).
Revised: 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).

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

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

Original: Appendix 2.1 for visualizations of the three-body systems.
Revised: [Table]

Original: On the top-right corner, we show the OTdistance matrix that is produced by our individual-based evaluation method on the test set.
Revised: [Table]

Original: The matrixcan clearly quantify the similarity between bodies,which tellsthe different typesof individuals.
Revised: can clearly quantify the similarity between bodies, tells types

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 responsesacross different neuron population groups and animals, and thus is used in previous methods foracross-animal decoding [44, 43, 13] and neural representation learning [34, 39].
Revised: 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].

Original: Mihi-Chewie is aspike sorted dataset where two rhesus macaques, Chewie (‘C’) and Mihi (‘M’), are trained to performa simplified reaching task to one of eight targets, while their neural activities in the primary motorcortex 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 neuronsfor 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 [54] containsthe primary motor and dorsal premotor cortices activities of a rhesus macaque named Jenkins (‘J’).
Revised: The Jenkins’ Maze dataset [64] contains the primary motor and dorsal premotor cortices activities of a rhesus macaque named Jenkins (‘J’).

Original: In this dataset, J reaches towards targets while avoiding the boundaries of maze that appears on thescreen.
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 ofcontinuous labels (movement velocity and hand position), we used this dataset as to examine if ourmodel 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 as to examine if our model creates an individual representation space that is rich enough to decode continuous targets.

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

Original: Experimental setup.
Revised: Experimental setup.

Original: The temporal transformer and the spatial transformer both have a depth of 2and 6 heads.
Revised: The temporal transformer and the spatial transformer both have a depth ofand 6 heads.

Original: We set the dimension of single neuron representation to be 16 through the transformertraining, and bottleneck it to be 1d when evaluating the activity representation, which makes theactivity 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 Adamoptimizer 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
Revised: We benchmark our models’ performance with a

Original: MLP, bi-directional GRU [55], NDT [35], and a supervised variant of the NDT that we implementedfor sequential data decoding.
Revised: MLP, bi-directional GRU [65], NDT [43], and a supervised variant of the NDT 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 alsotested EIT ’s performance when considering only non-sequential data and compared it with the recentself-supervised methods [34, 39] in Appendix 3.2.
Revised: We also tested EIT ’s performance when considering only non-sequential data and compared it with the recent self-supervised methods [42, 47] in Appendix 3.2.

Original: Decoding performance.
Revised: Decoding performance.

Original: As shown in Table 2, our model provides strong decoding performanceon various benchmarks in terms of both its classification of reach (Mihi-Chewie) and continuousdecoding 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 in terms of both its classification of reach (Mihi-Chewie) and continuous decoding of position and velocity in the regression task studied (Jenkins).

Original: For the classification taskon Mihi-Chewie, we followed [34, 39] and tested the model’s robustness under both shorter sequencesetting ( T = 2 ) and longer sequence setting ( T = 6 ), while for J’s Maze we evaluated our model on thefull sequence with a regression loss.
Revised: 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.

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

Original: Our results suggeststhat our model is competitive on these diverse decoding tasks from different neural populations inmultiple 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’sdecoding 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.

Original: Instead we train our model on one recording, freeze the weights, and apply the channel-invariantportion 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 performancein 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 outperformmany 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 temporaltransformer has learned information about the firing patterns of neurons that can be transferred acrosspopulations 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 (interms 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 theembedding of neurons in the training set (M1) had close neighbors in our test set (M2) that had sameoverall functional properties and 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) with similar tuning (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 betweenall 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 bytheir firing rate (bottom), we found that neurons already have a good overall global correspondencein terms of their latent embeddings.
Revised: 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.

Original: This provides some evidence that the learnt latent representationof individuals can accurately reflect how active the neurons are in terms of their overall firing rateand that neurons in both conditions can potentially be aligned.
Revised: 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.

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

Original: [Table]
Revised: [Table]

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

Original: Totest the models ability to transfer when the totalclass distribution shifts between the train and thetest, we tested the model when we fix the neuronsbut the targets are different between the train (2classes) and the test (all 8 classes).
Revised: 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).

Original: Again, we cantest generalization in this condition by training another linear decoder to decode a different set oflabels.
Revised: Again, we can test generalization in this condition by training another linear decoder to decode a different set of labels.

Original: Our results in Table 3 suggest that EIT canwork well when tested in new behavioral conditions and provides significant gaps over other approaches, even when the training data ishighly 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 collectivedynamics when modeling multi-variate time-series.
Revised: In this work, we proposed EIT as a model to learn representations of individual and collective dynamics in multi-variate time-series data.

Original: In our experiments on both synthetic many-bodysystems and real-world neural systems recordings, we demonstrated that our model not only providesstate-of-the-art decoding performance, but its permutation invariance (enabled by the separationof individuals and populations) led to excellent domain generalization performance.
Revised: 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.

Original: These resultssuggest that our model could open up a new way to decode complex time-series datasets, especiallyfor neural activities.

Original: While our model provides a novel point of view for building representations of collective andindividual dynamics, we note that there are also limitations: 1) treating individual systems separately,limits the  capacityof individuals (also discussed in [56]) as each individual has limited dimensionalityas in the final representation space.
Revised: 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: • 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.

Original: 2) our model incorporates individual dynamics in a deterministicway, while many individuals can actually be expressed with a probabilistic model (such as in [57]). Aside from addressing the above limitations, we also point out that our methods could be extended ina number of directions, including: 1) training the whole network in a self-supervised way, perhapsusing recent advances in transformers for vision [58, 59] and contrastive approaches for neural activity[34, 39]; 2) improving domain transfer through structured transport methods [60]. Checklist
Revised: 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. 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.

Original: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope? [Yes]
Revised: • EIT utilized versatile transformer encoders for both the individual and interaction modules.

Original: (b) Did you describe the limitations of your work? [Yes] (c) Did you discuss any potential negative societal impacts of your work? [N/A]
Revised: It is possible to replace the versatile 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.

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes] If you are including theoretical results...
Revised: Additionally, the individual module of EIT models individual dynamics in a deterministic way, but in certain cases (e

Original: (a) Did you state the full set of assumptions of all theoretical results? [N/A] (b) Did you include complete proofs of all theoretical results? [N/A]
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 [69]).

Original: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either inthe supplemental material or as a URL)? [Yes] The code and 
Revised: • 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.

Original: instructions will be included as a URL after the double-blinded review. (b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
Revised: (2) Exploring ways of improving domain transfer and individual alignment by using structured transport methods [72, 73].

Original: [Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes] We report the error bars in Appendix.3.
Revised: Impacts in neuroscience.

Original: (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)? [Yes]
Revised: Recently, foundation models have emerged as a key tool in natural language processing [74, 75] and vision [76].

Original: If you are  using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators? [Yes] (b) Did you mention the license of the assets?
Revised: These models are trained using large amounts of data from diverse sources, which have demonstrated an impressive ability to generalize to downstream tasks.

Original: [N/A] (c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
Revised: 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.

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating? [N/A] All datasets we used are publicly available (e) Did you discuss whether the data you are using/curating contains personally identifiable
Revised: Here we show that by decoupling learning from 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.

Original: information or offensive content? [N/A] If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if
Revised: Acknowledgements

Original: applicable? [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
Revised: This project was supported by NIH award 1R01EB029852-01, NSF award IIS-2039741, the NSF

Original: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation? [N/A]
Revised: 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.