Universal Physics Transformers: A Framework For Efficiently Scaling Neural Operators

Neural operators, serving as physics surrogate models, have recently gained increased interest. With ever increasing problem complexity, the natural question arises: what is an efficient way to scale neural operators to larger and more complex simulations -most importantly by taking into account different types of simulation datasets. This is of special interest since, akin to their numerical counterparts, different techniques are used across applications, even if the underlying dynamics of the systems are similar. Whereas the flexibility of transformers has enabled unified architectures across domains, neural operators mostly follow a problem specific design, where GNNs are commonly used for Lagrangian simulations and grid-based models predominate Eulerian simulations. We introduce Universal Physics Transformers (UPTs), an efficient and unified learning paradigm for a wide range of spatio-temporal problems. UPTs operate without grid-or particle-based latent structures, enabling flexibility and scalability across meshes and particles. UPTs efficiently propagate dynamics in the latent space, emphasized by inverse encoding and decoding techniques. Finally, UPTs allow for queries of the latent space representation at any point in space-time. We demonstrate diverse applicability and efficacy of UPTs in mesh-based fluid simulations, and steady-state Reynolds averaged Navier-Stokes simulations, and Lagrangian-based dynamics. Project page: https://ml-jku.github.io/UPT We follow the popular approach to approximate G via three maps [89] : The encoder E : U → R h1 takes an input function and maps it to a finite dimensional latent feature representation. For example, E could embed a continuous function to a chosen hidden dimension R h1 for a collection of grid points. Next, A : R h1 → R h2 approximates the action of the operator G, and D decodes the hidden representation, and thus creates the output functions via D : R h2 → V, which in many cases is point-wise evaluated at the output grid or output mesh. Particle vs. grid-based methods. Often, numerical simulation methods can be classified into two distinct families: particle and grid-based methods. This specification is notably prevalent, for instance, in the field of computational fluid dynamics (CFD), where Lagrangian and Eulerian discretization schemes offer different characteristics dependent on the PDEs. In simpler terms, Eulerian schemes essentially monitor velocities at specific fixed grid points. These points, represented by a spatially limited number of nodes, control volumes, or cells, serve to discretize the continuous space. This process leads to grid-based or mesh-based representations. In contrast to such grid-and mesh-based representations, in Lagrangian schemes, the discretization is carried out using finitely many material points, often referred to as particles, which move with the local deformation of the continuum. Roughly speaking, there are three families of Lagrangian schemes: discrete element methods [24] , material point methods [94, 12] , and smoothed particle hydrodynamics (SPH) [29, 62, 69, 70] . In this work, we focus on SPH methods, which approximate the field properties using radial kernel interpolations over adjacent particles at the location of each particle. The strength of SPH lies in its ability to operate without being constrained by connectivity issues, such as meshes. This characteristic proves especially beneficial when simulating systems that undergo significant deformations. Latent space representation of neural operators. For larger meshes or larger number of particles, memory consumption and inference speed become more and more important. Fourier Neural Operator (FNO) based methods work on regular grids, or learn a mapping to a regular latent grid, e.g., geometryinformed neural operators (GINO) [54] . In three dimensions, the stored Fourier modes have the shape h × n x × n y × n z , where h is the hidden size and n x , n y , n z are the respective Fourier modes. Similarly, the latent space of CNN-based methods, e.g., Raonić et al. [83], Gupta & Brandstetter [31], is of shape h × w x × w y × w z , where w x , w y , w z are the respective grid points. In three dimension, the memory requirement in each layer increases cubically with increasing number of modes or grid points. In contrast, transformer based neural operators, e.g., Hao et al. [33], Cao [18], Li et al. [53], operate on a token-based latent space of dimension n tokens × h, where usually n tokens ∝ n points , and GNN based neural operators, e.g., Li et al. [52], operate on a node based latent space of dimension