MMET: A Multi-Input and Multi-Scale Transformer for Efficient PDEs Solving

1 Introduction

The solution of complex partial differential equations (PDEs) has long been one of the great challenges in engineering. In the vast majority of cases involve elasticity, fluid dynamics or other physics fields, it is difficult to obtain analytical solutions for their governing PDEs. As an alternative, numerical methods such as the Finite Element Method (FEM) are usually used to approximate the result. However, these methods are computationally expensive and challenging to accelerate using parallel hardware.

Recent advances in the field of Artificial Intelligence (AI) have explored the use of machine learning techniques to approximate physical laws directly from PDEs [Raissi et al., 2019] or big data [Wang et al., 2021], and have demonstrated promising accuracy in Computational Fluid Dynamics (CFD) [Raissi et al., 2020; Mao et al., 2020], magnetohydrodynamics [Rosofsky and Huerta, 2023], and other fields. Nonetheless, traditional fully connected neural networks can typically solve only a single instance in a set of PDEs. Any changes in geometric shape, initial/boundary conditions, or material properties require retraining of the neural network. This makes the speed advantage of AI-based approach over traditional FEM a pseudo proposition, as the cost of training a neural network for a single instance can be comparable to or even exceed that of FEM.

Recently, generative pre-trained large-scale models based on the Transformer architecture [Vaswani et al., 2017] have shown strong zero-shot or few-shot generalization abilities in natural language processing [OpenAI, 2023], machine vision [Liu et al., 2024], and more scenarios. By extending this paradigm to the field of PDE-driven physics simulation, it becomes feasible to develop universal models that do not require expensive retraining for each instance. Several research efforts have attempted to fit PDEs using the Transformer architecture, targeting various aspects such as enhancing operator learning [Li et al., 2022; Shih et al., 2025], generalization [Santos et al., 2023], multi-scale meshes [Hao et al., 2023], simulation accuracy [Zhao et al., 2024], geometric model segmentation [Wu et al., 2024a], and improving the attention mechanism to enhance accuracy [Cao, 2021].

Despite these advancements, the above methods primarily address specific challenges within their respective domains. However, a universal paradigm that can systematically tackle all the key challenges within a unified framework remains absent. These challenges can be summarized as follows.

• Multi-Input Representation: A complete PDE solving problem can contain multiple inputs, such as differential operator, geometry domain, initial/boundary conditions, etc., which are in the form of variables or functions. The model must possess a robust and adaptable input encoding mechanism capable of efficiently encoding above diverse inputs.
• Multi-Scale Resolution: Different scenarios often require varying output resolutions, leading to inconsistencies in query inputs. However, Transformer-based models may alter their attention mechanisms as the input sequence changes, resulting in fluctuations in query accuracy under different resolutions.
• Computationally Expensive: Meshes in industrial scenarios often consist of tens to hundreds of thousands of node points, posing significant challenges to the cost and efficiency of attention mechanisms at such scales.

Based on these limitations, we propose the Multi-input and Multi-scale Efficient Transformer (MMET) architecture, which can be used to solve a class of general problems in real time for certain groups of PDEs. Our contributions can be summarized as follows:

• Proposing the Gated Condition Embedding (GCE) layer, embedding input functions and variables in units of spatio-temporal points, which can effectively encode complex input conditions while ensuring that the input sequence length does not increase significantly.
• Decoupling the query points and the geometry mesh into two sequences for the input of the encoder and decoder, respectively, enabling queries at any resolution.
• Reserializing and patching the mesh sequence using the Hilbert curve [Hilbert, 1935], significantly reducing the computational cost of the attention mechanism.

Experiments across various physical fields demonstrate that these innovative methods significantly enhance both the accuracy and computational efficiency of models in solving multi-input and multi-scale PDEs.

2 RelatedWork

2.1 Physics-Informed Neural Networks

Physics-Informed Neural Networks (PINNs) [Raissi et al., 2019] are a class of numerical simulation approaches based on machine learning aiming to approximate the solution of PDEs by incorporating physical laws directly into the loss function of the artificial neural network. This methodology can actually be regarded as a form of self-supervised learning rather than a network design, as the architecture of the neural network itself is highly flexible.

2.2 Operator Learning

The operator learning paradigm, originating from DeepONet [Wang et al., 2021], has recently replaced traditional neural network architectures as a new generalized PDE solver. Among these methods, the FNO [Li et al., 2021] stands out for its efficiency, leveraging linear projections in the Fourier domain to handle high-dimensional inputs. Extensions such as U-NO [Rahman et al., 2022] and F-FNO [Tran et al., 2023] improve multiscale representation or computational efficiency, further expanding the versatility of neural operators. On the other side, models like the GNO [Anandkumar et al., 2020], MGNO [Li et al., 2020], and Geo-FNO [Li et al., 2023b] have been proposed to enhance the perception of irregular geometric models by employing graph-based learning and geometric Fourier transforms. Several recent studies, such as 3D-GeoCA [Deng et al., 2024] and Geom-DeepONet [He et al., 2024], further enhance this capability by introducing techniques from the field of point cloud processing.

While operator learning methods demonstrate generalization ability in the face of variations in differential operators, challenges arise when dealing with other complex local variables such as geometries or boundary conditions. Without an efficient attention mechanism, these architectures often struggle to encode those intricate local features. It is worth mentioning that 3D-GeoCA [Deng et al., 2024] uses a backbone structure to encode the geometric features and boundary conditions of the problem domain. In their experiments, this method was also integrated with Point-BERT [Yu et al., 2022] and other advanced Transformers. However, its main purpose is to enhance the capacity of the model to perceive complex geometric configurations, while the ability to generalize across multiple types of inputs remains unexplored.

2.3 Transformer for PDEs

Several past works have explored the possibility of applying the Transformer architecture to the field of AI for PDEs. In which, HT-NET [Ma et al., 2022] takes advantage of the Swin Transformer [Liu et al., 2021] and the multigrid method to enhance the ability of multi-scale spatial perception. OFormer [Li et al., 2022] and IPOT [Lee and Oh, 2024] use Transformer to learn operators from meshes with varying resolutions. Recent studies, GNOT [Hao et al., 2023] and UPT [Alkin et al., 2024], attempted to enhance the multi-input generalization and multi-scale ability of the model. In addition, there are several studies devoted to improving the attention mechanism and increasing the accuracy of neural networks in fitting PDEs or operators [Cao, 2021; Li et al., 2023a; Xiao et al., 2024].

The vast majority of the above models are trained using labeled data, which are extremely expensive in real industrial scenarios. While PINNs have demonstrated considerable potential in reducing data dependency, these complex networks have more complex gradients and are harder to converge when using PDEs as part of the loss function [Rathore et al., 2024]. In this regard, PIT [Santos et al., 2023] and PinnsFormer [Zhao et al., 2024] explored direct training through physics-informed methods and have been successful in some small-scale cases. Nevertheless, they did not consider generalization to more complex input conditions.

In addition, most of these studies apply the attention mechanism to the whole domain and directly embed each sample point as a token. When the complexity of the geometric model increases, the number of sample points or meshes required may reach tens or even hundreds of thousands, making the sequence length unacceptable for the global attention mechanism. Transolver [Wu et al., 2024a] proposed a new physics-attention mechanism for segmenting geometric meshes to solve this problem, making it possible to solve complex geometric models under a fixed sequence length. Nevertheless, it relies on generating a weight map for each slice through a feedforward subnetwork, which lack global receptive fields to capture diverse input conditions. Therefore, a well-trained slicing subnetwork can only be effective for dynamic inputs within limited ranges. To address these varying conditions in a more universal manner, a predefined, non-learned slicing strategy should be considered.

3 Problem Setup

We consider the following PDE in a general form:

where u(x, t) is the expected value to be solved; N(u; λ) is the differential operator that describes the PDE, which is parameterized by physical parameters λ; Ω is the spatial domain that depends on the geometry; T is the max time of the field.

For the complete definition of a problem instance, the PDE should also have the following initial conditions:

and different types of boundary conditions:

where ∂Ω represents the boundary of the problem domain, and d, n, r are the Dirichlet, Neumann, and Robin boundary conditions, respectively.

To develop an ideal large-scale universal model for PDE solving, it is necessary to design a network characterized by the following mapping relationship:

where θ are the parameters of the model, the domain Ω and initial/boundary conditions g, d, n, r are highly complex and locally varying. The mapping f is required to accurately solve a set of instances with the above variables, which represents a general class of problems in the engineering field.

4 Our Method

4.1 Overall Architecture

Given the complexity of geometric features, a common approach is to discretize them into point cloud or unstructured mesh for feature extraction. In MMET, we employ two discretization strategies for the instance under consideration, as shown in Figure 2. First, we adopt the same mesh generation technique as used in classical FEM [Okereke and Keates, 2018]. The generated mesh serves as the input pseudosequence of the encoder to extract the global features. For the region of interest, the domain is further discretized into points to be queried, and is subsequently fed into the decoder.

In the encoder part of Figure 2, the input mesh pseudosequence comprises not only the spatial-temporal coordinates (x, t), but also the sampling values of the differential operator N, initial/boundary conditions g, d, n, r, and other necessary variables at this point. The coordinates of each point are directly fed into the positional encoding layer as position indices, and this layer can adopt a learnable feedforward network or fixed positional encoding as shown in PIT [Santos et al., 2023]. To handle complex input variables or functions,

(a) Pipeline of the GCE layer

(b) An example of the GCE layer to process different types of boundary conditions

Figure 3: Pipeline of the GCE layer and an example. The white square in the vector indicates that the value at that position is missing. Through dimension expansion and gating, this module can effectively extract features from different dimensions of input variables.

the innovative GCE layer is employed for effective embedding, which will be described in detail in Section 4.2. Subsequently, the sequence is reserialized using a Hilbert curve for spatial coherence and divided into patches of a specified size. Each patch is transformed into a token via a linear layer, which will be described in Section 4.3. The resulting token sequence is fed into the encoder for feature encoding.

In the decoder part of Figure 2, to enable querying at arbitrary resolutions, each query point is directly treated as a token and fed into the decoder. The self-attention layer typically used in the decoder of the classical Transformer [Vaswani et al., 2017] is removed to ensure that the output is not influenced by the construction of the query sequence itself. Therefore, the query point sequence only performs cross-attention with the output feature sequence of the encoder. This design offers significant flexibility, allowing query points to be sampled in any quantity and at any location within the domain Ω. Furthermore, if large number of points must be queried simultaneously or hardware resources are limited, the query sequence can be divided into batches and processed sequentially without affecting the final result.

In our model, all activation functions are replaced by the wavelet function proposed in PinnsFormer [Zhao et al., 2024] for better accuracy. The attention layer can be configured as classical dot-product attention [Vaswani et al., 2017], Galerkin attention [Cao, 2021], or linear attention [Hao et al., 2023], depending on the specific requirements.

4.2 Gated Condition Embedding

The previous approach GNOT [Hao et al., 2023] uses an architecture similar to Mixture-of-Experts (MoE) approach [Shazeer et al., 2017; Fedus et al., 2022] to encode different features, which may lead to extremely long input sequences caused by repeated encoding of different features at the same point. The GCE layer is a novel module we designed specifically to address this deficiency. As shown in Figure 3a, we embed different input functions in units of spatio-temporal points (x, t) rather than the sample values sequence of each function. Under this embedding approach, samples of different boundary conditions d, n, and r exhibit different dimensional properties depending on whether the boundary point corresponds to Dirichlet, Neumann, or Robin types.

As an example, in Figure 3b, we consider two boundary points A and B with Dirichlet and Neumann boundary conditions dA(xA, tA) and nB(xB, tB), respectively. The input vector of point A only needs to contain its sampled value dA, while point B may need to contain not only sampled value nB, but also the normal vector ⃗n of that boundary. However, neural networks require input vectors to have fixed dimensions. Naively masking missing features with zeros introduces ambiguity, as a fully connected network cannot distinguish between an input value that is genuinely zero and one representing missing data.

The GCE layer resolves this issue by mapping each type of input variable to a higher-dimensional space using a linear layer. After training, the original feature dimensions form a hyperplane in this higher-dimensional space that does not intersect the origin. If the input vectors of points A and B are expressed as vectors ⃗ dA ∈ Ra , ⃗nB ∈ Rb , the outputs of the proposed linear layers are given as:

where Wd ∈ R(a+1)×a and Wn ∈ R(b+1)×b are the weight matrices of the linear layers for Dirichlet and Neumann boundary condition inputs, respectively, and they are independent of each other; ⃗bd ∈ Ra+1 and⃗bn ∈ Rb+1 are the bias vectors of these layers.

A gate unit is then used to control whether the mapped features ⃗hA and ⃗hB should be integrated into the embedding vector. If this feature is missing, the value at that position in the embedding vector will be set to zero, representing the origin in the high-dimensional space. Importantly, this gate unit is parameter-free and relies solely on the presence or absence of the corresponding variables for a given point. Therefore, the embedding vectors of points A and B can be expressed as:

where ⃗oA ∈ Rb+1 and ⃗oB ∈ Ra+1 are zero vectors masked by the gate unit, which are used to pad missing inputs.

This design allows the model to reliably distinguish between missing inputs and inputs with a value of zero, ensuring a more robust and interpretable representation.

4.3 Reserialization and Patch Embedding

Patching the mesh sequence to a larger slice not only reduces its overall length but also enhances the model's performance on complex instances, as directly applying attention mechanisms to a large number of mesh points often hinders the model's ability to capture meaningful relationships [Wu et al., 2022]. To ensure generalization to different input conditions, we avoid using any learnable patch method such as those proposed in Transolver [Wu et al., 2024a].

Inspired by the work in point cloud processing, PTv3 [Wu et al., 2024b], we first reserialize the mesh points using the Hilbert curve [Hilbert, 1935], which is a type of space-filling curve designed to map high-dimensional discrete spaces into one-dimensional sequences while preserving spatial coherence. Compared with alternatives like Z-order curves or random raster scans, Hilbert ordering more consistently maps nearby nodes in 2D/3D space to adjacent positions in the serialized sequence, making it well-suited for downstream patching. This improves the likelihood that spatially related features are grouped within the same patch, which is especially beneficial for attention-based models operating under fixed memory budgets.

As shown in Figure 4, we take a 2D mesh as an example. Consider a node (x, y) within this unstructured mesh, linearly mapped into a square where each dimension is discretized into 2n units. The Hilbert curve recursively subdivides the square into four subregions, with the number of recursion determined by the order number n. Each subregion at a given order i is sequentially assigned a code bi ∈ [0, 3], which can be expressed as:

where Ri represents the rotation direction at order i, which depends on the starting orientation of the recursion. The index() function is visualized in Figure 4a, dynamically adjusting the rotation direction Ri to ensure that spatially adjacent coordinates maintain locality in the serialized sequence. The Hilbert code e for each coordinate can be computed as:

Each node in the original mesh is sorted according to their Hilbert code e to form a new 1D sequence. This sequence is then patched according to a fixed size. The vectors within each patch are concatenated and passed through a linear layer to extract patch features, which serve as individual tokens for the encoder.

This approach further benefits from FEM's adaptive meshing: regions with sharp gradients are a priori assigned denser meshes. By fixing the patch size, these regions generate more tokens, allowing the attention mechanism to prioritize critical areas. This design enhances MMET's capability in multiscale scenarios while preserving accuracy for complex geometries.

Third-Order Hilbert Curve

(b) Patching the mesh nodes as 1D sequence

Figure 4: Reserialize and patching of the mesh. Points with same color indicate the same patch. Regions with denser nodes are assigned more patches, as patch 1 and patch 2.

5 Experiment and Result

5.1 Benchmark Comparisons

In this section, we conduct benchmark experiments and compare our model to demonstrate its advantages over previous methods.

Experimental Setting. We select six representative datasets spanning elasticity, fluid mechanics, and thermodynamics. Here, A indicates a dataset with multiple geometric shapes, B indicates datasets with various boundary conditions, and C indicates datasets with multi-scale query resolution. The following is a brief introduction.

• Poisson: A simple 2D Poisson problem that is mainly used to verify the effectiveness of each baseline under physics-informed training.
• Darcy Flow [Takamoto et al., 2024]: The 2D Darcy flow problem from the PDEBench datasets, aims to learn the mapping from spatially varying diffusion coefficients to pressure distributions.
• Shape-Net Car (A) [Umetani and Bickel, 2018]: A multi-scale 3D aerodynamics problem used to predict the flow velocity v and surface pressure p of different cars. This dataset includes complex geometric features.
• Heat2d (A, B) [Hao et al., 2023]: A multi-scale heat conduction problem used to predict the temperature field T from input functions.
• Beam2d (B, C) [Bai et al., 2023]: This dataset solves the stress σx, σy, τxy, and displacement u, v of a rectangular cantilever beam when the end is subjected to a bending moment. We modify this case so that the bending moment varies dynamically within a certain range.

• HeatSink2d (A, B, C): A thermodynamic case designed to solve the internal temperature T of a 2D heat sink with variable height and surface temperature. This case relies solely on PDEs for physics-informed training, and without using labeled data.

Meanwhile, we selected another three recent Transformerbased SOTA methods as baselines, including PinnsFormer [Zhao et al., 2024], Transolver [Wu et al., 2024a], and GNOT [Hao et al., 2023]. For MMET, we also tested the encoderonly configuration on some of the datasets that do not need for multi-scale queries.

For all models, we employ AdamW [Loshchilov and Hutter, 2019] or L-BFGS [Liu and Nocedal, 1989] optimizers for combinatorial optimization. The model is trained on two NVIDIA H20 GPUs. The relative L2 error is used to evaluate all models, which has also been widely applied in previous works [Wu et al., 2024a; Xiao et al., 2024; Hao et al., 2023]. The error is computed as follows:

where ui is the prediction value output by the model, u′i is the ground truth, and N is the number of data in the test set.

For the Poisson, Shape-Net Car and Darcy Flow dataset, since there are no multiple boundary condition inputs to embed, we therefore replace the GCE layer by a simple feedforward network. For the HeatSink2d dataset, the error is expected to be larger than for others because training relies solely on PDEs, while the ground truth from Ansys Workbench (via FEM) also contains approximation errors.

Result and Discussion. The results in Table 1 demonstrate that our model consistently outperforms selected SOTA baselines in nearly all complex datasets, highlighting the effectiveness of MMET in addressing large-scale and dynamic physical fields. In contrast, PinnsFormer and Transolver lack effective encoding mechanisms for boundary conditions, rendering them incapable of handling problems with multiple input conditions. It is worth noting that on the Heat2d dataset designed by GNOT's team (which is characterized by multiinput, multi-geometry, and multi-scale), our model shows higher accuracy than GNOT. This further validates the capability of MMET in tackling intricate and diverse scenarios.

However, we also observe that the encoder-decoder structure in complete MMET converges more slowly and is sensitive to large learning rates. Therefore, for tasks that do not require arbitrary resolution queries, such as Shape-Net Car, an encoder-only variant is recommended to improve training speed, stability, and accuracy.

5.2 Multi-Scale Query Experiment

We designed a supplementary experiment to verify the zeroshot query ability of our model at different resolutions. For the Beam2d dataset, in the training stage, our model uses an unstructured mesh with 5404 nodes generated by Ansys Workbench. The query sequence is a regular point matrix with a resolution of [50×20]. In the inference stage, we keep the mesh input the same as in the training stage and test the performance of the query points at four different resolutions, as shown in Figure 6.

Experimental results indicate that decoupling the mesh from the query sequence enables our model to maintain fully consistent error performance across varying query resolutions. This demonstrates the distinct fitting capability of MMET within the continuous function space, comparable to that of traditional fully connected networks. In contrast, other Transformer-based models suffer from accuracy fluctuations as the input sequence length changes.

5.3 Ablation Experiment

In this section, we evaluate the impact of disabling the GCE layer and modifying the patch size in different datasets, and analyzing their effects on the performance of our model.

GCE Layer. To assess the significance of the GCE layer, we designed an experiment introducing ambiguity in boundary representation. This experiment is modified from HeatSink2d dataset, where a boundary is represented as a heat flux (Neumann) boundary in part of the dataset and as a zerovalue temperature (Dirichlet) boundary in the remainder. If a zero pad is used to pad the missing inputs, then the sampled characteristics of the two parts on this boundary would be exactly the same. We evaluated the performance of the model using three embedding approaches: a simple MLP, an MLP with an additional input node to encode the boundary type, and our GCE layer. The results of this experiment are presented in Table 2.

It can be seen that using a simple MLP layer results in complete failure, with errors exceeding 20% in both parts of the dataset. This is because the exact same input makes the MLP unable to effectively distinguish between different types of boundary conditions. Introducing a boundary type indicator into the input indeed enables the model to learn effectively. Nonetheless, our GCE layer achieves significantly higher accuracy. This result shows the superior performance of our method in resolving input ambiguities and dynamically encodes boundary conditions with different dimensions.

Avg L2 Error GPU Memory

Patch Embedding. In the Shape-Net Car and Beam2d datasets, we dynamically adjust the patch size in a certain range to evaluate the effectiveness of our patching method. Notably, a patch size of 1 is equivalent to disabling patch embedding. To ensure a fair comparison, the model scale and training configuration are maintained consistently in all settings. The best accuracy and GPU memory consumption are shown in Figure 7.

For the Shape-Net Car dataset, the experimental results indicate that our method surpasses the no-patch configuration when the patch size ranges from 2 and 3, while significantly reducing GPU memory consumption during training and inference. However, excessively large patch sizes lead to a noticeable decline in performance without obvious gains in efficiency. In contrast, for the Beam2d dataset, the model fails to converge with small patch sizes, indicating an inability to effectively extract information from a complex mesh. As the patch size increases, the convergence rate and final accuracy of the model improve progressively.

These findings highlight the effectiveness of our approach in handling complex meshes. The Shape-Net Car dataset, derived from 3D point clouds sampled uniformly across model surfaces, differs fundamentally from the Beam2d dataset, which consists of meshes generated by Ansys Workbench. The latter a priori incorporates physical information. This enables the model to prioritize regions with complex gradients more effectively. By employing fixed-size patches, our method implicitly assigns different weights to various regions of the geometry, thereby enhancing the efficiency and accuracy of the attention mechanism.

6 Conclusion

In this work, we proposed the MMET, a novel framework for real-time solving of generalized PDEs with multi-input and multi-scale resolution. The architecture incorporates key innovations, including the GCE layer and Hilbert curve-based reserialization and patching, enabling efficient handling of complex geometries and input conditions with scalability and accuracy. Experimental results across various benchmarks demonstrate the superior performance of MMET compared to SOTA methods, showcasing its potential for broad applications in complex physics-based simulations.

In the future, we plan to scale up the model, striving to establish it as a foundational pre-trained model for specific physics fields.

Acknowledgments

Yichen Luo is sponsored by the China Scholarship Council for PhD study at KTH Royal Institute of Technology, Sweden (No.202408320060). This work was partially supported by the National Natural Science Foundation of China (Grant No.72401232), the XJTLU External Research Fund under 2024-015, and LLL24ZZ-02-01 and LLL24ZZ-02-02 funded by the Liaoning Liaohe Laboratory.