A probabilistic framework for task-aligned intra- and inter-area neural manifold estimation
Authors: Edoardo Balzani
Jean-Paul G Noel
Pedro Herrero-Vidal
Dora E Angelaki
Cristina Savin
Venue: ICLR 2023 notable top 25% (2023)
Keywords: neuroscience
dimensionality reduction
probabilistic methods
inter-area interactions

Abstract: Latent manifolds provide a compact characterization of neural population activity and of shared co-variability across brain areas. Nonetheless, existing statistical tools for extracting neural manifolds face limitations in terms of interpretability of latents with respect to task variables, and can be hard to apply to datasets with no trial repeats. Here we propose a novel probabilistic framework that allows for interpretable partitioning of population variability within and across areas in the context of naturalistic behavior. Our approach for task aligned manifold estimation (TAME-GP) explicitly partitions variability into private and shared sources which can themselves be subdivided in task-relevant and task irrelevant components, uses a realistic Poisson noise model, and introduces temporal smoothing of latent trajectories in the form of a Gaussian Process prior. This TAME-GP graphical model allows for robust estimation of task-relevant variability in local population responses, and of shared co-variability between brain areas. We demonstrate the efficiency of our estimator on within model and biologically motivated simulated data. We also apply it to several datasets of neural population recordings during behavior. Overall, our results demonstrate the capacity of TAME-GP to capture meaningful intra- and inter-area neural variability with single trial resolution.

Published as a conference paper at ICLR 2023
A PROBABILISTIC FRAMEWORK FOR TASK- ALIGNED INTRA- AND AREA NEURAL MANIFOLD ESTIMATION
Edoardo Balzani, Jean Paul Noel,
Cristina Savin
Pedro Perrero-Vidall & Dora E. Angelaki
Center for Neural Science
Center for Neural Science
Center for Data Science
New York University New York, NY, 10003
New York University
New York, NY, 10003
{eb1 62, jpn5, pmh354 da93) @nyu.
cs5360enyu. edu
ABSTRACT
Latent manifolds provide a compact characterization of neural population activity and of shared co-variability across brain areas. Nonetheless, existing statistical tools for extracting neural manifolds face limitations in terms of interpretability of latents with respect to task variables, and can be hard to apply to datasets with no trial repeats. Here we propose a novel probabilistic framework that allows for interpretable partitioning of population variability within and across areas in the context of naturalistic behavior. Our approach for task aligned manifold estimation (TAME-GP) explicitly partitions variability into private and shared sources which can themselves be subdivided in task-relevant and task irrelevant components, uses a realistic Poisson noise model, and introduces temporal smoothing of latent trajectories in the form of a Gaussian Process prior. This TAME-GP graphical model allows for robust estimation of task relevatl variability in local population responses, and of shared co- oo-vaiabblltyy between brain areas. We demonstrate the efficiency of our estimator on within model and biologically motivated simulated data. We also apply it to several datasets of neural population recordings during behavior. Overall, our results demonstrate the capacity of TAME-GP to capture meaningful intra- and inter-area neural variability with single trial resolution.
INTRODUCTION
Systems neuroscience is gradually shifting from relatively simple and controlled tasks, to studying naturalistic closed-loop behaviors where no two observations "trial""""ls") are alike (Michaiel et al. 2020; Noel et al., 2021). Concurrently, neurophysiological techniques are advancing rapidly (Steven son & Kording 1011jAngotz et al. 2019 Boi et al. 2020) to allow recording from an ever-increasing number of simultaneous neurons (i.e., "neural populations") and across multiple brain areas. These trends lead to a pressing need for statistical tools that compactly characterize the statisties of neural activity within and across brain regions. Dimensionality reduction techniques are a popular tool for interrogating the structure of neural responses (Cunningham & Byron 2014). However, as neural responses are driven by increasingly complex task features, the main axes of variability extracted using these techniques often intermix task and nuisance variables, making them hard to interpret. Alternatively, dimensionality reduction techniques that do allow for estimating task-aligned axes of variability (Brendel et al. 201 (Semeoo al., 20199] Keeley et al. 2020 Glaser et al. [2020] Hurwitz et al. 2021), do not apply to communication between brain areas, and/or necessitate trial repeat structure that does not occur in natural behavior.
Here, we introduce a probabilistic approach for learning interpretable task-relevant neural manifolds that capture both intra- and inter-area neural variability with single trial resolution. Task Aligned Manifold Estimation with Gaussian Process priors (TAME-GP) incorporates elements of demixed
Webpage: https:/aagehakkiabnyy.ooog
Wehpage: http://www.cns.yyucom csavin

---

Published as a conference paper at ICLR 2023
PCA (dPCA: Machens (200)) Kobak al. (2016)) and probabilistic canonical correlation analysis (pCCA; Bach & Jordan (2005) into a graphical model that additionally includes biologically relevant Poisson noise. The model uses Gaussian Process (GP) prior to enforce temporal smoothness, which
allows for robust reconstruction of single-trial latent dynamics (see DDamianou et al. (2016) for a similar approach using Gaussian observation noise). We demonstrate the robustness and flexibility of TAME-GP in comparison to alternative approaches using synthetic data and neural recordings from rodents and primates during naturalistic tasks. This reveals TAME-GP as valuable tool for dissecting sources of variability within and across brain areas during behavior.
Related work. Dimensionality reduction is usually achieved by unsupervised methods that identify axes of maximal variability in the data, such as PCA. In neuroscience, this is often accompanied by additional smoothing over time reflecting the underlying neural dynamics (c.g.. Gaussian process factor analysis (GPFA) (Yu et al., 2008); see GP-LVM (Hk & Lawrence) 2009) for similar approaches outside of neuroscience). This low dimensional projection is followed by a post hoc interpretation of latents in the context of behavioral variables, often by visualization. Alternative approaches such as
dPCA (Machens 2010 Kobak al. 2016) explicitly look for axes of neural variability that correlate with task variables of interest (see also Zhou & Weil (2020) for nonlinear version). However, these require partitioning trials into relatively few categories, based on experimental conditions or behavioral choices and averaging within conditions. This makes them unusable in naturalistic tasks where a single trial treatment is needed. Similarly, SNP-GPFA (Keeley et al., 2020) can partition (multi-region) neural activity into 'shared signal* and 'private noise' components, but only using data with stimulus repeats. Under 'no-repeat' conditions, pCCA (Bach & Jordan 2005) can find subspaces of maximal cross-correlation between linear projections of task variables and neural responses (under gaussian noise assumptions), without the need for a priori grouping of trials by experimental condition choice. This approach also be applied for determining shared axes of co-variability across areas, an analog for communication subspaces (Semedo et al. (019). Nonetheless, its noise model assumptions are mismatched neural data. More fundamentally, PCCA only considers pairwise relationships, preventing a joint multi-area and task variables analysis. Overall, existing approaches come with practical limitations and do not directly address the routing of task-relevant information across brain areas.
TASK-ALIGNED MANIFOLD ESTIMATION WITH GP PRIORS (TAME-GP)
In its most general form, the graphical model of TAME-GP models set of spike-count population responses x from up to n different areas 2 together with task variable of interest y (Fig. [A). The neural responses are driven by a set of n + low-dimensional latent variables z(1) Specifically, the responses of neuron in area j arise as a linear combination of private latent variability z and shared latents (0)), which reflect task interpretable aspects of the underlying dynamics, with Poisson noise and an exponential link function:
x 2 (0.m))
Poisson (exp (w:(0.)z4 + W(G5) +00++10 +n10))
(I)
with parameters W (0/55) and h))
To make latents interpretable with respect to task variables y, we adapt probabilistic framing of CCA (Bach & Jordan 2005) to introduces dependencies between any of the latents ((k)), which could be private or shared across areas, and
p (y)l20)) =N (x: (z Cz +
v), with parameters C, d, V.
'See Appendix A.1 for background on probabilistic PCA, CCA and their relation TAME-GP. VVariables x y tensors with dimensions corresponding to 1) area-specific number of neurons/ task variable dimension, 2) time within trial, and 3) trial index. We make indices explicit only where strictly needed.
2

---

Published as a conference paper at ICLR 2023
A
Z
reurano
PEUTORS
X
X
X
time [sea]
ground tuuhh
Inuth TAME-GG
www--oreeuul
EM-ltaratio
[sed]
0.0
time [sec]
[sea]
Figure I: A. TAME-GP generative model. 2 denotes shared latent dimensions while ( denote private latents of the corresponding area i; denotes the task variables. activity and task variable from the TAME-GP graphical model. C. Model log-likelihood function of the EM iteration (left) and cross-validated lcave-one-neuron-on marginal likelihood as function of (0) dimension (right). D-F. Latent variables estimation for within model simulated data: ground truth latent factors and model posterior mean + 95% CI for three latent dimensions.
B. Example draws of spiking
Finally, we regularize all latents to be smooth over time, through the introduction of Gaussian Process prior, as GPFA (Yu et al. 2008).
200
GP
(t- t²)
(10) ()) ky 2i 22,0))
dee exp
(4)
2T(10 (1)
.0) with area and dimension specific hyperparameters t, "t, is the i-th component of the j-th latent at time t, and 6,P the Kronecker delta.
Putting these elements together results in a factorization of the joint distribution of the form, p (x (1::m) y, (Uxx) = IT P (z()) p (ylz(0)) II,, P (z)1) (2(") 2 .Tii general form allows for unified mathematical treatment of several estimation tasks of interest. We will detail key instances of this class that have practical relevance for neuroscience when presenting our numerical results below.
3 EM-BASED PARAMETER LEARNING
E-step Since a closed form solution of the posterior is not available (due the Poisson noise), construct Laplace approximation of the posterior P p (z x2)x, y, 0) s (z)x, y, 0) y,
N (z; -H-)), where is the MAP of the joint log-likelihood and H is its corresponding Hessian. Both of these quantities are estimated numerically.
The MAP estimate is obtained by gradient descent on the joint log likelihood. The gradient of the joint log likelihood w.r.t. the latents can be written as
200 logp (2, x)
Σ
V zit logn (2()) + Vesn logp (yelzi
120
1>0
2i logp
x
(3)
lz
2 )
t>
1>0
We group latents in spike counts in and 0 (w(0/23) h0), C. d. v 760) simplify notation.

---

Published as a conference paper at ICLR 2023
where 1 (1 M) refers to the trial number, explicit index omitted for brevity. For given trial,
expanding one term at the time we have
2inh logp 2 = k()
log yy2, (0) = CTT v -ye Czi (y-x))) -d)
V
2) logp 2,
1 zi
(5)
= w(ki)
X; exp (w(l.)))) + w(id))) + +b()))
where j > 0, k {0. and K() the GP-prior covariance matrix (Eq. 3). The corresponding
second moments are
V ztn logp (261) -K() j E (0:
0o logp (ylz( 10))
-CT
(8) (k) logp
(1)
( 2 (1) (Tk. diag [ex (w(0))))) + w
(0)
22 + hho))) w(h)))
with h k E [0.j]. Inverting the x dimensional Hessian matrix is cubic in = dj. where T is the trial length and dj denotes the dimensionality of latent 2(1), which restricts the number and dimensionality of latents in practice. The Hessian of the log likelihood is sparse but does not have factorized structure. Nonetheless, we can take advantage of the block matrix inversion theorem, to speed up the computation to O Σ, d)) (see Appendix A.2)) with additional improvements based on sparse GP methods (Wilson & Nickisch 2015) left for future work.
M-step Given the approximate posterior q found in the E-step, the parameters updates can be derived analytically for a few parameters, and numerically for the rest (see Suppl. Info. A.3 for details). The other observation model parameters are computed numerically by optimizing the expected -ikellidd under the posterior. In particular, for neuron i in population we have
(0)
C (w!10)) w-(hhi) ns)
Σ
Tei
h4 +
[w/(a>)
w.G.a)
M (1(]
- exp h4
e.1
(0) (
(0.0) (T,))
((.5)
+ (w:00>)
(((0.)) 2
Σ
Σ (i-)))
(5)
For each neural population, we jointly optimized the projection weights and the intercept of all neurons with full Newton scheme by storing the inverse Hessian in compressed sparse row (CSR) format (see Appendix A.4 for the gradient and Hessian of L).
The GP-prior parameters were also learned from data by gradient based optimization (using the limited-memory Froydee--Feeccer-GGnn (2),)) ), and optimize for x to enforce a positive time constant. We define k(()) E RTT' such that [K("]] ts
Shanno--hnoooo scheme (Virtanen et al. (0220). First, we
set (5) - -og((>
exp
(1 (1-)))
.)²)- The resulting objective function will
take the form, L (A(3) trace K (K())
1) E, 2i ( log [((') Gradients are provided in
(j)
(
Appendix A.5 together with the procedure for parameter initialization (Appendix A.6).
4 RESULTS
Latent reconstruction for within model data. To validate the estimation procedure, we first used a simulated dataset sampled from the TAME-GP graphical model, with predefined parameters. Specifi- cally, we simulated two neural populations x(1) and x2)) each with 50 units and one-dimensional task relevant variable y. We fixed the private latent factors (1)) and z x(2) to two dimensions, and that of the shared factor z (0) to one. The projection weights W and C, the intercept terms d and h1), the observation variance matrix o and the GP time constants of the factors were randomly assigned. The parameters were chosen such that the overall mean firing rate about 20Hz in both
4

---

Published as a conference paper at ICLR 2023
areas. We simulated spike counts at 50ms resolution for 200 draws from the process (which we will
refer to as 'trials' in analogy to experiments), each lasting 2.5 seconds (see example trial in Fig. [B]
Given this data, we assessed the ability of our EM-based estimator to recover its true latent structure 4
The marginal log likelihood saturated after a relatively small number of EM iterations (Fig. I
As basic test of our ability to determine the dimensionality of latents, we systematically varied
the dimensionality of the shared latent, while fixing the dimensions of 2 and (2) to their ground truth value of 2. We found that the best model fit was achieved at the ground truth task dimension 1, demonstrating that we are able to infer true latent dimensionality from data (Fig D-F))
Finally, we assessed the quality of the recovered latents in individual test trials. Due to known degeneracies, originally documented in linear gaussian latent models (Roweis & Ghahramani 1999), the latent factors in TAME-GP are identifiable up to an affine transformation of the latent space. To address this, we used Procustes (Schönemann) [966) realign the latent axes back the original space. The resulting posterior mean estimate of the latents show an excellent agreement with the ground truth factors ccros-vvalddate latents and ground truth, Fig. D-F), while the model predicted rates explained 98% of the ground truth firing rate variance. The ability to reconstruct ground truth structure for within model data persists when considering more than two areas with shared covariability (Suppl. Fig. (ST). Overall, these numerical tests confirm that EM provides a veridical estimation of ground truth latent structure for within distribution data.
linear regression R² of 0.99 between the MAP estimate of
Task-aligned latent reconstruction for simulated latent dynamical systems models. The simple graphical model of TAME-GP captures axes of neural variability of scientific interest, but is far from an accurate generative model for neural dynamics during behavior. To assess the ability of TAME-GP to extract underlying structure from complex and dutttf-distributtin neural data, we used latent dynamical systems models in which we can explicitly define the flow of information from external stimuli and between areas, in several scenarios of practical interest.
The first in silico experiment focuses on identifying axes of task-relevant variability in neural responses. As a simple test case, we modeled a single neural population with a 6d latent structure (Fig. 2 Two of the latent dimensions were task-relevant, driven by an observed temporally smooth external input Yr, while the other four dimensions were intrinsic to the circuit. The key distinction between this process and the TAME-GP model assumptions is that the observed task variable acts as an input drive to the underlying latent dynamics rather than mapping to the latents directly. The latent dynamics take the form of a multivariate AR(I).
Apr (Zpr.t 140 At + /2At dw, (0)
Zpr.+1
(6)
Zaz+1 = Air (Zr, y.) At + 22At dw! (1)
where Apr E R4x4 and Ar E R2x2 the private and task relevant dynamics, y: E R² and 14 E R4
inputs drawn from a factorized RBF kernel, and w."" (6) is independent white noise for 0, Given these latent dynamics, spikes are generated as described by the TAME-GP observation model with W IR 100066 and d E R2000 We adjusted the parameters as to cover several average population firing rates by regulating d, for a fixed number of trials (200) and fixed trial duration (5 seconds). For simplicity, we circumvent the hyperparameter selection step by assuming that all estimators have access to the ground truth latent dimensionality: TAME-GP assumed 2 shared and 4 private latents. Unsupervised methods (pPCA, P-GPFA) were tasked with extracting the main two axes of neural variability in the data, while the supervised methods (pCCA) estimated 2d latents that correlate with task variable y; the same alignment procedure was used in all cases.
Fig. 2B illustrates the latent dynamics estimated by TAME-GP, pPCA (Tipping & Bishop, 1999), P-GPFA (Hooram 2015), and PCCA (Bach & Jordan 2005 We quantify the latent space estimation accuracy by mean squared error, demonstrating that TAME GP captured the stimulus driven dynamics better than other methods (Fig. (F and Suppl. Fig. S2) P-GPFA showed a tendency to smor-h
which obscured most of the underlying fine timescale latent structure. PCA failed by focusing on main axes of variability irrespective of task relevance, while CCA estimates were visually less interpretable. Only PCCA and TAME-GP found projections that selectively encoded for Zςy with TAME-GP outperforming PCCA across conditions. Finally, TAME-GP maintained its ability to *Here and in all subsequent analyses 90% of the data used for training the model and 10% for testing.
5

---

Published as a conference paper at ICLR 2023
B
ground truth
TAME-GP
P-GPFA
PPCA
DODA
privalo
bark relower
-0. -01
-02
0 03
LD-1
LD-1
LD-1
look
LO6
Polexon notso
pip
"2d tatents
E
'6d latents
veriecblee
aplesi
Lasort moc
rate JH
TAME OP POPPA
fr. mome
10.1 [+2]
[Mg]
[swcc
Figure 2: Methods comparison for single area task manifold alignment. A. TAME-GP graphical model for single area (top) and schematic for data generating process (bottom). Zur denotes task relevant shared latent dimensions while 7pr denotes private task-irrelevant variability. B. Ground truth task relevant dynamics (green) and estimated low dimensional projection for TAME-GP (purple), GPFA (blue), pPCA (dark gray) and PCCA (light gray) Mean squared error between the true shared dynamics and the model reconstruction, mean + s.d. over 10-fold cross-validation. D. Example single trial firing rate reconstruction. E. Mean squared error between the true and reconstructed firing rate across conditions, mean + s.d. over 10 folds of the data.
recover the underlying structure even when the model assumptions do not match the data exactly, in particular when the effect of the latents were modeled be approximately additive (Suppl. Fig. [S3).
We also compared these methods in terms of their ability to predict the ground truth firing rate generating the observed spiking responses (total dimensions matching the ground truth 6). Both TAME-GP and P-GPFA showed stable and accurate firing rate reconstruction error across conditions (Fig. (D..). while the factorized linear gaussian methods (pPCA, PCCA) performed poorly. This may be due to the larger model mismatch, while additionally suffering from the lack of temporal smoothing, especially for low firing rates. Overall, TAME-GP was the only procedure that both captured the overall data statistics well and extracted accurate task-interpretable latents.
Assessing inter-area communication in simulated latent dynamical systems. the second set of numerical experiments, focused estimating low-dimensional communication sub-spaces across neural populations (Fig. AA). The ground truth data was again constructed using latent dynamical systems models, which now included two populations (Fig. 3B), where a low dimensional projection of the dynamics in one area, the sender, drive the dynamics of the other area, the receiver:
As (zs.t yt) + 221w (0)
28
7yh
= P.zs
Zzh.t) At 22tW (1)
ZR,++1
AR (ZR,C
where As E RR4x4 and An E R4x4 are the sender and receiver dynamics, y, and A are temporally smooth inputs drawn from independent GPs with factorized RBF kernels, P E R2x4 defines the shared submanifold projection, and w;" is independent white noise. These latents map into spikes as above. We simulated three average firing rate conditions and varied the ground truth number of shared dimensions, from one to three. We compared our method with the two most commonly used approaches to communication manifold: PCCA and Semedo's 's reduced-rank regression procedure for communication manifold estimation (Semedo et al.) (219)) (Fig. BC), as well as SNP- GPFA (Keeley ct al. 2020) (both with and without trial repeats, see Appendix A.7 and Suppl. Fig. S4).
TAME-GP (without task alignment) outperformed alternative approaches in terms of the recon- struction error of both ground truth firing rates (Fig. D, F) and shared latent dynamics (Fig- 3).. Furthermore, when testing the ability of different approaches to infer the dimensionality of the shared manifold through model comparison, the lcave-one-out likelihood saturated at the ground truth dimension for all simulations (Fig. ). and peaked at the correct dimension 75% of the times (Fig. 3 H). In contrast, the Semedo estimator tends to systematically overestimate the dimensionality of the shared manifold in this dataset.
Finally, we tested the general case in which we search for a communication subspace that aligns to task variable y. To do so, fit TAME-GP to the same dataset but assuming that y, is observed.

---

Published as a conference paper at ICLR 2023
A
sendarr
recelver
pleca ardie Elophia notver
pricte
HeabeT hajedary
siveed
Er. mume
TAME-GP
linear proj-
EAP
Poisson spkes
10:7 rate [Hz]
var.
y,
15 -0.5
-1.0 PODA
0.0
-1.0
-0.5
0.0
sender
Samedo
-0.5
00
10.7 rate (Hz)
receiver
LD-1
14
ground Truth Er.
lssvo ore
H
Drob.
time [sec]
4
dim
(estimated - -rue)) dim.
latent dim.
Figure 3: A. Schematic of communication subspace (left) and associated TAME-GP graphical model versions (right). B. Ground truth spike count generation process. C. Example shared latent reconstruction for TAME-GP (purple), PCCA (light grey) and, reduced rank regression (dark grey); ground truth in orange. D. Statistics for firing rate prediction quality. E. Statistics of shared dynamics reconstruction. F. Example reconstructions of the receiver firing rates compared the ground truth (green). G. TAME-GP leave-one-neuron-our log-likellhond for different ground truth shared manifold dimensionality (d=1,2,3) and increasing population rate from 5.1, 10.7, 15.9 Hz (respectively, dashed, dashed-dotted and continuous lines). Lines styles show different average firing rate conditions. H. Difference between estimated and true zu, dimensionality for TAME-GP (purple) and reduced rank regression (grey). I. Model fit quality as function of latent dimensionality for all estimators. Ground truth dimension d=2 (dashed line). Error bars show mean + s.d. over 10-folds of coos-validation.
We found again that TAME-GP has the best reconstruction accuracy, which saturates the ground truth dimensionality (d=2). These observations are consistent across firing rate levels (see Suppl. Fig. (Sp.. When fitting SNP-GPFA to simulated data in the case of precise stimulus repetitions and comparing it to TAME-GP, we find that both models are able to capture the latent space factorization. However, only TAME-GP works well in the case when latent dynamics vary across episodes, as would be the case during natural behavior (i.e. without stimulus repeats, see Suppl. Fig. [S4] Table ST and Appendix | (A.7 for details). Overall, these results suggest that TAME- MP can robustly recover meaningful sources of vo-variabylity across areas in range of experimentally relevant setups.
Mouse neural recordings during open-field exploration. As first validation of the method, we estimated the manifold formed by simultaneously recorded head direction cells (n = =33) in the anterodorsal thalamic nuclei (ADN) (Taube 1995) of a mouse exploring a circular open field.
These neurons are known to form circular manifold representing heading direction (Chaudhuri et al. 2019). Thus, they provide the opportunity to examine the ability of TAME-GP to recover the underlying structure of data for which we know the biological ground truth. Recorded responses were segmented in 10sec time series, discretized 20ms bins, and with a either a head -directtin aligned 2d latent manifold (Fig(A); private noise dimension d=5), with unsupervised methods PPCA and PGPFA, each with Tatent dimensionality d=2. All methods recovered the underlying circular structure of the heading representation to some degree (Bi We decoded head direction from the extracted 2d latents 6 and confirmed that TAME-GP preserved more information than pPCA, and comparable P-GPFA (Fig with an overall superior data fit quality relative to pPCA (Fig(4)) as assessed by the R² between model leave-one-neuroe neeeen-euoo-ouu firing rate predictions and the raw spike
'Surgeries and procedures were approved by the Institutional Animal Care and Use Committee Baylor College of Medicine and New York University and were in accordance with National Institute of Health guidelines, protocol number 18-1502.
Decodinng was performed using Lasso regression, with hyperparameter selection by 5-fold cross-validation.
7

---

Published as a conference paper at ICLR 2023
TAME-GP
PPCA
P-OPFA
C
&c
ACN
headdn
60 ED4
do LD4
JPCA
LD41
E
initel largel
monkey trapector
TAME-GP
P-GPFA
LDG
floor
distance
0.2
8+ PPCA
00
02
LD-1
LD-1
K
enguler havelled
recial
0.5
IT TAME--?
0. &
TAME-GP
20.1
0.1
2.1
GPFA
plc
.0.3
wt
lapal.eng Bogel mng wat the dia eyel hosd
5
0.5
0.0
03
in
dimensonailly
onaliy
Figure 4: Fitting TAME-GP to neural data. A. Graphical model for heading aligned mouse population responses in area ADN. Zur denotes heading related shared latent dimensions while Zpr denotes private task-irrelevant variability. B. Latent population dynamics, colored by time-varying heading, for various manifold estimators. C. Head direction decoding from 2d latents extracted with each method (by Lasso regression). Mean + standard deviation over 5folds. D. Scatter plot of seae-ooee-eeroon PPCA. Dots represent individual neurons. E. Schematic of the firefly task. Initial target location randomized and remains visible for 300ms. The monkey has use the joystick navigate to the internally maintained target position. F. Top view of example monkey trajectories; increasing contrast marks initial location the target (right, center, left). G. Within-area TAME-GP estimation aligned latent task variable: the distance travelled. H. Scatter plot of leave-one-neuroonool spike count variance explained for prteension-maacher TAME-GI and pPCA. Dots represent individual neurons. I. Single trial TAME- GP estimates of the task relevant dynamics, compared to J. those of P-GPFA. Trajectories are color golordgraded according to the initial angular target location (as in B). Lasso regression decoding of K. and L. linear distance travelled. TAME-GP decoding π² (purple) is based on 2d task relevant latent. P-GPFA R² (blue) estimates were obtained for a range of latent dimensions (1-10). M. Communication subspace estimation between MSTd and dIPFC. N. As H. for shared latent space. O. Lasso regression decoding of task relevant variables (sorted by their shared subspace information content) from the shared (orange) and private latents (green, red) estimated by TAME-GP. Mean R² s.e.m. estimated across 10 folds of the data.
spike count variance explained for dimension matched TAME-GP and
counts (Yu et al., 2008). Overall, these results confirm that the TAME-GP estimator can extract sensible coding structure from real data that does not exactly match the assumptions of the model.
Multi-area neural recordings in monkeys during VR spatial navigation Finally, we tested the ability of TAME-GP to find task aligned neural manifolds in challenging dataset characterized by a high-dimensional task space and lack of trial repeats. Specifically, monkeys navigate virtual reality by using a joystick controlling their linear and angular velocity to "catch fireflies" (Fig4e. F) Lakshminarasimhan et al., 2018). Spiking activity was measured (binned in 6ms windows, sessions lasting over 90min) and neurons in the two recorded brain areas (MSTd and dIPFC) showed mixed selectivity, encoding a multitude of task relevant variables (Noel et al. 2021). As a result, responses are high dimensional and unsupervised dimensionality reduction methods capture an hard to interpret mixture of task relevant signals in their first few latent dimensions.
We used TAME-GP to extract latent projections that align with the ongoing distance from the origin, decomposed in an angular and radial component (Fig. 4G). We set the task relevant latent (0) dimensions to two, matching the number of task variables. We verified the accuracy of the model by computing deave-one-neurone predictions and raw spike counts. The TAME-GP estimator systematically outperformed pPCA with matched number of latents by this metric (Fig. 4).. We also compared the latent factors found
firing rate predictions and calculating the R² between model
8

---

Published as a conference paper at ICLR 2023
by TAME-GI to those obtained by P-GPFA (Fig. A. J). For both variables, found that the task variables were better accounted for by a two-dimensional TAME-GP estimated latent than by up to 10 dimensional latent spaces extracted with P-GPFA (Fig. AK. L). A similar compression of the manifold was achieved in a separate dataset of monkey (pre-) (pre-)motor responses during sequential reaches (see [A.9 and Suppl. Fig. [S7)- This confirms that TAME-GP provides compact low dimensional account of neural variability with respect of task variables of interest.
Lastly, we probed the model's ability learn communication subspace (Fig. 4M) between MSTd and dIPFC, brain areas that are known to interact during this task (Noel et al. 202T). In this instance, we selected the number of shared and private latent dimensions by maximizing the leave-one-neurow out spike counts variance explained over grid of candidate values (see Suppl. Fig. and As before, we find that the TAME-GP reconstruction accuracy surpasses that of dimensionality-mathhi PPCA, for both MSTd and dIPFC NN). Since the shared manifold estimation was agnostic to task variables in this case, we used decoding from latent spaces to ask if the shared variability between these areas carried information about task variables known to drive single neuron responses in these areas. We found that the monkey's ' horizontal eye position, as well as latent task variables such the travelled distance the distance still remaining to target shared, as opposed to private, axes of variability (Fig. (p)). This recapitulates prior observations made at the single-cell level (Noel et al., 202T). Overall, the results demonstrate that TAME-GP can extract interpretable low-dimensional latents and shared neural subspaces from complex and high- nigh-dimensional datasets.
mostly accounted for in
5 DISCUSSION
Technological advances in systems neuroscience place an ever verrincceasingg premium the ability to concisely describe high-dimensional task-relevant neural responses. While sophisticated methods based recurrent neural networks are increasingly used for fitting neural responses PPandarinat ett al. (2018), the extracted dynamics are also not necessarily easy to interpret. Here we introduce TAME-GP, flexible statistical framework for partitioning neural variability in terms of private or shared (i.e., inter-area) sources, aligned to task variables of interest, and with single trial resolution. We show that our method provides compact latent manifold descriptions that better capture neural variability than any of the standard approaches we compared it against.
An important nuance that distinguishes various neural dimensionality reduction methods is whether the covariability being modeled is that of trial-averaged responses (i.e. stimulus correlations), residual fluctuations around mean responses (i.e. noise correlations) or combination of the two (total correlations). Since isolating either the signal or the noise correlations alone would require across trial averages, our approach models total correlations, time resolved within individual trials. This differentiates our shared variability estimates from the traditional definition of a communication subspace [Semedo et al. 2019), which uses noise correlations alone, while keeping some of its spirit. It also makes it applicable to datasets without trial repeats.
The model adapts the approach of pCCA as a way of ensuring that the extracted latents reflect axes of neural variability that carry specific task relevant information. This choice has appealing mathematical properties in terms of unifying the problems of finding interpretable axes and communication subspaces, but is not the most natural one in terms of the true generative process of the data. While behavioral outputs can be thought of as outcomes of neural activity as described by the TAME-GP graphical model, sensory variables act as drivers for the neural responses and should affect the latent dynamics, not the other way around. Hence natural next step will be to incorporate in the framework explicit stimulus responses, perhaps by taking advantage of recent advances in estimating complex tuning functions during naturalistic behavior (Balzani et al.) 2020).
would be interesting to explore the use temporal priors with more interesting structure, for instance spectral mixture kernels (Wilson & Adams, 2013), introducing prior dependencies across latent dimensions (de Wolff et al. or using non-reversible GP priors that better capture the causal structure of neural dynamics (Rutten et al.) 2020). More generally, the probabilistic formulation allows the ideas formalized by TAME-GP GP to be combined with other probabilistic approaches for describing stimulus tuning and explicit latent neural dynamics (Duncker et al., 2019 Glaser et al.) 2020 Duncker & Sahani 202T). Hence, this work adds yet another building block in our statistical arsenal for tackling questions about neural population activity as substrate for brain computation.

---

Published as a conference paper at ICLR 2023
Broader impact We do not foresee any negative consequences to society from work. Task aligned manifold extraction may prove useful in clinical applications, specifically for in- creasing robustness of BMI decoders by exploiting the intrinsic structure of the neural re- sponses. Code implementing the TAME-GP estimator and associated demos is available at http://hhw....cm/aaaaaaaaaaaaaaa
This work was supported by the National Institute of Health under the U19 research program (grant agreement number NIH U19NS118246).
Acknowledgements.
REFERENCES
Gian Nicola Angotzi, Fabio Boi, Aziliz Lecomte, Ermanno Miele, Mario Malerba, Stefano Zucca, Antonino Casile, and Luca Berdondini. Sinaps: An implantable active pixel sensor emos-probe for simultaneous large-scale neural recordings. Biosensors and Bioelectronics, 126:355-364, 2019.
Francis R Bach and Michael Jordan. A probabilistic interpretation canonical correlation analysis.
Technical report, 2005.
Edoardo Balzani, Kaushik Lakshninarasinhan, Dora Angelaki, and Cristina Savin. Efficient esti- mation of neural tuning during naturalistic behavior. Advances in Neural Information Processing Systems, 33:12604- 2200-12214, 2020.
Christopher M Bishop and Nasser M Nasrabadi. Pattern recognition and machine learning, volume 4.
Springer, 2006.
Fabio Boi, Nikolas Perentos, Aziliz Lecomte, Gerrit Schwesig, Stefano Zordan, Anton Sirota, Luca Berdondini, and Gian Nicola Angotzi. Multi-shanks sinaps active pixel sensor cmos probe: 1024 simultancously recording channels for high-density intracortical brain mapping. bioRxiv, pp. 749911, 2020.
Wieland Brendel. Ranulfo Romo, and Christian K Machens. Demixed principal component analysis. Advances in neural information processing systems, 24, 2011.
Rishidev Chaudhuri, Berk Gerçek, Biraj Pandey, Adrien Peyrache, and Ila Fiete. The intrinsic attractor manifold and population dynamics of canonical cognitive circuit across waking and sleep. Nature neuroscience, (199):1512-15555 2019.
John P Cunningham and M Yu Byron. Dimensionality reduction for large-scale neural recordings. Nature neuroscience, 17(11):1500-1509 2014.
Andreas Damianou, Neil Lawrence, and Carl Henrik Ek. Multi-view learning as a nonparametric nonlinear inter-battery factor analysis. arXiv preprint arXiv: 04439, 2016.
Taco de Wolff, Alejandro Cuevas, and Felipe Tobar. Mogptk: The multi-output gaussian process toolkit. Neurocomputing, 424:49-53, 2021.
Lea Duncker and Maneesh Sahani. Dynamics on the manifold: Identifying computational dynamical activity from neural population recordings. Current opinion in neurobiology. 70: 70:163-170, 2021.
Lea Duncker, Gergo Bohner, Julien Boussard, and Maneesh Sahani. Learning interpretable continuous-time models of latent stochastic dynamical systems. In International Conference on Machine Learning. pp. 1726-1733 PMLR, 2019.
Carl Henrik Ek and PHTND Lawrence. Shared Gaussian process latent variable models. PhD thesis, Citeseer, 2009.
Joshua Glaser, Matthew Whiteway, John P Cunningham, Liam Paninski, and Scott Linderman. Re- current switching dynamical systems models for multiple interacting neural populations. Advances in neural information processing systems, 18:14688144878 2020.
Nam Hooram. Poisson extension of gaussian process factor analysis for modeling spiking neural populations master's thesis. Department of Neural Computation and Behaviour, Max Planck Institute for Biological Cybernetics, Tubingen, 2015.
10

---

Published as a conference paper at ICLR 2023
Cole Hurwitz, Akash Srivastava, Kai Xu, Justin Jude, Matthew Perich, Lee Miller, and Matthias Hennig. Targeted neural dynamical modeling. Advances in Neural Information Processing Systems, 34:29379-29399 2021.
S.L. Keeley, M.C. Aoi, Y. Yu, S.L. Smith, and Pillow J.W. Identifying signal and noise structure in neural population activity with gaussian process factor models. NeurIPS, 34, 2020.
Dmitry Kobak, Wieland Brendel, Christos Constantinidis, Claudia E Feierstein, Adam Kepecs, Zachary F Mainen, Xue-Lian Qi, Ranulfo Romo, Naoshige Uchida, and Christian K Machens. Demixed principal component analysis of neural population data. Elife, 5:e 30e10989, 2016.
Kaushik Lakshminarasimhan. Marina Petsalis, Hyeshin Park, Gregory C DeAngelis, Xaq Pitkow, and Dora Angelaki. dynamic bayesian observer model reveals origins of bias in visual path integration. Neuron, 99(1):1942006 2018.
Christian K Machens. Demixing population activity in higher cortical areas. Frontiers in computa- fional neuroscience, 4: 4:126, 2010.
Angie Michaiel, Elliot TT Abe, and Cristopher Niell. Dynamics of gaze control during prey capture in freely moving mice. Elife, 9:e57458, 2020.
Jean-Paul Noel, Edoardo Balzani, Eric Avila, Kaushik Lakshminarasimhan. Stefania Bruni, Panos Alefantis, Cristina Savin, and Dora Angelaki. Flexible neural coding in sensory, parietal, and frontal cortices during goal- direceee virtual navigation. bioRxiv, 2021.
Chethan Pandarinath, Daniel J O'Shea, Jasmine Collins, Rafal Jozefowicz, Sergey D Stavisky, Jonathan C Kao, Eric M Trautmann, Matthew T Kaufman, Stephen Ryu, Leigh R Hochberg, et al. Inferring single-trial neural population dynamics using sequential auto-encoders. Nature methods, 15(10))805-815, 2018.
Matthew G Perich, Patrick N Lawlor, Konrad P Kording, and Lee E Miller. Extracellular neural recordings from macaque primary and dorsal premotor motor cortex during a sequential reaching task. https: //renss///rcnscccc org/ 2018.
Sam Roweis and Zoubin Ghahramani. A unifying review of linear gaussian models. Neural computation, 1(2):305-345, 1999.
Virginia Rutten, Alberto Bernacchia, Maneesh Sahani, and Guillaume Hennequín. Non-reversible gaussian processes for identifying latent dynamical structure in neural data. Advances in neural information processing systems, 33:9622-9633 2020.
Peter H Schönemann. A generalized solution of the orthogonal procrustes problem. Psychometrika, 31(1):1-10, 1966.
João D Semedo, Amin Zandvakili, Christian K Machens, M Yu Byron, and Adam Kohn. Cortical areas interact through a communication subspace. Neuron, (02(1))24992))) 2019.
Ian H Stevenson and Konrad P Kording. How advances in neural recording affect data analysis. Nature neuroscience. 14(2):199-12, 2011.
JS Taube. Head direction cells recorded in the anterior thalamic nuclei of freely moving rats. 15(1): 70-86, 1995. doi: 10. 10.1533/EURRSSC.550001
1995.
Michael E Tipping and Christopher M Bishop. Probabilistic principal component analysis. Journal of the Royal Statistical Sociery: Series B (Statistical Methodology), 61(3):611-62 622, 1999.
Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, JJ Carey, ilhan Polat, Yu Feng. Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientifie Computing in Python. Nature Methods, 17:261-272 2020. doi: 10.1038/s41592 192-019-066--22
11

---

Published as a conference paper at ICLR 2023
Andrew Wilson and Ryan Adams. Gaussian process kemnels for pattern discovery and extrapolation. In International conference on machine learning, pp. 1067-1075. PMLR, 2013.
Andrew Wilson and Hannes Nickisch. Kernel interpolation for scalable structured gaussian processes (kiss-gp). In International conference on machine learning, pp. 1775-1784. PMLR, 2015.
Byron Yu, John Cunningham, Gopal Santhanam, Stephen Ryu, Krishna V Shenoy, and Maneesh Sahani. Gaussian-proces factor analysis for low-dimensional single-trial analysis of neural population activity. Advances in neural information processing systems, 21, 2008.
Ding Zhou and Xue-Xin Wei. Learning identifiable and interpretable latent models of high- dimensional neural activity using pi-vae. Advances in Neural Information Processing Systems, 33: 7244,724, 2020.
A APPENDIX
A.1
BACKGROUND ON PPCA, PCCA, AND THEIR RELATION TO TAME-GP
A.1.1 CANONICAL CORRELATION ANALYSIS
Given a random vector æ. PCA aims to find linear transformation such that the components of the transformed vector are uncorrelated. In other words, it tries to find linear transformation that diagonalizes the co-variance matrix of the random vectors. Similarly, CCA starts from two random vectors #1 and Ø2 of dimensions m, and m2, and tries find two linear transformations U E IRm Xm and V E RR such that each component of U æ1 is correlated with a single component of Ø2. terms of correlation matrix, this corresponds to,
[pi if
corr (U (U. V x2) = =10
(8)
otherwise
where are called canonical correlations.
En Σ2
Σ12
Letting the joint empirical co-variance be E
Ema
turns out that CCA projections are the singular vectors of the correlation matrix re-scaled by the inverse square-root of the individual co-variances. Namely, if ü; 0; are the i-th singular vectors of the correlation matrix coir (æ1, x2) / / then the canonical vectors are 1/21 u, Eg 5-11/2g), ú.). The two projection matrices U and V are obtained by stacking the canonical vectors; is immediate to verify that Lm, U LNUU, Lmz V V 2 V and P VTVV, where P mi m2 diagonal matrix with diagonal entries the canonical correlations.
(U, Vi) = -ET1
Again, making the parallel with PCA, we know that the first PCA vector is the eigenvector of the em- pirical co-variance corresponding the largest eigenvalue and satisfies wi argmax cov(2)
||w||=1 Similarly, can be shown that the canonical vector corresponding to the largest singular value of the correlation matrix satisfies,
(U, Vi) argmax corr(u T -0 T ..))) || =1/e|l1
(9)
Finally the n-th canonical vector satisfies,
(Un, Vn) argmax corr(u 20,, TT.)) weul- WEVE
(10)
with 21 {u |ul|| 1, E (UI...-Un-1)) } and V {v ||a|||
1,
(V,, Vn-1) 1)]
12

---

Published as a conference paper at ICLR 2023
A.1.2 THE PROBABILISTIC INTERPRETATION OF PCA AND CCA
As first shown by Tipping and Bishop Tipping Bishop (199). PCA can be expressed in terms of the maximum likelihood solution of the following probabilistic latent variable model,
p(z) ~ 1) p(=)2) N(W= + pr, I,,
(11)
(12)
where / is the x D identity matrix, W is x D projection matrix, j E R. is an intercept term and o² a positive constant.
Similarly, Bach and Jordan Bach & Jordan (005) showed that the canonical directions emerge from the maximum likelihood estimates of a simple probabilistic model,
(D)
min(mi, m2)
(13)
21 2 ~ N(W2= + +1,, V1)
1 č 0
(14)
2/2 N(W2= μ2, V2)
1/2 >
(15)
where we use notation similar to that of equations [11 12 for the projection weights, the intercept and the identity matrix, while , and 1/ are generic positive semi-definite N x N matrices.
We will refer to these models as the probabilistic PCA and probabilistic CCA, or pPCA and PCCA.
Too better highlight the link between CCA and PCCA we report the MI. estimates of the PCCA model parameters,
W, EIUM
(16)
We = E22VM VM
(17)
V = Ê1 W.
(18)
2 E22 W.W,
(19)
1
AI N
Σ
(20)
P22
N
1
Σ
T2j1
(21)
where Mj are arbitrary D x D matrices such that M, P, the diagonal matrix of the canonical correlations, and V the canonical directions.
The posterior means and co-variances are given by.
E| Φ1 M (I1 P11)
(22)
E[ [2] = MºT V (x22 (U2)
(23)
cov (x(æ1) I - M,MTT
(24)
cov (z|x²) M2MM
(25)
M,
(T- -2)-1
(2-p2) -2)-
UT (m) pa))
2 [[22(11,022)
(26)
M2
(2)
VT (x2 p2)
M1
(1- 11-12)-1
(1-p2)-13
Mi
Cov (2|01, 222)
M2
(2-
(T- p2)12
M2
(27)
13

---

Published as a conference paper at ICLR 2023
It is important to notice that, independently of the M1 and M2 matrices, the observation projected into the D-dimensional subspace the canonical directions. See Beespishpp & Nasrabadi (2066) for similar argument bridging PCA and pPCA.
A.1.3 TAME-GP COMBINES AND EXTENDS THE PPCA AND PCCA GENERATIVE MODELS
The probabilistic interpretation of PCA and CCA, Eqs. (1 allows (1) extending the model to non-Gaussian observation noise, (2) replacing the normal prior over the latent with a smoothing GP-prior, and (3) combining the two graphical models a more general framework.
In particular, TAME-GP assumes shared latent factor 20)) with a GP prior that captures fine time scale correlations between some continuous task variables of interest (modelled as conditionally Gaussian) and the spike counts from multiple brain regions (modelled as conditionally Poisson). This approach extends the ideas of pCCA to the analysis of spike trains driven by smooth temporal dynamies. Further, we extended our graphical model by including additional area-specific latent factors 2 (GP-distributed). The projection associated with those factors aim specifically to capture the residual inter-area co- co-fucttations, in close resemblance to the role of the pPCA projection weights.
The general formulation of the TAME-GP generative model is given by Eqs.1-3 1-3 in the main text.
A.2 INVERTING THE HESSIAN THE JOINT LOG-LIKELIHOOD
The dimensionality of the individual latents and trial duration pose computational challenges for TAME-GP approximate inference. For each trial, evaluating the posterior covariance requires inverting the Hessian of the joint log-likelihood, of dimensionality D x D, where D = E, dj. dj
the dimension (i) and 7 is the number of time points of the trial (for simplicity. we assume all trials the same length here, but the implementation allows for variability in trial duration). Hence, a naive implementation of the posterior estimation would require O (D*) operations (the cost of inverting D-dimensional matrix). Nonetheless, the specific conditional independence assumptions of our model allow us to speed this computation by using the block matrix inversion theorem. In particular, if we define
V (h) V (2()) log logp(z, H
H has the following structure,
Hon
Ho1
Hog
Hon
Hg
H11
0
0
H
Hgz
0
H2
0
0
Hnn
therefore, can be inverted according to,
A CT (A (A-CTB-IC)
-A-CCTB-0Q-ICTT- B CA- B-44--CCACCC--000C
C
Hn
Hil
Huz
0
by setting A Hoo and B
and C :
computing B-1
H
Han
requires only inverting the block diicooaa elements, while (A CT DTTB-CC -'C) Hoo. achieving inversion of H in O Σ, d3) operations.
has the same size as
14

---

Published as a conference paper at ICLR 2023
A.3 PARAMETER UPDATE DETAILS
Eg| and D (kaA)
= =1, 21))
(1) (k)
Introducing the notation μ4
Mi 14
we have
-1
C
Σ
(T) Vil
Σ 1
Yt
Σ
(T)T
Σ
Σ
(0.0)
Σ
14
(0) (O)
Σ
0
1
(T)
TM
TM
1.1
1,1
1,8
1,t
d TM
Σ yr P.r-cE
C.EC Mi
(0)
(Eva
V TM
Σ yiy.
Vi O0T CT C++CE
vilk
(0) y?
+ + =aEx, Yr
1.1
1.
1,1
+C
Σ 22 (0.0) +umll"))
CT+ +
C
c
(0) p
a
d
Σ
mi
(0)T
tT
+ Tdd
1,
1,
where = 1: M and f = 1: : T" are trial and time within trial indices.
A.4 LEARNING THE POISSON OBSERVATION PARAMETERS
In order to learn the Poisson observation parameters we numerically maximize E, log(p(x, y, z(00) as a function of (0.0)), W(d) and h()) Our implementation follows a Newton scheme which requires both the gradient and the Hessian of the optimization objective.
In order to simplify notation, we fix a unit i from population and we set
(
M
pe
I(0.0)
Σg
(0,))
W(Os) T
Ft Ft (
h (.)))
where We Rdo+d, and h IR. The corresponding gradient and derivative will be,
OF. (,,
Σ
Tifli N+W T T.W (14 + E,W)
(28)
OW
1,t
OE [log(p(x, y,
24 - T H T.w Σ
(29)
Oh
1,t
0 E. [log(p(x y, 2,
N+W f E.W [m + E,W) (ML + E,W) + +E,]
(30)
DW2
8 E. [pog(p(x y, 0))
N+W #. TIW (14 + E,W)
(31)
OhoW
B [log(p(x, y, z
N+W m.t.
(32)
0h2
where = 1, M and = 1,.... T are the trial and time indexes respectively.
0 (W1)) h C. L.d,v,7l T,)
15

---

Published as a conference paper at ICLR 2023
A.5 LEARNING THE GP TIME CONSTANTS
GP hyperparameters (time constant) are learned by gradient based numerical optimization of the joint log-likelihood. Following the notation of the main text we set, X -log(2r;)). and we define
kernel K ( R - RTxT such that, (1) =exp (t
The objective function takes the form,
trace (x(0))(2(00)) (l (2)=(0)) 2(T)) log (2()) + const, Σ
Fgg [log(p(xx))) =00)
where j = O,.... n. is the latent factor, = 1,... M is the trial number and i = 1, d2 is the component of z 2(0), Using the chain rule we obtain,
Bg (oo((((x,y,z)) trace
IE [lo[[((x y, v
OK())
OK!1
ax(1)
with
EE [log(p(x,)) y,
1 Σ
(KK + (5) xx(3)~1E,(( E,1x (0)
ok! (5)
2
0
K
to
(t 112 exp (- - ())))
0A
A.6 PARAMETER INITIALIZATION
Factorized TAME. Before running EM on the full TAME, we obtain initial condition for the model parameters (all except the GP kernel hyperparameters) by means of rumning five iterations of EM for the temporally factorized version the model. In particular, we replace the GP-prior over the latents
with product of Gaussian normal distributions, i.e. p(z(²) II. p( (2()), ). and p(=)) (3) (0, 1).
Under this prior assumption the joint likelihood as a whole factorizes over the temporal axis (i.e. the observations are temporally independent given the latents). As consequence, the Hessian matrix of the joint pdf is sparse, and can be stored and inverted efficiently, allowing for the implementation of full Newton scheme to numerically optimize for the MAP estimate of the posterior ever latents
The EM-based optimization the factorized TAME also needs an initial choice for parameters. We found empirically that CCA-based heuristic works well for this purpose. Specifically, we set:
W(()) to the first du canonical directions V between the routee- mean-centered spike counts of population j. أاوو v Vx M4 and the task variables (My is the empirical mean of the rquare-rooted spikes).
W(id) as the first d, principal direction for the orthogonal complement of the counts w.r.t the canonical directions, $(5) x()) VTV This will initially enforce orthogonality
in the task relevant and private latent subspaces.
h (i) was set to the log of the empirical mean of the counts.
was set to the first dg canonical directions between and the square-rooted counts from all the neural populations, Y [y ()) :y:::.....))))
d was set to the empirical mean of and Y to the empirical covariance.
time constants. The initial GP time constants were drawn from a uniform random distribution 0.5.
A.7 COMPARISON OF TAME-GP AND SNP-GPFA
We compared our framework to that of SNP-GPFA Keeley et al. (2020). which identifies shared fluctuation between neural populations, under the assumption of trial repeats with common
16

---

Published as a conference paper at ICLR 2023
dtimulus-driven mean (corresponding dimensionality reduced peristimulus time histogram, or PSTH).
Briefly, the multi-area SNP-GPFA assumes that the spike counts of two areas, area A and area B, generated according to
Poisson (f (
WAA AA
X A,n7
(33)
B
BBB
[ B,n
with Y.A/B the spike counts of population A and for trial j. f the soft-max non-linearity; X A./B. are drawn from GP with factorized RBF covariance, which captures within area co- co-fluctuations for trial j; X corresponds to draws from another GP, which is shared across trials and populations, thus capturing the shared across area co-fluctuations.
We generated spike counts from the graphical model in figure S4A assuming fixed trial duration (necessary for the SNP-GPFAS in different conditions: 1) fixing the shared dynamics across trials (as in SNP-GPFA, figure S4B, top), or 2) varying the shared dynamics across trial (igure(S4BB bottom).
Specifically, for the first condition the counts followed (33 but replacing the non-linearity with an exponential. For the second case, the counts follow Poisson statistics of the form
WAR AA
0
X A.n
B
Poisson
exp
W
X
+
(34)
O
W BBB
x B,n
where we added trial dependency the shared Gaussian process factor.
We set the dimensionality the shared factor to 2, and of each private factors to 3. We simulated spike counts from two populations of 30 neurons for 50 trials, each having 100 time points with 0.05 second resolution. The average firing rate of both population was set to 10Hz. We fit the simulated spike counts with TAME-GP and SNP-GPFA for both conditions. The results show that TAME-GP captures the between area co- oo-luctuation in both scenarios while SNP-GPFA fails when the shared dynamics vary between trials, as expected by the model assumptions (Fig. S4C.E). We assessed the accuracy of the factorization of the spike-count variance by means of Lasso regression. In particular, we regressed the ground truth latents from the estimated latents of the different models, and quantified regression goodness-of offdness-offfttt in terms of cross volsdaaddd R (Fig. S4D.F.. We quantified the contribution of each latent factor to the regression terms of the magnitude the associated coefficients. Results (reported in Table ST) show that I) both models can factorize the variance when the shared dynamics are fixed across trials, with SNP-GPFA achieving cleaner decomposition (expected given that is a optimal
closer model of the true data generating process in this case); 2) TAME-GP achieves factorization when the shared latents vary across trials (as assumed by its generative model), while SNP-GPFA is unable find the appropriate decomposition. Overall, TAME-GP estimator proves more robust to deviations from its underlying model assumptions.
A.8 SELECTING THE NUMBER OF PRIVATE AND SHARED DIMENSIONS IN REAL DATA
We select the number of private and shared dimensions hyperparameters via grid search. A priori we set the maximum number of dimensions to be evaluated as the number of PCs needed to account for 80% the population variance (in this case, dimensions). Fig S6 shows estimates of model fit quality included in private and shared latents for the multi-area TAME-GP presented in 4l-.4 The results show well-behaved cross- validated R² landscape, with optimal dimensionalities (5, 5).
in real data by optimizing these
function of the number of dimensions
A.9 FITTING TAME-GP TO MONKEY PREMOTOR AND MOTOR RESPONSES DURING REACHES
We also tested our estimator on a publicly available dataset Perich et al. (2018) that records neural activity in premotor cortex (PMd) and primary motor cortex (M1) of macaques during sequential
17

---

Published as a conference paper at ICLR 2023
reaches (binned 10ms resolution). Specifically, the monkey controls on-screen cursor and is rewarded for moving that cursor to an indicated reach target, with multiple targets presented in trial. Since there minimal kinematic requirements for the reaching movements (e.g., very brief hold times), the monkey typically makes relatively smooth series reaches. As pPCA latent structure was very poor quality in this dataset, restricted our comparison between TAME- with task manifold aligned to screen position, and PGPFA, with latent dimensionality d=2 (Fig Visually, the latent structure extracted by TAME-GP seems to better capture the animal behavior, (in R terms @ how much information about the task variables can be linearly decoded from their respective latents for TAME-GE and PGPFA with variable latent dimensionality (Fig. [S7C, F). These results confirms that in this dataset as well, the TAME-GP task aligned manifold provides a compact account of neural variability with respect task variables of interest, which does not align with the overall axes of neural variability of the data. PGPFA needs substantially higher dimensional latent spaces (10d vs. 2d) to capture the same amount task-relevant neural variability.
asked
A.10 SUPPLEMENTARY FIGURES
B
2r
M
3h
LD [a.u
X
X
ground truth TAME-GP
2
time [sec]
time [sec]
time [sec]
time [sec]
Figure S1: Multi-area parameter reconstruction. TAME-GP generative model for three brain areas with shared interactions. B. D- Latent variables estimation for within model simulated data: ground truth latent factors and model posterior mcan 95% CI for all latent dimensions.
method
CCA
0.175
TAME-GP
0.150
0.100
0.075
5.5
10.1
23.1
Figure S2: Task-aligned latent dynamics reconstruction (extends Fig. 2C). Mean squared error between the true task relevant dynamics and model reconstruction based on the dimensional task relevant latent factor for CCA and TAME-GP, and the full 6 dimensional latent space for P-GPFA and PCA. In contrast, figure 2C shows the MSE based only on the first 2 principal latents for PCCA and P-GPFA; Error bars represent the mean + s.d. over 10-fold cross-validatio
*All decoding used I. Lasso regression, with 5-fold cross-folding for hyperparameter estimation.
18

---

Published as a conference paper at ICLR 2023
A
B
ground truth
TAME-GP
PPCA
POCA
private subssaace
1.00
0.5
-0.1
0.0
0.0
1.0
-0.2 0.0
subspace
LD-1
LD-1
LD-1
LD-1
"2d latents
10
ground futh
lask variables
LDS
ReLu
Poleson roise
spikes
pop.
latent mse
0.7
rata (H2]
pPCA
PODA
TAME-GP
04
Y,
time (sec]
Figure S3: Effects of model mismatch on latent structure estimation. A. TAME-GP graphical model for single area (top) and schematic for data generating process (bottom). LDS dynamics where transformed into firing rates through a ReLu non linearity (replacing the exponential used in Fig. additive instead of multipliative effects on the observed neural
2). that the latents have activity. B. Ground truth task relevant dynamics (green) and estimated low dimensional projection for TAME-GP (purple). pPCA (dark gray) and pCCA (light gray). Mean squared error between the true shared dynamics and the model reconstruction, mean + s.d. over 10-fold cross-folding. D.
Example single trial firing rate reconstruction.
A
rato 1-th nou. pOD
shared latent (constant trie)
privato latent
triel reppeated dynamics
W
exp
Z
Z
trials
w
rato 1-th nou.
sharad latent (varies per trial)
private Iatent
A
B
pial venyingg Byynan
W
B
TAME-GP
SNP-GPFA
D
2
&
0.0
sharnd
E
latent
1.0
&
0.0
share!
prisee
prinate
[sec]
[sec]
Ialent
Figure S4: Communication subspace estimation with SNP-GPFA and TAME-GP (extends Fig. 3). A. Scheme of the spike count generative model for trial repeated (top right) and trial varying (bottom right) shared dynamics. B.C. Ground truth shared dynamics (black lines) and model reconstructions (colored lines) for the trial repeated (B) and trial varying (C) conditions. D.E. Ground truth shared and private dynamics variance explained by model predictions for the trial repeated (D) and trial varying (E) conditions.; error bars represent mean standard deviation over 5-fold cross validation.
19

---

Published as a conference paper at ICLR 2023
15.9 Hz
TAME-GP Semedo 2019
10.7 Hz
5.1 Hz
CCA
0.8
0.6
t 0.4.
0.2
0.0.
2
3
4
latent dim.
Figure S5: Model fit of shared and task aligned dynamics. of the linear regression between the ground truth task aligned latent dynamics and the model MAP estimate for TAME (purple), PCCA (light grey) and reduced rank regression (dark grey). Extends fig. 3I the main text to multiple average firing rates.
PFC MST
0.04
0.02
private
R*
0.00
0.02
0.04
1
2
à
shared
Figure S6: Latent dimensionality selection, for the MSTd and dIPFC communication manifold analysis (extends fig. 41-K). Hcat-map of the mave-onn-neuron-ouu R2 of the spike count variance explained by a TAME-GP for different combination of shared and private latent dimensions. The upper bound dimensionality was set to the number of principal components needed explain 80% of the population spike count variance.
20

---

Published as a conference paper at ICLR 2023
A
position screen
TAME-GP
P-GPFA
TAME-GR
pos (cm)
MON NONN
ar
P.GPFA
PMd
posttion
pos, (cm)
latent dimensionality
E
position screen
TAME-GP
P-GPFA
F
0.6
Z
pos (cm)
10
2
0.4
0.2
0.0
5
18
position
pos. (cm)
latent dimensionality
Figure S7: TAME-GP manifold estimation for monkey premotor (PMd) and motor (M1) neural responses during reaching. A. Graphical model for PMd neural manifold aligned to 2d coordinates of hand on screen. B. Behavior and corresponding PMd latent population trajectories for four example individual reaches (colors), extracted with TAME-GE GE and P-GPFA. C. Lasso regression decoding of position: TAME-GP R² (purple) based on 2d task relevant latent. P-GPFA R (blue) estimates were obtained for range of latent dimensions (1- 10). D, E, F. Same as A, B, C for MI.
21

---

Published as a conference paper at ICLR 2023
A.11
SUPPLEMENTARY TABLE
Lasso results
model
sim type fixed
ground truth latent private A
model latent
[
SNP-GPFA
private
0.252458
private B
0.008377
shared
0.065025
private B
private A
0,001043
private
0.373415
shared
0.012722
shared
private
0.014438
private
0.026413
shared
0.665547
variable
private A
private
0.128838
private
0.256987
shared
0.044492
private B
private A
0.046482
private
0.332019
shared
0.114382
shared
private
0.058657
private
0.386729
shared
0.010308
TAME-GP
fixed
private A
private
0.222231
private
0.00472
shared
0.16166
private B
private A
0.016308
private
0.419728
shared
0.02365
shared
private
0.1016
private
0.006841
shared
0.476519
variable
private A
private A
0.268177
private
0.011 153
shared
0.032323
private B
private
0.003969
private
0.411102
shared
0.020695
shared
private
0.019493
private 0.005826
shared 0.658865
Table $1: Lasso regression coefficients, related to session A.7 Norm of the coefficients of the Lasso regression between the ground truth latent dynamics and the SNP-GPFA/ TAME-GP predicted latents. Lasso hyperparameters are set by grid search with 5-fold cross-folding procedure.
22


Reviews:
[{"review_id": "oe-rhW63N4", "reviewer": "Reviewer_RPBA", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a probabilistic approach for learning interpretable task-relevant neural manifolds that capture both intra- and inter-area neural variability with single trial resolution. Task Aligned Manifold Estimation with Gaussian Process priors (TAME-GP) incorporates elements of demixed PCA and probabilistic CCA into a graphical model that additionally includes biologically relevant Poisson noise. The model uses a Gaussian Process (GP) prior to enforce temporal smoothness, which allows for robust reconstruction of single-trial latent dynamics. Experiments using both synthetic data and neural recordings from rodents and primates during naturalistic tasks demonstrate the robustness and flexibility of TAME-GP in comparison to alternative approaches. ", "summary_of_the_review": "This is a solid piece of work. The paper has solid technical contributions which have been adequately demonstrated through extensive experiments and indepth discussions. The proposed TAME-GP could be very useful as solution/tool in practice for dissecting sources of variability within and across brain areas during behavior.", "strength_and_weaknesses": "Strengths:\nThe paper is clealy presented with excellent use of English. The technical contributions are significant. The pros and cons as well as the broader impact of the proposed method are discussed in depth.\n\nWeaknesses:\nIn Section 2, the notations may need to be explained a bit more.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "h3kTMwK3Vcx", "reviewer": "Reviewer_4RNV", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Neuroscientists are generating rich datasets with population recordings from multiple brain areas simultaneously while the animal is engaged in a complicated behavior. It urges for an analytical tool to extract interpretable information from these rich datasets. Latent variable models are popular approaches for generating a compressed summary of such data. However, the interpretability of the latent model is limited. In the current paper, the authors presented a novel latent variable model which combines dPCA and pCCA into a graphic model to learn task or stimulus-relevant latent variables of neuron population data. The author demonstrated the model successfully learns interpretable latent variables that could account for both inter- and intra- area variance. Overall, the current model achieved superior performance compared to several population models. The message of the paper is clear and the results are solid to me. ", "summary_of_the_review": "Overall, I recommend accept the current work as the authors have not only developed a new model, but also carried out thorough study on the model.  ", "strength_and_weaknesses": "Strength: \n1.\tThe authors have a good sense in demonstrate the interpretability and application of the new model using various simulation data. I really like the part about the communication subspace (figure 3). The fact that the current model achieved better performance compared to Semedo 2019 and Keeley 2020 for both simulation datasets with or without trial repeats is exciting. \n2.\tThe authors have compared the current model with multiple popular models in fitting both simulation and real datasets, which makes the overall results more convincing. \n\nWeakness: \n1.\tSome recent RNN-based latent models eg. LFADs and Oerich 2020, were overlooked in the current manuscript. It would be great to discuss those. \n2.\tIt is not clear to me whether such a model could generate novel knowledge or testable hypothesis about neuron data.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "TkYP6EdASZ3", "reviewer": "Reviewer_1fJi", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "In this submission, the authors proposed a method to partition neural response variability within and across brain areas.", "summary_of_the_review": "On the one hand, it’s a decent paper. I don’t see fatal errors in this paper. However, I also don’t feel the work to be too exciting, so perhaps not quite up to par for ICLR.  Overall, I consider this paper to be slightly below the border of acceptance. But I’d also be ok if this paper gets accepted. \n", "strength_and_weaknesses": "****Pros: \n- The paper is generally well-written and, for the most part, clear. \n- Understanding inter-area communications represents an interesting problem in neuroscience. This is generally a less studied topic, but has attracted more attention recently due to the increased experimental capacity to monitor the activities from multiple regions.  \n- The proposed method is pretty straightforward. \n- Several applications were considered, based on both simulations and neuroscience data.\n\n\n****Cons:\n- The way proposed in the paper to model the communication is not totally convincing. The signals across the different regions are only allowed to interact via a multiplicative modulation. It is unclear whether this is reasonable for neural systems. \n- The technical contribution presented in this paper is rather incremental. The model is largely linear and assembles multiple components from prior work. Most of the components are fairly standard. The main innovation seems to the interpretation of Eq (1) as a multi-region communication model (but see my next comment).\n- It is unclear what’s special about the inter-area communication. Eq. (1) could also be interpreted as a model for a single brain area with increased latent dimension. This needs to be further clarified.\n- The real neuroscience data applications are detailed but unconvincing. Maybe the authors could articulate better what new insights were revealed. This is particularly important given the rather incremental technical contribution. \n- The authors showed examples in which case the model could be recovered. Are there also cases for which the model components can not be recovered? This seems to be an important question for interpreting the output of the model, but it is studied in the paper.\n\n\nOther comments:\n- Practically, how to set the number of latent dimensions in each area?\n- The title should be changed. It is unclear what “INTRA- AND INTER-AREA NEURAL MANIFOLD” really means. I understand that “neural manifold” is good for branding, but that should come with the cost of sacrificing the accuracy and faithfulness of the title.\n- It is unclear how to think about the mouse ADN recording in the context of inter-area or intra-area communication. This needs a clarification. \n- There should be a more thorough treatment of the literature on using the task variable to guide the estimation of latent variables. A few papers that would be useful to discuss:\nSani, Omid G., et al. \"Modeling behaviorally relevant neural dynamics enabled by preferential subspace identification.\" Nature Neuroscience 24.1 (2021): 140-149.\n\nZhou, D., et al \"Learning identifiable and interpretable latent models of high-dimensional neural activity using pi-VAE.\" Advances in Neural Information Processing Systems 33 (2020): 7234-7247.\n\nHurwitz, Cole, et al. \"Building population models for large-scale neural recordings: Opportunities and pitfalls.\" Current Opinion in Neurobiology 70 (2021): 64-73.\n\n- Fig 1a. The model variables need to be better explained in the figure caption/legend.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "3PF1ZkQh0G", "reviewer": "Reviewer_7WyW", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "A method (TAME-GP) to capture task-related variance in neural firings observations is presented. The method has its roots on probabilistic canonical correlation analysis (pCCA), but proposes an extension for non-Gaussian noise. This is achieved by replacing pCCA’s prior that assumes a normal distribution with a Gaussian process as well as adding an explicit hidden (latent) layer z for relating inputs x (spike counts) and the output y (task variables).\nTwo verification efforts on synthetic data and two validation efforts on experimental data are provided. Moreover, of the latter two one is of segregational nature (local activities), and the other of integrational nature (communication between regions).\n", "summary_of_the_review": "+ Including the latent layer makes me wonder, how would hidden graphical models e.g. hidden Markov models or others sophisticated variants, perform compared to TAME-GP?\n+ Regarding the aforementioned weakness on the relation of the mathematical decisions having a physiological or functional counterpart, why n+1 latent variables? Why Gaussian processes? Why in the expectation step not using a transformation  e.g. Anscombe or Freeman-Tukey or other, rather than approximating the posterior again with a normal distribution? What would be the physiological interpretation of the model hyperparameters (for instance, but not limited to; the linear approximation Cz^j +d, the \\tau, etc)?\n+ Has Procrustes also been applied to the other methods being compared in Fig 2B? It seems that at least P-GPFA and pPCA are not correctly aligned?\n+ The order of presentation of some mathematical elements is confusing. For instance. Parameter K is introduced in pg 3 during the expectation step but not defined until the end of the maximization step pg 4.\n+ As one of the advantages of TAME-GP over pCCA, it is stated that pCCA only considers pairwise relations. This somewhat implies that TAME-GP handles relations of arity higher than 2 between regions. The k_j terms are clearly pairwise among latent z variables. It is not clear to me how this propagates to brain circuits of more than two regions. \n\nMinor details\n+ Cross-validation: Popular terminology but mathematically inaccurate. Suggestion: internal validity using cross-folding.\n+ In the discussion: Causality of the outcomes is claimed, but I can’t see how any of the traditional causal assumptions for any popular probabilistic causal definition (Spirtes-Glymour, Pearl, Suppe, etc) are shown or demonstrated here.\n+ Appendices A.1.1 and A.1.2 are possibly unneeded. References should suffice.\n", "strength_and_weaknesses": "Strengths\n+ The novel way to generalize pCCA to non-Gaussian noise model.\n+ Results do accompany; the new model not only captures more observed variance but also appears to better decode task-related information.\nWeaknesses\n+ It is unclear how the mathematical decisions match physiological and or functional rationales, which means that even if the method works, its interpretation remains cumbersome.\n+ Conciseness….or the lack of it.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%