SEA: State-Exchange Attention for High-Fidelity Physics Based Transformers

Authors: parsa  esmati
amirhossein  dadashzadeh
vahid  goodarzi
nicolas  larrosa
nicolò  grilli
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 03:34:40

Abstract: Current approaches using sequential networks have shown promise in estimating field variables for dynamical systems, but they are often limited by high rollout errors. The unresolved issue of rollout error accumulation results in unreliable estimations as the network predicts further into the future, with each step's error compounding and leading to an increase in inaccuracy. Here, we introduce the State-Exchange Attention (SEA) module, a novel transformer-based module enabling information exchange between encoded fields through multi-head cross-attention. The cross-field multidirectional information exchange design enables all state variables in the system to exchange information with one another, capturing physical relationships and symmetries between fields. Additionally, we introduce an efficient ViT-like mesh autoencoder to generate spatially coherent mesh embeddings for a large number of meshing cells. The SEA integrated transformer demonstrates the state-of-the-art rollout error compared to other competitive baselines. Specifically, we outperform PbGMR-GMUS Transformer-RealNVP and GMR-GMUS Transformer, with a reduction in error of 88% and 91%, respectively. Furthermore, we demonstrate that the SEA module alone can reduce errors by 97% for state variables that are highly dependent on other states of the system. The repository for this work is available at: https://github.com/ParsaEsmati/SEA

Main paper: 
 introduction: 
 Solving partial differential equations (PDE) has been a primary concern of many fields in science and engineering, including physics [Salvini et al., 2024] , chemistry [Grilli et al., 2020] , and material sciences [Grilli et al., 2018 , Esmati et al., 2024] . In many cases where a direct analytical solution of the PDE is impossible to obtain, numerical simulations are used. To solve these equations numerically, the domain is discretized into smaller cells using a discretization method such as finite element or finite volume methods. In some cases, the domain of interest is divided into millions of smaller elements forming large matrices representing the equation [Moukalled et al., 2016] . Solving these discretized equations typically follows an iterative technique, which can take days and sometimes weeks of runtime to converge to the full solution on the defined temporal domain [Liu et al., 2024 , Jiang et al., 2023] . In some cases to resolve some specific features of the system a fine mesh is required. The graphical representation of small droplet [Um et al., 2018] , cracking and brittle behaviour in thin components [Pfaff et al., 2014] , multiphase and turbulent flows in computational fluid dynamics (CFD) [Kochkov et al., 2021 , Thuerey et al., 2020 , Heyse et al., 2021a,b] are all instances of such scenario.Finer mesh however comes at the price of computational cost, which may not be feasible in industry settings. Consequently, there is a critical need for frameworks that can either bypass these detailed simulations or accelerate the solvers.Recent advances in sequential networks have shown promising results in estimating the state variables of dynamical systems over time. However, they have not yet been effectively utilized to bypass or accelerate numerical models [Yousif et al., 2022] . Key challenges include lengthy training times and steep gradients in rollout errors [Sun et al., 2023 , Han et al., 2022] . Additionally, many of these networks are task-specific, necessitating retraining for each unique test case [Li et al., 2020a] . If a model is to be integrated into a solver, steep error accumulation necessitates frequent retraining of networks to maintain efficacy when bypassing numerical solvers. This leads to unnecessary computational costs, making it challenging to integrate these models with solvers. Ideally, a model should be capable of learning the underlying physics and constitutive laws to minimize error gradients. Such a model could reduce retraining frequency and enable reuse through transfer learning approaches [Yosinski et al., 2014] .The current trend towards probabilistic models, such as diffusion models for science and physics, tends to overcome the rollout errors [Han et al., 2024 , Valencia et al., 2024] . However, regardless of their domain (science, video, etc.) , these models can produce unrealistic motions, especially for long sequence generations, and result in a temporal limitation [Weng et al., 2024] . Hence, the autoregressive generation seems unavoidable for longer sequences.Transformer-based architectures, in particular, are at the forefront of the autoregressive generations [Sun et al., 2023 , Zhao et al., 2023] . By drawing a parallel between the video and spatio-temporal PDE simulations, we intend to improve the unavoidable autoregressive rollout error for long sequence generation using transformers as the building block. Hence, inspired by the spatio-temporal crossattention mechanisms used in vision models [Lin et al., 2022 , Chen et al., 2021] , we propose a State-Exchange Attention (SEA) module for physics-based transformers, designed to mimic PDE by explicitly capturing variable coupling. The effective multidirectional information exchange between fields enables the correction of some fields by others, allowing the model to learn the complex features of the physical system. This work employs a Vision Transformer (ViT) like [Dosovitskiy et al., 2020] encoder for mesh embedding. The embedded mesh at each instance is then used in the temporal State-Exchange Attention (SEA) integrated Transformer to predict the system's future states.In summary, the contributions of this work include:1. Design and integration of a novel SEA module for physics-domain Transformer models.This module enables learning the underlying physics through a multidirectional information exchange process between the state variables.2. Assembly of a full ViT-SEA integrated framework that demonstrates state-of-the-art results in generating the complete sequence of the physical system given the initial sequence and specific time-invariant parameters representing the model.3. Comprehensive evaluation of SEA module, and ViT-SEA integrated transformer across different computational fluid dynamics cases, showing over 60% reduction in error in all cases compared to state-of-the-art models. 
related work: 
 In recent years, deep learning has led to major advancements in modeling physical systems, with contributions ranging from applying computer vision techniques to enhance the resolution of coarse meshes to incorporating physical symmetries and constraints through innovative modifications to learning objectives.The work of [Raissi et al., 2019] demonstrated the feasibility of directly incorporating physical information into the objective function, including initial conditions and necessary physical residuals.Other examples of such approaches include [Jeon et al., 2024 , Haghighat et al., 2021 , Yu et al., 2022] .While these techniques embed physical knowledge into the objective function, they lack generality, as the obtained parameters tend to be case-specific. The broader underlying physics is not fully captured, and only the physical symmetries specific to a particular case are addressed.Another class of methods is neural operators [Li et al., 2020b] . These methods generalize well across the PDEs they are trained on and are not case-specific. However, these models require data to be represented in higher dimensions to capture complex relationships. The application of an integral kernel in a high-dimensional representation of complex geometries with dense data points poses a significant computational challenge [Li et al., 2020a] . While neural operators can employ different architectures to reduce rollout error, they lack explicit mechanisms for integrating temporal data, unlike transformers, which handle sequence dependencies robustly through their attention mechanisms. This limitation can affect their effectiveness in applications that require sequential data processing.A notable trend in recent years involves the use of encoder-decoder pairs to process dynamical states in latent space. The work of [Wiewel et al., 2019] is an early example of this approach. In general, the input must first be encoded into the latent space while preserving the context. A sequential network, such as LSTM, GRU, or other variants, is then trained on the encoded data. Thus, these approaches typically require two components: an encoder-decoder pair and a temporal model. Another example of this approach is Mesh Graph Networks (MGN) [Pfaff et al., 2020] and Graph-Network-based Simulators (GNS) [Sanchez-Gonzalez et al., 2020] . The encoder modules of these models are based on graph networks. After processing nodes, edges, and features, they use a simple multi-layer perceptron (MLP) to compute the derivatives of the features over time and update the state using a forward-Euler scheme. While the encoder-decoder pair is a powerful tool, especially for unstructured mesh spaces, the lack of a robust time-stepping model significantly limits performance, with rollout error becoming a dominant issue.Current state-of-the-art models that demonstrate optimal performance on baseline datasets typically combine a graph network encoder with more advanced time-stepping algorithms. Notably, the study by [Han et al., 2022] employs a Graph Mesh Reducer (GMR) for encoding, along with a sequential time-stepping network and a Graph Mesh Up-Sampling (GMUS) decoder. This research explored various sequential models, including LSTM, GRU, and Transformers, with the latter achieving the lowest rollout error. Building on these principles, [Sun et al., 2023] introduced a modified version of GMR-GMUS, adding a RealNVP normalizing flow model to the time-stepping transformer. While adding a normalizing flow model does not yield a fully tractable model, it allows for the direct maximization of log-likelihood over the final data point, improving the overall objective and resulting in the most competitive baseline model reported thus far. However, these models still cannot fully address the rollout error in a systematic way that incorporates our physical understanding of the system.As a result, improving rollout errors and reducing training time remain significant challenges in this field. In this work, we demonstrate that both objectives can be addressed by separating the fields and allowing them to learn the inherent physical relationships and symmetries. To this end, we introduce the SEA module, implemented on top of a Vision Transformer (ViT) based mesh autoencoder.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: Section 5Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
problem statement: 
 Temporal generation of the states in a dynamical system from an initial condition is analogous to video generation tasks, where the process is conditioned on both the initial state and an external input. Similarly, the evolution of a dynamical system can be conditioned on known parameters, such as the Reynolds number in fluid flow and the system's initial condition. Given this similarity to autoregressive generative models like ART-V [Weng et al., 2024] , we formulate the temporal evolution as an autoregressive generation of states, conditioned on both the initial state and a time-invariant parameter.However, autoregression on the mesh space is challenging due to the large size of the elements and the number of variables stored on each element. Hence the mesh must be embedded into a manageable embedding, and our formulation becomes an autoregressive sequence generation in latent space.If we denote the fields with index i, and time with index t, then the stored data on field i at time t is presented by X i t . Given the proposed framework with the SEA module, we intend to encode the field groups separately. The field group refers to fields with the same dimensions (e.g., velocity in the xand y-directions belong to the same group, while pressure lies in a separate group). Let G represent the partition of fields into groups based on their dimensions (e.g., velocity, pressure). Each group is then encoded separately with the mesh encoder E, resulting in an encoded representation Z G k t at time t. The final encoded representation at time t is then given by:Z G k t = E Concat {X i t } i∈G k , ∀G k ∈ GWe now denote Z j t as the encoded representation of the group G k , where j indexes the group embedding space. This embedding is then used to formulate the autoregressive sequence generation in time where the initial condition Z j 0 and the time invariant parameter Θ (e.g., Reynolds number) are always available and used to condition the generation. We can formally express this problem as arg max Z j 1:T P (Z j 1:T | Z j 0 , Θ), where we aim to find the sequence of latent variables Z j 1:T that maximizes the conditional probability given the initial condition Z j 0 and the time-invariant parameter Θ. This is achieved through a pointwise estimation of the conditional probabilities in the continuous embedding space, with optimization performed by minimizing the L2 loss. Further detail on objectives and training is provided in Appendix C.The backbones commonly used to create embedding spaces in image and video models, such as Latent Diffusion Models (LDM) [Rombach et al., 2022] and Video Diffusion Models (VDM) [Ho et al., 2022] , cannot be directly applied to mesh data due to their inherent structured pixel inductive biases. Therefore, our proposed autoencoder must specifically overcome these biases for unstructured mesh space. Given that the temporal model employs a transformer backbone, the embedding must generate tokens compatible with the transformer's input requirements. These tokens are generated similarly to that of ViT [Dosovitskiy et al., 2020] . Following the approach used in ViT, the space is divided into multiple patches, with each patch containing a number of cells. Let Ω denote the domain in which our study is conducted, with dimension d. Assume that the domain is discretized into a set of nodes {x i } N i=1 , where each node x i represents a point in R d . To construct patches, we define a partitioning function f P : R d → {1, 2, . . . , M }, which assigns each node x i to one of M patches based on its coordinates.Assuming the boundaries between patches are showing with B = {b 1 , b 2 , . . . , b m } then the function f P is defined as follows:f P (x i ) =    1 if x i < b 1 , j if b j-1 ≤ x i < b j for 2 ≤ j ≤ m, m + 1 if x i ≥ b m .To address the challenge of irregular and unstructured meshes, which lead to varying numbers of nodes per patch, each patch is padded to align with the length of the largest patch. A padding value of 0 is used throughout the framework. Moreover, bias terms are excluded in the embedding process, and the Gaussian Error Linear Unit (GELU) activation function [Hendrycks and Gimpel, 2016] is applied to ensure that the padded elements do not influence the spatial encoding. To achieve spatially aware embeddings and coherent reconstructions, we apply a multi-head self-attention mechanism (MHSA). This padding and embedding strategy is illustrated in Figure 1 .The output generated by the Vision Transformer (ViT) embedding module is subsequently flattened and utilized as tokens within the temporal model. The complete process is thoroughly explained in appendix A.The State-Exchange Attention module is integrated into the temporal model in this work to enhance the autoregressive generation of sequences in time. The temporal model utilizes a transformer architecture to capture the temporal dependencies of the state variables. This transformer includes the adaptive layer-norm by [Peebles and Xie, 2023] and the rotational positional embedding (RoPE) as developed in [Su et al., 2024] and adapted by state of the art autoregressive image generation models [Lu et al., 2024 , Sun et al., 2024] . The adaptive layer-norm employed in this work is modified to take the continuous time invariant parameters as input and act as a secondary conditioning mechanism on the temporal model. Full detail of the temporal model is provided in appendix B. This work initializes a decoder block for each state variable in the given PDE For instance, the Navier-Stokes equations governing the fluid dynamics, requires the resolution of variables such as velocity and pressure each of these are assigned an expert decoder. The SEA module is designed to allow the exchange of information amongst these experts with cross attention. We further investigate other modes of information exchange in appendix D.The flow of information through the expert layers and the SEA module can be formulated by building on the well-known attention mechanism. For clarity, the terms regarding the positional embedding are omitted here. We start from the encoded groups of variables presented in section 3.1, denoted by Z j t . To simplify the notation and remove the explicit dependence on t, we represent the sequence of encoded variables across all time steps as a single matrix, Z j , which stacks the encodings of all time steps. The attention mechanism then reads:EQUATIONWhere K(Z j ), Q(Z j ), and V (Z j ) represent the key, query, and value matrices, and are obtained from the linear transformation of the input Z j by a set of learnable weights W K , W Q , and W V .Additionally d k represents the model dimension.Furthermore the adaptive layer-norm is demonstrated by 'AdaLN' and hence following a pre-norm convention the multihead self-attention (MHSA) becomes:(Z j ) SA = Z j + Attention AdaLN(Z j ) , ∀j ∈ G (2)Given the autoregressive task at hand all the attention mechanisms including the presented MHSA have causal mask to improve autoregressive generation.The field information flows through the information exchange module after the temporal self attention. This information exchange module is represented by the state-exchange attention in Figure 2 , where we allow the state variables to exchange relevant information with a causal cross attention mechanism. This is formulated as: (Z j ) SEA = (Z j ) SA + k∈G k̸ =j f SEA AdaLN((Z j ) SA ), AdaLN((Z k ) SA ) , ∀j ∈ G (3)Here, f SEA (•, •) represents our information exchange mechanism SEA. This module takes in two arguments: the first is the adaptive layer norm of the expert for which the attention is taking place, and the second is the adaptive layer norm from the expert to which the module is attending. These are represented by AdaLN(Z j ) and AdaLN(Z k ), where k and j are non-equal embedding field indices.To enhance efficiency during information exchange, we adopt a bottleneck mechanism inspired by expert-based architectures such as [Lee et al., 2024] . Introducing a bottleneck at this stage helps keep the model scalable by enabling selective information exchange in a lower-dimensional space. Another instance of such strategy is the Perceiver architecture [Jaegle et al., 2021] , where cross-attention creates a bottleneck for high-dimensional data from different modalities. In our case, we employ a simpler method, using a down-projection with learnable parameters, following the approach in [Lee et al., 2024] . If the down-and up-projection matrices for mapping to the bottleneck are denoted by W j d and W j u , respectively, the State-Exchange Attention mechanism is fully described by Equation 4. The arguments to this function are the AdaLN of (Z j ) SA and (Z k ) SA , as shown in Equation 3. However, for clarity in illustrating the equation, we use A j and A k to represent the input arguments here.EQUATIONIn the presented equation K(W j d A j ), Q(W j d A j ), and V (W j d A j ) represent the key, query and value matrices obtained form the down projection of the inputs.The conditioning of the generations on external parameters such as the Reynolds number is done using an indirect method of adaptive layer norm and a direct method of time invariant parameter injection module (TIPI). These components replace the more computationally intensive attention mechanism commonly used in physics domain autoregressive models. The TIPI component processes the timeinvariant parameters using a multilayer perceptron (MLP) with a GELU activation function. The MLP maps the time-invariant parameters Θ to the model's embedding dimension through learnable parameters. As shown in Equation 5, the information injector, represented by TIPI[•, •], injects these processed parameters into the current state by summing the model's embedded information with the upscaled time-invariant parameters.EQUATIONFigure 2 illustrates the architecture's detailed schematic.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [NA] .Answer: [No] Justification: The code is being processed and will be uploaded on github soon. Guidelines:• The answer NA means that paper does not include experiments requiring code. • The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
procedure: 
 In this section, we evaluate the proposed model on two benchmark datasets and compare its performance with other frameworks. First, the complete model is tested on the cylinder flow benchmark, a widely used dataset in computational fluid dynamics. This evaluation includes the comparison of the full model with recent physics domain autoregressive models.We then explore a multiphase case, where the model is tasked to resolve the interface by taking into account the fluxes caused by the velocities using the SEA module. In this section, only the temporal aspects of the model are varied (SEA module), while the ViT mesh autoencoder is fixed to isolate and eliminate the impact of encoding method on model's performance. To this end, separate decoders are assigned, one for velocity and another for volume fraction, similar to structure depicted in 2. The inclusion of the volume fraction allows us to study the extent of the model's capability to resolve interfaces and the effect of State-Exchange Attention on capturing the multiphase scenarios.For consistent comparison with state-of-the-art models [Sun et al., 2023 , Han et al., 2022] , relative mean squared error is used to quantify the errors. The model was trained on an A100 GPU for approximately 2 hours for both datasets. Furthermore, a consistent Transformer architecture was adopted in both cases, utilizing 1 layer and 8 attention heads. The embedding dimension of the model for the cylinder flow case was set to 1024, in line with the literature [Sun et al., 2023] , while a dimension of 2048 was used for the multiphase flow case to effectively capture the interface. Full details of the configurations and datasets are provided in Appendix G and F, respectively.The autoencoder used in the following experiments was trained exclusively with a reconstruction objective. The complete training procedure for this model is detailed in Appendix A. Reconstruction errors, compared to recent graph-based autoencoders, are presented in Table 1 . We assess our complete architecture using the 2D cylinder flow, a benchmark dataset employed by other leading baseline models. In this case, the Navier-Stokes equation governs the motion of the fluid throughout the domain. Consequently, the trajectory to track is the velocity and pressure. The mesh was initially tokenized at all time steps based on the explained ViT mesh encoder. These tokens were then fed through the SEA integrated model illustrated in Figure 2 . During the training, the entire sequence was fed to the network for each batch. During inference, the model was set to estimate the trajectory autoregressively, and hence, the test errors correspond to the total rollout error. For the cylinder flow dataset, the error was evaluated over the case with Reynold's number of 400. This case was chosen to keep consistent with the other competitive models.The recorded rollout error, and its comparison with other baseline models is illustrated in Figure 3 . In Figure 3a , we show that our model outperforms the established state-of-the-art models across the board. Specifically, our results indicate an average improvement of 88% and 91% in reducing the error in autoregressive sequence generation compared to the PbGMR-GMUS Transformer-RealNVP and GMR-GMUS Transformer architectures, respectively. Furthermore, our model outperforms both variants of MGN architecture with 99% and 98% improvement over the base MGN model and MGN-NI, respectively. Full detail of the errors are presented in Table 2 . [Pfaff et al., 2020] 25 778 136 313 GMR-GMUS Transformer [Han et al., 2022] 4.9 89 38 43.96 PbGMR-GMUS Transformer-RealNVP [Sun et al., 2023] The evaluated test cases were post-processed to generate a contour map for further observation of the learned patterns and potential areas of error. For consistency with the work in literature [Han et al., 2022 , Sun et al., 2023] the contour map of the case with Reynolds number of 400 is presented here. To fully explore our model's ability to capture complex flow features, such as the downstream vortex, we present visualizations at the 250th timestep, a point at which the Von Karman vortex street is fully developed, as shown in Figure 4 . This timestep was deliberately chosen to visualize complex dynamics that are often challenging for traditional models to capture.Evaluation of the presented model is extended to include a multiphase flow scenario, which complements the analysis presented in Section 4.3. The evaluated test case in this experiment corresponds to an immiscible collapse of two blocks of liquid due to density differential. In studies of multiphase flows, a critical aspect is the precise identification of fluid interfaces. This is accomplished by using a volume fraction state variable, denoted by α, which indicates the region occupied by the fluid. The value of α ranges from 0 in one fluid to 1 in the other, effectively distinguishing between the two phases. To capture the phase, an additional block is assigned to the volume fraction which can communicate with other fields (Velocity in this case). Given the minor variations of pressure, we discard this variable here and only focus on the importance of the field communication between velocities and volume fraction through SEA module.We investigate three possible variations of the Transformers to achieve this. First, we estimate the rollout error of the model with the SEA module. Second, we evaluate a basic model that encodes the fields into different latent spaces with no mode of information exchange. Finally, we assess a model that encodes all fields together into a single latent space, referred to as the Field Fusion Encoder (FFE) Transformer. This latter model corresponds to the Transformer architecture used in the PbGMR-GMUS Transformer-RealNVP and GMR-GMUS Transformer, as indicated in the provided results. The rollout error of these models are presented in Figure 5 . respectively. This represents approximately a 52% and 97% reduction in error with the integration of the SEA module, compared to the basic and FFE Transformers. Additionally, improvements of 48.5% and 40% are observed in the averaged velocity components.Further demonstration of the contour maps is provided in Figure 6 where the actual interface tracking capability of the SEA enhanced transformer can be observed. Further visual results on the velocities are provided in Appendix H. 
discussion: 
 The presented ViT-based mesh autoencoder and State-Exchange Attention (SEA) integrated Transformer module were evaluated through two different experiments on computational fluid dynamics (CFD) problems. A significant improvement in relative mean squared error was observed when the SEA module was deployed. During the cylinder flow evaluation, the full ViT-SEA integrated transformer framework achieved over 80% improvement compared to all competitive baselines. This improvement was accompanied by a lower gradient in the error, demonstrating a form of selfcorrection through information exchange between the velocity and pressure fields. The isolated SEA module was then tested on a multiphase case, resulting in a 97% improvement in the volume fraction compared to the field fusion encoder transformer, where the entire field is encoded into the same latent space. In the multiphase case, it was evident that the velocities exhibited a marginal error difference; approximately 40-50%; however, the volume fraction dominated the overall improvement. This is due to the significance of velocity in the displacement of the interface, whereas the volume fraction did not provide any valuable information to the velocity field. The error reduction observed with the deployment of the SEA module suggests that SEA module provides the necessary tools for the underlying physics of the governing to be captured.The presented module, however, may face challenges when scaling to equations involving a large number of state variables. For instance, in multiphase flows with more than two phases or in grain growth within materials where each grain is represented by a state variable, the model would require a corresponding number of transformers to operate in parallel. This could lead to inefficiencies as the number of variables increases significantly. 
conclusion: 
 The presented work introduces SEA, a novel module that enables the state variables of a physical system to exchange information within the transformer architecture. Transformers integrated with SEA demonstrated state-of-the-art performance, surpassing previously established benchmarks by other transformer-based models with improved autoregressive rollout error. The improved rollout error is an indication that the SEA module is enabling the model to learn the underlying physics. In future works, we will explore the scaling complexity of SEA to handle more state variables and its application in other domains. Furthermore, the integration of SEA with probabilistic models, such as diffusion models, will be investigated. experiments, the variational encoding does not lead to significant improvements, and the embedded space error the error (Refer to Appendix D). Therefore, we choose to proceed with pointwise encoding, as it consistently results in better reconstruction accuracy in the studied cases.The hyperparameters associated with this autoencoder include the number of patches, the MLP dimensions and architecture, the number of heads in the attention mechanism, the number of attention layers, and the embedding dimension. The attention block is configured similarly to the original ViT, with exact specifications provided in Appendix G. The number of patches is primarily determined by the number of cells and the mesh dimensions. In this work, we use a constant number of 64 patches, arranged as 8 patches along the x-direction and 8 patches along the y-direction. Consequently, the number of cells per patch depends on both the total number of cells and the mesh density. Table 3 summarizes this information.  
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [NA] Justification: Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 13. New Assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: The transformer used to capture temporal dependencies is designed with a decoder-only architecture, incorporating the SEA and TIPI modules. In this work, we develop and evaluate two variations of information exchange methods, namely exchange by addition and exchange via cross-attention, as shown in Figure 8 . Both methods operate within an information bottleneck, where the data is down-projected to retain only the relevant information required for exchange. To ensure improved gradient flow and prevent information loss due to this bottleneck, we introduce a residual connection from the output of the multi-head self-attention (MHSA) to the flow outside the SEA module. The down-projection within the bottleneck is set to a factor of 2, following the approach in [Lee et al., 2024] . The architecture referred to as the basic transformer does not incorporate any connections or information exchange between the decoders assigned to different variables.Additionally, the TIPI module is incorporated as a direct conditioning mechanism, while AdaLN serves as an indirect conditioning technique. In the final design, the primary TIPI module utilizes the addition-based approach, as shown in Figure 9 . However, we also explored injection via attention, which is commonly used as a conditioning technique in physics-based models.The adaptive layer normalization (AdaLN) technique, used as the indirect conditioning method, is adopted from [Peebles and Xie, 2023] . In this approach, conditioning is applied by replacing the standard layer normalization with scale and shift parameters derived from the time-invariant parameter Θ. To remain consistent with [Peebles and Xie, 2023] , we use an MLP to compute the scale and shift values. First, Θ is up-projected to the model dimension, followed by the application of the SiLU activation function. A subsequent layer is then used to obtain the final scale and shift parameters.An schematic of the AdaLN block is presented in Figure 10 . For the temporal model, the data is first structured similarly to the section A. After the encoding process, we obtain a data structure of [B*T, P, D] for each field group. To process different trajectories separately and avoid mixing, we separate the time from the batch and collapse the patch in the encodings. Hence, the final shape processed by the temporal model is [B, T, P*D], where B represents different trajectories, which we batch from and learn the temporal evolution.Autoencoder: The autoencoder was trained wtih the objective of the reconstruction error with L2 loss:L = E X∼D ∥X -X∥ 2 2 (6)Where X represents the model's output, and the expectation is done with respect to the data drawn from the batch.For a consistent comparison with the literature we provide the errors using relative mean squared error where we refer to this as relative reconstruction error, and relative rollout error for the autoencoder and temporal models respectively. The relative rollout MSE is formulated as:RelMSE = E X∼D N i=1 ∥X i -Xi ∥ 2 2 N i=1 ∥ Xi ∥ 2 2 (7)Here the summation is done on the cell dimension.Temporal model: The temporal model is trained using L2 loss in a teacher forcing manner, where the model generates the encoded field for the subsequent time step after observing the full trajectory. To accurately assess the autoregressive performance of the model, we conduct an autoregressive evaluation every 250 epochs, allowing us to track the learning history for the required autoregressive task. This evaluation includes monitoring both the encoded and decoded errors. We observe that, although the magnitude of variations in error is not perfectly aligned, the trends in the embedding space and decoded space errors follow a similar trajectory, peaking simultaneously. The following figure provides an example of the autoregressive learning curves for both the embedded space and the decoded space, using the ViT mesh autoencoder described in earlier sections. Provided curves in Figure 11 demonstrate the dynamic of the joint ViT-autoencoder and the temporal model. Informed by these curves, a non-variational encoding approach is adopted, and we assign 0 weights to the KL term used for the regularization of the latent space.Ablation of the information exchange method is performed using a vanilla architecture. The architecture incorporates layer normalization, a common practice in transformer models within the physics domain. We also initialize with absolute positional embeddings. Following our exploration of the vanilla transformer with the information exchange module, we extend the ablation studies to include the components of the main architecture. Specifically, we show that adding TIPI and AdaLN as conditioning mechanisms further reduces error and eliminates the need for an explicit attention mechanism for conditioning. To demonstrate the effect of the TIPI module, we investigate how different types of conditioning impact the error. This information is summarized in Table 5 , where "no conditioning" refers to the absence of TIPI, and "attention" indicates conditioning through attention mechanisms. Additionally, we evaluate whether the optimal placement of the TIPI module is before or after the SEA module.It is evident that all three types of conditioning significantly reduce the error across both datasets. Furthermore, our secondary conditioning mechanism, AdaLN, reduces the errors even further, resulting in the lowest recorded values. We compare the effect of using AdaLN versus not using AdaLN, along with the best errors obtained from Table 5 , as shown in Table 6 . The addition of the second conditioning mechanism yields similar results regardless of TIPI placement, effectively removing TIPI placement as a critical model parameter.The effect of RoPE embedding has been widely investigated in the literature, and thus, we do not explore this any further. However, in our experiments, the incorporation of RoPE significantly improved convergence within this specific domain. To gain deeper insight into the problem and the sources of error, we explore the errors addressed by our model and framework. In this analysis, we disregard the spatial and temporal discretization errors originating from the numerical solver and instead focus on the potential errors introduced by the autoregressive models grounded in physical simulations.The coupling of variables in PDEs is a common occurrence. Many physical systems are governed by sets of linear or nonlinear coupled PDEs. Examples include Maxwell's equations (describing electromagnetism), the Ginzburg-Landau equation (used in superconductivity and phase transitions), the Einstein field equations (in general relativity), and the Navier-Stokes equations (governing fluid dynamics).A common approach to modeling these equations involves embedding the variables into a shared latent space, where the system is resolved temporally in an autoregressive manner. For long trajectories, an autoregressive approach becomes essential, even for temporally-aware diffusion models. Consequently, it is critical to ensure that the embedded space closely resembles the real physical fields to maintain accuracy throughout the prediction process.Assume ψ 1 , and ψ 2 are two non-equal coupled field variables governed by the following advection equation, where f and g are any functions enforcing the coupling:EQUATIONSay Ψ is a variable that includes both ψ 1 and ψ 2 then the entropy of Ψ follows the following condition:H(Ψ) > max{H(ψ 1 ), H(ψ 2 )} (9)For the rate-distortion trade-off, this implies that, at the same rate, we can achieve lower field-wise distortion when the embedding is performed separately. In the case of two fields, achieving the same distortion requires doubling the rate to embed them together. This results in a higher-dimensional embedding space, which introduces two significant challenges for the temporal model. First, transformers scale quadratically with the embedding dimension, leading to greater computational complexity when the dimension increases, as opposed to using separate decoders for each field. Second, in higher dimensions, our temporal model begins to suffer from sparsity, resulting in poor generalization. Let the distortion error be denoted as ϵ D . Moreover, embedding the physical fields together may result in an uneven distribution of information from the fields, potentially weakening the coupling described in Equation 8.We can now demonstrate how any form of information exchange theoretically leads to lower error when separate embedding spaces are created. Based on the definition of mutual information, we have:EQUATIONwhere I(ψ 1 ; ψ 2 ) represents the mutual information between ψ 1 and ψ 2 , and H(ψ 1 ) is the entropy of ψ 1 . If f (ψ 2 ) is non-zero, then from the PDE, we know that I(ψ 1 ; ψ 2 ) is non-zero as well. This implies that the entropy of ψ 1 , conditioned on ψ 2 , is lower than the entropy of ψ 1 alone. A lack of such conditioning in the model introduces an additional error term due to increased uncertainty, which we refer to as the coupling error ϵ C . Our SEA module addresses this error by establishing a channel for information exchange between the fields.To summarise the errors addressed in our work:• Distortion error introduced the from joint embedding of the variables ϵ D .• Coupling error introduced by lack of information exchange or unjust compression of information due to joint embedding ϵ C .The dataset used in this work consists of two fluid mechanics simulations motivated by physical phenomena: flow around a cylinder and the mixing of immiscible fluids. The governing equations for these problems are the conservation of mass (continuity equation) and the conservation of linear momentum (Navier-Stokes equation), which were numerically solved using the finite volume method implemented in the open-source software OpenFOAM. For the multiphase flow simulation, an additional advection-diffusion equation was solved to track the volume fraction of the phases. The results of these simulations were labeled and used as the ground truth in our experiments. The continuity and momentum equations for incompressible fluid flow are given by:EQUATION)where ν is the kinematic viscosity, ρ is the density, p is the pressure, u is the velocity vector, and g = 9.81 m/s 2 is the acceleration due to gravity.The simulation of flow around a 2D cylinder was conducted using icoFoam, a transient incompressible solver for laminar flows in OpenFOAM. To ensure comparability with the literature, we generated 70 trajectories at different Reynolds numbers, with 60% used for training and 20% for validation. The input flow variables for this test case are the velocity components in the xand y-directions (u, v) and pressure (p).In the second test, the mixing of two fluids in the liquid phase with different densities was simulated using the twoLiquidMixingFoam solver. The flow was modeled as transient, incompressible, and isothermal, with the dynamics of the mixture being the primary focus of this test case. The model input parameters included the velocity components and the volume fraction, which separates the fluids and ranges between 0 and 1. We generated 40 trajectories for this case, with a similar ratio for training, validation, and testing as in the cylinder flow case.Table 7 details the configurations for the ViT-based mesh autoencoder and the temporal model, including key hyperparameters such as the number of patches and TIPI-specific settings. These include scaling mechanisms (MLP, Gaussian Fourier projection, linear) and injection methods (addition, attention). The bottleneck down-projection, adapted from Lee et al. [2024] , is introduced as an information exchange-specific hyperparameter. Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 


Decision: Reviewers appreciate how the paper adapts vision transformers to PDEs with a "state exchange attention (SEA)" mechanism and a "time invariant parameter injection (TIPI)" component. I personally find the paper to be frustratingly lacking in key details, and I strongly encourage the authors to address these concerns in the camera ready version. 

1. The authors do a decent job of motivating what they want to accomplish with SEA, but the mathematical description in eq. (3) doesn't make sense. Index $j$ appears on the right-hand side but not the left, and there's no specification of how results are reduced over $j$.
2. The so-called "time invariant parameter injection (TIPI)" component is not well motivated, and again the mathematical description in eq. (5) doesn't make sense: a mysterious "information injector" $f_{IJ}(\mu)$ is introduced without explanation. The variable $\mu$ is never even defined!
3. The captions for Figure 1 and Figure 2 are insufficient to describe what is being shown. The text sometimes refers to Figure 5 when it is meant to refer to Figure 3. The tick labels are illegible in Figures 4-6.
4. Table 1 presents the authors' method in bold even when it does not achieve the lowest error (specifically for $u$).
5. The authors use \citet when they should use \citep.
6. The paper needs to be edited for spelling mistakes, grammatical errors, and redundant phrases. For example, "We can formally formulate this problem..." (line 125) and "To process and prepare the spatial information..." (line 130).

Reviews: review: paper_summary: paper_summary: The study proposes an approach for autoregressive spatiotemporal estimation of a dynamical system state, e.g. solution of time-dependent PDE. In particular, a Vision Transformer (VIT) model is adapted to solving PDEs where each state variable is tokenized similarly to tokenizing an image and a novel State-Exchange Attention (SEA) module allows interaction and fusion of these tokens through cross-filed attention. This enables state variables in the system represented by tokens to exchange information and capture the physical relationships and symmetries between fields. The approach is evaluated on the problem of flow past a cylinder benchmark showing improvement of the error against current neural networks approaches.
strength_weakness: strength_weakness: W1. It is claimed that the rollout errors are effectively controlled by the SEA model. There's very little explanation/derivations of the source of these errors. It would be helpful to define these errors rigorously and also if possible to show (graphically or rigorously) how SEA model allows for better control of the error.

W2. Ablations of SEA model components and how their contribution to the overall accuracy of solution are missing. These could inform the critical parts of SEA and better characterize SEA.

W3. While generalization seems plausible, evaluation was performed on only one basic dataset from which introduces uncertainty in how the approach generalizes.
questions: questions: Q1. What is the computational complexity (time, storage) of SEA and how does it compare to other approaches and PDE simulation?

Q2. Does SEA apply to high dimensional state variables, e.g. 3D flows or ND state variables? If yes, it would be important to include in the manuscript (appendix) practical steps of extending the model to higher dimensions.
limitations: limitations: The technical contributions of this work are subtle since ViT are well known as powerful models and the SEA component seems to be sensible and inherits structure from prior approaches. This could be considered as a limitation, however, on the other hand, the impact of this work lies in the application of solving dynamical equations and simulating dynamical processes and appears to be significant advancement. The ability to implement existing components in such a way that they will be effective is not a straightforward endeavor. Hence, ablation results, discussion of computational complexity and more explanation of the error would be vital in this work. Also, addressing Weaknesses and Questions above would improve clarity, generalization and evaluation of this promising work.
paper_summary: paper_summary: The authors present a novel and interesting approach to physics-domain transformer models that exchanges information between state variables and demonstrates strong performance improvements relative to SOTA models on hard fluid dynamics problems. The paper is well written and the results compelling.
strength_weakness: strength_weakness: #### Major concerns:

1. Perhaps I missed this somewhere, but will code be made available with this paper? It seems too much to require readers to re-implement this architecture to use it.

2. The supplement is helpful but could contain further experimental details, further details on the layers (and their trainability), how projections are performed, and more complete explanation of results and metrics.

3. If training the models takes just 15-20 minutes, it would be nice to see the effects of the embedding dimensionality and simulation parameters explored more thoroughly.

4. The authors point out that a limitation might be scaling to a large number of state variables. Could they present an example on a problem with more than two state variables so that the "multidirectional" component is more than bidirectional, and perhaps add some thoughts on how this might scale? 

#### Minor concerns:
There are a number of places the notation could be clarified and some typos, for example:
By line:
- 126: argmax optimization variable is not specified (presumably z_1:T)
- 159: the X variable is never defined, and may be superfluous here... maybe just use z_i^k? 
- 166: T should be italic?
- A number of variables throughout seem to be used without introduction.
questions: questions: - How does the padding work in the case of graphs with highly variable node density (for example on an adaptive mesh with much denser representation near the cylinder as is common in practice)?
limitations: limitations: The authors have touched on potential limitations related to number of state variables but could discuss this in more detail, as many problems of interest contain more than two state variables.
paper_summary: paper_summary: The paper introduces the State-Exchange Attention (SEA) module, a new cross-attention module that enables information exchange between state variables in physics-domain transformer modules. The authors evaluate the performance of this module in both single-phase and multi-phase 2D fluid settings.
strength_weakness: strength_weakness: 1. Scalability Concerns: The current architecture is relatively small, using a transformer with only 1 layer and 8 attention heads. I am curious how the model will scale at larger architecture sizes. Moreover, as the authors themselves note, there are also concerns about how the model will scale with larger number of state variables. 
2. Diversity of experiments: The authors evaluate performance on one instance of single-phase and one instance of multi-phase flow, both in two dimensions. How does the model fair in 1D and 3D settings? How does it fair in a wider variety of fluid mechanics problems (like those presented in PDEBench/PDEArena)? 

Minor comments:
- It seems like the general gist of paragraphs on lines 179-186 and lines 187-192 are almost identical. Was this paragraph repeated by accident?
questions: questions: 1. Are there ablation studies available for single-phase flow (as is sort of provided for multi-phase flow). In other words, how does a basic transformer fair in this setting? 
2. How resource-intensive is training this model vs. traditional transformer architectures? Some indication of this would be very helpful.
limitations: limitations: See weaknesses.
paper_summary: paper_summary: This paper presents a novel neural network architecture (SEA) for solving partial differential equations (PDEs) in physical problems. This architecture effectively utilizes information exchange between multiple fields and the conservation quantities of physical systems to correct model predictions, achieving smaller errors and improved generalization. This has significant value for industrial applications, as it can greatly reduce the time cost for solving specific problems.
strength_weakness: strength_weakness: Further exploration of the model's generalization ability is needed, such as evaluating model errors under more complex boundary conditions. Additionally, presenting the training and inference costs of the model would provide a better assessment of its potential for industrial applications.
questions: questions: See cooments
limitations: limitations: yes

Comments: title: Authors Rebuttal Response
Thank you.
comment: Thanks for your responses. They help resolve my questions.
I'd like to thank the authors for providing an informative rebuttal and additional results in response to comments that I made in my review.  The authors clarified the questions and comments that I made and I would like to encourage the authors to include the results in the provided pdf and their discussion in the revision and also rigorous definition of the involved errors along with succinct explanation as the authors provide in the rebuttal to W1.
Thank you for your thoughtful response and for the additional details in the provided PDF. I am satisfied with the authors comments and additional details. I will keep my score.
We are satisfied with the authors response and have increased our score accordingly.