Characterizing control between interacting subsystems with deep Jacobian estimation

We thank the reviewer for their constructive feedback and questions about the paper, and address them here.

Dense writing and boundaries of the method

We rewrote section 3 to clarify key intuitions and moved technical details to the appendix for improved narrative flow. If accepted, we will use the extra page to space out text and figures.

To better define the method’s boundaries and clarify our claims, we substantially revised the limitations section (see also reviewer 1):

A future challenge for JacobianODEs is partially observed dynamics. Recent work has identified that it is possible to learn latent embeddings that approximately recover the true state from partial observation [108-112]. Jacobian-based dynamics learning has been performed in a latent space [61, 113], yet it is unclear whether this translates to accurate Jacobian estimation, given the challenges related to automatic differentiation and Jacobian estimation presented here.

We also note that reachability estimates depend sensitively on several factors: the alignment of cross-subsystem interactions and within-subsystem dynamics, the eigenvectors and eigenvalues of within-subsystem Jacobians, and the way activity propagates within each subsystem (see appendix C.1). While our method reliably captures broad trends in reachability over time, fine-grained, time point-specific comparisons should be interpreted with caution.

Finally, although JacobianODEs scale well to moderately high-dimensional systems, their performance in systems that are orders of magnitude larger than those considered here (e.g., recordings of thousands of neurons via calcium imaging) remains to be tested. In many practical settings, this may not be a limitation: neural representations during tasks often exhibit intrinsic dimensionalities comparable to those studied here, enabling JacobianODEs to operate in a reduced dimensionality. Notably, most models converged in under 45 minutes on a single GPU, suggesting favorable scaling. Future work should explore the viability of the method in higher dimensions, and assess whether dimensionality reduction strategies (such as latent state models or low-rank Jacobian approximations) can further improve scalability and accuracy.

The model of the RNN should be shown in the main text

We agree that this was an oversight. We now introduce the RNN architecture directly in section 5, as suggested.

Suggested analysis of more complex RNN architectures, due to potentially low-dimensional Jacobian estimation

TL;DR: The Jacobians were not low-rank. While both RNN activity and Jacobians lie on low-dimensional manifolds relative to the 128D state space, they match dimensionality observed in neural data, as the reviewer noted. The 64D Lorenz 96 system had much higher intrinsic dimensionality, supporting our method’s effectiveness on complex systems.

The reviewer astutely notes that RNN activity is often confined to low-dimensional manifolds, potentially simplifying Jacobian estimation. We address this point through the following analyses.

Jacobian rank: We first note that both the RNN hidden weight matrix and nearly all Jacobians had full rank (128; ~0.1% had rank 127). To assess the effective rank, we computed the participation ratio (PR), a singular value-based measure of intrinsic dimensionality. The RNN weight matrix had a PR of 123, and the Jacobians had a mean PR of ~106 (std 0.7).

Intrinsic-data dimensionality: As the reviewer notes, the data itself may be low-dimensional. Since the Jacobian is a function of the trajectory data, the Jacobian manifold’s intrinsic dimension is at most that of the trajectory manifold. To estimate the intrinsic dimensions, we pooled time points from a large sample of trajectories. For RNN state vectors, we computed PR directly; for Jacobians, we flattened them into $n^2$-dimensional vectors before computing PR. We found the RNN trajectory PR to be ~13.5, and the Jacobian PR to be ~6. Given the 128D extrinsic space, this confirms both RNN states and Jacobians lie on lower-dimensional manifolds.

Notably, prior work estimated the dimensionality of prefrontal cortex delay activity at 24 using ~4000 neurons, falling to ~6 when subsampled to 100 neurons (Rigotti et al., The importance of mixed selectivity in complex cognitive tasks. Nature, 2013). These values were interpreted as high-dimensional. Thus, the dimensionalities we observe are consistent with biological systems performing complex tasks.

Furthermore, in the 64-dimensional Lorenz 96 system, the trajectory PR was ~47.5, and the Jacobians had a PR of ~50, which we believe to be high-dimensional enough to demonstrate the effectiveness of our method on complex systems.

We note that in larger systems, approximating the Jacobian as low rank may be more efficient. Rather than computing each of the $n^2$ outputs directly, we can constrain the Jacobian to be low-rank by designing the output as a low-rank combination of vectors (e.g., it learns to output a low-rank singular value decomposition).

We agree that applying our method to more complex architectures (e.g., LSTMs) and tasks is an important direction for future work. For now, we believe our results sufficiently demonstrate the method’s ability to capture rich, high-dimensional dynamics.

As suggested, we added a discussion of dimensionality to the limitations section, as well as to the appendix. We thank the reviewer for prompting this clarification of an important aspect of the method.

How the Jacobian at a given time point can be approximated without repetition

We appreciate this question and recognize it was not fully addressed. The loop closure loss provides the primary mechanism as it constrains the Jacobian’s form beyond fitting observed trajectories. We have rewritten the loop closure section as follows (and see reviewer 2 discussion):

The Jacobian captures how perturbations to the system will propagate along any direction in state space. Estimating it purely from dynamics constrains only the direction of the flow, leaving the full solution underdetermined. To address this, we again exploit the fact that each row of the Jacobian is a conservative vector field. Specifically, we note that for any piecewise smooth loop $\mathcal{C}\_{\text{loop}}$, we have $\left \Vert \oint\_{\mathcal{C}\_{loop}} \mathbf{J} ds \right \Vert_2 = 0$. Thus, by integrating along loops that contain directions orthogonal to the system's dynamics (and penalizing the deviation from zero), we encourage the estimated Jacobians to capture information about other directions in state space (see appendix A.3 for full technical details). To ensure broad coverage of tangent space directions, we form loops from concatenations of line integrals between randomly selected data points. This strategy samples diverse directions from the tangent space while remaining easy to compute. This self-supervised loss term builds on the loss introduced by Iyer et. al. [63]. It constrains $\hat{\mathbf{J}}^{\theta}$ to satisfy both the dynamics and conservativity. This improves Jacobian estimation accuracy significantly (see appendix C.4 for ablation studies).

Observation noise also helps the model learn responses in directions beyond the primary flow. We added small observation noise to each batch during training on the task-trained RNNs. As noted in the “Jacobian reconstruction quality” section: “Noise encourages the model to explore how perturbations around the observed trajectory evolve, which is crucial for learning accurate Jacobians in high-dimensional systems.”

We hope this addresses the reviewer’s question and are happy to clarify further.

Gramian discussion is too quick

We added the Gramian equation into the main paper, and agree it is a key quantity of interest. We also expanded the discussion of this quantity as follows:

Since Jacobians are time-dependent, we obtain a linear time-varying representation of control dynamics along the trajectory. This enables computation of time-varying reachability ease, capturing how readily each area drives the other toward novel states [56-58]. Reachability is quantified via the reachability Gramian, a matrix defining a local metric in the tangent space. For the above control system capturing the influence of area B on area A, the time-varying reachability Gramian on the interval $[t\_0,t\_1]$ is defined as $\mathbf{W}\_r(t\_0, t\_1) \triangleq \int\_{t\_0}^{t\_1}\mathbf{\Phi}(t\_1, \tau)\mathbf{B} (\tau)\mathbf{B}^T(\tau)\mathbf{\Phi}^T(t\_1, \tau) d\tau$ where $\mathbf{\Phi}$ (computed from $\mathbf{J}^{A \to A}$) denotes the state-transition matrix of the intrinsic dynamics of subsystem A without any input (i.e., $\delta\mathbf{x}^A(t) = \mathbf{\Phi}(t, t_0)\delta\mathbf{x}^A(t\_0)$), and $\mathbf{B}(\tau) = \mathbf{J}^{B \to A}(\mathbf{x}(\tau))$ [56]. The Gramian $\mathbf{W}\_r(t\_0, t\_1)$ is symmetric and positive semidefinite for every $t\_1 > t\_0$ [56]. Each eigenvalue of the reachability Gramian quantifies how easily the target subsystem can be driven along its corresponding eigenvector. Thus, the trace of the reachability Gramian reflects average ease of reaching new states. The minimum eigenvalue indicates the lowest ease of control (i.e., the highest cost), and its corresponding eigenvector denotes the most difficult direction to control. For further details see appendix B.1.

Typo on line 150

This has been fixed.

Misleading title

Thank you for the comment. We’ve considered adding more neuroscientific context, and would appreciate your thoughts on this alternative: Deep Jacobian estimation for characterizing control between interacting subsystems with applications to neuroscience.

Transition to equation (1)

Equation 1 has been moved to better separate it from Figure 1.

Relocating the related work section

We moved the related work section (section 6) to follow immediately after the introduction.

We thank the reviewer for their helpful comments and feedback, and are glad that they found the method to be well-motivated, interesting, and theoretically strong. We now address the questions and comments posed by the reviewer.

Section 2.1 is out of place

We have moved section 2.1 to the analysis of the task-trained RNN (section 5), which we agree is a more conceptually aligned location.

Main text/appendix split & loop construction explanation

To clarify the narrative flow, we moved “Constructing the time derivative” and technical loop closure details to the appendix. In their place, we provide a more intuitive discussion of loop closure in the main text (see next point).

Formulation of the loop closure loss and methods of constructing loops

TL;DR: The loop closure loss encourages the Jacobian to learn about directions in the tangent space that are orthogonal to the flow of the dynamics. We construct loops from concatenated lines between random points on the data manifold to encourage uniform sampling of these directions.

We agree that the loop closure loss deserves more explanation. We have revised the main text as follows:

The Jacobian captures how perturbations to the system will propagate along any direction in state space. Estimating it purely from dynamics constrains only the direction of the flow, leaving the full solution underdetermined. To address this, we again exploit the fact that each row of the Jacobian is a conservative vector field. Specifically, we note that for any piecewise smooth loop $\mathcal{C}\_{\text{loop}}$, we have $\left \Vert \oint\_{\mathcal{C}\_{loop}} \mathbf{J} ds \right \Vert_2 = 0$. Thus, by integrating along loops that contain directions orthogonal to the system's dynamics (and penalizing the deviation from zero), we encourage the estimated Jacobians to capture information about other directions in state space (see appendix A.3 for full technical details). To ensure broad coverage of tangent space directions, we form loops from concatenations of line integrals between randomly selected data points. This strategy samples diverse directions from the tangent space while remaining easy to compute. This self-supervised loss term builds on the loss introduced by Iyer et. al. [63]. It constrains $\hat{\mathbf{J}}^{\theta}$ to satisfy both the dynamics and conservativity. This improves Jacobian estimation accuracy significantly (see appendix C.4 for ablation studies).

As mentioned, we have also moved the technical details to the appendix.

Other choices (e.g., forward and backward loops along observed trajectories or arbitrary circles) were considered, but our method offers a strong tradeoff between tractability and uniformity in sampling. We have added a discussion of these other potential choices for loops to the appendix as well.

We appreciate the reviewer’s suggestion, which helped clarify a core component of our method.

Images of trajectories from sample systems

We agree that the images of the trajectories in their current forms may not be optimally informative. As suggested, we will run tSNE, as well as UMAP, on the data and add the most visually salient / illustrative representations to Figures 4, 5, and S2. We have also added time-series plots of the top Principal Components in the appendix to highlight the temporal unfolding of each system.

Limited spacing, table captions, and small figure information

Regarding spacing, as mentioned, we have moved many details to the appendix, which should create more space. If accepted, we will further space out the text and figures to fill out the extra allotted page.

We have updated the following table captions and figure information to make it easier to read:

• More detail has been added to the caption for Table 1, Table S1, and Table S6
• Table S2 and S3 have been made larger
• The text in Figure 6 has been made larger

Delayed reachability curve in Figure 6B

TL;DR: The reachability Gramian is highly sensitive, and takes into account the alignment between cross- and within-area dynamics, as well as how within-area dynamics propagate over time. Small estimation errors can thus cause temporal offsets. We conjecture that small estimation errors in the within-area eigenvalues are the source of the observed delay, and now include a discussion of this phenomenon in our limitations.

As the reviewer astutely notes, in Figure 6B, cognitive-to-visual reachability peaks slightly late, while visual-to-cognitive hits a minimum slightly early. This likely stems from the sensitivity of the reachability Gramian, which arises from a complex interplay of components. These include the alignment of the cross-area interaction terms with the eigenvectors of the within-area dynamics, the corresponding eigenvalues of each of the eigenvectors, and how activity propagates within each area over successive time points. This phenomenon can be observed in the results in appendix C.1. Mathematically, consider the definition of the reachability Gramian currently in appendix B.1, in which the $\mathbf{\Phi}(t_1, \tau)$ term captures how within-subsystem dynamics propagate between times $\tau$ and $t_1$, and the multiplicative term $\mathbf{\Phi}(t_1, \tau)\mathbf{B}(\tau)$ captures how the directions of interaction between the subsystems aligns with the within-subsystem dynamics. The Jacobian estimation for the task-trained RNN, while robust considering the Jacobian consisted of a $128 \times 128$ matrix at each time point, was not perfect (we achieved an $R^2$ of 0.68 with 5% observation noise, as reported). Thus, for such a sensitive measure, even small Jacobian estimation errors can lead to timing deviations.

We can therefore conjecture that the source of this delay might have to do with small estimation errors in the eigenvalues of the estimated within-area dynamics matrices over time. Broadly speaking, when considering control over short time windows, (moderately) unstable eigenvalues in the target subsystem yield higher reachability, as activity can more easily be driven outward along their corresponding eigenvectors. As mentioned, the exact reachability values will depend sensitively on many factors, but the eigenvalues should provide a clue. For instance, the average predicted maximum eigenvalue in the visual area peaks later than the ground truth, which may explain the cognitive-to-visual delay. In the cognitive area, a transient dip in the average predicted maximum eigenvalue could cause the early minimum in visual-to-cognitive reachability.

Crucially, the delays in the reachability curves do not obscure the core result of the figure. On average, the cognitive-to-visual reachability tended to peak later in the delay than the visual-to-cognitive reachability. We furthermore note that this bias does not extend to the actual Jacobian estimation, for which the average residual is very close to zero, indicating that the JacobianODEs constitute an unbiased estimator.

This analysis highlights a limitation of our approach that we did not previously address, which relates to the interpretability of the time-resolved reachability curves. As demonstrated, reachability is a highly sensitive measure. While our method reliably captures broad trends in reachability over time, fine-grained, time point-specific comparisons should be interpreted with caution. We have added this to the discussion of limitations in the paper, and we thank the reviewer for encouraging this clarification.

Training on other laboratory tasks (e.g. motor control)

We agree that applying our method to laboratory tasks involving motor control, where control-theoretic frameworks are well established, would be highly compelling. This is a promising direction for future work, and we appreciate the suggestion.

I thank the authors for their conceptual clarifications, and I appreciate their effort in revising the relevant parts of their submission. I believe the core message of the submission will flow smoother now, I will continue to reccomend the acceptance of this work.

We thank the reviewer for their thoughtful feedback, and appreciate that the reviewer found our theoretical framework to be both elegant and broadly applicable. We address here the questions posed by the reviewer.

Comparison to echo state networks

TL;DR: Echo State Networks (ESNs) achieve comparable short-term prediction on chaotic systems, but fail to recover accurate Jacobians. This stems from structural limitations in mapping ESN dynamics back to the source system.

The reviewer aptly suggested a comparison to echo state networks (ESNs, or reservoir networks) in the case of chaotic dynamical systems. ESNs have been shown to be able to forecast the Lorenz 63 system for a particular choice of hyperparameters (Pathak, J. et. al. Using machine learning to replicate chaotic attractors and calculate Lyapunov exponents from data. Chaos, 2017). We note that Pathak et. al. (2017) implements the same method as Pathak et. al. (2018) referenced by the reviewer. We implemented the ESN architecture and hyperparameters from Pathak et al. (2017) for the Lorenz 63 system, which involved a reservoir of 300 nodes. Our implementation of the Lorenz system was simulated with a sampling interval of ~0.015 s, very similar to the 0.02 s used in the paper. Thus, for consistency of comparison with JacobianODEs, we fit ESN models to data sampled every 0.015 s.

We trained the ESNs on the same volume of data (with 1% observation noise) as the other models, and swept the regularization parameter $\beta$ to ensure a good fit. To ensure stable internal dynamics, we used a 150-step warm-up before generating predictions. The resulting ESNs predicted the Lorenz system well, achieving 10-step test MSE of 0.03—comparable to JacobianODEs (0.07) and NeuralODEs (0.04).

However, obtaining the Jacobian of the original Lorenz system from the ESN presents challenges. While the Jacobian of the internal reservoir dynamics (a 300×300 matrix) can be uniquely computed and was used in prior work to estimate Lyapunov exponents (Pathak et. al. 2017), recovering the Jacobian of the original system requires inverting the non-square readout matrix (Banerjee, A. et. al. Using machine learning to assess short term causal dependence and infer network links. Chaos, 2019). As shown in Banerjee et al. (2019), this inversion yields a non-unique solution, making such Jacobians ill-suited for system identification, and are used in the paper "only for causality estimation purposes".

We estimated these Jacobians nonetheless, using a long (7200-step) warm-up and true state input for stability. The resulting Jacobians had a Frobenius norm error (averaged over five random seeds) of $250 \pm 0.3$, compared to $3.3 \pm 0.2$ for JacobianODEs and $8.7 \pm 0.3$ for NeuralODEs. Thus, while ESNs can match predictive accuracy, they are significantly less accurate in capturing the system’s Jacobian, due to architectural limitations.

We have added a discussion of this comparison into the appendix of the paper, and we thank the reviewer for suggesting this valuable benchmark.

Stating limitations more explicitly

We have extended the limitations section to address limitations surrounding sensitivity of reachability and scalability, in addition to partial observability. The revised section is as follows:

A future challenge for JacobianODEs is partially observed dynamics. Recent work has identified that it is possible to learn latent embeddings that approximately recover the true state from partial observation [108-112]. Jacobian-based dynamics learning has been performed in a latent space [61, 113], yet it is unclear whether this translates to accurate Jacobian estimation, given the challenges related to automatic differentiation and Jacobian estimation presented here.

We also note that reachability estimates depend sensitively on several factors: the alignment of cross-subsystem interactions and within-subsystem dynamics, the eigenvectors and eigenvalues of within-subsystem Jacobians, and the way activity propagates within each subsystem (see appendix C.1). While our method reliably captures broad trends in reachability over time, fine-grained, time point-specific comparisons should be interpreted with caution.

Finally, although JacobianODEs scale well to moderately high-dimensional systems, their performance in systems that are orders of magnitude larger than those considered here (e.g., recordings of thousands of neurons via calcium imaging) remains to be tested. In many practical settings, this may not be a limitation: neural representations during tasks often exhibit intrinsic dimensionalities comparable to those studied here, enabling JacobianODEs to operate in a reduced dimensionality. Notably, most models converged in under 45 minutes on a single GPU, suggesting favorable scaling. Future work should explore the viability of the method in higher dimensions, and assess whether dimensionality reduction strategies (such as latent state models or low-rank Jacobian approximations) can further improve scalability and accuracy.