Nonlinear multiregion neural dynamics with parametric impulse response communication channels
Authors: Matthew Dowling
Cristina Savin
Venue: ICLR 2025 Spotlight (2025)
Keywords: neural dynamics
multiregion
variational inference
Primary area: applications to neuroscience & cognitive science
TL;DR: a multiregion neural population dynamics model capable of extracting nonlinear neural population dynamics with communication channels between regions parameterized by their impulse response.

Abstract: Cognition arises from the coordinated interaction of brain regions with distinct computational roles. Despite improvements in our ability to extract the dynamics underlying circuit computation from population activity recorded in individual areas, understanding how multiple areas jointly support distributed computation remains a challenge. As part of this effort, we propose a multi-region neural dynamics model composed of two building blocks:  _i)_ within-region (potentially driven) nonlinear dynamics and _ii)_ communication channels between regions, parameterized through their impulse response. Together, these choices make it possible to learn nonlinear neural population dynamics and understand the flow of information between regions by drawing from the rich literature of linear systems theory.  We develop a state noise inversion free variational filtering and learning algorithm for our model and show, through neuroscientifically inspired numerical experiments, how the proposed model can reveal interpretable characterizations of the local computations within and the flow of information between neural populations.  We further validate the efficacy of our approach using simultaneous population recordings from areas V1 and V2.

Published as a conference paper at ICLR 2025
NONLINEAR MULTIREGION NEURAL DYNAMICS WITH PARAMETRIC IMPULSE RESPONSE COMMUNICATION CHANNELS
Matthew Dowling & Cristina Savin Center for Neural Science
New York University {md6276, www.wneeny.nucnnonoooo
ABSTRACT
Cognition arises from the coordinated interaction of brain regions with distinct computational roles. Despite improvements in our ability to extract the dynam- ics underlying circuit computation from population activity recorded in individual areas, understanding how multiple areas jointly support distributed computation remains challenge. As part of this effort, we propose multi-region neural dynamics model composed of two building blocks: i) within-region (potentially driven) nonlinear dynamics and äij communication channels between regions, pa- rameterized through their impulse response. Together, these choices make pos- sible learn nonlinear neural population dynamics and understand the flow of information between regions by drawing from the rich literature of linear systems theory. We develop state noise inversion free variational filtering and learning algorithm for our model and show, through neuroscientiffcally inspired numerical experiments, how the proposed model can reveal interpretable characterizations of the local computations within and the flow of information between neural pop- ulations. We further validate the efficayy of our approach using simultaneous population recordings from areas VI and V2.
INTRODUCTION
Perception, choice and action engage neural circuits distributed across the brain (Chen et al. 2024 Khilkevich et al., 22024, Noel et al.) 2024 Pinto et al.) 2022 Machado et al. 2022 Ebrahimi et al, 2022). Despite technological advances that facilitate recording from multiple, anatomically distinct, populations of neurons (Steinmetz et al., (021), understanding neural computation at the level of multiple interacting populations remains statistical and theoretical challenge. Making progress requires new theoretical frameworks describing how global computations arise from multiple inter- acting circuits (Perich & Rajan 2020). each with potentially complex local nonlinear dynamics, and new statistical tools that extract such structure directly from recorded neural activity during behavior.
One prominent of approaches for measuring interarea interactions based on neural data focus on communication subspaces (Semedo et al., 2019). Rather than modeling local circuit dynamics ex- plicitly, these approaches aim to partition population response variability into 'private' dimensions, that are local to an area, and 'shared' dimensions reflecting the flow of information across areas. In its simplest form, this partitioning formalized as low-rank regression or canonical correlation analysis, for directional or undirectional communication, respectively (Semedo et al.) 2020). Ad- ditionally, Gaussian Process (GP) priors for the latents enforce temporal regularities, and explicitly model features like communication delays (Gokcen et al. 2022 2024), frequency and phase delays (Li et al. 2024 or additional task-relevant covariates (Balzani et al. 2023).
Building upon decade of progress in latent state estimation from neural population activity (Panin- ski et al. Cunningham & Yu 2014:) Duncker Sahani 2021), other approaches directly model the dynamics within each areas and the interactions between them. The simplest such models use linear dynamical systems (LDS) for capturing within area dynamics. For instance, gLARA (group
"Center for Data Seience, NYU

---

Published as a conference paper at ICLR 2025
latent autoregressive analysis) (Semedo et al. 2014) assumes that within- and between- popula- tion dynamics are both governed by LDSs More recently, the state-space representation of finitely differentiable GPs (Li at. (022)) blurs the distinction between GP-prior based communication sub- spaces and L.D methods, although these are mainly leveraged for efficient inference. Oftentimes, multi-area approaches can be seen special cases of single area models, with additional parameter constraints. For instance, Glaser et al. (2020) adapts recurrent switching linear dynamical systems (rSLDS) to construct multi-population sticky iSLDS (mp-srSLDS) of neural dynamics. This al- lows for nonlinear within area dynamics and instantaneous linear information flow between them. The most complex multi-area model is MR-SDS muulti-region switching dynamical systems), which uses neural networks to parametrize arbitrary nonlinearities for within area dynamics and across ar- eas communication and uses switching to capture global transitions between such nonlinear systems to model behavioral state "Karniol-Tambowt et al. 2022). Closest to the circuit level, multi-region recurrent neural networks (RNNs) can be fit directly to single neuron responses Perich et al. (2020), which provides direct current estimates but leaves understanding the low dimensional dynamical systems structure of the solution to post-hoe investigation. Overall, different approaches provide
different trade-offs between flexibility and interpretability (see Appendix Table [SI]. None of the existing methods fully reflect the nature of distributed computation as formalized in current circuit
level theories (Bredenberg et al. 2024 Langdon al. 2023
Misić & Sporns 2016).
Here we develop a probabilistic generative model that accounts for the nonlinear nature of neural dy- namics and characterizes communication between regions using channels that are parameterized by their impulse response blending expressive nonlinear region specific dynamics with interpretable characterizations of the flow of information between regions. Our major methodological contribu- tions include i) the generative model of latent neural dynamics that combines node-specific nonlin- ear dynamical systems parameterized by deep neural networks with linear communication channels between regions, which we term MRDS-IR (for MultiRegion Dynamical Systems with Impulse Response communication channels) ii) an end-to-end variational methodology, using a state-noise inversion free filtering algorithm, streamlining the treatment of approximate inference in state-space graphical models with hybrid stochastic/detemmiiitii transitions. Through several ncuroscientifi- cally inspired numerical experiments including integration, gating of information flow and rhythmic timing, we demonstrate the use of our approach to make sense of the underlying computation behind observed multi- polliiipoolltionn neural responses. We also show that approach reveals meaningful features of neural activity in joint population recordings from visual areas VI and V2.
2
MODELING MULTI-AREA NEURAL DYNAMICS DURING BEHAVIOR
2.1 BACKGROUND
State-space models. State-space graphical models provide a principled framework for data driven learning of neural population dynamics (Paninski et al., 2010). For a single neural population, recorded neural activity, yt E R. time, t, is modeled as reflecting a lower-dimensional popula-
tion latent state, Z; E R2, which evolves as a dynamical system parameterized by 0,
Ze = fo fz(z1-1, C2) + we (latent process) yt Z ~ p(ye z:) (observation model)
(1)
where w. ~ (O0 Q). and c, denotes (optional) inputs/stimuli.
Generalizing this formalism to simultaneous recordings from K regions, a natural choice is to par- tition the latent space into K groups of latent variables, with population responses in any given depending only on the latents of that region (Gokcen et al. 2022 Li et al.,
2024 Karnool-Tambbur
z where yf! and zi are the population activity and latent state associated with region k. respectively. Within this common structure, different multi-region models make different choices for the functional form of the latent space and the dependencies linking latents across regions. This structure determines not only the model's expressiveness but also its ability to capture crucial aspects of neural population dynamics. One important such feature is the latency in communication between regions, reflecting the time delays inherent in signal propagation, which is absent in most process models (Glaser et al.) finite state his- tory (Semedo et al.] 2014). To provide a more flexible framework for modeling signal propagation between regions, we consider principles from linear system theory.
al. 2022, Semedo et al. 2014). p(y, (k) z (2) ....,
= (R)
2020 Karniol-Tamii Tambour et al. 2022) or realized by introducing dependence
2

---

Published as a conference paper at ICLR 2025
Characterizing communication channels via their impulse response. The general premise of our approach for modeling communication channels is that sigmal propagation between regions can be well approximated by sufficiently expressive linear filters, allowing for propagation delays and temporal filtering, e.g. preferential transfer of information in a specifie frequency band (Bastos et al., 2013), while keeping the model tractable. Two fundamental concepts for understanding linear system are i) its impulse response and ii) its transfer function; they offer complementary perspectives on the system's 's input to output map, characterizing information flow in both time and frequency (Kailath 1980 Chen 1984 Brockett 2015). Consider an Ni dimensional input signal, u, driving a linear system to produce an Now dimensional output signal, Xt- The impulse response of the system, h, is a Nout x Ni dimensional matrix, with entry (i,j) given by the [xe]; output when [ the unit impulse. By superposition, X and u, are related by convolution so that,
X
Σ u,
(2)
T=-o
An alternative characterization, more suited to understanding frequency-don -omain properties, is the transfer function, which for discrete discrete-time systems is the Z-transform of the impulse response,
H(2)
Σ 2- h
(=
where the transfer function, H(z), also an Noun x N dimensional matrix whose (i,j) entry characterizes how frequency content changes from input dimension to output dimension i. If x(2) and u(z) are the Z-transform of x₁ and u, respectively, then in the Z-domain they can be related by x(2) = H(z)u(z). Importamtly, if the entries of Hzz) are all rational in and the degree of the denominator exceeds that of the numerator, then a miite-dimensionl realization of that system can be implemented by an LDS This means that for any strictly proper transfer function satisfying those properties, there exists tuple (A, B, C) that parameterize an LDS,
X = Cy whose impulse response and transfer function match h, and H(2) respectively, and can be written in terms of the LDS parameters as,
Yt = A 1 + Bu,
(4)
(5) Consequently, impulse response descriptions of communication channels can be directly incorpo- rated into state-space model descriptions of multi-region neural dynamics. For understanding tem- poral characteristics such as delays in communication channels, the impulse response can provide an informative description; how information may be attenuated amplified at different frequencies is better understood through the transfer function.
h = CA '
H(2) = C(2I A)
2.2 THE MRDS- IR GENERATIVE MODEL
We consider region specific latent states driven by their own recurrent dynamics subject to filtered content of other region's 's latent state history, with dynamics of the form,
(k)
+ Σ Hz.r(2102 + (c( + w((h)
fa (Z
where GR maps region specific stimuli/inputs to the latent space, and Ha,e transforms the latent state history of region l into an input to region k acting as directed and causa 3 channel that controls the transmission of information between regions. This space structure mathematically general, with many existing multi-region neural dynamics models in the literature special cases.
We model channels between regions, Ha.r, és as linear filters parameterized by their impulse response, as explained above, which allows us to build a fully Markovian representation in higher dimen- sional state-space (Aström & Wittenmark 2013). We structure the latent state-space according to the following coupled difference equations,
zi = (z((z > CR.
(k)
(K.6) + GRc (k) + +w
(K)
π (k.f) = A
(k.f) + (L,
(8)
'There exist infinite state-space model realizations of minimum state dimension (Rosenbrock 1970
2 Hence the lack of the D matrix that may be familiar in more general treatment (Chen 1984) "Causal" is used in the sense of systems theory, to mean that current outputs do not depend on future inputs.
3

---

Published as a conference paper at ICLR 2025
where 200)) are La dimensional region specific states, wt (0, Q+) 7i are LeMa,e di- mensional states of the channel from e to k, AI. is Mx,e x LęMx, BE, is MyMe, x L and Cx.e is L x Le We parameterize (-) as deep neural networks capable of learning highly nonlinear transition operators (specifically, we use the minimally gated unit (Heck & Salem (017)) similar Schimel et al. (2021).
(8)
(k.E)
). The channel parameters, (Ag,e, Bx,e: Ck.e). are learned alongside the other generative model parameters, and are constrained so that channel dynamics remain stable.
To streamline notation, we define r(+)) an extended latent state containing the latent state vectors of the LDS that processes messages entering node k, 1) (k,k+1)
( = % (1.1)
Pe (R,--1)
(R.K)
72
Yi
For further brevity we also define st == (z) (())), and adopt the notational convention that variables without a superscript represent the concatenation of all variables with that name,
24 == % Zi
(K)
r. (
r(K)
S = (zz, I+)
So that the full latent dynamics model, po (St
8t 8t-1) = N(s mo Q) factors as,
(s 8t-1) = II Po (z, (4:)
2 1 (())
(k)
II
(77 17-1 1--1,2
(k.0) (k.0)
(9)
where 8(.) is the Dirac delta function, and appears consequence of communication channel dy- namics having no noise component. More than notational brevity, this representation also simplifies the algebraic complexity of developing an efficient message passing algorithm for posterior infer- ence; since, as we discuss shortly, state-noise inversion free algorithm can be developed, allowing us formulate deterministic transitions as degenerate Gaussian distributions, or delta measures.
For the observation model, similar to other latent variable models of multiregion communication, each region's 's instantaneous activity made dependent only on the latent variables associated with that region. This leads to factorized likelihood, which in the linear Gaussian or Poisson GLM (generalized linear model) case we parameterize each region as,
ply:.) zz (y D&Z, (k) dg. R2))
(k)
(k)
(10)
ply; (R)
zz (1) Poisson (y exx((Dx2 (R) + dx))
(11)
An important concept worth noting, is that because one regions' latent variables do not instan- taneously affect another region's s activity, causal message passing algorithms should also not use observations of one region update the filtering belief of another. For linear and Gaussian observa- tion models, this structure arises naturally, but for amortized approximate inference (full details in Appendix D) we make sure adhere this principle when causally updating our filtered beliefs.
Parameterizing channels. The recurrent dynamics of between channel filters are parameterized as real representations of a diagonal matrix with M complex-conjugate roots,
Apt1 = diag LL II
be. I
IIL, -bk,,arI IL, IR.
(12)
Increasing M increases the order of the linear filter and makes it possible learn linear filters with increasingly nuanced frequency responses (as a result of adding additional pole-zero struc- tures) (Stoica al., 2005). Constraining Ax.i to be diagonal might first seem like a restrictive choice, however, with BR. and C2, free, this parameterization is able to capture any rational trans- fer function with greatest common denominator of order M and no repeated poles (Aokil 2013), and is the basis for Gilbert's method of constructing minimal realizations (Gilbert 1963). Additionally, considering that diagonal matrices are dense in the space of square matrices (Golub & Van Loan, 2013), it is not possible to learn non-trivial Jordan block structures through gradient descent without enforeing those structures.
We parameterize the complex conjugate roots of each block using their representation in polar co- ordinates and enforce stability during optimization through clipping a root's radius exceeds 1; however, alternative that allows for unconstrained optimization would be the stable exponential parameterization introduced in Orvieto et al. (2023). For the readout/readin matrices, Cx. and Bg, their parameters optimized without any additional constraints or structure. While we make this choice for practical simplicity, more sophisticated a priori pole-zero specifications could be intro- duced by considering the sparsity structure of these matrices (Kailath 1980 Kay 1988).
4

---

Published as a conference paper at ICLR 2025
Inference and end-to-end learning. While the choice of nonlinear dynamics for local popula- tion activity is well motivated by the inability of linear dynamics models to capture key features of neural computation, such attractor structure (Khona & Fiete, 2022), this choice also renders the exact posterior intractable - necessitating approximate inference. We approach this problem using an end-to-end variational inference methodology (Blei et al. (017). so that gradients of the evidence lower bound (ELBO),
[(() = > Egg [logp (y: C
Zi (4)
[KLL. ((5)
po (o(s))
| 1 1
(13)
can be used to optimize the parameters of an approximation g(s; St-1) = p(s: Vt), and parameters of the generative model (derivation of the ELBO in Appendix B). A Mante-Carlo
approximation of the ELBOO, suitable for gradient based optimization, can obtained by recursively sampling S g(s sن_1) ,) from the conditional variational approximation. For efficient sampling amenable to the reparameterization trick (Kingma & Welling 2014), we parameterize the variational conditional as the product of Gaussian potential, depending Y2, and the prior transition model by setting,
g(s 8; 8;-1) = (s: mije-1, P(tt-) x Φly: (x)) x po (s-1))
(14)
with (St 5 = Mss mg(s:-1), Q). This parameterization, inspired by conjugate potential amortized inference networks such as the structured variational autoencoder VAV) (Johnson al.) 2016), forces the backpropagated gradients of the ELBO to traverse through the latent dynamics model- an important component for learning meaningful dynamics models capable of long horizon
forecasts Karl et al. 2017 Klushyn et al.) 202T). Each Gaussian potential has the form,
Φly: B2) X exp (k(y))* S; +
(15)
When the observation model is not conjugate k(-) and K(-) are parameterized by neural networks whose parameters are learned maximizing the ELBO, but in the case of a linear and Gaussian obser- vation model, p(y, St) Ny(x Ds;, R), have optimal closed form solutions,
k(y:) = D RR-
K(y, (y.)) = D DDRR R
(16)
Now, given a sample 5/- forming the conditional Gaussian approximation statistics,
m -- my (s( )+ + Qge
P Q QK,
QK,)- Q
(17)
with g: = kg K, + K, QK,)-KK (Qk; - ms 1)). we can sample st ~ q(s sy 1)-
Having formulated the recursive belief updates without requiring inversion of the state-noise further allows us to treat hybrid/stochastic latent transitions similarly (Appendix (A). Proceeding with this recursion until time T produces a completely differentiable trajectory sampled from the series of causally constructed beliefs which can be used to evaluate the ELBO. We offer more in depth dis- cussion of the approximate filtering algorithm in Appendix D there, we also cover in greater detail how block structures appearing in k, and K, due to the multi-region observation model and deter- ministic channel transitions can be exploited to reduce the computational complexity of inference.
RESULTS
3.1 MRDS-IR RECOVERS GROUND TRUTH DYNAMICS
We first validated the efficacy of our inference algorithm by simulating a synthetic dataset with matched generative model structure (Fig. [A]. The synthetic three region system was crafted so that each area's 's recurrent dynamics are characterized by slow dynamic structures believed to play important roles for neural computation (Fig. D from left right a stable limit cycle produced by van der Pol's oscillator, stable spiral, and ring attractor), with limited set of connections between them, themselves modeled as linear with predefined impulse response functions (Fig. [[C and D. black). From this system we generate 1000 trials of 200 time points each and project each region's 's latent state to 100 dimensional observation space via a linear gaussian likelihood. We fitted a MRDS-IR model with three nodes and all-to-all connections between them, and a linear gaussian observation model to this simulated data, and assessed whether the estimated model could i) recover individual region dynamics ii) recover the linear filters between channels and iii) correctly identify whether a channel was 'open' or 'closed.'
Examining the estimated flow fields of each of the regions (Fig. TB, bottom row), one can see that our estimator was able to learn the true autonomous dynamical systems structure, up to expected model invariances(axis rotation, and re-scaling). Fig. IC shows the ground truth and recovered impulse
5

---

Published as a conference paper at ICLR 2025
A
60
60
(s(
C
true
model
E -
(U)
637.8
15
1.
437.0
2-
R pred. (1)
leareed
1C
0.00
time
20
40
zero pole
Figure| 1: Ground truth model recovery. A) Structure of the ground truth data: three nodes and the flow of information between them. B) Single node population dynamics for the (top) ground truth data and (bottom) learned dynamical system after fitting the model; trajectories from the autonomous dynamics in light gray. C) Learned and ground truth impulse responses for the open/used (left) and closed/unused (right) channels. D) Pole-zero plot of the impulse response channel (black: ground truth, red: estimates). E) negative ELBO and predicted neural responses R² for MRDS-IR (ours) and the linear (LN) and nonlinear (NL) baseline models, see text for details.
responses of the (left) open channels and (right) closed channels. We found that the recovered impulse responses match in periodicity, although the units of amplitude are arbitrary. Importantly, the model learned to prune inactive channels by setting their amplitude close to zero 4 The pole- zero plot in Fig. DD confirms that the estimated active channels have frequency responses matching ground truth. Reassuringly, we found that our model -which matches the true data statistics- is revealed as the better fit in model comparison using either the ELBO or the R² of latent trajectory predictions for a forecast horizon of 150 time points regressed to ground truth examples(Fig. [ where the latter metric helps to assess prediction capability of the learned dynamics. This is not a given, simpler model could in principle fit the data better (due to finite data, fewer parameters, and a smaller inductive bias), but we confirmed that our model fitting procedure can still recover a fuller description of the underlying multi-region population activity.
3.2 REVERSE ENGINEERING DISTRIBUTED COMPUTATION IN AN INTEGRATION TASK
Next, we tested the ability of MRDS-IR to reveal the principles underlying multi-region neural com- putation in a distributed temporal integration task (Fig. 2 which requires long time scales in the dynamics and gating of information flow between regions. Rather than engineering a multi-region dynamical system computation directly, we chose to use trained RNNs and reverse engineered their function using either the ground truth trained RNN parameters or the MRDS-IR corresponding esti- mates. We also included CURBD (Perich et al.) 2020) baseline comparison. Unlike the previous experiment, here there is a model mismatch between ground truth and the MRDS-IR estimator. Since CURBD uses RNNs for multi-region dynamics, it provides particularly stringent compari- son, but our explicit input conditioning which links the underlying dynamics to task meaning may help MRDS-IR to better identify the underlying computations behind the measured neural activity.
Concretely, the simulated circuit used three regions, each with a low-rank RNN architecture (Mas-) trogiuseppe & Ostojic 2018 Beiran al., 2023) that low-dimensional dynamics could be easily visualized and compared. The activity in each region k evolves as
(A) (1- A 1y + yr
W (
+ Σ W (0) + G&c,
(k) +c(()) (R) )
(18)
EGEt
(8) is input to region k read out linearly by Gg. We = M,NT and WE. M ke
where ot are low-rank within/between population weight matrices, of ranks and 2, respectively, E, ~ (O o'I). Each RNN region had 128 neurons with tanh nonlincarities, and linear readouts for region-specific outputs.
The task requires each of the three nodes to integrate their respective inputs (which constant over the course of a trial), with their computation gated by delays and region-specific go cues. Fig. 2 shows one particular example trial. At the beginning of each trial, continuous 'cue 1 and
Loosely defined, the closed channel amplitude is much smaller than local signals in the receiving area, although this can be precisely validated by model comparison.

---

Published as a conference paper at ICLR 2025
A
C
go 1
go 3
x fixed pl.
go input
comm. channols
cue
RNN
estim. dynamics estim.trajectory
RNN n²
MRDS-IR
CURBD
go1
cue 2
0.95
0.63
100
121
0.63
0.03
wWw go
3
0.88
0.82
1eg
que 2
gD
g0
R
CUC
19gio0 3
9 S
nouron 75
200
100
OLHBD
que sum
1.6
200
forociast
Figure 2: A distributed working memory circuit. A) A diagram of circuit connectivity. B) Struc- ture of single trial. Cues and determine speed of temporal integration; go signals start integra- tion in the corresponding region; computation in region 2 is gated by state of region 1 with temporal delay. C) Ground truth versus inferred dynamics comparison. (left) RNN autonomous dynamics and and example trajectories colored by the sum of the cues received by population 1 and 2: (middle) the dynamics learned and (right) corresponding single trajectories estimated by our MRDS-IR. D) Estimated inter-area communication compared to ground truth RNN signals and CURBD; numbers indicate goodness of fit measured by the R to single trial RNN currents. E) Comparison of MRDS- IR and CURBD in terms of ability to predict single trial future neural responses over forecast horizon up to 225 bins (R² of model predicted observations matched to true data, left) with one example neuron predictions for each region (right); dashed grey line marks beginning of forecasting window: black line shows ground truth neuron activity.
'cue 2' become active- with their amplitude proportional to a value, randomly chosen from the set (10, 25, 50, 75, 90), indicating different speeds of integration, determining the number of time bins taken for a linear transition from '-1' to '1' of the respective region's output. The 'go 1' event (transient, binary) starts the integration process in region 1. Then, 25 time bins after that signal saturates, region 2's activity should start ramping with speed determined by 'cue 2.' Finally, when "go is received, region 3 should start ramping with a speed given by the sum of 'cue 1' and 0.8 x 'uue 2". The RNN is trained by BPTT with a supervised objective to match this target mapping.
Reflecting the two-state nature of the outputs and the working memory aspects of the task, the low-dimensional dynamics learned by the RNN (Fig. c left column), have several fixed points that operate alongside the supplied input to produce the desired trajectories. Using neural activity from this RNN (all units) we then trained three region, all-to-all connected MRDS-IR, with 2- dimensional local latent states, L1 = L2 = 2, order M 1 filters for the channels, and a linear Gaussian observation model with diagonal noise. Fig. 2C (middle column) shows the phase portraits of the corresponding dynamical systems recovered by our model for each region. Similar to the ground truth RNN, the latent trajectories inferred by the MRDS-IR, Fig. D (right column), driven by a combination of fixed point structures and the inputs provided by the 'go' signals. For example, focusing Fig. [ the trajectories of both the RNN (left panel, bottom) and MRDS-IR (right panel, bottom), in region 3, after receiving the 'go 3' signal, are slingshot in the direction of the 'go 3' readout vector, with the autonomous dynamics then reducing their momentum according to the sum of cues for that trial.
When comparing the estimated communication ground truth inter-region RNN currents (cue- conditioned across trial averages), we also find a good match (Fig. 2D). As their RNN counterparts, the estimated channels show ramping signals that separate across the stimuli conditions and saturate early on in the trial. Similar to the trained RNN, the relative scale of channel contents also remains in proportion;
particular, the trained RNN chose to weakly use the to connection, something which the estimator also picks up Importantly, the match to RNN communication is substantially tighter than the CURBD baseline, trained on the same data with default pyerparaaeetess (Fig. 2D.
7

---

Published as a conference paper at ICLR 2025
go
QUe
go
Un
cue
off HWWMMM 200 bins
01
113
time
OFF
E
[[Hi (aB)
f = 0.017
1= 0.010
I = 0.007
MADS-IR CCA (0.95)
25 cehanaoo
AAAA
WMM
(0.79)
1 11mmmmm wm m0 time
-60
125 bins
100 bins
f
.2
Figure 3: Rhythmic timing task. A) A two-node version of rhythmic timing task: cue communi- cates target frequency, go the onset of the oscillation, with additional on-off switch gating outputs from region 2. B) Structure of an individual trial. C) two orientation angles showing input con- ditioned latent states produced by the learned autonomous dynamics of the MRDS-IR color marks different frequency cue values. D) Phase portrait of the dynamics for region 2 with overlayed latent trajectories, color coded by the value of the gating input. E) Recovered 2D latent state of region (top) and the output of the channel from region 1 to region 2 (bottom) for 3 example frequency conditions. F) Comparison to CCA, R² mee-uureeed match to ground truth above. G) Magnitude of Each line represents one com-
the frequency response for the channel from region 1 to region munication dimension between region to region 2 (since region 1 is 3-dimensional, region 2 is 2-dimensional, and a second order filter was used, hence 6 filters total).
compare R²). Moreover, the closed channels estimates close zero for method, but not for CURBD (Suppl. Fig S2), so the baseline would draw incorrect conclusions about the nature of inter-region communication in this dataset. Although it is not easy to derive the low-dimensional dynamical system structure of the multi-region dynamical system estimated by CURBD for a flow fields comparison, can directly estimate the quality of the solution in terms of the ability of the estimated dynamics to predict neural responses (Zhao & Park 20200; Hernandez et al. (2018) increasingly large horizons into the future (Fig. 2B). By this metric, our approach shows better ability to match ground truth RNN dynamics across forecast horizons and slower degradation of predictability over long time scales, reflecting a more accurate understanding of the dynamical system structure of the system.
3.3 MULTIREGION RHYTHMIC TASK WITH OUTPUT GATING
While the integration task involves interesting inter-area interaction structure, does not have a direct task equivalence to a known neuroscience experiment. For more direct neural equivalence, we considered a variant the rhythmic timing task introduced in Zemlianova et al. (2024). which is motivated by rhythmic timing experiments in primates (de Lafuente et al.] (022) (Fig. 3A). The oscillatory nature of the computations involves gives us an opportunity to not only expose a different form of local nonlinear dynamics than the ones considered before, but also allows
focus on the
spectral structure of the communication subspace between areas.
In the original RNN implementation of this task context/stimulus inputs indicate the desired fre- quency of a sinusoid (via a 'cue") to be linearly read out from the population activity after a 'go' trigger. For a multi-region variant, we assumed that the RNN passes its outputs to a second region, which aims to reproduce the same sinusoid, except, the second region has an "off""ff' switch - so its output should follow the same sinusoid region 1 produces when the switch is on, and output a value of 0 when the switch is off, simple form of behavioral output gating. Fig. 3 shows an example trial and the expected linear readout of the neural state for each region; each trial in the training data is chosen to have random period ranging from 60 to 150 time bins in intervals of 15. As in the previous experiment, the RNN was trained by BPTT with each region having 128 neurons, rank recurrent weights, and rank 1 cross region weights. We then fit an MRDS-IR model using latent dimensions L = 3 and Lg = 2 and order M = 2 channels with a linear and Gaussian observation model with independent noise for each region's 128 neurons.
8

---

Published as a conference paper at ICLR 2025
In Fig. BC, input conditioned samples from the learned dynamics model for region 1, show how the go imput propels the latent state into a plane of state-space producing oscillations at particular frequency- similar to computational mechanisms hypothesized to reproduce precise time inter- vals (Beiran et al. 2023 Zemlianova et al.) 2024). The learned dynamics of region 2 reveal latent states for on and off switch conditions that occupy a distinct region of state-space (Fig. BD. The communication channel from region 1 to 2 learned by MRDS-IR has clear periodic structure re- flecting the target periodicity (Fig. 3E,, which better matches ground truth RNN signals compared to CCA baseline (Fig. 7.. Finally, a detailed characterization of the channel frequency response (Fig. 33) reveals a passband at low frequencies (demarcated by the shaded region) which spans the range of frequencies that generated the data, and another passband at higher frequencies, presumably to allow for fast transitions that might occur at the time of the go signal if the gating signal 'on.' This suggests that our method can extract interesting spectral structure for inter-area communication.
3.4 V1/V2 RECORDINGS
Our final experiment asked whether our method can reveal neural correlates of cross-region com- munication from real neurophysiological recordings. To examine this, we considered simultaneous neural recordings taken from areas VI and V2 of macaque monkey as it observed gratings of different orientations on screen (Zandvakili & Kohn, 2015). This dataset has been used in other studies to examine the efficacy of intraregional models of neural signaling (Gokkenn et al., 2024] Li al. 2024), making it natural testbed for comparison. Like Li et al. (2024) and Gokkenn et al. (2024), we used spiking activity from session 106r001 126. To most directly model this data we used the Poisson GLM (generalized linear model) form of the likelihood (Weber & Pillow| 2017). To fa- cilitate comparison, the latent space dimensionalities follow those used in Li et al. (2024), and were chosen to be L1 = 3 and L2 2, and channels were parameterized all-to-all connectivity between regions. Since it is not clear how to best encode stimulus identity to match the biophysical circuit structure, chose to inputs that remained constant throughout the trial and had an amplitude proportional to the stimulus ID as a first pass.
M = 2 order filters, with
Fig. 4 shows the stimulus-averagee latent trajectories for area V1 (top) and area V2 (bottom) extracted by the model. A key feature of these latents is their oscillatory structure, reflecting the periodicity of the drifting grating experimental stimulus. This is apparent in single trial trajecto- ries (not shown), and remains visible in trial averages, despite potential across- trial fluctuations in phase. Moreover, different stimuli map into distinguishable regions of the latent space, reflecting the encoding of different stimuli. In Fig. TB. we examine the V1 inferred trajectories more closely; we performed linear discriminant analysis (L.DA) (Bishop) 2006) in order to find a projection that maximized the variance between trajectories with different stimulus ids and minimized variance be- tween trajectories with the same stimulus id. The dimension of highest variance explained, shown in Fig. 43 (bottom) shows that trajectories can be separated according to their stimulus ID; while the other two dimensions, Fig. 43 (top), account for oscillations in latent space.
Investigating the interaction between the areas revealed similar oscillatory structure (Fig. qual- itatively consistent with previous communication subspace based estimates (Gokcen et al. 2024). Interestingly, the amplitude of these oscillations differed markedly between feedforward and feed- back information flow. As expected for sensory-driven initial communication: the V1 to V2 signal was most prominent immediately after stimulus onset, whereas the V2 V1 communication slowly ramped up over the course of the trial, as one would expect from a top-down cognitive refinement of local information. Overall, while not in themselves surprising given the details of the experiment, the model results confirm that the extracted within area dynamics recover interpretable, known fea- tures of visual neural activity in V1/V2, which cannot be readily identified with alternative methods (see DLAG estimates in Fug
Finally, we used the same data to quantitatively compare our approach to contemporary multi- region probabilistic methods. Specifically, we fit separately MRDS-IR, controls LN and NL. Sec. C.2). and MRM-GP (Li et al., (2024) to each of the 8 stimuli conditions. To keep the models equal grounding, we did not make use of stimulus conditioning in MRDS-IR and used linear Gaussian observation models for each region after square-rooting the spike counts to make their statistics more Gaussian Yu et al. 2009). Then, test set of 50 trials, we computed the MSE of a held-out set of neurons whose firing rate was inferred from a separate set of held-in neurons. Following a similar procedure Li et al. (2024), we used 10 random partitions of held-in/held-out
9

---

Published as a conference paper at ICLR 2025
A
B
D
R
stim1
V2
comm
4
0.1
21
FF
22
time
V1
0.0
cure
DLAG MRM
time
200 ms
V1
E
MSE
OuTS
mim-gp
LN
NL
stim id
L V2
30
8
200 ms
time
250
stim
2
Figure 4: Distributed computation in monkey areas V1/V2. A) Stimulus averaged trajectories of the latent state associated with areas V1 (top) and V2 (bottom) B) Projection of the latent state of area VI using LDA, where the most significant dimension (bottom) shows trajectories separate in space across time depending on the stimulus ID and (top) oscillations explained in the other two dimensions. C) Inferred messages transmitted across channels for one example stimulus, the darker set of lines show the two-dimensional signal transmitted while the thin line shows their spectral envelope. D) Signal amplitude in feedforward (FF. VI to V2) and feedback (FB, V2 to V1) messages transmitted across channels averaged across all 8 stimulus conditions. E) MSE of held-out neuron predictions for each 8 stimuli for each method, averaged over 3 random seeds. F) R² of held-out neuron firing rate predictions (session 1071003p143, stimulus 1).
neurons, where 10% of neurons are held-out, and then reported the average across partitions. Across stimuli, MRDS-IR consistently achieved the smallest errors across models (Fig. [... (fg.(4t). Predicted re- sponses a single stimulus condition also proved competitive when compared to DLAG (Fig 4 F)
using published pre-pro pre-procsssing and R² of held-out neuron activity as a metric (Gokcen et al. 2022).
Overall these results confirm that our approach is competitive to and generally outperforms state of the art methods on real multi-region data.
DISCUSSION
Despite accumulating evidence of the complex distributed nature of across-area interactions during behavior, how different aspects the process are orchestrated across circuits remains poorly under- stood. To bring closer process model of the distributed circuit computations that give rise to observed neural activity, we proposed a new probabilistic generative model which combines non- linear dynamics that capture potentially complex underlying local computations with an easy interpret linear filtering model inter-area communication. This unique combination of features allows for an approximate inference algorithm that uses variational principles optimize model parameters and possible inference network parameters in an end-to-end fashion through maximiza- tion of the ELBO. Modeling channels between regions using impulse response descriptions of their ouput- relationship makes it possible to use mature and principled tools from linear systems theory to understand their properties. Across series of mootvvly datasets, our results demonstrate the utility of MRDS-IR in helping to determine the causal structure of multi-area com- putation, and recover known properties of inter-area communication in early sensory processing.
While we have allowed the model to consider all possible patterns of inter-area interactions, allowing the data to determine their relevance for explaining the observed activity, scenarios with limited data or additional known interaction structure it may be useful to enforce additional constraints on the architecture. Future work could incorporate further regularization promoting sparse connectivity structure, for example through automatic relevance determination group lasso penalties Yuuan & Lin 2006). Another possible future direction would be exploring more structured specifications of the readout/readin matrices of the dynamics and communication channels to equip the generative model with better inductive biases.
in Gokcen al. (2024), or
ACKNOWLEDGEMENTS
MD and CS were supported by the National Institutes of Health (NIH) under NIH Award ROINS127122 and NIH NINDS iRFINSI27122-01. We thank the anonymous reviewers and area chair for their helpful feedhack.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Masanao Aoki. State space modeling of time series. Springer Science & Business Media, 2013.
Karl J Aström and Bjöm Wittenmark. Computer-controlled systems: theory and design. Courier Corporation, 2013.
Edoardo Balzani, Jean-Paul G Noel, Pedro Herrero- Vidal,, Dora Angelaki, and Cristina Savin. A probabilisticc framework for task-aligned intra-and inter-area neural manifold estimation. In The Eleventh International Conference on Learning Representations, 2023.
Andre Bastos, Julien Vezoli, and Pascal Fries. Communication through coherence with inter-areal delays. Current opinion in neurobiology. 31:173-180, 2015.
Manuel Beiran, Nicolas Meirhaeghe, Hansem Sohn, Mehrdad Jazayeri, and Srdjan Ostojic. Para- metric control of flexible timing through low-dimensional neural manifolds. Neuron, 111(5): 739-753, 2023.
Christopher M. Bishop. Pattern recognition and machine learning (information science and statis- fics). Springer--erllg. Berlin, Heidelberg, 2006. ISBN 0387310738.
David M. Blei, Alp Kucukelbir, and Jon D. McAuliffe. Variational inference: A review for statisti- cians. Journal of the American statistical Association, 112(518):859-877) 877, 2017.
Colin Bredenberg, Cristina Savin, and Roozbeh Kiani. Recurrent neural circuits overcome partial inactivation by compensation and re-learning. Journal of Neuroscience, 44(16), 2024.
Roger W Brockett. Finite dimensional linear systems. SIAM, 2015.
Chi-Tsong Chen. Linear system theory and design. Saunders college publishing, 1984.
Susu Chen, Liu, Ziyue Aiden Wang. Jennifer Colonell, Liu D Liu, Han Hou, Nai-Wen Tien, Tim Wang, Timothy Harris, Shaul Druckmann, et al. Brain-wide neural activity underlying memory- guided movement. Cell, 187(3):676-69)) 2024.
John Cunningham and Byron Yu. Dimensionality reduction for large-scale neural recordings. Nature neuroscience, (7(11):1500-1555 2014.
Victor de Lafuente, Mehrdad Jazayeri, Hugo Merchant, Otto Gracfa-Garibay, Jaime Cadena- Valencia, and Ana Malagón. Keeping time and rhythm by replaying a sensory-motor engram. bioRxiv, pp. 2022-01, 2022.
Lea Duncker and Maneesh Sahani. Dynamics on the manifold: Identifying computational dynamical activity from neural population recordings. Current opinion in neurobiology, 70:163-170, 2021.
Sadegh Ebrahimi, Jérôme Lecoq, Oleg Rumyantsev, Tugce Tasci, Yanping Zhang, Cristina Irimia, Jane Li, Surya Ganguli, and Mark Schnitzer. Emergent reliability in sensory cortical coding and inter-area communication. Nature, 605(7911):713-725 2022.
Elmer G Gilbert. Controllability and observability in multivariable control systems. Journal the Society for Industrial and Applied Mathematics, Series A: Control, 1963.
Joshua Glaser, Matthew Whiteway, John Cunningham, Liam Paninski, and Scott Linderman. Re- current switching dynamical systems models for multiple interacting neural populations. Ad- vances in neural information processing systems, 33:1 4867-14478, 2020.
Evren Gokcen, Anna Jasper, João D Semedo, Amin Zandvakili, Adam Kohn, Christian K Machens, and Byron M Yu. Disentangling the flow of signals between populations of neurons. Nature Computational Science, 2(8):512-525 2022.
Evren Gokcen, Anna Jasper, Alison Xu, Adam Kohn, Christian K Machens, and Byron M Yu. Uncovering motifs concurrent signaling across multiple neuronal populations. Advances in Neural Information Processing Systems, 36, 2024.
Gene H Golub and Charles F Van Loan. Matrix computations. JHU press, 20133

---

Published as a conference paper at ICLR 2025
Joel C Heck and Fathi M Salem. Simplified minimal gated unit variations for recurrent neural net- works. In 2017 JEEE 60th International Midwest Symposium on Circuits and Systems (MWSCAS), pp. 1596. IEEE, 2017.
Daniel Hernandez, Antonio Khalil Moretti, Ziqiang Wei, Shreya Saxena, John Cunningham, and Liam Paninski. Nonlinear evolution via spatially-depende linear dynamics for electrophysiol- ogy and calcium data. arXiv preprint ftfi::/811.0245555 2018.
Matthew J Johnson, David Duvenaud, Alexander Wiltschko, Sandeep R Datta, and Ryan P Adams. Structured VAEs: Composing probabilistic graphical models and variational autoencoders. arXiv preprint
2:2016, 2016.
Thomas Kailath. Linear systems, volume 156. Prentice-Han Malti Englewood Cliffs, NJ, 1980.
Maximilian Karl, Maximilian Soelch, Justin Bayer, and Patrick van der Smagt. Deep variational Bayes filters: Unsupervised learning of state space models from data. In International Con- ference on Learning Representations, 2017. URL https:: http:://ppeevvviveee net. / forum?
id HyTgHL5xg
Orren Karniol-Tambone David M Zoltowski. E Mika Diamanti, Lucas Pinto, David W Tank, Car- los D Brody, and Jonathan W Pillow. Modeling communication and switching nonlinear dynamics in multi-region neural activity. bioRxiv, pp. 2022-09, 2022.
Steven M Kay. Modern spectral estimation. Pearson Education India, 1988.
Mohammad Khan and Wu Lin.
Variational Inference Converting Varia- tional Inference in Non- ConjuConee Models to Inferences in Conjugate Models. In Aarti Singh and Jerry Zhu (eds.), Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, volume 54 of Proceedings of Machine Learning Research, pp. 878-887. PMLR, 20-22 Apr 2017. URL htt.:tttsss/prrcceirrggg/.. httpr//gww.......r.reee.e.re.e..
wwwsss/v5//eet 7a.
Mohammad Emtiyaz Khan and Didrik Nielsen. Fast yet simple natural-gradient descent for varia- tional inference in complex models, 2018.
Mohammad Emtiyaz Khan and Hávard Rue. The bayesian learning rule. Journal of Machine Learn- ing Research, 24(281):1-46, 2023.
Andrei Khilkevich, Michael Lohse, Ryan Low, Ivana Orsolic, Tadej Bozic, Paige Windmill, and Thomas D Mlsic-Flogel. Brain-wide dynamics linking sensation to action during decision- making. Nature, pp. 2024.
Mikail Khona and Ila Ficte. Attractor and integrator networks in the brain. Nature Reviews Neuroscience, (311):7744,66,, 2022.
Diederik P Kingma and Max Welling. Auto Auto-Encoding variational bayes. ln International Conference on Learning Representution, May 2014.
6114
Alexej Klushyn, Richard Kurle, Maximilian Soelch, Botond Cseke, and Patrick van der Smagt. Latent matters: Learning deep state-space models. Advances in Neural Information Processing Systems, 34:10234410220 10245, 2021.
Rahul G. Krishnan, Uri Shalit, and David A. Sontag. Structured inference networks for nonlinear state space models. In AAAI Conference on Artificial Intelligence, 2016. Utthttss:://est
semant
org/Corrpsspppop 2D:2901055
Christopher Langdon, Mikhail Genkin, and Tatiana A Engel. A unifying perspective on neural manifolds and circuits for cognition. Nature Reviews Neuroscience, 24(6):363-377 2023.
Weihan Li, Chengrui Li, Yule Wang, and Anqi Wu. Multi-region markovian gaussian process: efficient method to discover directional communications across multiple brain regions. arXiv preprint arXiv:2402.02686 2024.
Timothy A Machado, Isaac V Kauvar, and Karl Deisseroth. Multiregion neuronal activity: the forest and the trees. Nature Reviews Neuroscience, 23(1 (0683334337 2022.
12

---

Published as a conference paper at ICLR 2025
Francesca Mastrogiuseppe and Srdjan Ostojic. Linking connectivity, dynamics, and computations in low-rank recurrent neural networks. Neuron, 99(3):609 623, 2018
Bratislav Mišié and Olaf Spors.. From regions connections and networks: new bridges between brain and behavior. Current opinion in neurobiology, 40:1-7, 2016.
Jean-Paul Noel, Edoardo Balzani, Cristina Savin, and Dora E Angelaki. supported by dynamic reconfiguration single unit functional connectivity prefrontal cortex of male macaques. Nature communications, 151):8573, 2024.
beliefs are
Antonio Orvieto, Samuel Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pas- canu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International Conference on Machine Learning, pp. 26670-26698. PMLR, 2023.
Liam Paninski, Yashar Ahmadian, Daniel Gil G. Ferreira, Shinsuke Koyama, Kamiar Rahnama Rad, Michael Vidne, Joshua Vogelstein, and Wei Wu. A new look at state-space models for neu- ral data. Journal of Computational Neuroscience, 1961-2):07 August 2010. ISSN
1573-6873. doi: 10.1007/s10827-0090017799 U s10827- 009- 0179-X
URL ottp:::///.cii.orr///o..oooo
Matthew G Perich and Kanaka Rajan. Rethinking brain-wide interactions through multi-region 'network networks' models. Current opinion neurobiology, 65:146-151, 2020.
Matthew G Perich, Charlotte Arlt, Sofia Soares, Megan E Young, Clayton P Mosher, Juri Minxha, Eugene Carter, Ueli Rutishauser, Peter H Rudebeck, Christopher Harvey, al. Inferring brain- wide interactions using data-constraing recurrent neural network models. BioRxiv, pp. 2020-12, 2020.
Lucas Pinto, David W Tank, and Carlos D Brody. Multiple timescales of sensory-evidence accumu- lation across the dorsal cortex. Elife, 11:e70263, 2022.
Prajit Ramachandran, Barret Zoph, and Quoc V Le. Searching for activation functions. arXiv preprint arXiv: 1710.0911. 2017.
Howard Harry Rosenbrock. State-space and multivariable theory. (No Title), 1970.
Marine Schimel, Ta-Chu Kao, Kristopher T Jensen, and Guillaume Hennequin. ilqr-vae: control- based learning of input-driven dynamics with applications to neural data. bioRxiv, pp. 2021-10, 2021.
Joao Semedo, Amin Zandvakili, Adam Kohn, Christian K Machens, and Byron M Yu. Extracting latent structure from multiple interacting neural populations. Advances in neural information processing systems, 27, 2014.
João Semedo, Amin Zandvakili, Christian Machens, Yu Byron, and Adam Kohn. Cortical areas interact through a communication subspace. Neuron, 102(():249-259, 2019.
João Semedo, Evren Gokcen, Christian Machens, Adam Kohn, and M Yu Byron. Statistical methods for dissecting interactions between brain areas. Current opinion in neurobiology. 65; 59-69, 2020.
Nicholas A Steinmetz, Cagatay Aydin, Anna Lebedeva, Michael Okun, Marius Pachitariu, Marius Bauza, Maxime Beau, Jai Bhagat, Claudia Böhm, Martijn Broux, et al. Neuropixels 2.0: A minia- turized high-density probe for stable, long-term brain recordings. Science, 3726653)). 2021.
Petre Stoica, Randolph L Moses, al. Spectral analysis of signals, volume 452. Pearson Prentice Hall Upper Saddle River, NJ, 2005.
Martin J. Wainwright and Michael I. Jordan. Graphical models, exponential families, and variational inference, Foundations and Trends in Machine Learning, 1 1(1-2): 2008 ISSN 1935-8237. doi: 10.1561/2200000001. URL httt:://dx.oooo http//xxxoooo.oooooooooooo http://pp.ooooo.o...ooooo//////. 1t11
13

---

Published as a conference paper at ICLR 2025
Alison Weber and Jonathan W Pillow. Capturing the dynamical repertoire of single neurons with generalized linear models. Neural computation, 2912)33200-3222 2017.
Byron M. Yu, John P. Cunningham, Gopal Santhanam, Stephen I. Ryu, Krishna V. Shenoy, and Maneesh Sahani. Gaussian-process factor analysis for low-dimensional single-trial analysis of neural population activity. Jourmal of Neurophysiooloy, 102(1))61-6635, jul 2009. doi: 10.1 10.1152/ jn.90941.2008
1wtp:e.11111.... 90941 2008
PMID: 19357332.
Ming Yuan and Yi Lin. Model selection and estimation in regression with grouped variables. Journal of the Royal Statistical Society Series B: Statistical Methodology, 68(1):49-61 2006.
Amin Zandvakili and Adam Kohn. Coordinated neuronal activity enhances corticocortical commu- nication. Neuron, 844)88277- 2015.
Klavdia Zemlianova, Amitabha Bose, and John Rinzel. A recurrent neural network for rhythmic timing. bioRxiv, pp. 2024-05, 2024.
Yuan Zhao and Memming Park. Variational online learning of neural dynamies. Frontiers in computational neuroscience, 14:71, 2020.
14

---

Published as a conference paper at ICLR 2025
A STATE-NOISE INVERSION FREE FILTERING
To maintain strict interpretation of channels between nodes as linear filters that deterministically process information from one region propagating to another, their recurrent dynamics do not have state-noise. For this reason, it is desirable to have filtering algorithm that does not require inversion
of the state-noise term. We accomplish this by parameterizing g(z, 2t-1) R p(z, 2;-1; as,
g 7t-1) Φ(y+) po (Zt 7z-1)
(19)
where (yy:) 2) is Gaussian potential in terms of Zc, so can be written as,
((y)) = exp (k.- z, K2z,
(20)
This allows sampling trajectories causally in time by exploiting conjugacy and recursively finding,
Pipe Pip-1 = Q ( + K, QK, -K, Q
(21)
m = m + Qg.
(22)
(23)
and then drawing sample z conditioned on the previous sample zi-1- Here, ge, is quantity we define as the residual information, and is given by,
ge k, (I K, (QK))-kK (Qk; mo
(24)
where we used the Woodbury identity to write,
m, t-1 = P 1/1-1 (Q (mo(2;11) ki)
(25)
= -Q- QK, (I K, QK,) -2 K QQ -a0000mm((a000 +.)
(26)
= .mo (2 + Qg.
(27)
With these expressions, the KL term in the ELBO,
DxLL (a) 11)- = [( me m (2)-1) Q (m - (Z(-1))
(28)
+ tr P P + log Q log P Pitt-1 L]
(29)
can be efficiently evaluated by rearranging and simplifying terms to find (K P(j-1) = (g/ Qg. tr(YK (KKTk) ) log rr
(30)
where
rr = (I+ K. QK,)
(31)
B ELBO
We use a causal amortized variational filtering procedure that operates by drawing samples recur- sively from approximations, ((; 22-1) ~ p(z: Zt ye). To bound the log-marginal likelihood using causally constructed approximations of the latent state, we first bound the one-step predictive
log-likelihood of observation through standard variational arguments
log plye 1 2 Equ [log p(y: 2,)] (g(z 1:11)llp(z z)-1))
(32)
= E(y: 2z-1)
(33)
where the bound is tight when (Zz z = p(z: -()). Using this bound, we proceed by applying log Ep(z.-1 (ex=1)(exp(-) to both sides,
log Ep ilv1. 1) [p(y: 21-1)] = p(y y1a-1)
(34)
2 Epta- (yn-1) [E(y; z1-1)]
(35)
allowing us to construct the following lower bound on the log-marginal likelihood of yı:r as,
log p(yr()
log p(y; V):-1)
(36)
2 E Epzz. (Ys--)) [E(y; 2x-1)]
(37)
This quantity depends on the intractable filtering distribution, so we approximate using the varia- tional approximation,
logp(y))r) Z Σ Ep (yin-s)) [Buge- [log p(y:
2z)1 - (g(z z; z:-1) p(zz 2;-1)]]
Σ By [log p(y; zz)] Eq..
z (z z(-1)
(38)
L(4)
(39)
Now, although this is fundamentally an approximation to the lag-marginal likelihood, we retain a
strict lower bound on the data log-marginal likelihood, since, shown in Krishnan al. 2016
15

---

Published as a conference paper at ICLR 2025
DLAG V1/V2 communication latents
MMM
time
Figure S1: Example of DLAG across-region latent variables for session 1071003p143 for stimulus 1; each column represent samples of the three-dimensional across-area latent trajectory and the direction of the arrow indicates the direction of signal propagation as inferred by DLAG. These do not show natural segregation into a FF-like early transient and a ramping up FB-like amplitude.
any variational approximation factorizing forward in time as ((2.r)) be used to lower bound log p(y1:)) as,
( II ((z
2; 2t-1), 1), can
F()) Egg [log ply: 2c]] gg [Dx (g(zz Z(-1)ll po 2;-1))]
(40)
the bound being tight when g(z, 2x-1) p(z, 1, 1t.T).. This means that, conditional varia-
tional approximations that are designed without using future data will lead to a looser bound. How- ever, the benefit is that the causal structure of the generative model is enforced during inference.
C MULTIAREA METHODS COMPARISON
Below quantify existing probabilistic models of multiple populations of neurons.
Table S1: Tabulating existing multipopulation models. Our model situates itself uniquely, allowing for nonlinear region specific dynamics but interpretable communication channels between regions.
method
single node latents
communication channel
conditions stimulus
reference
ours
nonlinear
impulse response
gLARA
LDS
lincar, autoregressive
no
Semedo et
2014
mp-srLDS
switching LDS
lincar, instant
dLAG
GP
GP (time-delay)
Bo
Gokeen d
2024
MRM-GP
GP (frqquency/phave delays)
no
4024
mr-s&LDS
switching nonlinear
switching nonlincar instantancous
yes
2022
CURBD
RNN
linear instantancous
B0
Perich ct al.
2020
C.1 DLAG V1/V2 COMMUNICATION LATENTS
We wondered whether qualitative evidence of feedforward/feedhack signaling similar to those vealed by MRDS-IR IR appeared naturally in the latent variables recovered by DLAG. In Fig. ST] we plot examples of those recovered latents suggesting those signaling characteristics absent.
C.2 LN/NL MODELS
To asses whether the introduction of temporal processing in the communication channels and lo- cal nonlinear dynamics actually led substantial improvements in estimation quality over simpler model choices, we considered two simpler models in which we removed the history dependence in the channels while preserving the local nonlinear dynamics (nonlinear, NL model) and additionally swapped nonlinearities for linear local dynamics (linear, LN model; equivalent to multi-region LDS), resulting dynamics of the form.
z
(M)
fr 1 +
(8) CC,
w(t))
(NL model)
Egk
(k) = F + - C2.e2 2,
w!
(LN model)
where fg and F5 are learnable nonlinear/linear recurrent dynamics functions.
16

---

Published as a conference paper at ICLR 2025
RNN
ours
CURBD
+ 2
1 < 3
<-1
2 <
100 bins
Figure S2: Inferred channels from MRDS-IR and CURBD as compared to the ground truth RNN currents. Three of the RNN inter-region connections are absent in the architecture and a forth (3 <- 1) is present but weakly functional, as a result of the training procedure. Our approach recovers the right scaling for all messages, whereas CURBD infers comparable size currents for all possible pairs of regions, missing the true underlying network communication structure.
C.3 MRDS-IR V1/V2 CHANNELS
Since the ELBO is not convex in terms of model parameters, we wondered whether the feedback- /feedforward structure of latent channels inferred by MRDS-IR was a feature that could be repro- duced across different seeds. In Fig. $3 we show 8 separate seeds and the channel contents inferred by each of those models sorted left to right by decreasing ELBO values on held out data notably, similar inferred structure is found across majority of seeds.
C.4 INTEGRATION TASK CHANNELS
In Fig. S2
show all 6 channels of inferred by MRDS-IR and CURBD compared the channels of the ground truth RNN that solve the task; examining the figure, we note that CURBD infers that the closed channels, not used by the RNN, have similar activity levels compared to the channels that were used by the RNN to solve the task.
D FILTERING UPDATES
Recalling that,
logp(y: 2z-1)
2 Exg Oly xc]] (xr.(((z:)) 1z-1)llpa(2) z(-1))
(41)
2r wright & Jordan 2008).
we write po (2,
1) and (2, 2x-1) through their exponential family representation WWin-
po [z) z 1) l((z)) exp (Z)( exp
A(Ao(Zi-1)))
(42)
<(z; 2t = l((z)) exp t(2:) - A(\rlc-1))
(43)
where h is the base measure, A(-) is the log- log-partition function, t(zz) are the sufficient statistics, and (z) with Ay--1 natural parameters of the conditional prior and approximation respectively.
Then, through the Bayesian learning rule (Khan & Rue 2023 Khan & Lin 2017 Khan & Nielsen 2018). the optimal variational approximation has natural parameters given by,
A l (2(2(1) V Maje-
E Ey,,,, [logp(y: 22)]
(44)
For linear and Gaussian models with p(y: 21)
N(zz, R), the RHS term can be solved for in
closed form, and after some algebra we get,
P-1
Q K.K.T
(45)
f|t-1
m t- P (Q a--moo kG)
(46)
where the terms natural parameter updates, k, and K,, are given by
k. D R Ryyy.
(47)
K,K, D 'D However, when the observation model is nonlinear available due to the implicit nature of the solution. In this case, we amortize inference by using neural networks to parameterize conjugate potential updates as in al. (2016); for example,
(48)
Gaussian closed form updates aren't 't
17

---

Published as a conference paper at ICLR 2025
VI-U2
Anto V2-V1
200 ms
decreasing ELBO
Figure S3: The stimuli averaged power content of communication channels for the area V1/V2 data across 8 different random seeds.
by directly parameterizing.
(2: 22-1) x Φ(y; Zz)) pg 7t-1)
(49)
where Φb Yt (kg, K2) parameterized by deep neural network. The parameters of the inference network can be learned alongside the generative model parameters through end-to-end learning using the ELBO objective. While for non-conjugate inference network could be parameterized in a completely black-box fashion, optionally, we can use the form of ideal implicit solutions to design better inference network. For example, when the observation model is Poisson and p(y: z[) = Poisson(y: exp(Dz, + d)), then the implicit and intractable optimal parameter updates are are given by,
non-Gaussian observations, the
k = (y: r, diag(r,)Dmo) Dmo(Z)11)
(50)
K,K, = D T diag(r;) )D
(51)
where, r exp (Dmo (z 1 diag (DQD T) + d). The way we amortize inference, inspired by these optimal updates, is to set,
ar NN(y.)
(52)
exp (at b)
(53)
k, D (y: r.)
(54)
K,K, = DT tiag((;;)
(55)
For our experiments analyzing recordings from areas V1 and V2, we use an MLP architecture with 128 hidden units and Swish nonlincarity (Ramachandran et al. 2017) to parameterize the inference network.
Importantly. close examination of these updates shows that only latent variables read out by the likelihood will have their latent state belief updated (those unread variables can be associated with fictitious 0 columns of D, and therefore the associated elements of k, and K, would be 0). In the context of multiregion models where one region's 's latent variables are not read out by another regions observation model, those other regions observations should play no part in updating belief if we causally filter data. Similarly, the belief over channel states is never updated directly through the observed data; intuitively this makes sense, showing our beliefs of the channel states can only be formed indirectly through our beliefs over the read out latent variables.
probabilistic
Furthemore, recalling the recursive belief updates are,
m m + Qg:
(56)
P P = Q QK, K, QK,) and that S; = (zz, rτ) represents the extended state, so that the blocks of Q associated with F zero, we see that for the deterministic states their conditional covariance 0, and m
(57)
mo expected.
Gaussian likelihood example. To gain some more intuition about the sparsity structure of the com-
plete generative model, consider the following generative model with linear and Gaussian observations,
examine its block structure in more depth. For illustrative purposes,
po(s; st-1) II Po (z(
(k)
z (b) r(() II (1.0)) 1 (-)
(h.0)
(k.6) (
(58)
E
ply 1z (k) = -y( De RR)
(59)
18

---

Published as a conference paper at ICLR 2025
where again, S; (Ze, [t). Now, if construct the block matrices,
Dj
0
0
0
R
D2
0
R2
D
(60)
D
R
where the last block column of D has as many columns as I3 has entries, then the observation model could also be written as,
ply: s.) Nyy, Dsr, R)
(61)
So that natural parameter updates inherit similar block structure,
D y D2R, Yi (2)
D
R D1
D2 R D2
k, = D DD
K,K, D diag
(62)
D RoRgy' (x)))
D R
DR
We can determine construct similar block parameter representations of the latent dynamics by sidering block operator dynamics and block state-noise, so that we can write.
(s( 8 = A(ss mo Q)
(63)
by setting,
(z (1) + E Ci, (1.0)
fz (2) + E1yz C2, (2.8)
Q Q2
mø(s:)
fx (z. (z( Σttk CK.E
Q = diag
(64)
A 1
(1.2)
B B
(2)
QK
AK (K.K1) K (rr-1((K]) (K-1)
Using these block representations makes more clear that the sampling procedure used for approx- imate inference only requires carrying out expensive linear algebraic operations with block matrices smaller than the full latent state dimension.
19


Reviews:
[{"review_id": "EKtqQat7G5", "reviewer": "Reviewer_SJFx", "summary": "This work introduces a new method for recovering the multi-area communication patterns in task-trained RNNs and real neural recordings. The authors validate the approach in both cases, introducing a new technically solid methodology to the field.", "questions": "I believe I understood most of the work adequately. I have some minor questions\n\n- Is Eq. (4) correct? Should the left hand side be C \\times γ? \n\n- Is the term \"neuroscientifically inspired\" a rigorous description? I am not sure if I would use it in such an otherwise rigorous study.\n\nTo address my weaknesses, you can take the following actions:\n\n- To address the audience problem, I would perhaps limit the jargon on linear system theory unless absolutely necessary to explain the findings. For example, the sentence \"Here we develop a probabilistic generative model that accounts for the nonlinear nature of neural dynamics and characterizes communication between regions using channels that are parameterized by their impulse response – blending expressive nonlinear region specific dynamics with interpretable characterizations of the flow of information between regions.\" came out of nowhere for me and took me a while to digest. Also, minor, but there wasn't much motivation built for this method in the introduction. Perhaps, you could add more insights into what is missing in current approaches and how this modeling approach will address those gaps.\n\n- I believe both Figs 2 and 3 should have the CCA and CURBD baselines. You can use one for CCA and another for CURBD.\n\n- Figure 4 is interesting, and shows that the modeling approach applies to real recordings. For a score higher than borderline, I would expect the authors to show a unique biological prediction with this model, and not a post-diction that explains already existing phenomena. \n\nEdit: I believe the authors have provided sufficient evidence to convince me that this work satisfies the criteria for publication at ICLR. Specifically, the new CURBD and CCA experiments have convinced me for the significance of this new approach and I look forward to using it in my own research.\n\n[1] Perich, M. G., Arlt, C., Soares, S., Young, M. E., Mosher, C. P., Minxha, J., ... & Rajan, K. (2020). Inferring brain-wide interactions using data-constrained recurrent neural network models. BioRxiv, 2020-12. (and Perich, M. G., & Rajan, K. (2020). Rethinking brain-wide interactions through multi-region ‘network of networks’ models. Current opinion in neurobiology, 65, 146-151.)\n\n[2] Ebrahimi, Sadegh, et al. \"Emergent reliability in sensory cortical coding and inter-area communication.\" Nature 605.7911 (2022): 713-721.\n\n[3] Vinograd, Amit, et al. \"Causal evidence of a line attractor encoding an affective state.\" Nature 634.8035 (2024): 910-918.\n\n[4] Musall, Simon, et al. \"Pyramidal cell types drive functionally distinct cortical activity patterns during decision-making.\" Nature neuroscience 26.3 (2023): 495-505.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- Simulation experiments are sound and provide strong evidence that the proposed approach is valid. Hence, this work is technically quite solid, satisfying the criteria for technical soundness.\n\n- The topic is of broader interest to the ML and neuroscience (or NeuroAI) communities, satisfying the criteria for interest but not significance (see below).", "weaknesses": "I believe there are three key weaknesses, which are the main reason behind more low score. If you can address these, I am happy to revise my final decision.\n\n## Audience problem\n\nThe paper is of broader interest to the community, but still suffers from a targeted audience problem. As a neuroscientist, I found the presentation to be too complicated for how intuitive and simple the actual method is. The mathematically heavy jargon may not help attract the biologists/experimental neuroscientists to read this paper. Yet, the paper has branded its audience as neuroscientists, which I do agree. Please see my suggestions below for how to address this.\n\nRelatedly, several relevant citations are missing. Crucially, CURBD [1], I believe there are far more groups in neuroscience working on these problems that deserve to be cited, a few well known papers include [2-4], though this is not an exhaustive list. I trust the authors will be able to balance papers from diverse groups in their revisions. \n\n## Missing baseline models\n\nI would have loved to see CURBD [1], and a simple CCA based analysis [2] as baseline models that should be compared to. Former uses a generative RNN model whereas latter does not model dynamics and simply learns co-fluctuation modes. I understand the appeal of \"causal\" modeling, but without proof that the current model captures the actual causal structure of the brain (as was done in [3] with intervention experiments), I do not believe these claims are substantiated. They do not necessarily need to be for a publication, instead it would suffice if the authors show how their model captures new and unique insights compared to these existing models. For example, I suspect CCA will create spurious correlations between two regions that are not directly connected. As for CURBD, it assumes a particular form for the nonlinearity and will therefore likely be less expressive. It would be great if these aspects are clearly fleshed out.\n\n## Missing new insights\n\nThough there is an application to neural recordings (Fig. 4), I am not convinced that there is a significant insight here. Oscillations are expected by the nature of the data, as also argued by the authors. The directionality is definitely an interesting prediction, but there is no clear actionable item for the biologists to follow up. Please see [1,2] for example works that have provided such actionable items such as computational predictions regarding the global communication modes between unsuspected brain regions.", "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": "", "extra_scores": {}}, {"review_id": "rWfU5hT9Mw", "reviewer": "Reviewer_UHzC", "summary": "This paper offers a new method for modeling multiregion neural data that combines a state space model with nonlinear local dynamics and emissions with linear, time-lagged cross-region communication channels parameterized by their impulse response. Parameterizing the communication channels in this way enables analysis of the time lagged and frequency components of communication across regions, offering high expressivity and interpretability. A structured variational inference approach is proposed for the model. The paper shows the effectiveness of the model and its interpretability through a series of simulation experiments and application to a multiregion neural data set.", "questions": "NA", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This is an excellent paper with a sophisticated, well motivated technical approach, that provides a novel contribution to the literature in an important and growing area of computational neuroscience - multiregion neural data analysis. The paper is well written, and does a good job of providing scientific context and literature review, and the model and technical approach are presented in a clear, instructive way. The proposed approach in many ways combines the benefits of DLAG (2022) and MR-SDS (2024) --- the proposed parameterization of the time-delayed communication channels allows interpretable time and frequency decomposition, while maintaining the expressivity benefits of a nonlinear state space modeling approach. The structured inference approach is elegant. Simulations are well motivated and their analysis and presentation are done in a careful and accessible manner despite their complexity. The application to a previously published and analyzed dataset provides novel insights.", "weaknesses": "The paper is very strong, but a few notes:\n\n- Comparisons -- the paper would benefit from additional comparisons with other time lagged communication models (DLAG) and related models with nonlinear communication (MR-SDS). \n\n- Given the sophistication of the structured inference approach, it would be nice to show and quantify that running the fit model in generative mode yields trajectories consistent with the inferred latents. \n\n- Figure 1 - comparing elbos across models can be problematic for a number of reasons, it would be helpful to have an interpretable measure such as r^2 on the inferred latents and reconstructed observations in the model vs NL/LN cases. Perhaps more importantly, and related to the previous point, comparing generative accuracy under these models should show a strong benefit of the additional expressivity of the full generative model; since it’s a simulation, generating more data should not be a problem.\n\n- It may be useful to comment on the limitations of the approach’s recovery of local vs communicated dynamics. In region1 of the RNN simulation, for example, the slingshot dynamics in the true RNN are not present in the local dynamics and appear to be induced by the incoming communication. By contrast, the learned flow fields for region1 in the model include the slingshot dynamics. It may be worth commenting on this, and / or visualizing the induced communication dynamics in this case.\n\n- Figure 3c and F: in figure 3c, 2 communication channels are shown, corresponding to the two orders of the filter, while in figure 3f, 6 dimensions are shown. The text explains that 3 and 2 dimensions are used for region 1 and 2 respectively, so 6 communication dimensions in all and an order 2 filter. It may help to include a brief summary of this in the figure caption to help readability. \n\n- Small typos: Line 291 - missing the word “are”. Line 462 -- these latents not this latents.", "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": "", "extra_scores": {}}, {"review_id": "AsGTNZRr7D", "reviewer": "Reviewer_TZvp", "summary": "This paper introduces MRDS-IR, a probabilistic generative model for capturing the dynamics and interactions of multiple neural populations. The model combines expressive nonlinear dynamics within each population with interpretable linear communication channels between them, parameterized by their impulse response functions. The authors develop an efficient variational inference and learning algorithm that handles the model's hybrid transitions without requiring state noise inversion. Through simulations and application to real neural data, they demonstrate MRDS-IR's ability to recover meaningful characterizations of the computations within and information flow between neural populations, providing a valuable tool for understanding distributed neural computation.", "questions": "- In Eq.10, $\\mathbf{D}_k$ and $\\mathbf{R}_k$ should be explicitly defined.\n- Do the red numbers (0.15, 0.0, -0.15) in Figure 1C represent the max/min amplitude of the impulse responses? Why do they differ from the black numbers (0.4, 0.0, -0.4), given that the model impulse responses and true impulse responses have similar amplitudes?\n- How can we explain the difference in 1->2 communication shown in Figure 2D? The estimated impulse response appears to differ from that of the RNN.\n- In Figure 3C, how should we interpret the statement, ''it is clear that the communication channel from region 1 to region 2 learned by MRDS-IRI transfers the frequency content of region 1''?\n- If the amplitude of a learned impulse response between channels \\(i\\) and \\(j\\) is close to zero, does this indicate that there is no communication between these channels?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- Propose a method to learn nonlinear neural population dynamics and the flow of information between regions.\n\n- Propose an end-to-end variational methodology.", "weaknesses": "- The writing could be improved, as there are several areas that seem unclear to me. Please refer to the questions for specific details.\n\n- The comparison is currently only with MRM-GP, a model with linear dynamics. A better choice would be to compare with a nonlinear model like MR-SDS, though this would require additional effort since public code for MR-SDS is not available.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)