Learning Manifold and Itô Dynamics with Branched Neural Rough Differential Equations

摘要

Neural rough differential equations (NRDEs) learn continuous-time dynamics from irregularly sampled sequences by encoding the input path with signature features, providing robustness to discretisation and sampling irregularity. However, existing NRDEs implicitly rely on algebraic identities that can fail in two important settings: *stochastic dynamics* interpreted in the Itô sense, and *dynamics evolving on manifolds* where curvature renders the effect of repeated derivatives order-dependent. In this work, we propose Branched Neural Rough Differential Equations (B-NRDEs), a unified framework that replaces geometric signature features with tree-based (branched) rough-path lifts, yielding models that remain well-defined under Itô noise and on manifolds. Building on these branched lifts, an Itô-consistent training objective is introduced via the branched signature kernel. We provide an efficient, autodifferentiable package *Stochastax* for computing branched (log-)signatures and solving (manifold) RDEs. Across various applications, including rough Bergomi volatility modelling, sim-to-real $\mathrm{SO}(3)$ dynamics forecasting, and SPD covariance dynamics, B-NRDE shows consistently strong results.