LASS-ODE: Scaling ODE Computations to Connect Foundation Models with Dynamical Physical Systems

TL;DR

LASS-ODE introduces a scalable approach to connect foundation models with physical systems by using locally linear ODE representations and a shared structure hub, enabling efficient, accurate, and generalizable dynamic predictions.

Contribution

It proposes a novel token representation respecting locally linear ODE evolution and a shared structure hub for knowledge sharing, scaling ODE computations for foundation models.

Findings

Enables zero-shot generalization across diverse ODE systems

Pretrained on 40GB of ODE trajectories for strong in-domain performance

Accelerates ODE integration while maintaining accuracy

Abstract

Foundation models have transformed language, vision, and time series data analysis, yet progress on dynamic predictions for physical systems remains limited. Given the complexity of physical constraints, two challenges stand out. $(i)$ Physics-computation scalability: physics-informed learning can enforce physical regularization, but its computation (e.g., ODE integration) does not scale to extensive systems. $(ii)$ Knowledge-sharing efficiency: the attention mechanism is primarily computed within each system, which limits the extraction of shared ODE structures across systems. We show that enforcing ODE consistency does not require expensive nonlinear integration: a token-wise locally linear ODE representation preserves physical fidelity while scaling to foundation-model regimes. Thus, we propose novel token representations that respect locally linear ODE evolution. Such linearity substantially accelerates integration while accurately approximating the local data manifold. Second, we introduce a simple yet effective inter-system attention that augments attention with a common structure hub (CSH) that stores shared tokens and aggregates knowledge across systems. The resulting model, termed LASS-ODE (\underline{LA}rge-\underline{S}cale \underline{S}mall \underline{ODE}), is pretrained on our $40$GB ODE trajectory collections to enable strong in-domain performance, zero-shot generalization across diverse ODE systems, and additional improvements through fine-tuning.