SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Mitchell Ostrow, Adam Eisen, Leo Kozachkov, Ila Fiete
Department of Brain and Cognitive Sciences, MIT
{ostrow,eisenaj,leokoz8,fiete}@mit.edu
How can we tell whether two neural networks utilize the same internal processes for a particular computation? This question is pertinent for multiple subfields of neuroscience and machine learning, including neuroAI, mechanistic interpretability, and brain-machine interfaces. Standard approaches for comparing neural networks focus on the spatial geometry of latent states. Yet in recurrent networks, computations are implemented at the level of dynamics, and two networks performing the same computation with equivalent dynamics need not exhibit the same geometry. To bridge this gap, we introduce a novel similarity metric that compares two systems at the level of their dynamics, called Dynamical Similarity Analysis (DSA). Our method incorporates two components: Using recent advances in data-driven dynamical systems theory, we learn a high-dimensional linear system that accurately captures core features of the original nonlinear dynamics. Next, we compare different systems passed through this embedding using a novel extension of Procrustes Analysis that accounts for how vector fields change under orthogonal transformation. In four case studies, we demonstrate that our method disentangles conjugate and non-conjugate recurrent neural networks (RNNs), while geometric methods fall short. We additionally show that our method can distinguish learning rules in an unsupervised manner. Our method opens the door to comparative analyses of the essential temporal structure of computation in neural circuits.
Comparing neural responses between different systems or contexts plays many crucial roles in neuroscience. These include comparing a model to experimental data (as a measure of model quality, Schrimpf et al. [2020]), determining invariant states of a neural circuit (Chaudhuri et al. [2019]), identifying whether two individual brains are performing a computation in the same way, aligning recordings across days for a brain-machine interface (Degenhart et al. [2020]), and comparing two models (e.g. to probe the similarity of two solutions to a problem, Pagan et al. [2022]). Current similarity methods include  $R^2$  via Linear Regression, Representational Similarity Analysis, Singular Vector Canonical Correlation Analysis, Centered Kernel Alignment, and Procrustes Analysis (Schrimpf et al. [2020], Kriegeskorte et al. [2008], Raghu et al. [2017], Williams et al. [2021], Duong et al. [2022]). Crucially, these all compare the geometry of states in the latent space of a neural network.
Identifying shared neural computations between systems calls for methods that compare similarity at the level of the fundamental computation. In the brain, computations are instantiated by the emergent dynamical properties of neural circuits, such as fixed points, invariant and fixed point manifolds,
limit cycles, and transitions between them (Amari and Arbib [1977], Hopfield [1984], Hopfield and Tank [1986], Seung [1996], Zhang [1996], Hahnloser et al. [2000], Burak and Fiete [2009], Wang [2002], Sussillo and Abbott [2009], Churchland et al. [2012], Chaudhuri et al. [2019], Vyas et al. [2020], Khona and Fiete [2022]. Geometric similarity metrics fall short in two manners when they are applied to dynamical systems. First, response geometries can differ due to sampling differences (Chaudhuri et al. [2019]) or other nonfundamental causes even though the dynamics are topologically equivalent. In such cases, measures of state-space geometry do not fully capture the core similarities between two systems, a false negative situation. Conversely, systems may exhibit distinct dynamics over similar geometries in the state space (Galgali et al. [2023]). In this case, geometric measures may fail to distinguish two distinct systems, a false positive situation. Therefore, we suggest that dynamical similarity should be the preferred lens through which to compare neural systems.
In dynamical systems, a key notion of similarity is called topological conjugacy: the existence of a homeomorphism that maps flows of one system onto flows of the other. When two systems are conjugate, they have the same qualitative behavior. Given two dynamical systems  $f: X \to X$  and  $g: Y \to Y$  with invertible mapping  $\phi: X \to Y$ , conjugacy is defined as:
$$
g \circ \phi = \phi \circ f \tag {1}
$$
Here, we develop a data-driven method called Dynamical Similarity Analysis (DSA), which combines work from two previously-unrelated fields, machine learning for dynamical systems and statistical shape analysis. In brief, DSA returns a similarity metric describing how two systems compare at the level of their dynamics. First, DSA uses the insight that appropriate high-dimensional embeddings can be used to describe a highly nonlinear system in a way that is globally linear. The features that enable this transformation are referred to as Koopman Modes (Mezić [2005]) and can be identified via the Dynamic Mode Decomposition (DMD) (Snyder and Song [2021], Brunton et al. [2021]). Subsequently, DSA employs a statistical shape analysis metric to compare vector fields within the Koopman Mode embeddings, thereby assessing the similarity between the systems' dynamics. This shape analysis measures how far a pair of systems are from being homeomorphic. Unlike most spatial shape analyses, which require averaging over multiple noisy trials in the same condition to arrive at one vector for each condition-time point (but see Duong et al. [2022]), small amounts of noise create a rich representation of neural dynamics in the estimated dynamics matrix. The importance of flows in response to intrinsic (natural) or induced perturbations in capturing dynamics has been highlighted in Seung [1996], Yoon et al. [2013], Chaudhuri et al. [2019], Galgali et al. [2023], as their evolution over time can uniquely identify dynamic structure.
Our results demonstrate that DSA effectively identifies when two neural networks have equivalent dynamics, whereas standard geometric methods fall short. DSA therefore has great relevance for neuroscience, providing another method to validate the fit of a model to the neurophysiology data. It additionally opens the door to novel data-driven analyses of time-varying neural computation. We expect our method to be interesting to computational neuroscientists due to its theoretical richness, valuable to deep learning researchers as a theoretically backed and rigorously justified tool through which to interpret and compare neural networks, and useful for experimental neuroscientists due to its simplicity of implementation and ease of interpretation.
Contributions We develop a general method, DSA, for comparing two dynamical systems at the level of their dynamics. To do so, we introduce a modified form of Procrustes Analysis that accounts for how vector fields change under orthogonal transformations. We apply our method in four test cases, each demonstrating novel capabilities in neural data analysis: (a) Our method can identify dynamical similarities underlying systems with different geometries, which shape metrics cannot. (b) Conversely, DSA can distinguish systems with different dynamics despite similar geometries, which shape metrics cannot. (c) We show that our method empirically identifies topological similarities and differences between dynamical systems, as demonstrated by invariance under geometric deformation and variance under topological transformation. (d) We demonstrate how to use DSA to disentangle different learning rules in neural networks without supervision, by comparing how representations change across training.
$^{1}$  https://github.com/mitchellostrup/DSA/
2.1 Theoretical Background: Koopman Operators and Dynamic Mode Decomposition (DMD)
A common approach to modeling nonlinear dynamics is to approximate them with linear systems, as those are more tractable and interpretable. One standard method involves finding fixed points of the dynamics, and then locally linearizing the system around each fixed point, to determine whether infinitesimal changes diverge from or converge to the fixed point. Another is to describe nonlinear dynamics via locally linear models, a separate one for each point in state space. These methods do not generate a description of the global dynamics. The former method also requires injecting perturbations, which may not be possible for experimental systems and pre-collected datasets.
The field of dynamical systems theory suggests a different approach, on which a new class of data-driven methods are based. Specifically, Koopman operators theoretically embed nonlinear systems into infinite-dimensional Hilbert space, which permits an exact and globally linear description of the dynamics (Koopman [1931]). Practical applications use the Dynamic Mode Decomposition (DMD) (Schmid [2010]) to create a finite-dimensional embedding that approximates the Koopman operator (Brunton et al. [2016a, 2017], Snyder and Song [2021], Dubois et al. [2020]). The DMD identifies a linear transition operator for the dynamics:  $X(t + \Delta t) = AX(t)$ , where  $X$  comprises a data matrix (consisting of functions of the observed states of the dynamical system). When applied to a sufficiently rich  $X$ , the DMD can capture global dynamical structure. A key challenge of the method is in constructing an appropriate space for  $X$  which is closed under the operation of  $A$  (a Koopman invariant subspace, Brunton et al. [2016b]). We employ the Hankel Alternative View of Koopman (HAVOK, Brunton et al. [2017], Arbabi and Mezić [2017]), which constructs  $X$  as a function of the delay-embedding of the observed variables (a Hankel Matrix), such that one column of  $X = [-x(t) - x(t - \tau) - \ldots - x(t - p\tau) - ]^T$ , where  $p$  is the number of included delays and  $\tau$  is the time-step between each measurement. The delay-embedding approach eliminates the difficulty of learning a nonlinear transformation of the data, and improves estimation capabilities over partially-observed systems (Kamb et al. [2020]). The latter capability arises from Takens' delay embedding theorem, which states that a delay embedding is diffeomorphic to the original system—that is, a partially-observed system can be fully reconstructed using a sufficient number of delays (Takens [1981]). However, our method could be used with a range of embedding algorithms in place of delay embeddings.
2.2 Dynamical Similarity Analysis (DSA)
Now, we introduce Dynamical Similarity Analysis, a general method that leverages advances in both Statistical Shape Analysis and Dynamical Systems Theory to compare two neural systems at the level of their dynamics. Simply put, we use Koopman operator theory to generate a globally linear embedding of a nonlinear system, via delay embeddings as in HAVOK, then use DMD to define the linear transition operators in this space. Finally, we apply a novel extension of the Procrustes Analysis shape metric to the resultant DMDs from two systems. This set of steps constitutes DSA.
Suppose we have two dynamical systems defined by the following equations:
$$
\dot {\mathbf {x}} = f (\mathbf {x}, t) \quad \mathbf {x} \in \mathbb {R} ^ {n} \quad \dot {\mathbf {y}} = g (\mathbf {y}, t) \quad \mathbf {y} \in \mathbb {R} ^ {m} \tag {2}
$$
We sample data from each system  $X \in \mathbb{R}^{c \times k_1 \times t_1 \times n}$  and  $Y \in \mathbb{R}^{c \times k_2 \times t_2 \times m}$ . Here,  $c$  indicates the number of conditions observed, which in neuroscience and deep learning refer to different input stimuli or contexts.  $k_i$  indicates the number of trials per condition, and  $t_i$  indicates the number of time steps observed per trial. Note that the observations must be sampled uniformly in time (Takens [1981]). Next, we produce delay-embedded Hankel tensors with a lag of  $p$ :  $H_x \in \mathbb{R}^{c \times k_1 \times (t - p - 1) \times np}$  and likewise for  $H_y$ . Using the Hankel Alternative View of Koopman (HAVOK) approach (Brunton et al. [2017]), we flatten all but the last dimension of  $H$  and fit reduced-rank regression models with rank  $r$  to the eigen-time-delay coordinates of the data, where the target is the coordinates of the next time step  $(V_{x,r}')$ :
$$
V _ {x, r} ^ {\prime} = A _ {x} V _ {x, r} \text {w h e r e} H _ {x} ^ {T} = U _ {x} \Sigma_ {x} V _ {x} ^ {T} \text {a n d} V _ {x, r} = V _ {x} [:, 1: r ] \tag {3}
$$
Only the rank of the HAVOK models explicitly need to be the same so that the two DMD matrices can be compared. Note that the eigen-time-delay coordinate formulation is equivalent to PCA whitening (Supplementary Information Section 12), which Williams et al. [2021] highlights to be important. The predictive capability of our model (via MSE,  $R^2$ , AIC, or others) can be used to identify the optimal rank  $r$  and number of delays  $p$ . In our experiments, hyperparameters for the DMD were chosen to ensure both sufficient capability to fit the dynamics in question as well as tractable computation times. Fitting the HAVOK model to each system returns our DMD matrices:  $\mathbf{A}_x, \mathbf{A}_y \in \mathbb{R}^{r \times r}$ , which we compare using the modified Procrustes Analysis algorithm detailed next.
Figure 1: Schematic and example of DSA. a. Schematic of the DSA pipeline. b. Schematic of attractor reconstruction with delay embeddings. c. Schematic of the Koopman Operator linearization over time-delay coordinates. d,e. PCA trajectories (black lines) of a GRU (left) and Vanilla RNN (right) trained to produce sine waves of a range of frequencies (Sussillo and Barak [2013]). Colored dots indicate center points of the limit cycles, colored by frequency of the output wave. f. Comparison of measured dissimilarity between the networks in d,e with (from left to right) Procrustes on the condition-averaged trajectories, Procrustes on the delay-embedded condition-averaged trajectories, DSA on the original data, and DSA on the delay embedding. Lower is more similar. Each metric uses angular distance, which ranges from 0 to  $\pi$ .
Procrustes Analysis over Vector Fields Procrustes Analysis is a shape metric which compares the geometric similarity of two datasets after identifying an optimal orthogonal transformation to align them (Dryden and Mardia [2016]). To use it in DSA, we must modify the notion of orthogonal transformations to be compatible with linear systems. Recall that the Procrustes metric solves the following minimization problem (Williams et al. [2021]):
$$
\mathrm {d} (\mathbf {X}, \mathbf {Y}) = \min  _ {\mathbf {C} \in O (n)} | | \mathbf {X} - \mathbf {C Y} | | _ {F} \tag {4}
$$
Where  $O(n)$  is the orthogonal group  $(C^T C = I)$ ,  $F$  indicates the Frobenius norm, and  $X, Y$  are two data matrices. A distance metric that is invariant to orthogonal transformations captures the intuitive notion that there is no privileged basis across neural networks.
However, vector fields, represented by the DMD matrices  $\mathbf{A}_x$ ,  $\mathbf{A}_y$  above, transform differently than vectors  $\mathbf{X}$  and  $\mathbf{Y}$  (see Supplementary Information Section 8 for an illustration). The invariance we seek is that of a conjugacy, or change of basis, between linear systems, which is reflected in the similarity transform,  $CA_xC^{-1}$ . Note that this is a special case of equation 1 over linear systems. Thus we introduce a novel similarity metric, which we call Procrustes Analysis over Vector Fields:
$$
\mathrm {d} \left(\mathbf {A} _ {x}, \mathbf {A} _ {y}\right) = \min  _ {\mathbf {C} \in O (n)} | | \mathbf {A} _ {x} - \mathbf {C A} _ {y} \mathbf {C} ^ {- 1} | | _ {F} \tag {5}
$$
The angular form of this distance can be written as:
$$
\mathrm {d} \left(\mathbf {A} _ {x}, \mathbf {A} _ {y}\right) = \min  _ {\mathbf {C} \in O (n)} \arccos  \frac {\mathbf {A} _ {x} \cdot \left(\mathbf {C} \mathbf {A} _ {y} \mathbf {C} ^ {- 1}\right)}{\left| \mathbf {A} _ {x} \right| _ {F} \left| \mathbf {A} _ {y} \right| _ {F}} \tag {6}
$$
where  $\mathbf{A}_x\cdot (\mathbf{C}\mathbf{A}_y\mathbf{C}^{-1})$  is computed as  $\mathrm{Tr}(\mathbf{A}_x^T\mathbf{C}\mathbf{A}_y\mathbf{C}^{-1})$ . In the Supplementary Information Section 9, we prove an equivalence between DSA and the 2-Wasserstein distance between the two DMD matrices' eigenspectrum, in special cases. This connection highlights the grounding of DSA in Koopman operator theory to measure conjugacy: when two dynamical systems are conjugate, their Koopman eigenspectra are equivalent (Proof in Supplementary Information 10, Budiszic et al. [2012], Lan and Mezić [2013]). We derive equation 5 from first principles and prove that it is indeed a proper metric in the Supplementary Information (6); the proper metric property is crucial for various downstream machine learning pipelines such as clustering or classification (Williams et al. [2021]).
Equations 5 and 6 are nonconvex optimization problems similar to those in Jimenez et al. [2013], although we solve them with gradient-based optimization. Across our experiments, we found that the metric asymptotes in roughly 200 iterations using the Adam optimizer (Kingma and Ba [2014]) with a learning rate of 0.01. To optimize over the orthogonal group,  $O(n)$ , we utilize an extension of the Cayley Map, which parameterizes  $SO(n)$  (Supplementary Information 5.1). In our experiments, we compare the response of DSA with direct application of Procrustes to the data matrices as it is the most analogous shape metric. If only the identification of conjugacy is desired, metric properties can be relaxed by optimizing  $C$  over the larger group of invertible matrices,  $GL(n)$ .
To demonstrate the viability of our method, we identified four illustrative cases that disentangle geometry and dynamics, and additionally performed an ablation analysis of DSA. First, a set of neural networks that are different by shape analyses but are dynamically very similar; second, a set of neural networks that are similar by shape analyses but are dynamically are very different; third, a ring attractor network whose geometry we can smoothly and substantially deform while preserving the attractor topology; finally, a line attractor network whose topology is transformed into a ring attractor by adding periodic boundary conditions.
3.1 Ablation Study
Importantly, both the reduced-rank regression and delay-embedding in DSA are required to isolate dynamical structure. We demonstrate this in Fig. 1f, where we compare the hidden state trajectories of an RNN and GRU solving the Sine Wave task (Sussillo and Barak [2013]) using the standard Procrustes Analysis (via the netrep package, Williams et al. [2021]), delay-embedded Procrustes Analysis (an ablation), DSA with no delay embedding (another ablation), and DSA score. To be exhaustive, we also computed a similarity metric that compares purely temporal information, defined as the angular Procrustes distance over the power spectrum of the RNN's hidden state trajectories. On the RNNs in Fig. 1, this metric reported 1.41. Only DSA identifies the systems as highly similar, indicating that both the delay embedding and the DMD are necessary to abstract geometric particulars from each system.
3.2 DSA captures dynamical similarity between RNNs despite differences in geometry.
We trained a collection of 50-unit RNNs on the well-known 3-bit Flip Flop task from Sussillo and Barak [2013]. In this task, an RNN is presented with a time series from three independent channels, each sporadically occurring pulses of either -1 or 1. The network's objective is to continuously output the last pulse value for each channel until it is updated by a new pulse (hence the Flip-Flop name). Following Maheswaranathan et al. [2019], We varied architecture (LSTM, GRU, UGRNN, Vanilla RNN), activation function (ReLU, Tanh), learning rate (.01,0.005), and multiple different seeds across each set of hyperparameters, ultimately collecting a dataset of 240 networks. We only saved a network if it solved the task with MSE less than 0.05. After training, we tested all models on the same 64 input sequences and extracted the recurrent states, upon which we apply Procrustes and DSA in parallel to compare networks.
Our RNNs have different geometry but equivalent attractor topology (Maheswaranathan et al. [2019]). Fig.2a. displays the trajectories of individual trials in a sample network in the first three principal
components, which capture on average  $95.9\%$  of the total variance. The computational solution to this task is characterized by eight stable fixed points at the vertices of a cube that correspond to the eight possible unique memory states  $(2^{3})$ . Importantly, this cubic structure is preserved across all networks, even though particular geometric aspects of the trajectories may differ. Standard shape analyses differentiate architectures and activation function by geometry, which we replicated in Fig. 2b (architecture) and in the Supplementary Information (activation, 13).
We fit individual HAVOK models with 75 delays and rank 100 to sampled trajectories from each network. For computational efficiency, we reduced the dimensionality of the network to 10 using Principal Component Analysis (PCA) before applying HAVOK. In Supplementary Information 14, we demonstrate that HAVOK's predictivity is robust to hyperparameter variation. Because the variance explained from the later PCs is negligible, they can be disregarded while preserving the original dynamics. This improves HAVOK's fit to the data and reduces the memory and time complexity of DSA. For each pair of networks, we compared these states directly using Procrustes Analysis and on the DMD matrices using Procrustes Analysis on Vector Fields. This yields two dissimilarity matrices (DMs), which we reduced to two dimensions using Multidimensional Scaling (MDS, Kruskal [1964]). We plot the results in this similarity space in Fig. 2, where each point is a trained network. As expected, the networks cluster by architecture (Fig. 2b) and activation function (Supplementary Information) when Procrustes is applied, indicating that aspects of their geometry differ. However, the networks intermingle under DSA (Fig. 2c, inset zoomed in), and have pairwise dissimilarities close to zero, indicating that their underlying dynamics are equivalent, which is correct (Maheswaranathan et al. [2019]). We quantified these differences by applying a linear SVM to the dissimilarity matrices and plotted the test accuracy in Fig. 2d. This further underscores the differences between Procrustes and DSA in Fig. 2b and c, as DSA's accuracy is close to chance, whereas Procrustes achieves high accuracy on both labels.
Figure 2: Same dynamics, different shape, only identified as similar with DSA. Dots indicate a single trained network. Analysis of RNNs trained on the 3-bit FlipFlop task. a. PCA trajectories of a single trained RNN. b. MDS projection of the dissimilarity matrix (DM) computed across condition-averaged hidden states between each RNN (architecture indicated by color). In this view of the data, RNNs cluster by architecture. c. MDS embedding of the DM generated from DSA of the same set of RNNs as in b. Here, RNNs do not cluster by architecture. d. Classification test accuracy of condition-averaged Procrustes and DSA similarity spaces on both architecture and activation labels. Dotted lines indicate chance.
3.3 DSA identifies dynamical differences despite geometric similarity