Explain Like I’m Five: Using LLMs to Improve PDE Surrogate Models with Text

The heat equation models a diffusive process and is given below in equation 1, where we are predicting the temperature distribution at each time step.

We generate data on a simulation cell given by: $\Omega=\left[-0.5,0.5\right]^{2}$ on a 64x64 grid with boundary conditions sampled from $\partial\Omega\in\left\{Neumann,Dirichlet,Periodic\right\}$. The values $v$ for the Neumann and Dirichlet boundary conditions are sampled uniformly: $v\in\mathcal{U}(-0.1,0.1)$. All four walls have the same boundary condition type and value for a given simulation. Initial conditions are chosen from three different distributions: exponential field given by $f(x,y)=\exp\left(100(x+y)^{2}\right)$, sum of sine and cosine given by $f(x,y)=\sin\left(c_{1}\pi x\right)+\cos\left(c_{2}\pi x\right)$, and product of sines given by $f(x,y)=\sin\left(c_{1}\pi x\right)\sin\left(c_{2}\pi y\right)$ for $c_{1},c_{2}\in\left\{2,4,6,8\right\}$. Lastly, our diffusion coefficient was sampled randomly according to $\beta\in\mathcal{U}\left(0.005/\pi,0.02/\pi\right)$. We simulated each trajectory for 2 seconds on a 64x64 grid. These distribution of diffusion coefficients was chosen so the diffusive dynamics behaved on approximately the same scale as the advection dynamics of Burgers Equation given below.

Burgers equation models shock formation in fluid waves, given below in equation 2, where we are predicting the [height], $u$ at each time step.

System specifications for the Heat equation are identical to Heat equation, given above, with the exception of the advection term. Our advection coefficient is sampled according to $\alpha_{x},\alpha_{y}\in\mathcal{U}\left(-0.5,0.5\right)$.

Third, we generate data from the Incompressible Navier-Stokes equations in vorticity form, given below in equation 3.

The Navier-Stokes data generated here follows the setup from Lorsung et al. (2024). Our simulation cell is given by: $\Omega=\left[0,1\right]^{2}$ with periodic boundary conditions. Our viscosity is sampled according to $\nu\in\left\{10^{-9},2\cdot 10^{-9},3\cdot 10^{-9},\ldots,10^{-8},2\cdot 10^{-8},\ldots,10^{-5}\right\}$, and our forcing term amplitude is sampled according to $A\in\left\{0.001,0.002,0.003,\ldots,0.01\right\}$. The initial condition is given by a Gaussian random field. Data generation is done on a 256x256 grid that is evenly downsampled to a 64x64 grid for this work.

Lastly, we take the Shallow-Water data set from PDEBench(Takamoto et al., 2023b). In this case, we are predicting the water height $h$ at each time.

Our simulation cell is: $\Omega=[-2.5,2.5]^{2}$ on a 64x64 grid, with Neumann boundary conditions, with 0 gradient on the boundary.

For each combination of initial conditions and boundary conditions, 900 samples were generated with randomly generated coefficients. For Navier Stokes, we did the vorticity formulation, from FNO, and varied viscosity and forcing term amplitude, with random gaussian field initial condition. 30 second simulation. Also used PDEBench’s shallow water data set. These are all single-channel. In our experiments, B refers to incorporating boundary condition information through text, C refers to incorporating coefficient information through text, and Q refers to incorporating qualitative information through text.

To explore how well our LLMs are able to incorporate different amounts of text information, we describe each of the Heat, Burgers, and Navier-Stokes equations with varying levels of detail. We will build a complete sentence description here, with all possible sentence combinations given in appendix A. At a base level, we can describe the basic properties of each governing equation, as well as identifying the equation. For Burgers equation, we generally have stronger advection forces, so we describe it as:

All of our basic system descriptions come from Wikipedia entries for their respective equations(Wikipedia, 2024a; b; c; d). Note: Burgers’ system is first order quasiliner hyperbolic only in the inviscid case, where we are working in the viscous case. The second sentence should instead read: Burgers equation is a second order partial differential equation.Additional testing has shown in appendix G that this misnomer does not affect results. Next, boundary condition information is added, in this case Neumann boundary conditions:

Third, we can add operator coefficient information:

In this case, the advection term has a coefficient of $\alpha_{x}$ in the x direction, $\alpha_{y}$ in the y direction, and the diffusion term has a coefficient of $\beta$.

Lastly, we can add qualitative information. The aim of this is to capture details about the system that are intuitive to practitioners, but difficult to encode mathematically. In this case, we have an advection dominated system:

This system is advection dominated and does not behave similarly to heat equation. The predicted state should develop shocks.

Our complete sentence description is passed into a pretrained LLM that is used to generate embeddings. These embeddings are then used as conditioning information for our model output.

The backbone of our multimodal surrogate model is the FactFormer(Li et al., 2023b) and our framework is given in figure 1. FactFormer was chosen because it provides a fast and accurate benchmark model. We add our system description as conditioning information both before and after factorized attention blocks. FactFormer is a neural operator that learns a functional $\mathcal{G}_{\theta}$ that maps from input function space ${\mathbb{A}}$ to solution function space ${\mathbb{U}}$ as $\mathcal{G}_{\theta}:{\mathbb{A}}\to{\mathbb{U}}$, with parameters $\theta$. In our case, we are specifically learning the functional conditioned on our system information ${\bm{s}}$. For a given input function $a$ evaluated at points ${\bm{x}}$, we are leaning the the operator given in equation 5. That is, our network learns to make predictions for our solution function $u$ also evaluated at points ${\bm{x}}$.

Sentences are first passed through an LLM to generate an embedding. The embedding is combined with the data through cross-attention, seen in figure 2, where our sentence embedding is the queries and the data embedding is the keys and values. The conditioned embedding is then added back to our data embedding, given below in equation 6. Our data embedding is embedding with convolutional patch embedding, mapping our data to a lower dimension as $f_{data}:\mathbb{R}^{b\times h\times w\times h_{FF}}\to\mathbb{R}^{b\times p\times h_{FF}}$, where $b$ is our batch size, $h_{FF}$ is our FactFormer embedding dimension, and $p$ is the number of patches, which is defined by kernel size and convolutional stride. Our sentence embeddings are projected to higher dimensional space to match our data embedding dimension, $f_{sentence}:\mathbb{R}^{b\times h_{LLM}}\to\mathbb{R}^{b\times p\times 1}\to\mathbb{R}^{b\times p\times h_{FF}}$ through two successive MLPs, where $h_{LLM}$ is the LLM output dimension. Once we have the data embedding ${\bm{z}}_{data}=f_{data}({\bm{a}})$ and our sentence embedding ${\bm{z}}_{sentence}=f_{sentence}({\bm{s}})$, for data sample ${\bm{a}}$ and sentence description ${\bm{s}}$, we calculate the multimodal embedding ${\bm{z}}_{multimodal}$ using multihead attention(Vaswani et al., 2017) given below in equation 6.

${\bm{z}}_{multimodal}$ is then returned from our multimodal block and used in the FactFormer architecture. We use cross-attention here instead of self-attention so that our sentence embeddings can be used to learn useful context for our data embeddings. After cross-attention, our multimodal embedding is upsampled back to the data embedding dimension using deconvolution(Zeiler et al., 2010) $f_{deconv}:\mathbb{R}^{b\times p\times h_{FF}}\to\mathbb{R}^{b\times h\times w\times h_{FF}}$ given in equation 7:

This multimodal block is agnostic to both the data-driven backbone as well as the LLM. We compare Llama 3.1 8B(AI@Meta, 2024) and all-mpnet-base-v2 from the Sentence-Transformer package(Reimers & Gurevych, 2019) as our pretrained LLMs. Llama 3.1 8B was used because of its good performance across a wide variety of benchmarks, as well as its small size allowing us to use only a single GPU to generate sentence embeddings. The word embeddings are averaged for each sentence to provide a single embedding that can be used in our projection layer. all-mpnet-base-v2 was used because it is designed to generate embeddings from sentences that are useful for tasks like sentence similarity. In all of our experiments, the LLM is frozen and a projection head is trained. This significantly improves training time by allowing us to generate the sentence embeddings before training, and avoiding expensive gradient computations for our LLMs.

First, we evaluate next-step predictive accuracy, where we take one frame of data and use it to predict the next across our entire temporal window. We see in figure 3 that our multimodal FactFormer significantly outperforms baseline FactFormer when using SentenceTransformer for our Heat, Burgers, and Shallow Water data sets, both with and without transfer learning, where we have small improvement on our Navier-Stokes benchmark when doing standard training. Transfer learning was done because Zhou et al. (2024a) found it to be an effective pretraining strategy. Our SentenceTransformer FactFormer had an average of 21.4% less error across all of our data sets for standard training and average of 20.6% less error than baseline FactFormer when using transfer learning. Additionally, we see smaller improvement over baseline results when using Llama as our LLM, with some notable increases in error on the Shallow Water and Navier Stokes data sets.

Second, we perform autoregressive rollout starting from our initial condition for 40 steps. This evaluates how well our model is able to temporally extrapolate beyond our training horizon. In this case, we use the models trained in section 5.1 with no additional fine-tuning. Accumulated rollout error is presented in table 1. We see that our SentenceTransformer FactFormer has lowest accumulated error for four of five data sets both for standard training and transfer learning, and outperforms baseline across all data sets. Our Llama FactFormer also outperforms baseline across all of our data sets except Shallow Water with standard training. For SentenceTransformer Factformer, we have an average reduction in accumulated error of 64.2% and 54.2% across all of our data set for standard training and transfer learning, respectively. For Llama FactFormer, we have an average reduction in accumulated error of 6.2% and 45.3% for standard and transfer learning, respectively. When removing Shallow Water for Llama Factformer with standard training, we have an average reduction in error of 63.1%.

We note that accumulated autoregressive rollout error for our Shallow Water benchmark is much higher for our transfer learning experiments than the standard training. This instability been noted before in Lorsung et al. (2024) when training errors are very low. In our case, training error for the Shallow Water data set is approximately an order of magnitude lower than other data sets. Error plots are given in figure 8 in Appendix C.

Using the same training setup as in section 5.1, we benchmark our multimodal models against baseline as we increase the amount of data for the Heat and Burgers data sets, given in figure 4, with scaling for the remaining datasets given in Appendix E. For the Heat and Burgers data sets, we see that our multimodal approach consistently outperforms baseline for our both large and small data settings, both with SentenceTransformer and Llama, and for both standard training and transfer learning. We add more data the relative improvement increases due to the plateau in our base models.

Lastly, we perform an ablation study based on text information, to determine which components of our text description are most helpful during the learning process. We compare our baseline model against our multimodal model with SentenceTransformer using 4,000 samples for each equation. Both baseline and multimodal models were trained with transfer learning. Results for next-step prediction are given in figure 5, and autoregressive accumulated rollout errors are given in 2. In both experiments, we see a correlation between richness of sentence description and performance, but that every level of text information improves upon baseline model with transfer learning. Additional error plots and results are given in appendix C.