Structure-Preserving Learning: Koopman, Transformers, and Hamiltonian Dynamics

This document outlines the theoretical framework for diffSimu-AI, integrating deep learning with geometric mechanics. The approach bridges the gap between the expressivity of sequence models (Transformers) and the rigorous conservation laws of Hamiltonian mechanics.

1. From NLP to Physical Dynamics

Recent advances in scientific machine learning (e.g., Geneva & Zabaras, 2021) have established a compelling parallel between residual self-attention in Natural Language Processing and multi-step explicit integration schemes (such as Runge-Kutta or Adams-Bashforth).

The Residual-ODE Analogy

A Transformer layer with a residual connection updates the latent state \(h_l\) to \(h_{l+1}\):

This is structurally identical to a forward Euler discretization of an Ordinary Differential Equation (ODE) \(\dot{h} = f(h, t)\) with step size \(\Delta t = 1\):

• The Analogy: The attention mechanism \(\text{Attention}(h_l)\) approximates the vector field (velocity) describing the system's evolution.
• The Limitation: While standard attention mechanisms can approximate complex dynamics, they lack intrinsic physical guarantees. They do not inherently conserve energy (Hamiltonian) or phase space volume (Symplecticity), often leading to non-physical drift in long-horizon simulations.

2. End-to-End Architecture (Koopman → Transformer → Geometric Solver)

To reconcile data-driven modeling with physical structure, the diffSimu-AI architecture rests on three pillars:

1. Koopman Embedding: Mapping non-linear state dynamics into a regularized space of observables.
2. Transformer Dynamics: Leveraging attention mechanisms to learn the scalar Energy field (Hamiltonian) and complex interactions, rather than directly predicting the next state.
3. Geometric Mechanics: Imposing symplectic and holonomic constraints to ensure conservation and stability via dedicated integrators.

Interfaces (inputs/outputs)

• Input state: \(x_n=(q_n,p_n)\), optionally with context \(u_n\) (boundary conditions, actuation, multi-body coupling, multi-physics signals).
• Koopman embedding: \(z_n=\psi(x_n)\) (learned encoder; optional decoder back to \(x\) for reconstruction losses).
• Energy model: a Transformer maps \((z_n,u_n)\) to \(H_\theta\) (or directly to \(\nabla_q H_\theta\)), and gradients are obtained either by automatic differentiation or direct prediction.
• Geometric solve + update: the residual/constraint solve returns \((a_n,\lambda_n)\), and a constrained symplectic integrator updates \((q,p)\) to produce \(x_{n+1}\).

Related work / Inspirations

The following references illustrate components that inspire the diffSimu-AI decomposition (latent embedding → energy model → structure-preserving integration). They are not meant as a strict blueprint, but as anchor points for the design space.

Geometric RO-HNN (reduced-order Hamiltonian neural networks)

Learning Hamiltonian dynamics at scale: a differential-geometric approach (2025)

Figure: Flowchart of the forward dynamics of a Geometric RO-HNN.

The architecture (blue) builds a constrained auto-encoder to learn the latent coordinate space \(\check{q}\). The dynamics (brown) are learned via a geometric HNN composed of SPD networks for mass/inertia matrices and an MLP for potential energy \(H_\theta\). The time evolution (right) is handled by a symplectic integrator, preserving the structure of the learned Hamiltonian system.

Transformers for physical dynamics

Transformers for modeling physical systems (2022)

Figure: Transformer-based structure-preserving dynamics architecture.

Koopman auto-encoders for Hamiltonian dynamics

Figure: Koopman auto-encoder (Koopman AE) for Hamiltonian dynamics.

Notation (used throughout)

• \(q\): generalized configuration coordinates (on the configuration manifold \(\mathcal{Q}\)).
• \(p\): generalized momenta, so the full Hamiltonian state is \(x = (q, p)\).
• \(z = \psi(x)\): Koopman observables / embedded state used as model input.
• \(h\): Transformer hidden states / tokens (used in the residual-layer analogy; distinct from \(z\)).
• \(H_\theta(q,p)\): learned unconstrained Hamiltonian (a scalar energy field); forces come from its gradients \(\nabla_q H_\theta\), \(\nabla_p H_\theta\).
• \(\Phi(q)=0\): holonomic constraints; \(G(q)=\partial \Phi/\partial q\) is the constraint Jacobian and \(\lambda\) are Lagrange multipliers.

3. Koopman Embedding: Learning Observables

Instead of learning dynamics directly on the full Hamiltonian state \(x=(q,p)\), we employ Koopman theory to project the state into a space of observables \(z=\psi(x)\).

The Koopman Operator

For a nonlinear dynamical system \(x_{t+1} = F(x_t)\), the Koopman operator \(\mathcal{K}\) is an infinite-dimensional linear operator acting on the space of scalar-valued observable functions \(\psi\):

• Linearization: While the state dynamics \(F\) are nonlinear, the dynamics of observables are linear: \(\psi_{t+1} = \mathcal{K} \psi_t\).
• Observables vs. Coordinates: Observables are functions of the state rather than just position/velocity vectors.

Types of Learned Observables

In diffSimu-AI, we do not hand-craft observables but learn them to capture specific physical structures:

1. Spatial Observables (Geometric Invariants): Geometric features extracted from the mesh topology (e.g., via Equivariant GNNs). They encode localized symmetries, strain invariants, or curvature, ensuring that the learned physics respects the underlying Euclidean symmetries.

2. Temporal Observables (Memory): States encoding the trajectory history. Mechanistically, these are analogous to Delay Embeddings (Takens' theorem) or "Time Cells" in neuroscience. They allow the Transformer to resolve non-Markovian dynamics by effectively reconstructing the phase space from partial observations.

3. Spectral Observables (Dynamic Modes): Learned spectral basis functions of the Koopman operator. These isolate invariant subspaces or specific time-scales (slow vs. fast modes), allowing the model to separate global rigid-body motion from high-frequency deformations.

4. The Dual Role of Attention

Although the symplectic integrator handles the temporal evolution, the Transformer remains the core "brain" of the system, responsible for computing a scalar energy (the Hamiltonian \(H_\theta\)) and/or its gradients from which forces are derived.

Why Attention is Essential

Even with an explicit integrator, the Transformer is critical for:

1. N-Body & Global Interactions: The mechanism of Self-Attention is mathematically ideal for modeling interactions in complex systems (fluids, soft tissues, molecular dynamics). It allows any part of the system to "attend" to any other, effectively computing the energy \(H_\theta(q, p)\) arising from long-range couplings.

2. Learning Unknown Physics: For materials with complex constitutive laws (e.g., biological tissues), exact analytical formulas for energy are often unknown. The Transformer acts as a Universal Approximator to learn the true shape of the Hamiltonian from data.

3. Dimensionality Reduction: By operating on the latent observables space \(z\), the Transformer computes physics in a compressed representation, enabling the simulation of massive systems much faster than classical FEM methods.

Self-Attention: The Potential Learner

Self-attention acts on the state of observables to compute the internal energy distribution. It functions as the "physicist" that analyzes the configuration to determine the potential field.

Cross-Attention: Hamiltonian Ports

Cross-attention couples the internal state with external factors. Physically, this corresponds to Hamiltonian ports in a Port-Hamiltonian System (PHS) framework. Cross-attention handles:

• External Forces & Boundary Conditions: Injections of energy or constraints from the environment (e.g., collisions, motor commands).

• Multi-Body Interactions: Couplings between distinct physical bodies.

• Multi-Physics: Interactions between different physical domains (e.g., thermal, electromagnetic).

The cross-attention mechanism effectively modifies the energy landscape in real-time based on external context.

5. Geometric Mechanics & Constraints

A fundamental principle of this architecture is that neither Koopman embeddings nor Transformers inherently guarantee symplectic structure. To ensure physical fidelity, we must ground the learning process in Geometric Mechanics.

Manifolds and Phase Space

The system evolves on a Configuration Manifold \(\mathcal{Q}\) (e.g., the space of all valid mesh deformations). The state of the system is a point \((q, p)\) in the Cotangent Bundle \(T^*\mathcal{Q}\) (Phase Space), which naturally carries a symplectic structure.

Constrained Hamiltonian Dynamics

Real-world systems are rarely unconstrained. They are subject to Holonomic Constraints \(\Phi(q) = 0\) (e.g., incompressibility, joints, contacts) that define a submanifold \(\mathcal{C} \subset \mathcal{Q}\).

The dynamics on this constrained manifold are governed by the Lagrange-d'Alembert Principle, leading to a system of Differential Algebraic Equations (DAEs) of index 3:

Where:

• \(H_\theta\): The Unconstrained Hamiltonian learned by the Transformer (representing internal potential energy and inertia).

• \(G(q) = \partial \Phi / \partial q\): The Jacobian of the constraints.

• \(\lambda\): The Lagrange Multipliers (constraint forces) that strictly enforce \(\Phi(q)=0\).

The Hybrid Split

This formulation reveals the division of labor in diffSimu-AI:

• The Neural Network learns the Physics (\(H_\theta\), \(\nabla H_\theta\)). It predicts the "trial" evolution based on energy.

• The Geometric Solver enforces the Constraints (\(\lambda\)). It projects the flow onto the valid manifold \(T\mathcal{C}\).

This ensures that even if the neural network makes approximation errors, the system never violates geometric validity (interpenetration, joint breakage).

6. Learning objectives

Training diffSimu-AI typically combines representation learning (Koopman AE), dynamics supervision (rollouts), and physics/structure regularization (energy, constraints, SPD).

Forward pass (conceptual)

Given current state x_n=(q_n,p_n) and context u_n:
  z_n = Encoder(x_n)                       # Koopman embedding
  H_theta = Transformer(z_n, u_n)          # scalar Hamiltonian (or predict grad_q H_theta)
  gradH = grad_{q,p} H_theta              # autodiff or direct prediction
  (a_n, lambda_n) = SolveResidual(q_n,p_n, gradH, u_n)   # Newton-Krylov / Schur complement
  x_{n+1} = SymplecticUpdate(x_n, a_n, lambda_n)
  xhat_n = Decoder(z_n)                   # optional (reconstruction)

Loss (typical decomposition)

Let \(x=(q,p)\) and \(z=\psi(x)\). A common objective is a weighted sum:

• Koopman AE reconstruction (optional, if using an encoder/decoder):

  \[ \mathcal{L}_{\text{recon}} = \sum_{n}\|\mathrm{Decoder}(z_n) - x_n\|^2. \]

• Rollout supervision (choose state-space or force-space supervision):

  \[ \mathcal{L}_{\text{roll}} = \sum_{n}\sum_{k=1}^{K}\|\hat{x}_{n+k} - x_{n+k}\|^2 \quad\text{or}\quad \sum_{n}\|\widehat{\nabla_q H_\theta}(x_n) - \nabla_q H^\star(x_n)\|^2. \]

• Energy drift penalty (for conservative/unforced segments):

  \[ \mathcal{L}_{\text{energy}} = \sum_{n}\sum_{k=1}^{K}\left|H_\theta(\hat{x}_{n+k}) - H_\theta(\hat{x}_{n})\right|. \]

• Constraint violation penalty (mostly for robustness, diagnostics, or when constraints are only approximately enforced):

  \[ \mathcal{L}_{\text{cons}} = \sum_{n}\|\Phi(\hat{q}_n)\|^2. \]

• SPD regularization (only if learning inertial terms such as \(M_\theta(q)\)):

  \[ \mathcal{L}_{\text{spd}} = \sum_{n}\left\|M_\theta(q_n) - M_\theta(q_n)^\top\right\|^2 \;+\; \text{(penalty for non-positive eigenvalues)}. \]

7. Integration with Sofax

This abstract geometric formulation maps directly to the concrete implementation of the Residual Operator in Sofax.

The Geometric Residual

The core object in Sofax is the residual \(\mathcal{R}\), which measures the deviation from the constrained symplectic flow, with unknowns the generalized acceleration \(a\) and multipliers \(\lambda\). For a discrete time step, finding the next state amounts to finding the root of:

Here \(M(q)\) is the (possibly learned) mass/inertia matrix, \(f_{ext}\) are external generalized forces, and \(f_{unc}\) denotes unconstrained predicted forces (before constraint projection).

subject to \(\Phi(q_{n+1}) = 0\).

Hybrid Solver Pipeline

The convergent pipeline for diffSimu-AI integrates the neural potential into the Newton-Krylov solver:

1. Koopman Embedding: The state \((q, p)\) is lifted to observables \(z\).
2. Hamiltonian Transformer: The network predicts energy gradients (e.g., \(\nabla_q H_\theta\)), providing the constitutive terms (stiffness, hyperelasticity) for the residual.

3. Schur Complement Solve: Sofax solves the saddle-point problem to find the accelerations \(a\) and multipliers \(\lambda\) simultaneously:

   \[ \underbrace{(G M^{-1} G^\top)}_{\text{Compliance}} \lambda = \underbrace{\frac{\Phi}{\Delta t^2} + G M^{-1} f_{unc}}_{\text{Violation + Prediction}} \]

   This step effectively projects the learned dynamics onto the constraint manifold.

4. Symplectic Step: The state is updated using a constrained symplectic integrator (e.g., RATTLE/SHAKE) to preserve phase space volume.

Note on Diffrax (JAX) symplectic solvers

For unconstrained or split Hamiltonian systems (canonical \((q,p)\) form), the JAX ecosystem also provides ready-to-use symplectic ODE solvers via Diffrax. In particular, Diffrax documents a dedicated family of symplectic methods (e.g., semi-implicit Euler / leapfrog-style solvers) and the expected term/state structure for Hamiltonian dynamics: see Diffrax — Symplectic methods.

End-to-end differentiability

In this pipeline, everything is differentiable by construction: the Koopman encoder/decoder, the Transformer energy model, and the time-stepping machinery. Gradients can flow through the geometric step either by differentiating the unrolled iterations (when stable) or via implicit differentiation of the residual solve, enabling end-to-end training and inverse problems with both models and solvers in the loop.

By embedding the learned Hamiltonian into the Sofax residual, we achieve a simulation that is expressive (capturing complex material behavior via AI) yet exact (respecting kinematic constraints via the solver).

8. Conclusion

By separating the learning of operators (via Koopman and Transformers) from the integration of dynamics (via Symplectic Solvers), diffSimu-AI creates a system where machine learning provides expressivity while mathematical structure guarantees physical validity.

The ML learns the dynamics within a specific latent structure, ensuring physical fidelity by construction.

Glossary

• \((q,p)\): generalized configuration and momenta (Hamiltonian coordinates).
• \(x=(q,p)\): full Hamiltonian state.
• \(z=\psi(x)\): Koopman observables / embedded state.
• \(h\): Transformer hidden states / tokens (internal representation of the sequence model).
• \(H_\theta\): learned unconstrained Hamiltonian (scalar energy field).
• \(\Phi(q)=0\), \(G(q)\), \(\lambda\): holonomic constraints, constraint Jacobian, and Lagrange multipliers (constraint forces).
• \(a\): generalized acceleration (unknown of the residual solve).
• \(\mathcal{R}\): Sofax residual operator defining the implicit constrained step.