Transformers on Manifolds

NeurIPS 2025 Workshop NeurReps Submission76 Authors

Keywords: Geometric deep learning, transformers, manifolds, Lie groups, projected dynamical systems, flow matching, constrained learning, neural ODEs

Abstract: Many prediction tasks require outputs to lie on structured constraint sets (e.g., spheres, rotation groups, Stiefel/PSD manifolds), yet standard transformers provide no mechanism to enforce feasibility. We study \emph{geometry-aware transformers}, derived from projected dynamical variants of neural ordinary differential equations (NODEs), that respect constraints by construction. First, we introduce projected transformer layers that apply analytical projections (e.g., $\mathbb S^{d-1}$, Stiefel, PSD-trace) between blocks. Second, for Lie-group targets we propose an \emph{exponential transformer} that learns Lie-algebra coefficients and updates intrinsically via the matrix exponential. Third, when no closed-form projection exists, we learn a differentiable projection with \emph{flow matching} that acts as a data-driven retraction onto the (unknown) manifold. On synthetic dynamical systems over $\mathbb S^2$ and $SO(3)$, we evaluate Euclidean accuracy (MSE) and feasibility (norm/determinant statistics). Explicit projections and intrinsic exponential layers achieve \emph{perfect} feasibility, with exponential updates particularly strong on $SO(3)$ and show that our models have the best performance.

Submission Number: 76