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

TL;DR

LASS-ODE addresses the challenge of scaling ODE-based dynamics within foundation models to multi-system, heterogeneous physical processes. It introduces token-wise locally linear ODEs for scalable trajectory decoding and a Common Structure Hub (CSH) to enable cross-system knowledge sharing via intra-/inter-system attention, trained on a large corpus of $~40$ GB of ODE trajectories. Key innovations include time-aware tokenization with RBF modulation, channel-independent data handling, and a piecewise-linear ODE decoder that reduces computational overhead while preserving physical fidelity. Empirically, LASS-ODE achieves strong in-domain and zero-shot generalization across diverse ODE systems and benefits further from LoRA fine-tuning; the approach lays groundwork for scalable, physics-informed foundation modeling in dynamical systems and suggests extensions to SDEs and control tasks.

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.

Figures (16)

• Figure 1: The proposed LASS-ODE framework.
• Figure 2: (a). Good zero-shot result. (b). Bad zero-shot result. (c). Fine-tuned result. (d). Loss curve in fine-tuning.
• Figure 3: The t-SNE visualization and heatmap.
• Figure 4: Ablation results.
• Figure 5: Inference time.