Disentangling Recurrent Neural Dynamics with Stochastic Representational Geometry

Authors: david  lipshutz
amin  nejatbakhsh
alex h willliams
Venue: ICLR_cc_2024_Workshop_Re-Align
Published: 2024-02-08 04:50:24

Abstract: Uncovering and comparing the dynamical mechanisms that support neural processing remains a key challenge in the analysis of biological and artificial neural systems. However, measures of representational (dis)similarity in neural systems often assume that neural responses are static in time. Here, we show that stochastic shape metrics (Duong et al., 2023), which were developed to compare noisy neural responses to static inputs and lack an explicit notion of temporal structure, are well equipped to compare noisy dynamics. In two examples, we use stochastic shape metrics, which interpolates between comparing mean trajectories and second-order fluctuations about mean trajectories, to disentangle recurrent versus external contributions to noisy dynamics.

Main paper: 
 introduction: 
 Biological and artificial neural networks represent external stimuli and actions in high-dimensional feature spaces. While different neural systems utilize distinct feature spaces, the geometry of representations in these spaces can be similar across systems (Dwivedi & Roig, 2019; Chung & Abbott, 2021) , hinting that underlying computational principles may be shared.Numerous distance measures between neural representations have been utilized, including: Representational Similarity Analysis (RSA, Kriegeskorte et al., 2008) , Centered Kernel Alignment (CKA, Kornblith et al., 2019) , Procrustes shape distance (Williams et al., 2021) , and individual matching of neurons (Li et al., 2015; Khosla & Williams, 2023) . Nearly all of these methods consider neural responses that are static and non-stochastic. That is, neural systems are idealized as a deterministic functions f : S → R N that map from a stimulus space S to N -dimensional feature space (where N is the number of neurons in the system).But many biological and artificial neural systems are neither deterministic nor static. In many brain regions, the mean neural response is often smaller than variance across trials (Goris et al., 2014) , and shared trial-to-trial fluctuations (so called "noise correlations") are thought to be a crucial determinant of a circuit's signalling capacity (Averbeck et al., 2006) . Moreover, neural responses unfold dynamically over time, with rich interplay of feedforward and recurrent interactions (Vyas et al., 2020) . Of course, many artificial network models also contain components that are recurrent (e.g. RNNs) or stochastic (e.g. variational autoencoders).Ignoring the stochastic and dynamical aspects of neural circuits may be a justified simplification in certain settings, but it is easy to imagine situations where traditional representation (dis)similarity measures are blind to important details. These deficiencies were separately addressed in two recent works by Duong et al. (2023) and Ostrow et al. (2023) . The former proposed stochastic shape distances (SSDs) to geometrically quantify differences in trial-to-trial noise across networks without addressing recurrent dynamics. The latter utilized Koopman operator theory to develop Dynamical Similarity Analysis (DSA), which aims to quantify similarity in dynamical elements of flow fields without an explicit focus on noise statistics.Here, we investigate the extent to which these two principles-stochasticity and recurrent dynamics-are interrelated in the context of quantifying representational similarity. The relation between the two has already been remarked upon in the broader literature. A recent paper by Galgali et al. ( 2023) demonstrates how stochastic deviations from an average dynamical trajectory ("residual dynamics") can be used to infer a dynamical flow field under certain conditions. It is also well understood that driving a deterministic system with noise (or some other "persistently exciting" signal) is necessary to infer its governing dynamics ( Åström & Torsten, 1965; Mareels, 1984; Green & Moore, 1986) . We show that SSDs, which lack any explicit notion of dynamics, can nonetheless be used to quantify differences arising from recurrent neural interactions. 
results: 
 We focus on SSDs that account for the mean and covariance of neural responses and ignore higherorder moments.1 Specifically, consider network A with mean responses {µ M } in an N -dimensional feature space across M conditions or "landmarks." The Wasserstein α-SSD between A and a similarly defined network B is:EQUATIONwhere B(•, •) is the Bures distance between positive semidefinite matrices, O(N ) is the set of orthogonal matrices, and 0 ≤ α ≤ 2 determines the relative weight placed on differences between the means or covariances of the neural responses. In the original conception of SSDs, each landmark corresponds to a different network input-i.e., m is an index over stimulus variables, such as natural images fed into a deep network. Below we will use m to index over timebins or dynamic behavioral states. Equation ( 1) is described in more detail in section 2.4 of Duong et al. (2023) .We begin by showing that eq. ( 1) can be used "off-the-shelf" to distinguish between different recurrent dynamics that give rise to the same trial-average trajectories. To demonstrate, we ran simulations of three network dynamics following the construction by Galgali et al. (2023) and later leveraged by Ostrow et al. (2023) . In these simulations, the recurrent dynamics are distinct-the first system is bistable, the second implements line attractor dynamics, and the third implements point attractor dynamics. However, the time varying input drive is tuned to render the trial-average responses identical. Figure 1A shows simulated dynamics under two conditions (yellow vs. blue), corresponding to a binary decision outcome (see Galgali et al., 2023, for context) . The top panels show single-trial dynamics (semi-transparent trajectories) and the trial-average. The bottom panels illustrate the trial-average (dots) and the marginal, across-trial covariance (ellipses) at individual time points. Full simulation details are provided in Appendix A.It is clear from inspection that the three systems can be distinguished from their covariances, but not their trial averages. To confirm, we computed SSDs using eq. ( 1) treating timebins in distinct conditions as "landmarks." Concretely, we have M = CT landmarks where T is the number of timebins in a trial and C = 2 is the number of conditions (blue vs. yellow trajectories). The SSD with α = 2 fails to distinguish the three cases, as expected (Figure 1B , mean-only distance). In contrast, the SSD with α = 0 or α = 1 successfully distinguishes the three systems (Figure 1B , mean + covariance and covariance-only distance). This reproduces one of the main capabilities of DSA (Fig. 3 in Ostrow et al., 2023) . Indeed, we qualitatively reproduce their finding that the line and point attractor systems are most similar (i.e. lower distance relative to the bistable system).Figure 1 shows a simple example where the shape of trial-to-trial "noise" can capture differences in recurrent dynamics, even when input-driven dynamics are tuned to make firing rate trajectories as similar as possible. In this section, we provide a theoretical explanation for this result which holds for any stochastic differential equation that takes the form:EQUATIONwhere x(t) is an N -dimensional vector of neural activities at time t, f is a sufficiently regular (e.g., Lipschitz continuous) function governing time-invariant recurrent dynamics, u : R → R N defines a deterministic time-varying input drive to the system, g : R N → R N ×N is a sufficiently regular function that defines the noise structure, and w(t) is an N -dimensional standard Brownian motion.The stochastic process {x(t)} defines a family of distributions on R N parameterized by t. Define the first two moments (mean and covariance) of these distributions as:EQUATIONEQUATIONWe drop the dependence on t below for brevity. Using standard methods from stochastic calculus (see e.g. Särkkä & Solin, 2019) , one can show that m(t) and P (t) evolve according to the differential equationsdm dt = E[f (x)] + u, dP dt = E[f (x)(x -m) T ] + E[(x -m)f (x) T ] + 1 2 E[g(x)g(x) T ].These differential equations provide a concrete explanation for the success of SSDs in Figure 1 . Specifically, we see that the evolution of the mean trajectory m(t) depends on the input drive u(t), whereas the evolution of the covariance matrix P (t) does not explicitly depend on u(t). These equations suggest that SSDs (with α = 2) are a useful metric for isolating differences in recurrent dynamics, while deterministic shape distances (α = 0) are useful for comparing the joint effect of recurrent dynamics and input drive. We caution, however, that this does not guarantee that SSDs will always distinguish between networks with different nonlinear recurrent interactions.To further elucidate the relationship between the covariance trajectories and the recurrent dynamics, it is instructive to consider the linear setting with additive noise-i.e., when f (x) = Ax and g(x) = √ 2L for N × N matrices A and L. In this case, the evolution of the covariance P (t) is governed by the linear differential equation:dP dt = AP + P A T + LL T .In particular, it is clear that the evolution of the covariance matrix depends only on the recurrent dynamics matrix A and noise matrix LL T , not the time-varying input drive u(t).As neuroscientists study more complex and naturalistic behaviors, it has become increasingly common for experimental setups to have less trial to trial structure (Williams & Linderman, 2021) . How can we disentangle recurrent neural dynamics between networks in such setups? As an example, suppose we have recordings from place cells in two mice while they are freely exploring an environment. Since the two mice will traverse unique trajectories, comparing time-aligned representations may not be a relevant measure for disentangling recurrent interactions. Instead, a more natural approach is to compare representations conditioned on position.We trained 3 classes of RNNs with 64 hidden units to estimate the position l(t) of an animal on a circular track (Figure 2A ). Each RNN received noise corrupted inputs (l(t) + ϵ l w l (t), v(t) + ϵ v w v (t)), where v(t) is the velocity of the animal and w l (t), w v (t) are standard Gaussian random vectors (Figure 2B ). The first (second, third) class of RNNs was trained with ϵ l = 0 (ϵ l = 0.5, ϵ l = 1) and ϵ v = 1 (ϵ v = 0.5, ϵ v = 0). Intuitively, since the first class receives noiseless position inputs, the RNN should simply output its position inputs (position networks). On the other hand, the third class has access to noiseless velocity inputs, so the RNN should integrate these inputs to estimate the position (velocity networks). Finally, we expect the second class learns an intermediate strategy (mixed networks). Therefore, we expect that the 3 classes of RNNs will have distinct recurrent interactions that can be distinguished by comparing their conditional responses.For each class, we trained the model with and without BatchNorm in their hidden-to-output layer. Since BatchNorm normalizes hidden state activations in each batch, we expect the networks trained with BatchNorm to have near zero trial averages at each time point. Therefore using BatchNorm is akin to adversarial examples in Fig. 1 where the trial averages are forced to be similar across models. Hence we expect SSD with α = 2 to fail distinguishing between 3 classes of RNNs.To confirm these intuitions, we computed SSDs between the networks using eq. ( 1) by treating the true position as "landmarks." Specifically, we partitioned the ground truth position l(t) into M bins and treated each bin as a landmark. We then computed the mean and covariance of a networks response conditioned on each landmark and compared their SSDs. Consistent with our intuition, we observed that for BatchNorm RNNs, the SSD with α = 2 fails to distinguish the position, velocity, and mixed networks. While the SSD with α = 0 successfully distinguishes the 3 classes of RNNs. This is expressed by the clustering of network types in the MDS embedding of the RNNs in Figure 2E and by the block structure of the SSD matrices in Figure 2F .Conversely, we observed that for RNNs trained without BatchNorm, SSD with both α = 0 and α = 2 distinguishes between different classes of RNNs (Figure 2G ). This suggests that both trial averages and noise correlations contain information about the recurrent strategies used by the networks. Details of the RNN training are provided in Appendix A.2.Figure 2 shows that we can leverage conditional covariance structure to distinguish recurrent interactions in a network when there isn't explicit trial structure. To better understand this, let z(t) be a stochastic process that encodes the behaviorally relevant variables (e.g., position or velocity). Here we assume z(t) is scalar valued, but the set-up readily extends to the case of a vector-valued process.Consider the following stochastic differential equation modulated by z(t):dx(t) = f (x(t))dt + u(z(t))dt + g(x(t))dw(t),where u : R → R N is the function that maps the behavioral variable to the dynamics and all other variables are as in eq. ( 2). The covariance P (t) evolves according to the differential equation:dP dt = E (f (x) + u(z))(x -m) T + E (x -m)(f (x) + T + 1 2 E g(x)g(x) ,which now includes an explicit on the time-varying input to the system u(z(t)) because z(t) itself is a stochastic process that is not independent of x(t). One way to obtain an equation that does not include the input to the system is by conditioning on the process z = {z(t) ≥ 0}:EQUATIONP (t, z) := E (x(t) -m(t, z))(x(t) -m(t, z)) T |z .In this case, the explicit dependence on the input u(z(t)) drops out of the evolution equation:dP (z) dt = E[f (x)(x -m(z)) T |z] + E[(x -m(z))f (x) T |z] + 1 2 E[g(x)g(x) T |z],which suggests that conditioning on z is useful for isolating the recurrent contribution to the dynamics. In the above, we have conditioned on the entire process z, so comparing two networks' responses requires that the behavioral processes have similar trajectories. However, if we further assume that the pair (x, z) is a stationary process-i.e., the joint distribution of (x(t), z(t)) does not depend on time t-then the conditional mean and covariance, defined by m(z) := E[x|z] and P (z) := E (x -m(z))(xm(z)) T |z , do not depend on time t and can be used to compare representations from networks with different trajectories as in the previous section. 
discussion: 
 We have shown that SSDs can be used to disentangle noisy dynamic systems with different recurrent interactions. As shown by Ostrow et al. (2023) , we find that the mean trajectories of networks with different recurrent interactions can be quite similar due to external inputs to the system. However, we find that comparing stochastic fluctuations about the mean trajectories can distinguish between systems with differing recurrent interactions, as demonstrated by (Galgali et al., 2023) . By using the full parametric form of SSDs (with α interpolating between 0 and 2), we can obtain a richer notion of the similarities and differences, quantified in terms of metrics, in recurrent interactions between networks with noisy dynamics.There are many important directions for future investigation. For example, the SSDs used in this work only capture first and second-order moments, which is optimal when comparing Gaussian processes; however, they are not guaranteed to capture differences between non-Gaussian processes.Another important avenue to explore is the relationship between SSDs measures of (dis)similarity discussed here and DSA measures introduced by Ostrow et al. (2023) . SSDs are metrics designed to capture similarities in geometric structure, whereas DSA has been shown to capture similarities in topological structure. Teasing apart these details is critical to enriching our understanding of the dynamics at play in neural systems.Karl-Johan Åström and Bohlin Torsten. Numerical identification of linear dynamic systems from normal operating records. IFAC Proceedings Volumes, 2(2): 96-111, 1965. 2nd IFAC Symposium on the Theory of Self-Adaptive Control Systems, Teddington, UK, September 14-17, 1965.A EXPERIMENTAL DETAILS A.1 THREE NETWORK DYNAMICSWe modeled three network dynamics-a bistable switch, a line attractor, and a point attractor-that were analyzed in (Galgali et al., 2023) and then in (Ostrow et al., 2023) . To simulate the trajectories, we used the code from the Jupyter notebook developed by Ostrow et al. (2023) , which can be found within the following Github repository: https://github.com/mitchellostrow/DSA. To compute the SSDs, we used the code from the following Github repository: https://github. com/ahwillia/netrep. The dynamics and simulation details are described below.The bistable switch is described by the 2-dimensional system of nonlinear stochastic differential equations with additive (degenerate) noisedx(t) = [f (x(t)) + 1ν i ]dt + σdw(t),where f (x 1 , x 2 ) := (ax 3 1 + bx 1 , cx 2 ), 1 = (1, 1) is the vector of ones, ν cond is a conditioned dependent constant input-drive that is applied to both coordinates, w(t) is a 2-dimensional standard Brownian motion and a, b, c, σ are constants. The parameters a, b, c are randomly sampled at the outset as follows: a ∼ Unif(-5, -3), b ∼ Unif(4, 7), c ∼ Unif(-4, -2). Depending on the condition, ν cond is set to -0.1 or 0.1 and drives the system to one of the stable fixed points. The noise coefficient σ is 5. The system is simulated using the Euler-Maruyama method:x n+1 = x n + f (x n )∆t + 1ν i ∆t + σ∆w nwith step size ∆t = 0.01 and ∆w n are i.i.d. with ∆w n,i ∼ N (0, ∆t) for i = 1, 2.The line attractor and point attractor are 2-dimensional linear stochastic differential equations of the form dx(t) = [Ax(t) + u(t)]dt + σdw(t),where w is a 2-dimensional standard Brownian motion and σ is as in the previous example. For the linear attractor, the recurrent dynamics matrix is A = V ΛV -1 , where Λ = diag(-1, 0) andV = 1 1 1 0 .For the point attractor, the recurrent dynamics matrix is A = diag(-0.5, -1). As in previous example, each linear system is simulated using the Euler-Maruyama method with In each case, the input drive u(t) is adversarially chosen to minimize the difference between the mean trajectory of the linear system with the mean trajectory of the nonlinear bistable switch. In particular, let xBS n denote a conditional mean trajectory of the bistable switch (i.e., conditioned on u = 1 or u = -1) and let xn denote the corresponding condition mean trajectory of the linear system. Then the input drive for the linear system u n and the mean trajectory for the linear system can be recursively defined by x0 = xBS 0 andu n = xBS n+1 -xn ∆t -A xn , xn+1 = xn + [A xn + u n ]∆t.The synthetic dataset was simulated to model the behavior of a mouse running around a circular track. The task is motivated by neuroscience findings on the existence of a ring attractor in the representations of entorhinal neurons and RNN models trained to track the position of animals (Low et al., 2021) . To generate the position and velocity of the fictitious animal in time, we first sampled a parameter θ t from a Brownian motion for 600 time steps with a drift of 0.2. The 2D position at each time step was calculated according to l t = (cos θ t , sin θ t ). The 2D velocity was then computed as the discrete differential of the position v t = l t -l t-1 . Each RNN has 4 inputs and 2 outputs and is tasked to estimate the position of the animal denoted by lt . Therefore we minimize the following loss function L(ϕ) = t,i ∥l t,i -lt,i ∥ 2 2 where t corresponds to the time point t and i denotes the trial i. Each network has 64 hidden units and follows the classical rate network dynamics:h t+1 = σ(W h h t + W i (l t + ϵ l t , v t + ϵ v t ) T + ϵ h t ), lt = W o h t , ϕ = {W h , W i , W o }where ϵ l t , ϵ v t , ϵ h t are i.i.d. noise terms added to the position input, velocity input, and hidden units respectively. The standard deviation of these noise terms determines how much the network can rely on each of the two input terms. For 5 velocity networks, we set the noise standard deviation of the position and velocity inputs to 1 and 0 respectively. We do the opposite for 5 position networks. For 5 mixed networks with injected noise with standard deviation 0.5 to both position and velocity inputs. The hidden noise standard deviation is always set to 0.1.For each class of RNNs (position, velocity, and mixed) we trained two separate models, with and without BatchNorm in their hidden-to-output layers.h t ← h t -E[h batch t ] Var(h batch t ) , lt = W o h tSince BatchNorm normalizes hidden state activations to have zero mean, it intuitively mimics an adversarial training strategy for limiting the information of the trial averages.To train the networks, we generated batches of size 64 and ran 1000 iterations of backprop through time using SGD optimizer with a learning rate of 0.05. All the experiments were done using Pytorch on CPU and every RNN takes about 3 minutes to train.An alternative SSD based on energy distance (instead of Wasserstein distance) may be used to investigate differences in higher-order moments (seeDuong et al.,2023). But we leave this possibility to future work. 
 
 Appendix:  



Reviews: review: title: title: Review of "Disentangling Recurrent Neural Dynamics with Stochastic Representational Geometry"
main_review: main_review: **Summary**

The paper shows that recent advances in shape metrics, specifically stochastic shape distances (SSDs), can distinguish different underlying dynamics even when the trial-averaged trajectories are similar. It analyses the interplay between stochasticity and recurrent dynamics and how the former can aid in comparing the latter. The authors provide simple experiments and interesting theoretical analyses to demonstrate this.

**Strengths**
* The paper is well-written and provides interesting insight on how SSDs, despite not explicitly modelling dynamics, can still be used to distinguish them.
* The numerical experiments are intuitive and comprehensive.

**Questions and suggestions**
* It would be nice if the authors could comment on whether SSDs could be used to distinguish different learning rules, i.e., dynamics of learning itself, similar to Ostrow et al. (2023).
* Perhaps showing some of the failure modes of these SSDs in distinguishing different networks would aid readers in understanding where best to use or not use these metrics.
* As the authors state, it would be interesting to explore the connections between DSA (Ostrow et al., 2023), which performs a modified Procrustes alignment of DMD matrices, and the SSDs which align network responses based on their moments. 

Overall, I believe that this work would be of great interest to the workshop's participants, and I recommend that the paper be accepted.
strength_weakness: strength_weakness: The submission demonstrates how recent advances in shape metrics can distinguish different dynamics. This builds upon a line of work that is highly relevant to the workshop, with several impactful future directions. Furthermore, the submission is well-written and provides a good mix of theoretical and empirical analyses.
title: title: The paper presents novel applications of stochastic shape distances (SSDs; Duong et al., 2023)  to disentangle recurrent neural dynamics
main_review: main_review: **Overall summary:**

The paper is very well written and its subject is relevant to the community of this workshop. The experiments are simple yet effective in supporting the paper's claims, which are also supported by theoretical arguments that, to the best of my understanding, seem correct. A suggestion is to add a bit more details for the experiments in section 2.1 (or at least add them in the appendix) so the reader doesn't have to refer to past works for something that is of key importance in the present paper.

**Typos:**

- Possible error in this sentence? “The first class of RNNs was trained on noiseless position inputs and noisy velocity inputs, whereas the second class of RNNs was trained on noiseless velocity inputs and noisy position inputs (Figure 2B)”
strength_weakness: strength_weakness: The paper is very well written, the topic is of interest to the community, and the claims are supported both by empirical results and by theoretical arguments. It's a high quality workshop paper.
title: title: SSDs as an improved measure of neural dynamic (dis)similarity
main_review: main_review: Summary: 
The authors demonstrate that the Stochastic Shape Distance, a measure of distance between two gaussian distributions is useful for disentangling neural dynamics which come from fundamentally different recurrent systems, but otherwise may be considered identical under other distance metrics. Specifically, first they show (both empirically and analytically) that by considering across-trial covariance they can disentangle systems with the same trial averaged trajectories. Secondly, they show how even without trial structure, the conditional covariance can be leveraged to a similar effect. 

Strengths:
- (Clarity) The paper is well written, the idea is thoroughly explored and well situated within the literature. 
- (Correctness) The experiments and analysis appear to justify the papers main claims sufficiently.
- (Novelty) The proposed application of the very recently introduced method is novel to the best of my knowledge.
- (Relevance) There is certainly relevance to designing accurate measures of dynamical similarity, and this is a good step in that direction.

Weaknesses: 
- The experimental details are quite minimal making the paper not very self-contained, with frequent references to external works for clarity. For an outstanding paper, it would be beneficial for the authors to include all these details in the appendix.
- The method does appear helpful in the situations the authors constructed, however it is not clear to what extent this method will remain beneficial in natural settings or if this if the presented examples are pathological. For example, does the SSD really capture meaningful variation between trial-to-trail differences in natural data, or are higher-order moments necessary as the authors allude to? 

Minor: 
- (Typo) Middle of Page 5, "One way to obtain an equation that does not include the input to the system is by conditioned on ..." --> "... is by conditioning on..."
strength_weakness: strength_weakness: The paper's topic is very timely and clearly in-line with the workshop goals. It includes analytical and empirical justification, and a new viewpoint on dynamic similarity. There are no flaws with the paper to my view other than the potential that more could be done to improve the method.