CViT: Continuous Vision Transformer for Operator Learning

## Abstract

Abstract Operator learning, which aims to approximate maps between infinite-dimensional function spaces, is an important area in scientific machine learning with applications across various physical domains. Here we introduce the Continuous Vision Transformer (CViT), a novel neural operator architecture that leverages advances in computer vision to address challenges in learning complex physical systems. CViT combines a vision transformer encoder, a novel grid-based coordinate embedding, and a query-wise cross-attention mechanism to effectively capture multi-scale dependencies. This design allows for flexible output representations and consistent evaluation at arbitrary resolutions.
We demonstrate CViT’s effectiveness across a diverse range of partial differential equation (PDE) systems, including fluid dynamics, climate modeling, and reaction-diffusion processes. Our comprehensive experiments show that CViT achieves state-of-the-art performance on multiple benchmarks, often surpassing larger foundation models, even without extensive pretraining and roll-out fine-tuning. Taken together, CViT exhibits robust handling of discontinuous solutions, multi-scale features, and intricate spatio-temporal dynamics. Our contributions can be viewed as a significant step towards adapting advanced computer vision architectures for building more flexible and accurate machine learning models in the physical sciences.

## 1 Introduction

Neural operators (Kovachki et al., ) have emerged as a powerful class of deep learning models for scientific machine learning applications, particularly in building efficient surrogates for expensive partial differential equation (PDE) solvers. These models aim to learn mappings between infinite-dimensional function spaces, enabling rapid acceleration of complex simulations in fields such as fluid dynamics and climate modeling (Azizzadenesheli et al., ; Serrano et al., ). Popular architectures like Fourier Neural Operators (FNO) (Li et al., ; Tran et al., ) and DeepONet (Lu et al., ; Wang et al., ) have shown promising results, supported by universal approximation guarantees (Chen & Chen, ; Lanthaler et al., ; Kovachki et al., ; De Ryck & Mishra, ) and interpretations as learning low-dimensional manifolds in function spaces (Seidman et al., ; ).

However, the architectural design of current neural operators has largely been motivated by mathematical intuition about the structure of PDE solution spaces, rather than leveraging recent advances in deep learning architectures. This approach, while theoretically sound, may limit the performance and flexibility of these models when applied to complex, real-world physical systems.

In parallel, the field of computer vision has seen significant advancements in neural field models, which represent continuous functions parameterized by neural networks (Stanley, ). A key innovation in this domain has been the development of conditioned neural fields, which allow for auxiliary inputs to modulate the field’s behavior (Xie et al., ). These models have demonstrated remarkable effectiveness in tasks such as representing high-resolution images and 3D scenes (Mildenhall et al., ; Müller et al., ).

Motivated by the success of conditioned neural fields in computer vision, we introduce the Continuous Vision Transformer (CViT), a novel neural operator architecture that bridges the gap between advanced computer vision techniques and scientific machine learning. Our main contributions can be summarized as follows:

- •
Continuous Vision Transformer (CViT). We introduce CViT, a neural operator that combines a vision transformer encoder with a novel, trainable grid-based coordinate embedding in its decoder. This design enables CViT to effectively capture multi-scale spatial dependencies and allows for continuous querying at arbitrary resolutions.
- •
Theoretical Insights into Lipschitz Constant and Spectral Bias. We present novel insights into the impact of coordinate embeddings on the Lipschitz constant of neural fields and neural operators. This analysis supports the choice of the proposed grid-based coordinate embedding and its efficacy in mitigating spectral bias, a common limitation in training neural operators .
- •
State-of-the-art Performance. Despite its architectural simplicity, CViT achieves state-of-the-art results on challenging benchmarks in climate modeling, fluid dynamics and reaction-diffusion processes. It often outperforms larger models while using fewer parameters and without requiring extensive pretraining.

Taken together, we explore the connection between operator learning architectures and conditioned neural fields, offering a unifying perspective for examining differences among popular operator learning models (see Appendices B and C). This perspective not only enhances understanding of existing architectures but also provides a general framework for designing future neural operators. By bridging the gap between computer vision and scientific machine learning, CViT paves the way for improved surrogates of complex physical systems, with potential applications ranging from climate modeling to engineering design.

## 2 Background and Related Work

The emergence of neural operators has sparked significant advancements in scientific machine learning, particularly in the realm of solving partial differential equations (PDEs) and modeling complex physical systems. Building upon the foundational architectures mentioned in the introduction, recent developments have focused on enhancing the expressiveness and efficiency of these models. For instance, extensions of the Fourier Neural Operator (FNO) (Li et al., ) have explored factorized representations (Tran et al., ) to reduce computational complexity while maintaining performance. Similarly, wavelet-based approaches (Gupta et al., ; Tripura & Chakraborty, ) have been proposed to capture multi-scale features more effectively. Most existing neural operator architectures can be interpreted as conditioned neural fields (Xie et al., ) with specific choices of base fields and conditioning mechanisms, see Appendices B and C for a detailed discussion.

Concurrently, the success of vision transformers in computer vision tasks has inspired their adaptation to scientific machine learning problems. The self-attention mechanism central to transformers offers a powerful tool for capturing long-range dependencies in spatial and temporal data, a crucial aspect in many PDE-governed systems. However, adapting vision transformers to the continuous nature of physical problems poses unique challenges, particularly in handling irregular domains and ensuring output consistency across different resolutions.

In the specific context of operator learning, several transformer-based approaches have emerged to address these challenges. OFormer (Li et al., ) introduced a novel way to embed continuous input functions and output queries into a transformer architecture. This approach demonstrated the potential of transformers in handling the inherent continuity of physical problems. GNOT (Hao et al., ) further extended this idea by incorporating graph structures, allowing for more flexible representation of input functions and improved handling of irregular domains. A significant advancement in this direction came with DPOT (Hao et al., ), which introduced a denoising pre-training strategy. This approach enables the model to learn from diverse PDE datasets, significantly enhancing its generalization capabilities.
However, these methods often require extensive computational resources for training and may struggle with maintaining consistency when evaluating at arbitrary resolutions.

Despite these advancements, several key challenges remain in the field of operator learning. These include efficiently handling high-dimensional data while capturing long-range dependencies, ensuring output consistency and accuracy across different resolutions, and balancing model expressiveness with computational efficiency.
Addressing even a subset of these challenges requires novel architectural designs that can leverage the strengths of both neural operators and vision transformers while overcoming their respective limitations. This sets the stage for our proposed Continuous Vision Transformer (CViT), which aims to tackle specific challenges related to resolution independence and efficient capturing of multi-scale features in regular grid-based problems, combining the continuous nature of neural operators with the powerful representation capabilities of vision transformers.

## 3 Continuous Vision Transformer (CViT)

Figure: Figure 1: Continuous Vision Transformer (CViT) Architecture: CViT consists of the following components: (1) Spatio-temporal patch embeddings to extract localized features. (2) A temporal aggregation module based on the Perceiver architecture, which captures temporal correlations to compresses tokens along the time axis. (3) A Transformer encoder that captures multi-scale spatial dependencies via self-attention layers. (4) A novel grid-based positional encoding scheme for query coordinates, allowing for flexible output representation and interpolation. (5) A cross-attention decoder that integrates information from the input function with query coordinates.
Refer to caption: extracted/5899740/cvit_figures/cvit_arch.png

While neural operators have shown promising results in learning mappings between function spaces, their architectural design has not fully leveraged recent advancements in deep learning. Popular architectures like FNO (Li et al., ), DeepONet (Lu et al., ), and NoMaD (Seidman et al., ) (see Appendix C) primarily rely on simple building blocks such as Fourier layers or fully connected networks, which may not effectively capture complex spatial dependencies and multiscale features present in many physical systems.

Motivated by the success of vision transformers in computer vision tasks, we introduce the Continuous Vision Transformer (CViT), a novel neural operator architecture that combines the strengths of vision transformers and continuous function representations. CViT aims to learn more powerful and flexible representations of time-dependent physical systems by leveraging a vision transformer encoder to capture complex spatial dependencies in input functions, a trainable grid-based positional encoding for flexible representation of query coordinates, and a cross-attention mechanism to integrate input function information with query coordinates.

Taken together, the overall design of CViT’s architecture draws inspiration from the analysis of existing neural operators presented in Appendix C, while leveraging the power of vision transformers to enhance expressiveness.

### 3.1 Architecture Description

The CViT architecture consists of three main components: a vision transformer encoder, a grid-based positional encoding, and a cross-attention mechanism.

#### Patch Embeddings.

The vision transformer encoder takes as input a gridded representation of the input function $u$, yielding a 2+1 dimensional spatio-temporal data tensor $\mathbf{u}\in{\mathbb{R}}^{T\times H\times W\times D}$ with $D$ channels. We patchify our inputs into 3D tokens $\mathbf{u}_{p}\in{\mathbb{R}}^{T\times\frac{H}{P}\times\frac{W}{P}\times C}$ by tokenizing each 2D spatial frame independently, following the process used in standard Vision Transformers (Dosovitskiy et al., ). We then add trainable 1D temporal and 2D spatial positional embeddings to each token:

$$ $\displaystyle\mathbf{u}_{pe}=\mathbf{u}_{p}+\text{PE}_{t}+\text{PE}_{s},\quad \text{PE}_{t}\in{\mathbb{R}}^{T\times 1\times 1\times C},\quad\text{PE}_{s}\in {\mathbb{R}}^{1\times\frac{H}{P}\times\frac{W}{P}\times C}.$ $$

#### Temporal Aggregation.

To reduce computational cost, we use a temporal aggregation layer based on the Perceiver architecture (Jaegle et al., ; Alayrac et al., ). Specifically, we learn a predefined number of latent input queries $\mathbf{z}\in{\mathbb{R}}^{1\times C}$ by applying a cross-attention module that maps the flattened visual features $\mathbf{u}_{f}\in{\mathbb{R}}^{(\frac{H}{P}\times\frac{W}{P})\times T\times C}$ to a latent tensor $\mathbf{z}_{agg}\in{\mathbb{R}}^{(\frac{H}{P}\times\frac{W}{P})\times 1\times C}$ as

$$ $\displaystyle\mathbf{z}^{\prime}$ $\displaystyle=\hat{\mathbf{z}}+\operatorname{MHA}\left(\operatorname{LN}\left( \hat{\mathbf{z}}\right),\operatorname{LN}\left(\mathbf{u}_{f}\right), \operatorname{LN}\left(\mathbf{u}_{f}\right)\right),$ $\displaystyle\mathbf{z}_{agg}$ $\displaystyle=\mathbf{z}^{\prime}+\operatorname{MLP}\left(\operatorname{LN} \left(\mathbf{z}^{\prime}\right)\right).$ $$

Here, $\mathbf{z}$ is initialized by a unit Gaussian distribution and $\hat{\mathbf{z}}\in{\mathbb{R}}^{(\frac{H}{P}\times\frac{W}{P})\times T\times C}$ is obtained by tiling the latent query $\mathbf{z}$ $\left(\frac{H}{P}\times\frac{W}{P}\right)$ times. Besides compressing tokens over the time axis, this module also enables the model to handle inputs with a variable number of time steps.

#### Transformer Encoder.

We then process the aggregated tokens $\mathbf{z}_{agg}$ using a sequence of $L$ pre-norm Transformer blocks (Vaswani et al., ; Xiong et al., ),

$$ $\displaystyle\mathbf{z}_{0}$ $\displaystyle=\operatorname{LN}(\mathbf{z}_{agg}),$ $\displaystyle\mathbf{z}_{\ell}^{\prime}$ $\displaystyle=\operatorname{MSA}\left(\operatorname{LN}\left(\mathbf{z}_{\ell- 1}\right)\right)+\mathbf{z}_{\ell-1},$ $\displaystyle\ell=1\ldots L,$ $\displaystyle\mathbf{z}_{\ell},$ $\displaystyle=\operatorname{MLP}\left(\operatorname{LN}\left(\mathbf{z}_{\ell} ^{\prime}\right)\right)+\mathbf{z}_{\ell}^{\prime},$ $\displaystyle\ell=1\ldots L.$ $$

#### Cross-Attention Decoder.

To enable continuous evaluation of outputs, we design a novel and efficient coordinate embedding to capture the fine-scale features of the target functions. Specifically, we create a uniform grid $\{\mathbf{y}_{ij}\}\subset[0,1]^{2}$, for $i=1,\dots N_{x}$ and $j=1,\dots N_{y}$, along with associated trainable latent grid features $\mathbf{x}\in{\mathbb{R}}^{N_{x}\times N_{y}\times C}$. For a query point $y\in{\mathbb{R}}^{2}$, we compute a Nadaraya-Watson interpolant over grid latent features:

$$ $\displaystyle\mathbf{x}^{\prime}=\sum_{i=1}^{N_{x}}\sum_{j=1}^{N_{y}}w_{ij} \mathbf{x}_{ij},\quad w_{ij}=\frac{\exp\left(-\beta\|y-\mathbf{y}_{ij}\|^{2} \right)}{\sum_{ij}\exp\left(-\beta\|y-\mathbf{y}_{ij}\|^{2}\right)}.$ (1) $$

Here $\beta>0$ is a hyperparameter that determines the locality of the interpolated features. It is important for determining the smoothness of the interpolation function. Specifically, larger values of $\beta$ yield more localized weight distributions $w_{ij}$, resulting in a higher-frequency interpolant that captures finer-scale variations. Conversely, smaller $\beta$ values produce a smoother interpolant by averaging over a broader neighborhood of points.

The interpolated grid feature $\mathbf{x}_{0}=\mathbf{x}^{\prime}\in{\mathbb{R}}^{1\times C}$ is used as query input to a transformer decoder. This decoder uses the output of the vision transformer encoder $\mathbf{z}_{L}$ as keys and values in a cross-attention mechanism:

$$ $\displaystyle\mathbf{x}^{\prime}_{k}$ $\displaystyle=\mathbf{x}_{k-1}+\operatorname{MHA}\left(\operatorname{LN}\left( \mathbf{x}_{k-1}\right),\operatorname{LN}\left(\mathbf{z}_{L}\right), \operatorname{LN}\left(\mathbf{z}_{L}\right)\right),$ $\displaystyle k=1\ldots K,$ $\displaystyle\mathbf{x}_{k}$ $\displaystyle=\mathbf{x}^{\prime}_{k}+\operatorname{MLP}\left(\operatorname{LN }\left(\mathbf{x}^{\prime}_{k}\right)\right),$ $\displaystyle k=1\ldots K.$ $$

Finally, a small MLP network projects the output to the desired output dimension.

#### Theoretical Insights.

The theoretical analysis of Lipschitz constants for different coordinate embeddings (as detailed in Appendix E) provides crucial insights that inform the design choices in CViT and have broader implications for neural operator architectures. Our findings reveal that the choice of coordinate embedding significantly influences the decoder’s capacity to learn complex, high-frequency functions.

For conventional MLP decoders with linear coordinate embeddings, the Lipschitz constant is approximately 1 at initialization, which does not enhance the network’s ability to capture fine-scale features, thereby leading to prediction susceptible to spectral bias (Rahaman et al., ). In contrast, random Fourier features can increase the Lipschitz constant through two mechanisms: by increasing the standard deviation of the sampling distribution or by increasing the network width. This observation not only explains the effectiveness of Fourier features in addressing spectral bias, but also highlights a potential trade-off between network size and the frequency range that can be captured.

Our proposed grid-based coordinate embedding in CViT offers a more direct and controllable approach to increasing the network’s Lipschitz constant. By adjusting the interpolation parameter $\beta$, we can tune the model’s ability to capture high-frequency components in the target functions. This provides a clear advantage as it allows for more precise control over the model’s spectral properties without necessarily increasing model complexity.

#### CViT vs other Transformer-based Approaches.

To further highlight the key innovations of CViT, we compare it with other popular transformer-based operator learning methods:

OFormer (Li et al., ):
OFormer uses an encoder-decoder structure. Its encoder combines a point-wise MLP and stacked self-attention blocks to process input functions, while employing random Fourier Features (Tancik et al., ) for query coordinate encoding. The model generates latent embeddings through cross-attention at sampled query locations.
In the decoder, OFormer uses recurrent MLPs for dynamic propagation instead of traditional masked attention. Unlike CViT, it doesn’t use ViT’s patch embedding to tokenize input functions, potentially limiting its ability to capture complex spatial dependencies.

GNOT (Hao et al., ): GNOT is a transformer-based framework for learning operators that addresses the challenges of multiple input functions and irregular meshes. It designs a heterogeneous normalized attention layer, providing an encoding interface for various input functions and additional prior information. Additionally, GNOT features a geometric gating mechanism, described as a soft domain decomposition approach to multi-scale problems. Unlike CViT, GNOT uses an plain MLP to process query coordinates, which may potentially suffer from spectral bias (Rahaman et al., ).

DPOT (Hao et al., ): DPOT is a transformer-based architecture designed for pretraining on multiple PDEs. It introduces a new auto-regressive denoising strategy and a novel model architecture based on Fourier attention. Unlike OFormer, GNOT, and CViT, DPOT does not explicitly take query coordinates as inputs, which implies that it cannot be evaluated at arbitrary locations.

MPP (McCabe et al., ): MPP is a large-scale PDE surrogate model that directly employs a Vision Transformer as its backbone. Similar to DPOT, it does not take query coordinates as inputs, thereby losing flexibility in diverse output representations.

## 4 Experiments

We compare the CViT against popular neural operators on three challenging benchmarks in physical sciences. In addition to demonstrating CViT’s performance against strong baselines, we also conduct comprehensive ablation studies to probe the sensitivity on its own hyper-parameters.

#### CViT model setup.

We construct CViT models with different configurations, as summarized in Table 1. For all experiments, unless otherwise stated, we use a patch size of $8\times 8$ for tokenizing inputs. We also employ a decoder with a single cross-attention Transformer block for all configurations. While we tested decoders with multiple layers, they did not yield improvements. The grid resolution is set to the spatial resolution of each dataset. The latent dimension of grid features is set to 512, and if it does not match the transformer’s embedding dimension, we align it using a dense layer. Besides, we use $\beta=10^{5}$ to ensure sufficient locality of the interpolated features. These choices are validated by extensive ablation studies, as illustrated in Figure 4. Full details of the training and evaluation procedures are provided in Appendix F.1.

**Table 1: Details of Continuous Vision Transformer model variants.**
| Model | Encoder layers | Embedding dim | MLP width | Heads | # Params |
| --- | --- | --- | --- | --- | --- |
| CViT-S | 5 | 384 | 384 | 6 | 13 M |
| CViT-B | 10 | 512 | 512 | 8 | 30 M |
| CViT-L | 15 | 768 | 1536 | 12 | 92 M |

#### Baselines.

We select the following methods as our baselines for comparisons. DeepONet (Lu et al., ): One of the first neural operators with universal approximation guarantees. NoMaD (Seidman et al., ): A recently proposed architecture leveraging nonlinear decoders to achieve enhanced dimensionality reduction. FNO
(Li et al., ): An efficient framework for operator learning in the frequency domain. FFNO (Tran et al., ): An improved version of FNO utilizing a separable Fourier representation, which reduces model complexity and facilitates deeper network structures. UNO (Ashiqur Rahman et al., ): A U-shaped neural operator with FNO layers.
$\textbf{U-Net}_{\textbf{att}}$ (Gupta & Brandstetter, ): A modern U-Net architecture using bias weights and group normalization, and wide ResNet-style 2D convolutional blocks, each of which is followed by a spatial attention block.
U-F2Net (Gupta & Brandstetter, ): A U-Net variant where the lower blocks in both the downsampling and upsampling paths are replaced by Fourier blocks, each consisting of 2 FNO layers with residual connections.
Dilated ResNet (Stachenfeld et al., ): A ResNet model that adapts filter sizes at different layers using dilated convolutions, providing an alternative method for aggregating global information.
GK-Transformer (Cao, ) / OFormer
(Li et al., ) / GNOT (Hao et al., ): Various transformer-based architectures for operator learning, with different designs of their encoder/decoder and attention mechanisms.
DPOT (Hao et al., ) / MPP (McCabe et al., ): Transformer-based networks with patch embedding as in ViT, and pre-trained on extensive PDE datasets using a rollout loss.

For all baselines, we report evaluation results from the original papers when applicable. When not applicable, we train and evaluate each model following the suggested settings in the respective paper. Details on the implementation of baseline models are provided in Appendix F.2.

### 4.1 Main Results

Here we provide our main results across three challenged benchmarks. The full details on the underlying equations, dataset generation and problem setup for each case are provided in the Appendix; see Section F.3, F.4 and F.5, respectively.

#### Advection of discontinuous waveforms.

The first benchmark involves predicting the transport of discontinuous waveforms governed by a linear advection equation with periodic boundary conditions. We make use of the datasets and problem setup established by Hoop et al. (de Hoop et al., ). This benchmark evaluates our model’s capability in handling discontinuous solutions and shocks in comparison to popular neural operator models.

Figure: Figure 2: Advection of discontinuous waveforms. Prediction (red dashed line) versus ground truth (blue line) for the worst-case example in the test dataset, for: DeepONet; NoMaD; FNO; CViT. Also reported is the associated Total Variation (TV) error ($\downarrow$).
Refer to caption: x1.png

Figure 2 illustrates the test sample corresponding to the worst-case prediction of each model we compared. We observe that CViT is able to better capture the discontinuous targets and yield the most sharp predictions across all baselines. This is also reflected in both the relative L2 and Total Variation error metrics; see Appendix Table 5 for a more complete summary, including details on model configurations, as well as average, median and worst case errors across the entire test dataset.

#### Shallow-water equations.

The second benchmark involves a fluid flow system governed by the 2D shallow-water equations. This describes the evolution of a 2D incompressible flow on the surface of the sphere and it is commonly used as a benchmark in climate modeling. Here we adhere to the dataset and problem setup established in PDEArena (Gupta & Brandstetter, ), allowing us to perform fair comparisons with several state-of-the-art deep learning surrogates.

**Table 2: Performance for 5-step roll-out predictions in the shallow-water equations benchmark.**
| Model | # Params | Rel. $L^{2}$ error ($\downarrow$) |
| --- | --- | --- |
| DilResNet | 4.2 M | 13.20% |
| $\text{U-Net}_{\text{att}}$ | 148 M | 5.68% |
| FNO | 268 M | 3.97% |
| U-F2Net | 344 M | 1.89% |
| UNO | 440 M | 3.79% |
| CViT-S | 13 M | 4.47% |
| CViT-B | 30 M | 2.69% |
| CViT-L | 92 M | 1.56% |

Table 2 presents the results of CViT against several competitive and heavily optimized baselines.
Our proposed method achieves the lowest relative $L^{2}$ error.
It is also worth noting that the reported performance gains are achieved using CViT configurations with a significantly smaller parameter count compared to the other baselines. This highlights the merits of our approach that enables the design of parameter-efficient models whose outputs can also be continuously queried at any resolution. Additional visualizations of our models are shown in Figure 3 and 10.

In addition, here we perform extensive ablation studies on the hyper-parameters of CViT. The results are summarized in Figure 4. As shown in Figure 4 (a), we observe that a smaller patch size typically leads to better accuracy, but at higher computational costs. Moreover, Figure 4 (b) demonstrates that the model performance is significantly influenced by the type of coordinate embeddings used in the CViT decoder. We compare our proposed grid-based embedding against a small MLP and random Fourier features (Tancik et al., ). The grid-based embedding achieves the best accuracy, outperforming other methods by up to an order of magnitude. This result is in agreement with the theoretical insights presented in Appendix E.

We also investigate the impact of the resolution of the associated grid. Figure 4 (c) reveals that the best results are obtained when the grid resolution matches the highest resolution at which the model is evaluated. As another direct empirical evidence of theoretical results in Appendix E,
Figure 4(d) shows the CViT model’s sensitivity to the hyper-parameter $\beta$ used for computing interpolated features; too small $\beta$ values can degrade predictive accuracy. Finally, in Appendix Figure 9, we observe minor improvements by increasing the latent dimension of grid features, or the number of decoder attention heads. Hence
the overall model performance remains robust to variations in these hyper-parameters.

Taken together, these findings not only justify our design choices for CViT but also provide general insights for the development of neural operators. They suggest that future research in this field should pay close attention to the interplay between spatio-temporal tokenization, coordinate representation, and the inherent multi-scale nature of physical systems. Such considerations could lead to even more effective architectures for a wide range of scientific machine learning tasks.

Figure: Figure 3: Shallow water benchmark. Representative CViT rollout prediction of the vorticity field, and point-wise error against the ground truth.
Refer to caption: extracted/5899740/cvit_figures/swe/swe_vor_pred.png

Figure: Figure 4: Ablation studies for CViT on the shallow-water equations benchmark. Convergence of test errors for: (a) different patch sizes; (b) different coordinate embeddings; (c) different resolutions of latent grids; (d) sensitivity on interpolating the CViT latent grid features, as controlled by the $\beta$ parameter. Results obtained using CViT-L with $16\times 16$ patch-size, varying each hyper-parameter of interest while keeping others fixed.
Refer to caption: x2.png

#### Navier-Stokes equations.

Our third benchmark corresponds to predicting the evolution of a 2D buoyancy-driven flow, depicting a smoke volume rising in a bounded square domain. The underlying governing equations are the incompressible Navier-Stokes equations coupled with a passive scalar transport equation. For a fair comparison with various neural operators and other state-of-the-art deep learning models, we use the dataset generated by PDEArena (Gupta & Brandstetter, ), and follow the problem setup reported in Hao et al. (Hao et al., ).

Our main results are summarized in Table 3, indicating that CViT outperforms all baselines we have considered. Notably, the previous best result (DPOT-L (Fine-tuned)) is achieved with a much larger transformer-based architecture, which involves pretraining on massive diverse PDE datasets followed by fine-tuning on the target dataset for 500 epochs. Here we also note that DPOT was pretrained and fine-tuned using a rollout loss, which directly optimizes for the multi-step prediction task on which the models are evaluated. In contrast, CViT was trained using a one-step loss without any rollout finetuning. This difference in training objective significantly favors DPOT for rollout evaluations across all benchmarks.
Furthermore, CViT exhibits remarkable parameter efficiency. CViT-S, with only 13M parameters, achieves accuracy comparable to the largest pretrained model (1.03B). Additionally, CViT shows good scalability; as the number of parameters increases, performance improves. These findings strongly support the effectiveness our approach in learning evolution operators for complex physical systems.
Additional results and sample visualizations are provided in Figure 5, as well as Figure 11 and 12 in Appendix.

**Table 3: Relative $L^{2}$ error of rollout predictions on the Navier-Stokes and Diffusion-Reaction benchmark.**
| Model | # Params | Navier-Stokes | Diffusion-Reaction |
| --- | --- | --- | --- |
| FNO | 0.5 M | 9.12 % | 12.00 % |
| FFNO | 1.3 M | 8.39 % | 5.71 % |
| GK-T | 1.6 M | 9.52% | 3.59 % |
| GNOT | 1.8 M | 17.20 % | 3.11 % |
| Oformer | 1.9 M | 13.50 % | 1.92 % |
| DPOT-Ti | 7 M | 12.50 % | 3.21 % |
| MPP-S | 30M | - | 1.12 % |
| DPOT-S | 30 M | 9.91 % | 3.79 % |
| MPP-L (Pre-trained) | 400 M | - | 0.98 % |
| DPOT-L (Pre-trained) | 500 M | 7.98 % | 2.32 % |
| DPOT-L (Fine-tuned) | 500 M | 2.78 % | 0.73 % |
| DPOT-H (Pre-trained) | 1.03 B | 3.79 % | 1.91 % |
| CViT-S | 13 M | 3.75 % | 1.13 % |
| CViT-B | 30 M | 3.18 % | 1.11 % |
| CViT-L | 92 M | 2.35 % | 0.68 % |

Figure: Figure 5: Navier-Stokes benchmark. Representative CViT rollout predictions of the passive scalar field, and point-wise error against the ground truth.
Refer to caption: extracted/5899740/cvit_figures/ns/ns_u_pred.png

#### Diffusion-Reaction equations.

Our final benchmark focuses on learning the density fields $\mathbf{u}=(u,v)$ governed by a coupled 2D diffusion-reaction equation (DR). This PDE, commonly used in chemistry, models substances reacting and diffusing across a spatial domain, illustrating chemical reaction processes and resulting in diverse spatial patterns. For a fair comparison, we use the dataset from PDEBench (Takamoto et al., ) and follow the problem setup described by Hao et al. ().

Table 3 presents our main results. Consistent with previous findings, CViT-L achieves the best performance across all baselines. CViT models outperform other approaches with comparable or fewer parameters, demonstrating good scalability and parameter efficiency.
Representative rollout predictions are provided in Figure 6 and 13 (see Appendix). Our predictions show excellent agreement with the corresponding ground truth.

Interestingly, we note that the solution of Diffusion-Reaction equations appears simpler than that of the Navier-Stokes equation dataset, resulting in uniformly lower test errors for various baselines in the DR benchmark.
In particular, OFormer outperforms GNOT in both examples, despite having comparable parameters. We postulate that this performance gain is due to the use of random Fourier features in OFormer, facilitating more efficient learning of complex functions. This observation further demonstrates the importance of coordinate embedding in designing neural operators and indirectly highlights the effectiveness of our proposed grid-based embedding.

Figure: Figure 6: Diffusion-reaction Benchmark. Representative CViT rollout prediction of the density field $u$, and point-wise error against the ground truth.
Refer to caption: extracted/5899740/cvit_figures/dr/dr_u_pred.png

## 5 Discussion

#### Summary.

This work introduces the Continuous Vision Transformer (CViT), a novel neural operator architecture that leverages advances in computer vision to address challenges in learning complex physical systems. CViT combines the strengths of vision transformers and continuous function representations, achieving state-of-the-art accuracy on challenging benchmarks in climate modeling, fluid dynamics and reaction-diffusion processes. Our approach demonstrates the potential of adapting advanced computer vision techniques to develop more flexible and accurate machine learning models for physical sciences. The key innovations of CViT include: (a) a vision transformer encoder for capturing multi-scale spatial dependencies in input functions, (b) a trainable grid-based positional encoding for flexible representation of query coordinates, and (c) a query-wise cross-attention mechanism ensuring consistent evaluation at arbitrary resolutions. These design choices ensure that CViT maintains the desirable properties of a well-defined conditioned neural field (as discussed in Appendix B), while offering improved flexibility and performance (as discussed in Appendix E).

Our empirical results across various PDE benchmarks demonstrate that CViT often outperforms larger models with fewer parameters and without extensive pretraining. This highlights the merits of our approach in designing parameter-efficient models that can be continuously queried at any resolution. The broader impact of this work lies in its potential to accelerate scientific discovery through more efficient and accurate simulations of complex physical systems, with applications ranging from climate modeling to engineering design.

#### Limitations.

Despite its promising performance, the proposed CViT architecture can be further improved in multiple areas. Firstly, like most transformer-based models, CViT’s self-attention blocks scale quadratically with the number of tokens, posing computational challenges for high-dimensional or high-resolution data. Recent advancements in computer vision (Liu et al., ) are tackling these scalability issues, and these ideas can be transferred to our setting to improve CViT’s efficiency for larger-scale problems. Secondly, transformer-based models like CViT generally require large training datasets for optimal performance and may struggle in data-scarce scenarios without careful regularization. To address this, future work could explore leveraging CViT’s ability to reconstruct continuous output functions in conjunction with physics-informed loss functions (Wang et al., ; ) for more data-efficient training and improved generalization. Lastly, while our current experiments treat input functions as sequences of images, many practical applications in scientific computing involve complex geometric structures and diverse input modalities such as unstructured meshes or point clouds. Addressing these cases requires developing appropriate tokenization schemes, and recent work in (Alkin et al., ; Pang et al., ) has taken steps in this direction. Extending CViT to handle such diverse input modalities represents an important avenue for future research.

#### Future Work.

Building on the success of CViT, several promising directions for future research emerge. One key area is the exploration of more efficient attention mechanisms to improve scalability to higher-dimensional problems. Additionally, investigating the potential of pre-training strategies, similar to those employed in DPOT (Hao et al., ), could further enhance CViT’s performance and generalization capabilities across diverse PDE systems. Finally, extending CViT to handle 3D problems and developing techniques for multi-resolution training and evaluation could significantly broaden its applicability in scientific computing domains.

#### Broader Impact.

The broader impact of this work lies in its potential to fuel the development of resolution-invariant foundation models for accelerating scientific discovery and enable more efficient and accurate simulations of complex physical systems, with applications ranging from climate modeling to engineering design.

### Acknowledgements

We would like to acknowledge support from the US Department of Energy under the Advanced Scientific Computing Research program (grant DE-SC0024563). We also thank the developers of the software that enabled our research, including JAX Bradbury et al. (), Matplotlib Hunter (), and NumPy Harris et al. ().

### Ethics Statement

Our contributions enable advances in accelerating the modeling and emulation of complex physical systems. This has the potential to significantly impact fields such as climate science, fluid dynamics, and materials engineering by providing more accurate and efficient simulation tools. While we do not anticipate specific negative impacts from this work, as with any powerful predictive tool, there is potential for misuse. We encourage the research community to consider the ethical implications and potential dual-use scenarios when applying these technologies in sensitive domains.

### Reproducibility Statement

The code used to carry out the ablation experiments is provided as a supplementary material for this submission. If the paper is accepted, we plan on making our entire code-base publicly available on GitHub.

## Appendix A Nomenclature

Table 4 summarizes the main symbols and notation used in this work.

**Table 4: Summary of the main symbols and notation used in this work.**
| Notation | Description |
| --- | --- |
| Operator Learning |  |
| $\mathcal{X}$ | The input function space |
| $\mathcal{Y}$ | The output function space |
| $u\in\mathcal{X}$ | Input function |
| $s\in\mathcal{Y}$ | Output function |
| $y$ | Query coordinate in the input domain of $s$ |
| $\mathcal{G}:\mathcal{X}\rightarrow\mathcal{Y}$ | The operator mapping between function spaces |
| ${\mathcal{E}}:{\mathcal{X}}\to{\mathbb{R}}^{n}$ | Encoder mapping |
| ${\mathcal{D}}:{\mathbb{R}}^{n}\to{\mathcal{Y}}$ | Decoder mapping |
| Fourier Neural Operator |  |
| $\mathcal{F},\mathcal{F}^{-1}$ | Fourier transform and its inverse |
| $\mathcal{F}_{n},\mathcal{F}_{n}^{-1}$ | Discrete Fourier transform and its inverse truncated on the first $n$ modes |
| $K$ | Linear transformation applied to the $n$ leading Fourier modes |
| $W$ | Linear transformation (bias term) applied to the layer inputs |
| $k$ | Fourier modes / wave numbers |
| $\sigma$ | Activation function |
| Continuous Vision Transformer |  |
| $\text{PE}_{t}$ | Temporal positional embedding |
| $\text{PE}_{s}$ | Spatial positional embedding |
| MSA | Multi-head self-attention |
| MHA | Multi-head attention |
| LN | Layer normalization |
| $P$ | Patch size of Vision Transformer |
| $C$ | Embedding dimension of Vision Transformer |
| $\mathbf{x}\in\mathbb{R}^{N_{x}\times N_{y}\times C}$ | Latent grid features |
| $\beta$ | Locality of interpolated latent grid features |
| Partial Differential Equations |  |
| $\zeta$ | Vorticity of shallow water equations |
| $\eta$ | Height of shallow water equations |
| $\mathbf{u}$ | Velocity field |
| $P$ | Pressure field |
| $c$ | Passive scalar (smoke) |
| $\nu$ | Viscosity |
| Hyperparameters |  |
| $B$ | Batch size |
| $Q$ | Number of query coordinates in each batch |
| $D$ | Number of latent variables of interest |
| $T$ | Number of previous time-steps |
| $H\times W$ | Resolution of spatial discretization |

## Appendix B Neural Fields & Conditioned Neural Fields

Figure: Figure 7: The structure of a conditioned neural field with both global and local conditioning.
Refer to caption: extracted/5899740/figs/CNF.png

A neural field is a function over a continuous domain parameterized by a neural network (Stanley, ). Recently, the computer vision community has adopted neural fields to represent 3D spatial fields, such as neural radiance fields (NeRFs) (Mildenhall et al., ), spatial occupancy fields (Mescheder et al., ), and signed distance functions (Park et al., ). These fields represent functions defined over continuous spatial domains, with NeRFs additionally taking an element of a continuous ray space as input. In this paper, a point in the domain of a neural field is referred to as a *query point*, here denoted by $y$. Conditioned neural fields allow these fields to change based on auxiliary inputs $u$, often represented as latent codes $z$, without the need to retrain the entire field; see Figure 7. The latent code affects the forward pass by modifying the parameters of the base field through either global conditioning, where the forward pass is affected uniformly across all query points, or local conditioning, where the modifications depend on the specific query point.

#### Global Conditioning.

Global conditioning in a neural field refers to modifying the forward pass only by the input $u$ and in a uniform manner over the entire query domain. This can be achieved
through concatenation, where a global latent code is appended to the query vector and passed together through an MLP (Chen & Zhang, ; Liu et al., ; Jang & Agapito, ). Alternatively, a hypernetwork could accept the auxiliary input and determine the neural field’s parameter values based on this input (Ha et al., ; Rebain et al., ). However, modifying every parameter of the base field for different auxiliary inputs can be computationally expensive. For example, feature-wise Linear Modulation (FiLM) (Perez et al., ), which originates from conditional batch normalization (De Vries et al., ), has been proposed to modify only the sine activation function parameters in a SIREN network (Sitzmann et al., ). This technique has been used for global conditioning in methods such as $\pi$-GAN (Chan et al., ).

#### Local Conditioning.

Local conditioning occurs when the latent code $z$, which modifies the base field’s parameters, is a function of the input query $y$, i.e., $z=g(y)$. This results in query-dependent field parameters that can vary across the query domain.
A common approach to generate local latent codes is to process the image or scene with a convolutional network and extract features from activations with receptive fields containing the query point. This can be done using 2D convolutions of projected 3D scenes (Xu et al., ; Yu et al., ; Trevithick & Yang, ) or 3D convolutions (Peng et al., ; Liu et al., ; Chibane et al., ). Muller et al. (Müller et al., ) learn features directly on multiresolution grid vertices and create local features via nearest grid point interpolation for each grid. Chan et al. (Chan et al., ) create features along three 2D planes and interpolate features from projections onto these planes. DeVries et al. (DeVries et al., ) learn latent codes over patches on a 2D “floorplan” of a scene. Other methods divide the query domain into patches or voxels and learn a latent code $z$ for each patch or voxel (Chabra et al., ; Tretschk et al., ; Chen et al., ).

## Appendix C Neural Operators as Conditioned Neural Fields

Here we demonstrate that many popular operator learning architectures are in fact neural fields with local and/or global conditioning. While connections between architectures such as the DeepONet and the Fourier Neural Operator have been described previously (Kovachki et al., ), they are typically formulated under certain architecture choices which allow these architectures to approximate one another. In contrast, here we show that, by viewing these architectures as conditioned neural fields, they are explicitly related through their respective choices of conditioning mechanisms and base fields.

### C.1 Encoder-Decoder Neural Operators

A large class of operator learning architectures can be described as encoder-decoder architectures according to the following construction. To learn an operator ${\mathcal{G}}:{\mathcal{X}}\to{\mathcal{Y}}$, an input function $u\in{\mathcal{X}}$ are first mapped to a finite dimensional latent representation through the encoding map ${\mathcal{E}}:{\mathcal{X}}\to{\mathbb{R}}^{n}$. This latent representation is then mapped to a callable function through a decoding map ${\mathcal{D}}:{\mathbb{R}}^{n}\to{\mathcal{Y}}$. The composition of these two maps gives the architecture of the operator approximation scheme, ${\mathcal{G}}={\mathcal{D}}\circ{\mathcal{E}}$. The class of encoder-decoder architectures includes the DeepONet (Lu et al., ), the PCA-based architecture of (Bhattacharya et al., ), and NoMaD (Seidman et al., ).

These encoder-decoder architectures can be seen as neural fields with a global conditioning mechanism by viewing the role of the decoder as an operator which performs the conditioning of a base field, with a latent code derived from a parameterized encoder, $z={\mathcal{E}}_{\phi}(u)$.

#### DeepONet.

For example, in a DeepONet, the base field is a neural network $t_{\theta}(y)$ (trunk network) appended with a last linear layer, and the decoder map conditions the weights of this last layer via the output of the encoding map ${\mathcal{E}}_{\phi}(u)$ (branch network); see Figure 8(a) for a visual illustration.
In other words, the decoder map is a weighted linear combination of $n$ basis elements $t_{\theta}(y)$, where the weights are modulated by a global encoding of the input function ${\mathcal{E}}_{\phi}(u)$, yielding outputs of the form $s(y)=\langle{\mathcal{E}}_{\phi}(u),t_{\theta}(y)\rangle$.
The PCA-based architecture of (Bhattacharya et al., ) can be viewed analogously using as basis elements the leading $n$ eigenfunctions of the empirical covariance operator computed over the output functions in the training set.

#### NoMaD.

Seidman et al. (Seidman et al., ) demonstrated how encoder-decoder architectures can employ nonlinear decoders to learn low-dimensional manifolds in an ambient function space. These so-called
NoMaD architectures typically perform their conditioning by concatenating the input function encoding with the query input, i.e., $s(y)=f_{\theta}(y,{\mathcal{E}}_{\phi}(u))$; see Figure 8(b) for an illustration.

### C.2 Integral Kernel Neural Operators

Integral Kernel Neural Operators form another general class of neural operators that learn mappings between function spaces by approximating the integral kernel of an operator using neural networks. They were introduced as a generalization of the Neural Operator framework proposed by Li et al. (Li et al., ).
The key idea behind here is to represent the operator $\mathcal{G}:\mathcal{X}\rightarrow\mathcal{Y}$ as an integral kernel, $\mathcal{G}(u)(y)=\int_{\Omega}\kappa(x,y,u(x))dx$,
where $\kappa$ is a kernel function.
In practice, the kernel function $\kappa$ is parameterized by a neural network, and the integral is then approximated using a quadrature rule, such as Monte Carlo integration, Gaussian quadrature, or via spectral convolution in the frequency domain.

#### Graph Neural Operator (GNO).

A single layer of the GNO (Li et al., ) is given by

$$ $s(y)=\sigma\Big{(}Wu(y)+\frac{1}{|N(y)|}\sum\limits_{x\in N(y)}\kappa_{\phi}(x ,y,u(x))u(x)\Big{)},$ (2) $$

where $W\in{\mathbb{R}}^{d_{u}\times d_{s}}$ is a trainable weight matrix, and $\kappa_{\phi}(\cdot)$ is a local message passing kernel parameterized by $\phi$ (Li et al., ). As illustrated in Figure 8(c), the GNO layer is a conditioned neural field with local and global conditioning on the input function $u$. The base field uses a fixed positional encoding $\gamma(y)$ that maps query coordinates $y$ to input function values at neighboring nodes $x\in N(y)$, i.e., $\gamma(y)=u(x)/|N(y)|$, followed by a linear layer. The encoding width equals the maximum number of neighboring nodes, with zeros filling tailing entries for queries with fewer neighbors. The GNO layer is conditioned globally via the parameterized kernel $\kappa_{\phi}(\cdot)$ modulating the linear layer weights, and locally via a position-depended bias term $Wu(y)$ added as a skip connection. The layer outputs $s(y)$ are finally obtained by summing the linear layer outputs and the skip connection, then applying a nonlinear activation function $\sigma$.

#### Fourier Neural Operator (FNO).

The FNO introduced a computationally efficient approach for performing global message passing in GNO layers assuming a stationary kernel $\kappa_{\phi}(\cdot)$ (Li et al., ; Bruna et al., ). Leveraging the Fourier transform, a single layer of FNO (Li et al., ) is given by

$$ $s(y)=\sigma\Big{(}Wu(y)+\mathcal{F}_{n}^{-1}\left(K\mathcal{F}_{n}(u)\right)(y )\Big{)},$ (3) $$

where $W\in{\mathbb{R}}^{d_{o}\times d_{o}}$ and $K\in{\mathbb{R}}^{n\times n}$ are trainable weight matrices, ${\mathcal{F}}_{n}:L^{2}\to{\mathbb{C}}^{n}$ is the Fourier transform truncated to the first $n$ modes, and ${\mathcal{F}}_{n}^{-1}:{\mathbb{C}}^{n}\to L^{2}$ is the left inverse of this operator. We show that the FNO layer can be viewed as a conditioned neural field employing both local and global conditioning on the input function $u$. Without loss of generality, we demonstrate this for a one-dimensional query domain $y\in{\mathbb{R}}$. To this end, note that the convolutional part of the FNO layer (equation 3) can be re-written as (see Appendix D for a derivation)

$$ $\mathcal{F}_{n}^{-1}\left(K\mathcal{F}_{n}(u)\right)(y)=\sum_{j=1}^{n}\Re(K \hat{u})_{j}\cos(2\pi\langle k_{j},y\rangle)-\sum_{j=1}^{n}\Im(K\hat{u})_{j} \sin(2\pi\langle k_{j},y\rangle),$ (4) $$

where $\hat{u}=\mathcal{F}_{n}u$ denotes the truncated Fourier transform of $u$, and $\Re(z)$ and $\Im(z)$ denote the real and imaginary parts of a complex vector $z\in{\mathbb{C}}^{n}$. We may interpret this transformation of $u(y)$ as a linear transformation (which depends on $\hat{u}$) acting on the positional encoding,

$$ $\gamma(y)=[\cos(2\pi\langle k_{1},y\rangle),\sin(2\pi\langle k_{1},y\rangle), \dots,\cos(2\pi\langle k_{n},y\rangle),\sin(2\pi\langle k_{n},y\rangle)]^{T}.$ (5) $$

The resulting expression (equation 4) is then acted on with a position dependant bias term $Wu(y)$, followed by a pointwise nonlinearity $\sigma$. In this manner, we see that an FNO layer is a neural field with a fixed positional encoding (first $n$ Fourier features), a global conditioning of the weights by the Fourier transform of $u$, and a local conditioning acting as a bias term by $Wu(y)$; see Figure 8(d) for an illustration. This neural field interpretation also suggests an alternative way for implementing FNO layers by explicitly using equation 4 to compute the inverse Fourier transform, thereby allowing them to be evaluated at arbitrary query points instead of a fixed regular grid.

#### Extending to other neural operators.

In practice, GNO and FNO layers are stacked to form deeper architectures with a compositional structure. These can also be interpreted as conditioned neural fields, where the base field corresponds to the last GNO or FNO layer, while all previous layers are absorbed into the definition of the conditioning mechanism. Analogously, one can examine a broader collection of models that fall under the class of Integral Kernel Neural Operators, such as (Gupta et al., ; Tripura & Chakraborty, ; Cao et al., ; Fanaskov & Oseledets, ; Liu-Schiaffini et al., ).

Figure: Figure 8: Viewing common neural operators as conditioned neural fields. Base field and conditioning mechanisms for: (a) DeepONet; (b) NoMaD; (c) the GNO layer; (d) the FNO layer.
Refer to caption: extracted/5899740/figs/no-nf.png

## Appendix D The FNO layer as a conditioned neural field

###### Proposition D.1 .

Suppose that a single layer of the Fourier Neural Operator (FNO) (Li et al., ) is given by

$$ $s(y)=\sigma\Big{(}Wu(y)+\mathcal{F}_{n}^{-1}\left(K\mathcal{F}_{n}(u)\right)(y )\Big{)}.$ (6) $$

Then the integral kernel term can be expressed as

$$ $\mathcal{F}_{n}^{-1}\left(K\mathcal{F}_{n}(u)\right)(y)=\sum_{j=1}^{n}\Re(K \hat{u})_{j}\cos(2\pi\langle k_{j},y\rangle)-\sum_{j=1}^{n}\Im(K\hat{u})_{j} \sin(2\pi\langle k_{j},y\rangle),$ (7) $$

###### Proof.

Recall that the map ${\mathcal{F}}_{n}:L^{2}\to{\mathbb{C}}^{n}$ is the Fourier transform truncated to the first $n$ modes and ${\mathcal{F}}^{-1}:{\mathbb{C}}^{n}\to L^{2}$ is the left inverse of this operator,

$$ ${\mathcal{F}}_{n}^{-1}(z)(y)=\sum_{j=1}^{n}z_{j}e^{2\pi i\langle k_{j},y \rangle}.$ (8) $$

Recall Euler’s formula,

$$ $e^{2\pi i\langle k,y\rangle}=\cos(2\pi\langle k,y\rangle)+i\sin(2\pi\langle k, y\rangle).$ (9) $$

When the (complex) Fourier coefficients $z_{1},\ldots,z_{n}$ come from a real valued signal, the imaginary components in the reconstruction equation 8 cancel and we are left with only real values. In this case, if we decompose each $z_{j}$ into real and imaginary parts, $z_{j}=\alpha_{j}+i\beta_{j}$, then using equation 9 we rewrite equation 8 as

$$ ${\mathcal{F}}_{n}^{-1}(z)(y)=\sum_{j=1}^{n}\alpha_{j}\cos(2\pi\langle k_{j},y \rangle)-\sum_{j=1}^{n}\beta_{j}\sin(2\pi\langle k_{j},y\rangle).$ (10) $$

This allows us to view the convolutional part of the FNO layer equation 3 as

$$ $\mathcal{F}_{n}^{-1}\left(K\mathcal{F}_{n}(u)\right)(y)=\sum_{j=1}^{n}\Re(K \hat{u})_{j}\cos(2\pi\langle k_{j},y\rangle)-\sum_{j=1}^{n}\Im(K\hat{u})_{j} \sin(2\pi\langle k_{j},y\rangle),$ (11) $$

where $\hat{u}=\mathcal{F}_{n}u$ denotes the truncated Fourier transform of $u$, and $\Re(z)$ and $\Im(z)$ denote the real and imaginary parts of a complex vector $z\in{\mathbb{C}}^{n}$.

∎

###### Proposition D.1 .

###### Proof.

## Appendix E Theoretical Insights on Coordinates embeddings

Without loss of generality, we consider an MLP with scalar inputs and outputs. Specifically, let $x\in{\mathbb{R}}$ be the input coordinate, $g^{(0)}(x)=x$ and $d_{0}=d_{L+1}=1$. An MLP $f_{\mathbf{\theta}}(x)$ is recursively defined as follows

$$ $\displaystyle\mathbf{f}^{(l)}_{\mathbf{\theta}}(x)=\mathbf{W}^{(l)}\cdot \mathbf{g}^{(l-1)}(x)+\mathbf{b}^{(l)}\,,\quad\mathbf{g}^{(l)}(x)=\sigma( \mathbf{f}_{\theta}^{(l)}(x))\,,\quad l=1,2,\dots,L\,,$ (12) $$

with a final linear layer defined by

$$ $\displaystyle f_{\mathbf{\theta}}(x)$ $\displaystyle=\mathbf{W}^{(L+1)}\cdot\mathbf{g}^{(L)}(x)+\mathbf{b}^{(L+1)}\,,$ (13) $$

where $\mathbf{W}^{(l)}\in{\mathbb{R}}^{d_{l}\times d_{l-1}}$ is the weight matrix in $l$-th layer and $\sigma$ is an element-wise activation function. Here, $\mathbf{\theta}=\left(\mathbf{W}^{(1)},\mathbf{b}^{(1)},\ldots,\mathbf{W}^{(L+
1)},\mathbf{b}^{(L+1)}\right)$
represents all trainable parameters of the network. In particular, we suppose that the network is equipped with Tanh activation functions and all weight matrices are initialized by the Glorot initialization scheme (Glorot & Bengio, ), i.e., $\mathbf{W}^{(l)}\sim\mathcal{N}(0,\frac{2}{d_{l}+d_{l-1}})$ for $l=1,2,\dots,L+1$. Moreover, all bias parameters are initialized to zeros.

###### Definition E.1 .

A function $f:{\mathbb{R}}^{n}\rightarrow{\mathbb{R}}^{m}$ is called Lipschitz continuous if there exists a constant $L$
such that

$$ $\displaystyle\|f(x)-f(y)\|_{2}\leq L\|x-y\|_{2},\quad\forall x,y\in\mathbb{R}^ {n}.$ (14) $$

The smallest $L$ for which the previous inequality is true is called the Lipschitz constant of $f$ and will be denoted $Lip(f)$.

###### Theorem E.2 .

((Federer, ), Theorem 3.1.6]
If $f:\mathbb{R}^{n}\rightarrow\mathbb{R}^{m}$ is a locally Lipschitz continuous function, then $f$ is differentiable almost everywhere. Moreover, if $f$ is Lipschitz continuous, then

$$ $\displaystyle Lip(f)=\sup_{x\in\mathbb{R}^{n}}\left\|\mathrm{D}_{x}f\right\|_{ 2},$ (15) $$

where $\|M\|_{2}=\sup_{\{x:\|x\|=1\}}\|Mx\|_{2}$ is the operator norm of the matrix $M\in\mathbb{R}^{m\times n}$.

In particular, if $f$ is real valued (i.e. $m=1$), its Lipschitz constant is the maximum norm of its gradient $Lip(f)=\sup_{x}\|\nabla f(x)\|_{2}$ on its domain set.

By Theorem 1 and the chain rule, the Lipschitz constant of MLPs have an explicit formula (Virmaux & Scaman, ; Latorre et al., )

$$ $\displaystyle Lip\left(f_{\theta}\right)=\sup_{x\in\mathbb{R}^{n}}\left\|\prod _{i=L+1}^{2}\left(\mathbf{W}^{(i)}\cdot\operatorname{diag}\left(\dot{\sigma} \big{(}\mathbf{f}_{\theta}^{(i-1)}(x)\big{)}\right)\right)\cdot\mathbf{W}^{(1) }\right\|_{2}.$ (16) $$

In Eq. equation 16, the diagonal matrices $\operatorname{diag}\left(\dot{\sigma}\big{(}\mathbf{f}_{\theta}^{(i-1)}(x)\big
{)}\right)$ are difficult to evaluate, as they may depend on the input value $x$ and previous layers. Fortunately, most major activation functions are 1-Lipschitz. More specifically, these activation functions have a derivative $\dot{\sigma}(\cdot)\in$ $[0,1]$. Hence, we may replace the supremum on the input vector $x$ by a supremum over all possible values:

$$ $\displaystyle Lip\left(f_{\theta}\right)$ $\displaystyle\leq\sup_{\mathbf{z}^{(i-1)}\in[0,1]^{n}}\left\|\prod_{i=L+1}^{2} \left(\mathbf{W}^{(i)}\cdot\operatorname{diag}\left(\mathbf{z}^{(i-1)}\right) \right)\cdot\mathbf{W}^{(1)}\right\|_{2}$ (17) $\displaystyle\leq\prod_{i=1}^{L+1}\left\|\mathbf{W}^{(i)}\right\|_{2}.$ (18) $$

From the above derivation, we immediately obtain that the Lipschitz constant of a MLP is inherently tied to the weight matrices of its hidden layers. MLPs are known to suffer from spectral bias, struggling to learn high-frequency functions efficiently. One reasonable explanation is that $Lip(f)$ is bounded by the product of its parameter norms, which can only increase gradually during training by gradient descent. Consequently, higher frequency components of the target function are learned later in the optimization process (Rahaman et al., ). Random Fourier features (Tancik et al., ) have been proposed as a potential solution to this spectral bias issue. In the following theorem, we demonstrate that both random Fourier features and our proposed position embedding can substantially increase the network’s Lipschitz constant at initialization. This enhancement enables more efficient learning of complex, fine-scale signals from the early stages of training.

###### Theorem E.3 .

Given the definition of MLP in equation 12 and equation 13, the Lipschitz constant of the following coordinates embeddings can be bounded as follows

- (a)
For conventional MLPs, the coordinating embedding is trival and can be viewed as the first linear layer $\phi_{1}(x)=\mathbf{W}x+b$, where $\mathbf{W}\in\mathbb{R}^{d}$, $b\in\mathbb{R}$, and $d$ is the width of the network. Each element of $\mathbf{W}$ is independently sampled from $\mathcal{N}(0,\frac{1}{d})$. The expected Lipschitz constant of $\phi_{1}$ satisfies
$\displaystyle\mathbb{E}\operatorname{Lip}\left(\phi_{1}\right)=\frac{\sqrt{2}}
{\sqrt{d}}\frac{\Gamma((d+1)/2)}{\Gamma(d/2)}\approx 1-\frac{1}{2d}.$
(19)
- (b)
The coordinate embedding of random Fourier features is given by
$\phi_{2}(x)=\begin{bmatrix}\cos(\mathbf{W}x)\\
\sin(\mathbf{W}x)\end{bmatrix}$, where $\mathbf{W}\in{\mathbb{R}}^{d/2}$ and each element of $\mathbf{W}$ is independently sampled from $\mathcal{N}(0,s^{2})$, with $s>0$ being a hyperparameter. The expected Lipschitz constant of $\phi_{2}$ satisfies
$\displaystyle\mathbb{E}\operatorname{Lip}\left(\phi_{2}\right)=2s\cdot\frac{
\Gamma\left(\frac{d/2+1}{2}\right)}{\Gamma(d/4)}\approx s\sqrt{d}\left(1-\frac
{1}{8d}\right).$
(20)
- (c)
The proposed coordinate embedding is defined as:
$\displaystyle\phi_{\beta}(x)=\sum_{i=1}^{n}\omega_{i}(x)\mathbf{f}_{i}=\sum_{i
=1}^{n}e^{-\beta(x-x_{i})^{2}}\mathbf{f}_{i},$
(21)
where $\{x_{i}\}_{i=1}^{n}$ is a uniform grid in $[0,1]$, and $\mathbf{f}_{i}\in{\mathbb{R}}^{h}$ are bounded latent features independent of coordinates, satisfying $\|\mathbf{f}_{i}\|\leq M$ for some $m,M>0$. Let $h$ denote the spacing of the grid, i.e., $h=x_{i+1}-x_{i}$ for any $i\in{0,1,...,n-1}$.
For sufficiently large $\beta>0$ such that $\frac{1}{\sqrt{2\beta}}<h$, the Lipschitz constant of $\phi_{\beta}$ is bounded as:
$\displaystyle M\sqrt{2\beta}e^{-\frac{1}{2}}\lesssim\operatorname{Lip}\left(
\phi_{\beta}\right)\lesssim 2M\sqrt{2\beta}e^{-\frac{1}{2}}.$
(22)

###### Proof.

(a) Let $\mathbf{W}=(W_{i})_{i=1}^{d}$. Then we have

$$ $\displaystyle\operatorname{Lip}(\phi_{1})=\|\mathbf{W}\|_{2}=\sqrt{\sum_{i=1}^ {d}W_{i}^{2}}=\sqrt{\sum_{i=1}^{d}\frac{1}{d}Z_{i}^{2}}=\frac{1}{\sqrt{d}}\chi (d).$ (23) $$

Since $\mathbb{E}\chi(d)=\sqrt{2}\frac{\Gamma((d+1)/2)}{\Gamma(d/2)}$,

$$ $\displaystyle\mathbb{E}\operatorname{Lip}(\phi_{1})=\frac{1}{\sqrt{d}}\mathbb{ E}\chi(d)=\frac{\sqrt{2}}{\sqrt{d}}\frac{\Gamma((d+1)/2)}{\Gamma(d/2)},$ (24) $$

where $\Gamma(\cdot)$ is the Gamma function.

For large $d$, this can be approximated by:

$$ $\displaystyle\mathbb{E}\operatorname{Lip}(\phi_{1})\approx\frac{1}{\sqrt{d}} \sqrt{d}\left(1-\frac{1}{4d}\right)=1-\frac{1}{4d}.$ (25) $$

(b) We observe that

$$ $\displaystyle\operatorname{Lip}\left(\phi_{2}\right)=\sup_{x\in R}\left\|\phi_ {2}^{\prime}\right\|_{2}=\left\|\left[\begin{array}[]{c}\mathbf{W}\cos(\mathbf {W}\mathbf{x})\\ -\mathbf{W}\sin(\mathbf{W}\mathbf{x})\end{array}\right]\right\|_{2}\leq\sqrt{2 \sum_{i=1}^{d/2}W_{i}^{2}}=\sqrt{2}\|\mathbf{W}\|_{2}.$ (28) $$

Similar to the proof of (a), we have

$$ $\displaystyle\mathbb{E}\operatorname{Lip}(\phi_{2})\leq s\sqrt{2}\mathbb{E} \chi(d/2)=2s\cdot\frac{\Gamma\left(\frac{\frac{d}{2}+1}{2}\right)}{\Gamma\left (\frac{d}{4}\right)}.$ (29) $$

For large $d$, this can be approximated by:

$$ $\displaystyle\mathbb{E}\operatorname{Lip}\left(\phi_{2}\right)\approx s\sqrt{2 }\sqrt{\frac{d}{2}}\left(1-\frac{1}{2d}\right)=s\sqrt{d}\left(1-\frac{1}{2d} \right).$ (30) $$

(c) To estimate the Lipschitz constant of $\phi$, we apply Theorem E.2, which reduces the problem to estimating its first-order derivative:

$$ $\displaystyle\phi_{\beta}^{\prime}(x)=\sum_{i=1}^{n}\omega^{\prime}_{i}(x) \mathbf{f}_{i}.$ (31) $$

For any given $x$, there exists a unique $k$ such that $x\in[x_{k},x_{k+1}]$. We can then decompose $\phi^{\prime}(x)$ as follows:

$$ $\displaystyle\phi_{\beta}^{\prime}(x)=\omega^{\prime}_{k}(x)\mathbf{f}_{k}+ \omega^{\prime}_{k+1}(x)\mathbf{f}_{k+1}+\sum_{i\neq k,k+1}\omega^{\prime}_{i} (x)\mathbf{f}_{i}:=I_{1}+I_{2}+I_{3}.$ (32) $$

Let us first estimate $I_{1},I_{2}$. Since all $\mathbf{f}_{i}$ are bounded, we consider $g(x)=e^{-\beta x^{2}}$, as each $\omega_{i}(x)$ is a translation of $g$ by $x_{i}$. By differentiation, we obtain:

$$ $\displaystyle g^{\prime}(x)$ $\displaystyle=-2\beta xe^{-\beta x^{2}}$ (33) $\displaystyle g^{\prime\prime}(x)$ $\displaystyle=(4\beta^{2}x^{2}-2\beta x)e^{-\beta x^{2}}.$ (34) $$

Setting $g^{\prime\prime}(x)=0$ yields $x=\pm\frac{1}{\sqrt{2\beta}}$. Consequently, $g^{\prime}(x)$ achieves its maximum value of $\sqrt{2\beta}e^{-\frac{1}{2}}$ at $x=-\frac{1}{\sqrt{2\beta}}$ and its minimum value of $-\sqrt{2\beta}e^{-\frac{1}{2}}$ at $x=\frac{1}{\sqrt{2\beta}}$. Therefore, we can bound $\omega^{\prime}_{i}(x)$ as follows:

$$ $\displaystyle-\sqrt{2\beta}e^{-\frac{1}{2}}\leq\omega^{\prime}_{i}(x)\leq\sqrt {2\beta}e^{-\frac{1}{2}}.$ (35) $$

To estimate $I_{3}$, recall that we assume $\frac{1}{\sqrt{2\beta}}<h$. For $i\neq k,k+1$ and fixed $x$, we have $mh\leq|x-x_{i}|\leq(m+1)h$ for a unique integer $m$, and $\omega^{\prime}_{i}(x)$ decreases as $x$ moves away from $x_{i}$. Then:

$$ $\displaystyle|\omega_{i}^{\prime}(x)|\leq 2\beta mh\cdot e^{-\beta m^{2}h^{2}}.$ (36) $$

Therefore, we have

$$ $\displaystyle\sum_{i\neq k,k+1}\left|\omega_{i}^{\prime}(x)\right|$ $\displaystyle\leq 2\sum_{m=1}^{n}2\beta mh\cdot e^{-\beta m^{2}h^{2}}$ (37) $\displaystyle\leq 2\sum_{m=1}^{n}\int_{m-1}^{m}2\beta hx\cdot e^{-\beta h^{2}x ^{2}}dx$ (38) $\displaystyle=2\int_{0}^{n-1}2\beta hx\cdot e^{-\beta h^{2}x^{2}}dx$ (39) $\displaystyle=\frac{2}{h}\left[-\left.e^{-\beta hx^{2}}\right|_{0}^{n-1}\right]$ (40) $\displaystyle=\frac{2\left(1-e^{-\beta h^{2}(n-1)^{2}}\right)}{h}$ (41) $\displaystyle\leq\frac{2}{h}.$ (42) $$

Since $\|\mathbf{f}_{i}\|\leq M$,

$$ $\displaystyle\left\|\phi_{\beta}^{\prime}(x)\right\|_{2}$ $\displaystyle\leq M\left(\left|w_{k}^{\prime}(x)\right|+\left|w_{k+1}^{\prime} (x)\right|+\sum_{i\neq k,k+1}|\omega_{i}^{\prime}(x)|\right)$ (43) $\displaystyle\leq M\left(2\sqrt{2\beta}e^{-\frac{1}{2}}+\frac{2}{h}\right).$ (44) $$

Therefore, we conclude that

$$ $\displaystyle\operatorname{Lip}\left(\phi_{\beta}\right)\lesssim 2M\sqrt{2 \beta}e^{-\frac{1}{2}}.$ (45) $$

To obtain a lower bound for $\operatorname{Lip}(\phi_{\beta})$, we proceed as follows:
Without loss of generality, assume $\|\mathbf{f}_{2}\|_{2}=M$. Let $x^{*}=h-\frac{1}{\sqrt{2\beta}}$, where $|\omega^{\prime}_{2}|$ achieves its maximum. Then

$$ $\displaystyle\|\phi^{\prime}_{\beta}(x^{*})\|_{2}\geq\|\omega^{\prime}_{2}(x^{ *})\mathbf{f}_{2}\|_{2}-\|\omega^{\prime}_{1}(x^{*})\mathbf{f}_{1}\|_{2}-\sum_ {i\geq 3}\|\omega^{\prime}_{i}(x^{*})\mathbf{f}_{i}\|_{2}=J_{1}-J_{2}-J_{3}.$ (46) $$

We can bound each term as follows:

$$ $\displaystyle J_{1}=M\sqrt{2\beta}e^{-\frac{1}{2}},$ (47) $$

and

$$ $\displaystyle J_{2}\leq M(2\beta)\left(h-\frac{1}{\sqrt{2\beta}}\right)e^{- \beta\left(h-\frac{1}{\sqrt{2\beta}}\right)^{2}}.$ (48) $$

Note that $J_{2}$ goes to 0 as $\beta$ goes to infinity. So $J_{2}$ is bounded with respect to $\beta$.
For $J_{3}$, we similarly estimate:

$$ $\displaystyle J_{3}\leq M\sum_{m=1}^{n}\beta mh\cdot e^{-\beta m^{2}h^{2}}\leq \frac{\left(1-e^{-\beta h^{2}(n-1)^{2}}\right)}{h}\leq\frac{1}{h}.$ (49) $$

Thus, combining these estimates, we conclude that

$$ $\displaystyle\operatorname{Lip}(\phi_{\beta})\gtrsim M\sqrt{2\beta}e^{-\frac{1 }{2}}.$ (50) $$

∎

Our theorem yields key insights into the Lipschitz constants of various coordinate embedding methods. For trivial coordinate embeddings (linear layer MLP), the Lipschitz constant is approximately 1. This has no impact on the network’s global Lipschitz constant, resulting in the network remaining susceptible to spectral bias.

In the case of random Fourier features, the Lipschitz constant can be improved through two main approaches. First, by increasing the standard deviation of the Gaussian used for sampling the frequency matrix. Second, by increasing the network width. This observation implies a trade-off between network width and the frequencies used in the Fourier features.

For our proposed coordinate embeddings, the Lipschitz constant primarily depends on the parameter $\beta$ in the interpolation process. This dependence on a single parameter potentially offers more straightforward control over the Lipschitz constant compared to Fourier features.

These findings significantly impact neural field design and optimization using different coordinate embedding techniques. They provide a quantitative basis for comparing embedding methods’ capabilities, particularly in addressing the spectral bias inherent in standard MLPs. The straightforward control of the Lipschitz constant in our proposed method (via $\beta$) may offer advantages in model tuning and optimization.

###### Definition E.1 .

###### Theorem E.2 .

###### Theorem E.3 .

###### Proof.

## Appendix F Experimental Details

### F.1 Training and evaluation

#### Training recipe.

We use a unified training recipe for all CViT experiments. We employ AdamW optimizer (Kingma & Ba, ; Loshchilov & Hutter, ) with a weight decay $10^{-5}$. Our learning rate schedule includes an initial linear warm-up phase of $5,000$ steps, starting from zero and gradually increasing to $10^{-3}$, followed by an exponential decay at a rate of $0.9$ for every $5,000$ steps.
The loss function is a one-step mean squared error (MSE) between the model predictions and the corresponding targets at the next time-step, evaluated at randomly sampled query coordinates:

$$ $\displaystyle\text{MSE}=\frac{1}{B}\frac{1}{Q}\frac{1}{D}\sum_{i=1}^{B}\sum_{j =1}^{Q}\sum_{k=1}^{D}\left|\hat{s}^{(k)}_{i}(\mathbf{y_{j}})-s^{(k)}_{i}( \mathbf{y_{j}})\right|_{2}^{2},$ (51) $$

where $s^{(k)}_{i}(\mathbf{y_{j}})$ denotes the $k$-th variable of the $i$-th sample in the training dataset, evaluated at a query coordinate $\mathbf{y}_{j}$, and $\hat{s}$ denotes the corresponding model prediction. It is worth noting that our objective function is a one-step loss, which differs significantly from the rollout loss used in training DPOTs
All models are trained for $2\times 10^{5}$ iterations with a batch size $B=64$.
Within each batch, we randomly sample $Q=1,024$ query coordinates from the grid and corresponding output labels.

#### Training instabilities

We observe certain training instabilities in CViT models, where the loss function occasionally blows up, causing the model to collapse. To stabilize training, we clip all gradients at a maximum norm of 1. If a loss blowup is detected, we restart training from the last saved checkpoint.

#### Evaluation.

After training, we obtain the predicted trajectory by performing an auto-regressive rollout on the test dataset. We evaluate model accuracy using the relative $L^{2}$ norm, as commonly used in Li et al.(Li et al., ):

$$ $\displaystyle\text{Rel. }L^{2}=\frac{1}{N_{\text{test}}}\frac{1}{D}\sum_{i=1}^ {N_{\text{test}}}\sum_{k=1}^{D}\frac{\|\hat{s}^{(k)}_{i}-s^{(k)}_{i}\|_{2}}{\| s^{(j)}_{i}\|_{2}},$ (52) $$

where the norm is computed over the rollout prediction at all grid points, averaged over each variable of interest.

### F.2 Model Details

#### CViT.

For the 1D advection equation, we create a tiny version of CViT to ensure the model size is comparable to other baselines. Specifically, we use a patch size of 4 to generate tokens. The encoder consists of 6 transformer layers with an embedding dimension of 256 and 16 attention heads. A single transformer decoder layer with 16 attention heads is used.

For the shallow water equation and Navier-Stokes equation, the CViT configurations are detailed in Table 1. While we use a patch size of $8\times 8$ by default, we use a smaller patch size of $4\times 4$ when training CViT-L for the Navier-Stokes problem. We find that using a smaller patch size better captures the small-scale features in flow motion, leading to improved accuracy.

#### DeepONet.

For the DeepONet, the branch network (encoder) is an MLP with 4 hidden layers of 512 neurons each. The trunk-net (decoder) also employs an MLP with 4 hidden layers of 512 neurons.

#### NoMaD.

The encoder consists of 4 hidden layers with 512 neurons each, mapping the input to a latent representation of dimension 512. The decoder, also an MLP with 4 hidden layers of 512 neurons, takes the latent representation and maps it to the output function.

#### Dilated ResNet.

Following the comprehensive experiments and ablation studies in PDEArena (Gupta & Brandstetter, ), we adopt the optimal Dilated ResNet configuration reported in their work. This model has 128 channels and employs group normalization. It consists of four residual blocks, each comprising seven dilated CNN layers with dilation rates of [1, 2, 4, 8, 4, 2, 1].

#### U-Nets.

Based on comprehensive experiments and ablation studies in PDEArena (Gupta & Brandstetter, ), we select the two strongest U-Net variants reported in their work: $\textbf{U-Net}_{\textbf{att}}$ and U-F2Net. For both architectures, we use one embedding layer and one output layer with kernel sizes of $3\times 3$. Besides, we use 64 channels and a channel multipliers of $(1,2,2,4)$ and incorporate residual connections in each downsampling and upsampling block. Pre-normalization and pre-activations are applied, along with GELU activations. For $\text{U-Net}_{\text{att}}$, attention mechanisms are used in the middle blocks after downsampling.
For U-F2Net, we replace the lower blocks in both the downsampling and upsampling paths of the U-Net architecture with Fourier blocks. Each Fourier block consists of 2 FNO layers with modes 16 and residual connections.

#### UNO.

Following the comprehensive experiments and ablation studies in PDEArena (Gupta & Brandstetter, ), we adopt the optimal UNO of 128 channels.

#### FNO.

For the 1D advection equation, our FNO model employs 3 Fourier neural blocks, each with 256 channels and 12 Fourier modes. The output is projected back to a 256-dimensional space before passing through a linear layer that produces the solution over the defined grid.

For the shallow water equation, we employ the best FNO model from comprehensive experiments and ablation studies in PDEArena (Gupta & Brandstetter, ). Specifically, we use the FNO consisting of 4 FNO layers, each with 128 channels and 32 modes. Additionally, we use two embedding layers and two output layers with kernel sizes of $1\times 1$, as suggested in Li et al.(Li et al., ). We use GeLU activation functions and no normalization scheme.

For the Navier-Stokes equation, we follow the problem setup in DPOT (Hao et al., ), which differs from the setup in PDEArena (Gupta & Brandstetter, ). Therefore, we directly report the results from DPOT (Hao et al., ).

#### FFNO / GK-T / OFormer / GNOT / DPOT / MPP.

We directly report the results from DPOT (Hao et al., ). Detailed implementations can be found in their paper.

### F.3 Linear Advection Equation

We consider the one-dimensional linear advection equation

$$ $\displaystyle\frac{\partial u}{\partial t}+c\frac{\partial u}{\partial x}=0, \qquad x\in[0,1],$ $$

with periodic boundary conditions $u(0,t)=u(1,t)$ given an initial condition $u(x,0)=u_{0}(x)$ with the advection velocity $c=1$. To generate discontinuous initial conditions, the initial condition $u_{0}$ is assumed to be

$$ $\displaystyle u_{0}=-1+2\ \mathbb{1}_{\left\{\tilde{u_{0}}\geq 0\right\}},$ $$

where $\tilde{u_{0}}$ a centered Gaussian,

$$ $\displaystyle\widetilde{u_{0}}\sim\mathcal{N}(0,C)\text{ and, }\mathrm{C}= \left(-\Delta+\tau^{2}\right)^{-d}.$ $$

Here $-\Delta$ denotes the Laplacian on $D$ subject to periodic conditions on the space of spatial mean zero functions. Morevoer, $\tau$ denotes the inverse length scale of the random field and $d$ determines the regularity of $\widetilde{u_{0}}$. For this example, we set $\tau=3$ and $d=2$.

#### Dataset.

We make use of the datasets released by Hoop et al.(de Hoop et al., ). This dataset consists of 40,000 samples of discretized initial conditions on a grid of $N=200$ points.

#### Problem setup.

We follow the problem setup by Hoop et al.(de Hoop et al., ). Our objective is to predict the solution profile at time $t=0.5$. We compare the CViT model’s performance against several strong baseline neural operators: DeepONet (Lu et al., ), NoMaD (Seidman et al., ) and Fourier Neural Operators (Li et al., ). For the training and evaluation, we considered a split of 20,000 samples used for training, 10,000 for validation, and 10,000 for testing.

#### Training.

For all models, we use a batch size of 256 with a grid sampling size of 128, except for the FNO, which uses the full grid size of 200. The models are trained using the AdamW optimizer (Loshchilov & Hutter, ) with an exponential decay learning rate scheduler. The models are trained for 200,000 iterations, and the best model state is selected based on the validation set performance to avoid overfitting. The final performance is evaluated on the test set of 10,000 samples.

#### Evaluation.

Due to the discontinuous nature of the problem, the relative L2 norm may not properly reflect the performance of the models. Instead, a more suitable metric for this task is the total variation (Stein & Shakarchi, ), defined as

$$ $\displaystyle TV(f,g)=\int_{\Omega}||\nabla f|-|\nabla g||\mathrm{dx}.$ $$

The total variation quantifies the integrated difference between the magnitudes of the gradients of the predicted solution and the ground truth, providing a better indicator of a model’s ability to capture and preserve discontinuities in the solution.

Table 5 presents a summary of our results. We notice that the CViT model achieves the lowest mean, median and worst-case relative $L^{2}$ errors, outperforming the other architectures on this metric. Furthermore, when considering the total variation metric, the CViT model exhibits competitive performance for the mean and median total variation values, and achieves the lowest prediction error the worst-case sample in the test dataset.

**Table 5: Performance of neural operator architectures under the relative L2 and Total Variation metrics for the linear advection with discontinuous initial conditions benchmark.**
|  |  | Rel. L2 error (%, $\downarrow$) | Total Variation error ($\downarrow$) |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Model | # Params | Mean | Median | Worst-case | Mean | Median | Worst-case |
| FNO | 4.98M | 13.94 | 12.03 | 56.52 | 0.0406 | 0.0309 | 0.1910 |
| DeepONet | 2.20M | 13.77 | 12.20 | 39.64 | 0.0572 | 0.0453 | 0.2300 |
| NoMaD | 2.21M | 13.64 | 11.01 | 69.00 | 0.0326 | 0.0245 | 0.1732 |
| CViT | 3.03M | 11.70 | 10.68 | 35.44 | 0.0334 | 0.0274 | 0.1516 |

### F.4 Shallow-Water equations

Let $\mathbf{u}$ denote the velocity and $\eta$ interface height which conceptually represents pressure. We define $\zeta=\nabla\times\mathbf{u}$ and divergence $\mathcal{D}=\nabla\cdot\mathbf{u}$. The vorticity formulation is given by:

$$ $\displaystyle\frac{\partial\zeta}{\partial t}+\nabla\cdot(\mathbf{u}(\zeta+f))$ $\displaystyle=0,$ $\displaystyle\frac{\partial\mathcal{D}}{\partial t}-\nabla\times(\mathbf{u}( \zeta+f))$ $\displaystyle=-\nabla^{2}\left(\frac{1}{2}\left(u^{2}+v^{2}\right)+g\eta\right),$ $\displaystyle\frac{\partial\eta}{\partial t}+\nabla\cdot(\mathbf{u}h)$ $\displaystyle=0,$ $$

subject to the periodic boundary conditions. Here
$f$ is Coriolis parameter and $g=9.81$ is the gravitational acceleration, and $h$ is the dynamic layer thickness.

#### Data generation.

The dataset is generated by PDEArena (Gupta & Brandstetter, ) using SpeedyWeather.jl on a regular grid with spatial resolution of $192\times 96\left(h=1.875^{\circ},\Delta y=3.75^{\circ}\right)$, and temporal resolution of $\Delta t=48\mathrm{~{}h}$. The resulting dataset contains 6,600 trajectories with a
temporal spatial resolution $11\times 96\times 192$.

#### Problem setup.

We follow the problem setup reported in PDEArena (Gupta & Brandstetter, ). Our goal is to learn the operator that maps the vorticity and pressure fields from the previous 2
two time-steps to the next time-step. All models are trained with 5,600 trajectories and evaluated on the remain 1,000 trajectories. We report the relative $L^{2}$ errors over 5-steps rollout predictions.

#### Training.

For training CViT models, please refer to section F.1. For training other baselines in Table LABEL:tab:_swe, we follow the recommended training procedure and hyper-parameters presented in PDEArena (Gupta & Brandstetter, ) Specifically, we use the AdamW optimizer with a learning rate of $2\times 10^{-4}$ and a weight decay of $10^{-5}$. Training is conducted for 50 epochs, minimizing the summed mean squared error. We use cosine annealing as the learning rate scheduler (Loshchilov & Hutter, ) with a linear warm-up phase of 5 epochs. An effective batch size of 32 is used for training.

#### Impact of Grid Resolution.

In our experiments, we matched the resolution of the embedding grid to the dataset for convenience. However, our model’s performance is not significantly affected as long as the embedding grid roughly matches the dataset resolution.

To validate this, we conducted experiments with three different grid resolutions: $95\times 191,97\times$ 193 , and $128\times 200$. We performed these experiments using the shallow water equation example under the same hyper-parameter settings. The resulting relative $L^{2}$ errors are summarized below:

**Table 6: Swallow Water equation: Relative $L^{2}$ error for different grid resolutions.**
| Grid Resolution | $96\times 192$ | $95\times 191$ | $97\times 193$ | $128\times 200$ |
| --- | --- | --- | --- | --- |
| Relative $L^{2}$ Error | $1.563\times 10^{-2}$ | $1.572\times 10^{-2}$ | $1.568\times 10^{-2}$ | $1.562\times 10^{-2}$ |

The results demonstrate negligible variation in error across different resolutions. This empirical evidence supports our claim that CViT’s performance is robust to small variations in grid resolution, suggesting that precise matching between the embedding grid and dataset resolution is not critical for model efficacy.

Figure: Figure 9: Shallow water benchmark. Ablation studies on the latent dimension of grid features and number of attention heads in decoder.
Refer to caption: x3.png

Figure: Figure 10: Shallow water benchmark. Representative CViT rollout prediction of the pressure field, and point-wise error against the ground truth.
Refer to caption: extracted/5899740/cvit_figures/swe/swe_pre_pred.png

### F.5 Navier-Stokes equation

We consider the two-dimensional incompressible Navier-Stokes equations in the velocity-pressure formulation, coupled with a scalar field representing a transported particle concentration via the velocity field. The governing equations are given by:

$$ $\displaystyle\frac{\partial c}{\partial t}+\mathbf{u}\cdot\nabla c$ $\displaystyle=0,$ $\displaystyle\frac{\partial\mathbf{u}}{\partial t}+\mathbf{u}\cdot\nabla \mathbf{u}+\nabla p-\nu\nabla^{2}\mathbf{u}$ $\displaystyle=\mathbf{f},$ $\displaystyle\nabla\cdot\mathbf{u}$ $\displaystyle=0,$ $$

where $c$ is the scalar field, $\mathbf{u}$ is the velocity field, and and
$\nu$ is the kinematic viscosity. In addition, the velocity field is affected through an external buoyancy force term $\mathbf{f}$ in the $y$ direction, represented as $\mathbf{f}=(0,f)^{T}$.

#### Data generation.

The 2D Navier-Stokes data is generated by PDEArena (Gupta & Brandstetter, ), which is obtained on a grid with a spatial resolution of $128\times 128(h=0.25,\Delta y=0.25)$ and a temporal resolution of $\Delta t=1.5\mathrm{~{}s}$. The simulation uses a viscosity parameter of $\nu=0.01$ and a buoyancy factor of $(0,0.5)^{T}$. The equation is solved on a closed domain with Dirichlet boundary conditions $(v=0)$ for the velocity, and Neumann boundaries $\left(\frac{\partial s}{\partial x}=0\right)$ for the scalar field. The simulation runs for $21\mathrm{~{}s}$, sampling every $1.5\mathrm{~{}s}$ . Each trajectory contains scalar and vector fields at 14 different time frames, resulting in a dataset with 7,900 trajectories.

#### Problem setup.

We follow the problem setup reported in Hao et al. (Hao et al., ). We aim to learn the solution operator that maps the passive scalar and velocity fields from the previous 10 time-steps to the next time-step. The models are trained with 6,500 trajectories and tested on the remaining 1,300 trajectories. We report the resulting relative $L^{2}$ errors over 4-steps rollout predictions.

#### Evaluation.

We directly report the results as presented in DPOT (Hao et al., ). For comprehensive details on the training procedures, hyper-parameters, please refer to their original paper.

Figure: Figure 11: Navier-Stokes benchmark. Representative CViT rollout prediction of the velocity field in $x$ direction, and point-wise error against the ground truth.
Refer to caption: extracted/5899740/cvit_figures/ns/ns_vx_pred.png

Figure: Figure 12: Navier-Stokes benchmark. Representative CViT rollout prediction of the velocity field in $y$ direction, and point-wise error against the ground truth.
Refer to caption: extracted/5899740/cvit_figures/ns/ns_vy_pred.png

### F.6 Diffusion-reaction equation

The 2D diffusion-reaction equation is given by

$$ $\displaystyle\partial_{t}u=D_{u}\partial_{xx}u+D_{u}\partial_{yy}u+R_{u},$ $\displaystyle\partial_{t}v=D_{v}\partial_{xx}v+D_{v}\partial_{yy}v+R_{v},$ $$

where $u=u(t,x,y)$ is the the activator and $v=v(t,x,y)$ denotes the inhibitor. Here $D_{u}$ and $D_{v}$ are the diffusion coefficient for the activator and inhibitor, respectively, $R_{u}=$ $R_{u}(u,v)$ and $R_{v}=R_{v}(u,v)$ are the activator and inhibitor reaction function, respectively. The domain of the simulation includes $x\in(-1,1),y\in(-1,1),t\in(0,5]$.
The reaction functions for the activator and inhibitor are defined by:

$$ $\displaystyle R_{u}(u,v)=u-u^{3}-k-v,$ $\displaystyle R_{v}(u,v)=u-v,$ $$

where $k=5\times 10^{-3}$, and the diffusion coefficients for the activator and inhibitor are $D_{u}=1\times 10^{-3}$ and $D_{v}=5\times 10^{-3}$, respectively. The initial condition is generated as standard normal random noise $u(0,x,y)\sim\mathcal{N}(0,1.0)$ for $x\in(-1,1)$ and $y\in(-1,1)$.

#### Data Generation.

We use the dataset generated by PDEBench (Takamoto et al., ), where discretized into $N_{x}=512,N_{y}=512$ and $N_{t}=501$, as well as the downsampled version for the models training with $N_{x}=128,N_{y}=128$, and $N_{t}=101$, the spatial discretization is performed using the finite volume method, and the time integration is performed using the built-in fourth order Runge-Kutta method in the scipy package.

#### Problem Setup.

We follow the problem setup reported in Hao et al. (Hao et al., ). We aim to learn the solution operator that maps the activator and inhibitor fields from the previous time-steps to the next time-step. The models are trained with 900 trajectories and tested on the remaining 100 trajectories. We report the resulting relative $L^{2}$ errors over rollout predictions.

#### Training.

For training CViT models, please refer to section F.1. It is worth noting that for this specific task, we employ an initial learning rate of $5\times 10^{-4}$, different from the $10^{-3}$ used in other experiments.

#### Evaluation.

We directly report the results as presented in DPOT (Hao et al., ). For comprehensive details on the training procedures, hyper-parameters, please refer to their original paper.

Figure: Figure 13: Diffusion-reaction Benchmark. Representative CViT rollout prediction of the density field $v$, and point-wise error against the ground truth.
Refer to caption: extracted/5899740/cvit_figures/dr/dr_v_pred.png

### F.7 Computational Cost

All experiments were performed on a single Nvidia RTX A6000 GPU. Average training times varied between 5 hours and 60 hours, depending on the task, input resolution, model size, and patch size. The following table summarizes the training times for different models. We use a brief notation to indicate the model size and the input patch size: for instance, CViT-L/16 refers to the "Large" variant with a
$16\times 16$ input patch size.

**Table 7: Training times (hours) on a single Nvidia RTX A6000 GPU for different models and patch sizes ($8\times 8$, $16\times 16$, $32\times 32$) on benchmarks of Shallow water equation and Navier-Stokes equation.**
| Method | Shallow Water | Navier-Stokes |
| --- | --- | --- |
| CViT-S/8 | 15 | 9 |
| CViT-B/8 | 25 | 14 |
| CViT-L/8 | 57 | 28 |
| CViT-L/16 | 17 | 10 |
| CViT-L/32 | 9 | 6 |

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: - I wondered how will the proposed method be applied for more general vision tasks, especially for 3d vision. I wondered what are the authors’ opinions on this? What are the difficulties, such as computations, that might be difficult to apply? 

Please see the above section for detailed comments, too.
soundness: 3
strengths: - The proposed coordinate embedding for the query is interesting, and is effective in the continuous ViT learning, compared to MLP and RFF. Also, it is simple and easy to control through the interpolation parameter $\beta$.
- Discussions and analysis in the paper are extensive and interesting. Overall, the paper is well-organized and easy to follow.
- The proposed coordinate embedding could have the potential to apply to more general continuous learning beyond physical domains. Especially for the vision tasks, such as 3d geometries including point cloud representations, it is worth exploring.
summary: This paper proposes a continuous operator learning method using the vision transformer with a newly introduced coordinate embedding for the query value. The proposed method is based on the previous continuous operator learning method. The authors have provided analysis including comparisons with recent operator learning methods on learning different operator equations, ablation studies for different structure choices, discussions of the proposed method compared to recent arts, theoretical insights of the proposed query embeddings, etc.
weaknesses: - Why not do coordinate embedding for all query, key, value? Only doing coordinate embedding for query preferable in which way? Maybe some further analysis on this could be included to further highlight the effectiveness of the proposed coordinate embedding.
- The first question brought to the second one, that in the paper (appendix), the authors discussed the Lipschitz constant for different embeddings from linear embedding to random Fourier features (RFF), and to the proposed coordinate embedding, which is interesting and worth discussing. The authors mentioned that the proposed coordinate embedding is easier to tune since it is mainly based on one parameter $\beta$. However, for RFF, there is also only one parameter to control the spectral bias, and thus control the Lipschitz constant. I wondered if the authors could also provide an analysis on only using RFF for the proposed framework. 
- Figures 3, 5, and 6 lack notations to distinguish columns.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: * "2+1 dimensional spatio-temporal data tensor" but it is then written that $\boldsymbol u \in \mathbb R^{T\times H \times W \times D} *i.e.* four coordinates. 
* How separating the positional encoding (PE) between 1D temporal and 2D spatial is different than having one spatio-temporal (thus 3D) PE?
soundness: 3
strengths: This paper is well-written, easy to read, and technically sound. The proposed approach appears to scale more efficiently than the current baseline thanks to the perceiver-inspired architecture. The results seem also promising.
summary: This paper introduces a new neural operator architecture that tackles recent computer vision techniques (vision transformers, cross-attention). In particular, the CViT model -- for Continuous Vision Transformer for operator learning -- leverages the perceiver architecture to retrieve the representation of a query at a given continuous coordinate from spatiotemporal data. The relevance of the representations learned by CViT is evaluated on a diverse range of PDE such as fluid dynamics, climate modeling, and reaction-diffusion processes.
weaknesses: * No clear structure of the related work section. It would help the reader to add more justifications on why transformer approaches struggle with various resolutions. Even if I tend to agree with the conclusion, it is also not clear, why a novel architecture design is required to solve the limitations related to high dimensional data and long-range dependencies.  
* The transformer architecture in Fig. 1. does not give any insight and is well known. IMO, Authors could either remove it or replace it with the perceiver one. 
* It is unclear how many visual latent queries are used. From the figure and text, it seems that there is only one query which is then duplicated. In the perceiver architectures, these latent queries are learned, does this mean that only one query is learned? The stability of the model with respect to the number of queries could be presented in the ablations.  
* The ablation with respect to parameter $\beta$ is difficult to understand as, in Fig. 4, $\epsilon$ is not introduced. Is $\epsilon = \beta$? The value of $\beta$ seems extremely high ($10^5$) and should be compared to the input's resolution. Here, does eq. (1) reduces to a naive averaging with $\beta=10^5$? 
* Results on standard benchmarks would help judge the method's efficiency. Here, it seems that the authors show results only on selected equations from PDEBench or PDEArena.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: My main point of inclarity is the initialization of latent queries in the encoder. The paper states they are initialized with a unit Gaussian - is this a learned initialization, or are the queries sampled every time the model is learned? The latter would mean, for example, that we would expect to see slightly different outputs every time the model is run - is that the case? (I am not saying this is a bad thing btw, just want to understand). Also - how many query vectors are used? the paper mentions it's more than one, but the number is not supplied in the paper.

Related to the above - I am not sure I understood why the query vector z is tiled across space? if it's the same vector across all of space then the output of the query would be the same for every spatial location, am I wrong? or is there more positional encodings added there? It would be good to clarify this.
soundness: 3
strengths: I found this paper very interesting.

Originality:
I think this is a very good variation of Vision Transformers which significantly expands their capabilities. I found the interpolated readout mechanism of particular novelty.

Quality:
This is a well executed paper. There is ample experimental validation using several different datasets. The baselines chosen (as far as I can tell) are diverse and strong. There is a good set of ablation experiments and deep analysis of the method in the appendix.

Clarity:
The paper is well structures, well written and easy to read. All the nomenclature is well defined and I had no problem following the paper.

Significance:
While I'm not an expert in the field to me this seems like a significant contribution to the field which greatly expands the expressive power of vision transformers.
summary: This paper presents a variant of Vision transformers which allows handling of continuous outputs. The model works by first encoding an input image sequence into a set of latent vectors using cross attention (a la Perceiver style encoding). The latents compress the input over time. The resulting latents are fed through several transformer layers. The readout is the interesting part - a learnable grid of spatial queries is interpolated using Nadarya-Watson interpolation to produce queries which are used to cross attend into the output of the vision transformer. The resulting attention weighted value is used as the outputs - this is done for every output coordinate.
The method is shown to work on several interesting physics oriented datasets with several ablation experiments. More analysis is shown in the appendix.
weaknesses: All in all I think this is a strong paper, however:

* Other domains? I would have loved to see if the resulting method is applicable to other domain or training setups beyond L2 prediction of physical systems. Does it work, for example, on natural video? Would it be useful as a diffusion model back-bone?

* More analysis of cross attention - one of the more interesting parts is the use of cross attention, both in the encoding and decoding. It would be nice to see visualizations of cross attention patterns emerging in both directions. Do specific latents take on specific roles when decoding? are they responsible over different parts of space? time? both?

While I do not expect these to be addressed in the rebuttal I really do think these would maker the paper stronger.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - Does the grid resolution similarly impact computational efficiency as the patch size?
soundness: 3
strengths: - The author conducts extensive experiments, comparing several recent model architectures to demonstrate performance outcomes. The results are presented with additional clarity by including the number of parameters.

- The detailed discussion of limitations and future directions is appreciated. Notably, the computational inefficiency for high-resolution data should be addressed in future work.
summary: This paper presents a neural operator architecture inspired by computer vision to address spatio-temporal scientific problems. The proposed model is structured with a transformer encoder featuring spatio-temporal embedding and time-aggregation layers, a grid-based positional embedding, and a cross-attention decoder. The author employs multiple benchmark datasets and model architectures for comparative experiments.
weaknesses: - Since this method draws on the vision transformer, have comparisons been considered with other vision-based extensions for spatio-temporal applications, such as ViViT[1] and VMAE[2]?

- Regarding the datasets selected for experimentation, incorporating more complex, real-world cases could further substantiate performance, as some baseline architectures still yield results similar to those of the proposed model, despite the model size.

[1] Arnab, Anurag, et al. "Vivit: A video vision transformer." Proceedings of the IEEE/CVF international conference on computer vision. 2021.

[2] Tong, Zhan, et al. "Videomae: Masked autoencoders are data-efficient learners for self-supervised video pre-training." Advances in neural information processing systems 35 (2022): 10078-10093.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: Please see the weakness section.
soundness: 3
strengths: 1. It achieved state-of-the-art performance on various benchmark datasets with the network's great parameter efficiency.
2. Coordinate embedding is an interesting approach to injecting the coordinate information into the network.
3. The theoretical result related to coordinate embedding is quite informative and useful for the community.
summary: This paper explores the architectural design of transformers for operator learning methods. To efficiently implement continuous vision transformer architecture, they developed (or adopted) two techniques: perceiver resampler and coordinate embedding. The perceiver modules effectively aggregate temporal dimension, and the coordinate embedding enables the model to handle continuous input domains and improve accuracy. The proposed architecture achieved state-of-the-art results on various challenging PDE benchmarks while maintaining higher parameter efficiency.
weaknesses: 1. The proposed grid-based positional encoding may not be a desirable solution for high-dimensional PDEs. As the authors said, it requires high-resolution grids to achieve competitive performance, and this approach would not scale up to more than 3-d time-dependent PDEs, e.g., (3+1) NS equation. 

2. The proposed Nadaraya-Watson interpolant would not scale with high-resolution grids since it has to look at all grid features. I am curious the inference cost of the proposed approach according to grid resolutions. Is the inference speed as fast as the previous methods?
The grid resolution 128 (N_x) x 128 (N_y) x 512 (feature dim) is already > 9M, and CViT-S model size is 13M. Does it mean that most of the parameters are used for grids?

3. Does this scale well for high-resolution solutions? The authors mentioned that the grid resolution should be similar to the solution resolution. Then, 512 resolution means 512^3, which is roughly 134M. Could you show me how the proposed method can make good predictions for high-resolution solutions? e.g., testing 128 grid resolution model for predicting 512 solution resolution.

4. This is a rather suggestion than a question. Why not coordinate embedding for temporal dimension? currently, y \in R^2, but why not y \in R^3? Temporal resolution can be coarse in this case to maintain parameter efficiency.

5. How about other datasets? e.g., other PDEs in PDEBench and PDEArena?

[minor comments]
1. Numbering for all equations would be desirable during the review process.
2. How do you determine \beta = 10^5?
3. L171: is u_f a flattened version of what? u_pe? not mentioned in the paper.
4. L174: ‘tilling the latent query’, I am not sure “tiling” is a widely accepted term, I think it is a pytorch’s terminology.
5. Eq (1): why not just using linear interpolation?
6. Figure 3, 5, 6: what do different columns mean? Time evolution?
7. Figure 4: (d) \epsilon -> \beta?

Meta-review: The submission proposes a method for learning physics and phenomena expressed as PDEs and resorts to ViTs with a new coordinate embedding and an interpolated readout mechanism. Five reviewers appreciated the contributions, in particular technical soundness, good results, scaling behavior, discussions and analysis, high potential, sound evaluation with good baselines and ablations, clear writing. They were all unanimous on the significance of the contribution, and the AC concurs.

Decision: Accept (Poster)