TRASE-NODEs: Trajectory Sensitivity-aware Neural Ordinary Differential Equations for Efficient Dynamic Modeling

TL;DR

This work extends neural ordinary differential equations by embedding trajectory sensitivity to the control input into an augmented state, enabling joint learning of state and sensitivity via a memory-efficient adjoint method. The TRASE-NODEs framework defines an augmented ODE for $[\mathbf x, \mathbf s]$, where $\mathbf s=\partial \mathbf x/\partial \mathbf u$, and derives an augmented adjoint system to compute gradients with respect to model parameters, despite the added second-order derivative costs. Evaluation on a damped oscillator and inverter-based resources demonstrates markedly improved generalization and accuracy for unseen control inputs, with substantial gains in sensitivity prediction and smaller gains in highly nonlinear settings. The approach offers a data-efficient, control-oriented pathway for dynamic modeling in systems where trajectory sensitivity to inputs is crucial for robust control and stability analysis.

Abstract

Modeling dynamical systems is crucial across the science and engineering fields for accurate prediction, control, and decision-making. Recently, machine learning (ML) approaches, particularly neural ordinary differential equations (NODEs), have emerged as a powerful tool for data-driven modeling of continuous-time dynamics. Nevertheless, standard NODEs require a large number of data samples to remain consistent under varying control inputs, posing challenges to generate sufficient simulated data and ensure the safety of control design. To address this gap, we propose trajectory-sensitivity-aware (TRASE-)NODEs, which construct an augmented system for both state and sensitivity, enabling simultaneous learning of their dynamics. This formulation allows the adjoint method to update gradients in a memory-efficient manner and ensures that control-input effects are captured in the learned dynamics. We evaluate TRASE-NODEs using damped oscillator and inverter-based resources (IBRs). The results show that TRASE-NODEs generalize better from the limited training data, yielding lower prediction errors than standard NODEs for both examples. The proposed framework offers a data-efficient, control-oriented modeling approach suitable for dynamic systems that require accurate trajectory sensitivity prediction.

Figures (8)

• Figure 1: Comparisons between true and NODE-predicted state trajectories for a simple linear ODE $\dot x = -3x + u$. (a) Training results with both inputs of $u=1,2$ show an accurate model fit. (b) Testing results on unseen inputs of $u=3,4$ reveal a noticeable prediction mismatch. This example highlights that even for a simple ODE, the sample efficiency and generalizability issues could still be of concern for the data-driven NODEs. Our proposed TRASE-NODEs can improve the generalizability to unseen data thanks to the inclusion of trajectory sensitivity information.
• Figure 2: Proposed TRASE-NODEs training builds upon updating the NODEs framework. Standard NODEs components are marked by solid lines, while the updated by dotted lines. Forward pass only predicts the terminal state using an ODE solver, while backward pass consists of forming the adjoint variables and using it for gradient computation.
• Figure 3: Selected state trajectories for testing cases with $u = \{2,3,4\}$. Solid lines denote ground-truth trajectories, with dotted lines for predictions. TRASE-NODEs predictions in (a) displacement $x(t)$ and (b) velocity $v(t)=\dot{x}(t)$. NODEs predictions in (c) $x(t)$ and (d) $v(t)=\dot{x}(t)$.
• Figure 4: Selected sensitivity trajectories for testing cases with $u = \{2,3,4\}$. Solid lines denote ground-truth trajectories, with dotted lines for predictions. TRASE-NODEs predictions in (a) $s_x(t)$ and (b) $s_v(t)$. NODEs predictions in (c) $s_x(t)$ and (d) $s_v(t)$. Note that ground-truth sensitivity is identical across all test cases, as the analytical solutions for \ref{['eq:sxv']} do not depend on the input $u$.
• Figure 5: NMSE of TRASE-NODEs (blue) and standard NODEs (red) in logarithmic scale over the range of input $u$ for (a) displacement $x(t)$, (b) velocity $v(t)$, (c) displacement sensitivity $s_x(t)$, and (d) velocity sensitivity $s_v(t)$. The y-axis is plotted on a logarithmic scale to highlight differences across several orders of magnitude.