Equivariant Manifold Neural ODEs and Differential Invariants (2401.14131v2)

Abstract: In this paper, we develop a manifestly geometric framework for equivariant manifold neural ordinary differential equations (NODEs) and use it to analyse their modelling capabilities for symmetric data. First, we consider the action of a Lie group $G$ on a smooth manifold $M$ and establish the equivalence between equivariance of vector fields, symmetries of the corresponding Cauchy problems, and equivariance of the associated NODEs. We also propose a novel formulation, based on Lie theory for symmetries of differential equations, of the equivariant manifold NODEs in terms of the differential invariants of the action of $G$ on $M$, which provides an efficient parameterisation of the space of equivariant vector fields in a way that is agnostic to both the manifold $M$ and the symmetry group $G$. Second, we construct augmented manifold NODEs, through embeddings into flows on the tangent bundle $TM$, and show that they are universal approximators of diffeomorphisms on any connected $M$. Furthermore, we show that universality persists in the equivariant case and that the augmented equivariant manifold NODEs can be incorporated into the geometric framework using higher-order differential invariants. Finally, we consider the induced action of $G$ on different fields on $M$ and show how it can be used to generalise previous work, on, e.g., continuous normalizing flows, to equivariant models in any geometry.

• The paper introduces a geometric framework that leverages differential invariants to parameterize equivariant manifold NODEs.
• It demonstrates augmented NODEs as universal approximators for equivariant diffeomorphisms on various manifolds.
• The methodology extends NODEs to non-Euclidean domains, enhancing generative and discriminative models with symmetry constraints.

Introduction

The paper of neural ordinary differential equations (NODEs) has seen substantial progress, particularly since the proposal of NODEs in their modern form. NODEs originate from the perspective of dynamical systems as the limit of recurrence in infinitely deep networks. Notably, NODEs facilitate the exploration of continuous transformations in deep learning, yielding models with desirable properties, such as reversible transformations that are conceptually attractive for generative models. Recent advancements have extended NODEs to non-Euclidean domains through neural manifold ODEs (manifold NODEs), addressing applications where data exhibit symmetries under a group of transformations and necessitating equivariant NODEs.

Equivariant Manifold NODEs and Differential Invariants

This paper presents a rigorous geometric framework for equivariant manifold NODEs by drawing on classical Lie theory for symmetries in differential equations. This framework allows for the parameterization of equivariant vector fields, regardless of the specific manifold or symmetry group, by using differential invariants of the action of a Lie group on the manifold. In doing so, equivariant differential equations are parameterized. Furthermore, the paper introduces augmented manifold NODEs that embody universal approximators for equivariant diffeomorphisms on path-connected manifolds. Higher-order differential invariants permit the construction and parameterization of augmented NODEs within this geometric approach.

Related Work and Contributions

The paper situates itself among seminal works that introduced normalized flows as generative models constructed from simple distributions and NODEs acting on manifolds, and recent efforts that have achieved flows on Riemannian manifolds equivariant under subgroups of isometries. Extending beyond these contributions, the authors establish a comprehensive approach that can accommodate any symmetry group and manifold and address previous restrictions to conservative vector fields by relying on the theory of symmetries in differential equations. The universality of NODEs is a pivotal consideration, with previous studies demonstrating that augmenting NODEs with sufficient dimensions assures their ability to approximate diffeomorphisms. The investigation also appraises the critical task of efficiently parameterizing the space of vector fields on manifolds to articulate expressive models.

Augmentation, Universality, and Applications

The paper explores the construction of universal approximators of equivariant diffeomorphisms through augmented equivariant NODEs on tangent bundles. The differential invariants serve to parameterize these models, offering an insightful geometric understanding of augmentation. Beyond the mathematical underpinning, the authors discuss how NODEs can transform various geometric objects on manifolds, extending from scalar fields to vector fields. The work ensures that transforming a G-equivariant object with a G-equivariant diffeomorphism preserves the equivariance property. This aspect has applications in constructing models for various feature maps in machine learning that uphold symmetry principles, corresponding to 'physically meaningful constraints.

Conclusion

The devised geometric framework for equivariant manifold NODEs enriches the methodological tools available for machine learning applications involving symmetries. By harnessing the foundational principles of differential invariants in equivariant differential equations, the research marks a significant advancement in parameterizing and modeling equivariant flows on manifolds, suitable for a broad spectrum of symmetries and topological intricacies. The implications of such a framework are wide-ranging, from more expressive and accurate models in physical sciences to widespread applications in generative and discriminative machine learning tasks. Future work proposes to further test the capacity of NODEs across varying applications, as well as to explore the relationship between the geometric framework and conservation laws prevalent in many scientific domains.