R EAL - TIME I NFERENCE AND E XTRAPOLATION VIA A
D IFFUSION - INSPIRED T EMPORAL T RANSFORMER O PERATOR(D I TTO)

Oded Ovadia
Department of Applied Mathematics
Tel Aviv University
Tel Aviv, 69978, Israel
odedovadia@mail.tau.ac.il
Ahmad Peyvan
Division of Applied Mathematics
Brown University
Providence, RI 02912, USA
ahmad_peyvan@brown.edu

Vivek Oommen
School of Engineering
Brown University
Providence, RI 02912, USA
vivek_oomen@brown.edu

Eli Turkel
Department of Applied Mathematics
Tel Aviv University
Tel Aviv 69978, Israel
eliturkel@gmail.com

Adar Kahana
Division of Applied Mathematics
Brown University
Providence, RI 02912, USA
adar_kahana@brown.edu
George Em Karniadakis
Division of Applied Mathematics
Brown University
Providence, RI 02912, USAgeorge_karniadakis@brown.edu

A BSTRACTExtrapolation remains a grand challenge in deep neural networks across all application domains. We
propose an operator learning method to solve time-dependent partial differential equations (PDEs)
continuously and with extrapolation in time without any temporal discretization. The proposed
method, named Diffusion-inspired Temporal Transformer Operator (DiTTO), is inspired by latentdiffusion models and their conditioning mechanism, which we use to incorporate the temporal
evolution of the PDE, in combination with elements from the transformer architecture to improveits capabilities. Upon training, DiTTO can make inferences in real time. We demonstrate its
extrapolation capability on a climate problem by estimating the temperature around the globe forseveral years, and also in modeling hypersonic flows around a double-cone. We propose different
training strategies involving temporal-bundling and sub-sampling and demonstrate performance
improvements for several benchmarks, performing extrapolation for long time intervals as well aszero-shot super-resolution in time.
Keywords Scientific machine learning · Diffusion models · Transformers · Partial differential equations.

1

Introduction

The field of scientific machine learning (SciML) has been
growing rapidly in recent years, and many successful methods for modeling scientific problems using machine learning (ML) methods have been proposed [36, 28, 24, 27, 52].
Many tools designed for standard ML and data scienceproblems can also perform well on SciML tasks. Recent
innovations in the field of ML primarily originate from thedomains of natural language processing [14] and computer
vision [20]. Our work exploits and further develops themain idea of a recently proposed method called diffusion
models [15] (used in generative AI) for solving forward
partial differential equations (PDEs).

Solving time-dependent PDEs is an essential topic forthe scientific community. If the underlying mathematical operators that govern the temporal evolution of thesystem are non-linear and/or there are observational data
available, the task of assimilating and simulating suchprocesses using discretization-based numerical methods
can become increasingly challenging and computationally expensive. The burden associated with the traditional
numerical solvers is further increased when separate simulation runs become mandatory for every new initial condition. Moreover, in certain application domains, such as
autonomy and navigation or robotics, real-time inference isInternational Conference on Scientific Computing and Machine Learning 2024 (SCML2024)

Diffusion-inspired Temporal Transformer Operator (DiTTO)

required. SciML methods such as neural operators specif- ditioning in a U-Net-based neural operator. The DiTTO
ically address these issues by significantly reducing the method proposed here is a diffusion-inspired model [45].
associated computational costs [28, 24, 48, 3].
The common use of diffusion models involves a generative process used to create data samples. It incorporates
Several methods for solving PDE-related problems usinga Markov chain of diffusion steps, where in each stage
ML methods, and specifically transformers, have been proa different texture is added to the data sample. Usually,
posed [36, 28, 32]. We solve the forward problem of athe texture is noise, so new noise distributions are incretime-dependent PDE by training the neural operator to acmentally added in each step. The models have also been
curately estimate the state field variable at a later time fromused with other kinds of textures, for example creating cara given initial condition. The main focus of this work istoonish images from plain ones. Herein, we use a similar
forecasting a continuous in-time solution in real-time for aframework, but instead of conditioning on the noise distriplurality of initial conditions and extrapolating beyond thebution, we do so for the temporal evolution. We explore
training domain. Solving PDE-related problems involves
several implementations and training strategies, in addiseveral challenges including generalizations for differenttion to the diffusion models themselves. These enhanced
problem conditions and dependence on the physical domethods form the class of explored DiTTO models. Immain’s discretization. To tackle the first, we utilize tools
portantly, we demonstrate how this framework can be usedfrom the growing field of operator learning [24, 28], wherefor extrapolation, i.e., it can make accurate predictions for
we use learning techniques to map a function space to ansamples outside the time interval it was trained to handle.
other one. Thus, we learn a family of solutions of PDEs
corresponding to a family of initial conditions. For thesecond challenge, we propose a method that, while being
2 Methodology
dependent on the spatial discretization, is continuous in thetemporal evolution of the solution, which is a prominent
We approximate the time evolution of a PDE solution (forchallenge in solving dynamical systems.
ward process). Instead of incrementally adding noise to
Recent works [35, 47] have demonstrated the efficiency of the inputs, as done with diffusion models, we incremenU-Net-based architectures for modeling time-dependent tally evolve the PDE solution over time. We replace the
PDEs. However, the outputs of these U-Net-based archi- noise level parameter ε with the temporal variable t (seetectures are discrete in time. Gupta et al. [7] performed Appendix A.1, A.2 for more details). Then, we use thea systematic ablation study to analyze the significance of conditioning capabilities of diffusion models to learn the
Fourier layers, attention mechanism, and parameter con- relations between the initial condition, the PDE solution,

Figure 1: DiTTO architecture. The discretized initial condition u(x, 0) concatenated with the corresponding spatialgrid, and the desired time t ∈ R+ are the respective inputs to the U-Net and the time-embedding network comprising
DiTTO. The U-Net illustrated here consists of ResNet blocks with temporal conditioning, a Spatial-Attention block,and Channel-Attention blocks at 4 levels of coarseness and the corresponding residual connections across the samelevels. The ResNet block conditions the non-linear representations of u(x, 0) with respect to the temporal embedding
vector, f⃗(t), by performing element-wise multiplication across the channels. Spatial-Attention and Channel-Attention
blocks learn to extract correlations across space and channels, respectively.
2

Diffusion-inspired Temporal Transformer Operator (DiTTO)

and the time domain. After the training is complete, themodel can interpolate between the initial and final time,creating a numerical solution that is continuous in time.
We define its time evolution {xt ]} as the following process:{xt | t ∈ [0, tf inal ], xt := u(x, t)},randomly sub-sampling the trajectory timesteps considered
at each epoch (DiTTO-s), and 2) adopting the temporalbundling [2] for enhancing the forecast capabilities. These
extensions are further explained in the Appendix A.2.

(1)

where u is the solution of the differential equation we at- 3 Experiments and Results
tempt to solve. Using this notation, the operator learning
We have tested the proposed DiTTO approach on variproblem becomes:
ous PDEs. Full details regarding these experiments are
x0 −→ xt , ∀t ∈ [0, tf inal ],(2) given in the appendix. Here, we present partial results forxt ≈ G(x0 )(t), ∀t ∈ [0, tf inal ],(3) the hypersonics and climate cases in Figure 2. We compare DiTTO to another operator learning framework called
where G represents the surrogate operator, DiTTO, where Fourier neural operator (FNO) [24] when applicable.
the operator learning technique we employ is the diffusion
process. It is discrete, while (1) is continuous. We discretize {xt } by taking a partitioning {tn }Tn=0 of [0, tf inal ], 3.1 Hypersonic Flow
where 0 = t0 < t1 < . . . < tT −1 < tT = tf inal .
We train DiTTO to learn the inviscid airflow around a
The discrete process is then defined as {xn }Tn=0 , where
double-cone object flying at a high Mach number. The flow
xn := u(x, tn ). In PDE terms, given an initial condition
physics at hypersonic Mach numbers around the doublex0 , we approximate the analytic solution at a set of specificcone geometry features complex transient events, stationT
future time steps {tn }n=1 . In operator learning terms, weary bow shock at the leading nose of the cone, and themap a family of functions of the form x0 = u(x, 0) tointeraction of the oblique shock wave originating fromanother family of functions of the form u(x, t).
the leading edge with the bow shock formed around the
The role of the neural network in diffusion models is to upper part of the geometry of the cone. The shock layer
perform conditional denoising in each step. We repurpose is the narrow band between the wall of the double-cone
this network structure to solve a PDE-related problem. geometry and the bow shock wave. In the shock layer, due
Since x0 , x1 , . . . , xT are taken from the analytical solu- to the interaction of shock waves, a triple point forms andtion, there is no noise in this process. Therefore, there is generates vortical flow structures that move downstream.
no need for denoising. Thus, we replace the conditional This interaction is challenging to capture using numeridenoising operation with a conditional temporal evolution cal solvers. The details regarding data creation and the(x0 , tn ) −→ xn .
governing equations are given in the Appendix B.
Next, we describe the DiTTO architecture. The network receives two main inputs: the initial condition x0 = u(x, 0)
and a time t = tn . Recall that x0 is a d-dimensional tensor,
and t is a nonnegative scalar. For the temporal input t, we
use an embedding mechanism based on the original Transformer positional encoding [49]. Each scalar t is mapped
into a vector of size demb , and then passed through a simple multi-layer perceptron (MLP) with two linear layers
and a GELU [13] activation function.

We now present two experiments conducted for this hypersonic problem. First, we learn the temporal evolutionof the density field near the double cone. Specifically, we
train the neural operator to learn the mapping from theincoming horizontal velocity field to the time-dependent
density field around the double cone structure. The dataset
comprises only 61 trajectories corresponding to Mach numbers M ∈ [8, 10]. It is split into training, validation, andtesting datasets in the ratio 80:10:10. We use a cosine
annealing type learning rate scheduler starting from 10−3that decays to 0 during the training. We compare FNO, UNet, DiTTO and DiTTO-s. Complete results are reported
in the appendix. We observe in Figure 2A that DiTTO canaccurately resolve the vortices close to the surface of the
double cone.

For the spatial input x0 , we concatenate it with a discrete
spatial grid and provide it as an input to a U-Net [38]. We
use a U-Net variant common in many diffusion models,such as DDPM [15]. It follows the backbone of PixelCNN++ [40], a U-Net based on a Wide ResNet [53, 12].
A sketch of the spatio-temporal architecture is given in
Figure 1.
Additionally, we perform an experiment where insteadof conditioning on time as in most of the examples preThis architecture is not limited to a specific dimension. The
sented, we explore the ability to condition DiTTO on asame mechanism can be implemented for d-dimensional
different quantity: the Mach Number. We train the surroproblems. The only major difference is the usage of dgate operator to learn the mapping from the density field
dimensional convolutions for the relevant problem.
at a specific time step and M = 8 to the density field at
We extend DiTTO to develop three variants: DiTTO-s, the same timestep but at different Mach numbers in the
DiTTO-point and DiTTO-gate. DiTTO-point is a memory- range M ∈ [8, 10]. We compare DiTTO with other neural
efficient version, and DiTTO-gate incorporates a gated operators. Our results suggest that DiTTO serves as an
sub-architecture motivated by Runge-Kutta methods. We accurate surrogate conditioned on any scalar parameter
also demonstrate the efficacy of two training strategies - 1) and not necessarily time.
3Diffusion-inspired Temporal Transformer Operator (DiTTO)

3.2

Climate Modelingaccumulation of errors. The contour plots illustrate themeasured (left) and predicted (right) global temperature
profiles in Spring, Summer, Autumn, and Winter. DiTTOClimate models are complex, involving non-linear dy- reduces the growth of the relative L2 error across the 5
namics and multi-scale interactions between multiple years of extrapolation with a mean of 0.014.
variables. Describing climate behavior using PDEs is
not a straightforward task. Consequently, there aremany different numerical models and ML-based models Acknowledgements[16, 1, 18, 33, 22, 30] that have been developed to modelclimate. Climate-related problems are challenging to apThis work was supported by the Vannevar Bush Faculty
proximate accurately at a low computational cost. Our
Fellowship award (GEK) from ONR (N00014-22-1-2795).
goal is to efficiently learn the temporal evolution of the surIt is also supported by the U.S. Department of Energy,
face air temperature across the globe, and make accurate
Advanced Scientific Computing Research program, under
forecasts by extrapolating beyond the training period.
the Scalable, Efficient and Accelerated Causal Reasoning
The results in Figure 2B demonstrate that DiTTO is able to Operators, Graphs and Spikes for Earth and Embedded
extrapolate for long time intervals without any substantial Systems (SEA-CROGS) project, DE-SC0023191.

(A) Modeling high-speed flow past a double-cone.

(B) Five-year climate forecast (extrapolation).

Figure 2: Two examples of DiTTO experiments.
Fig 2A: Modeling high-speed flow past a double-cone. a) The image of the double-cone model installed in the LENS
XX tunnel [39]. b) The histogram illustrates the sensitivity of (i) DiTTO-point, (ii) DiTTO, and (iii) FNO to different
noise levels γ. c) Temporal evolution of the density field (ρ) corresponding to test Mach number (M ) = 9.18 aroundthe double cone geometry predicted by DiTTO-point and its comparison with the solver simulation. d) Additional
illustration of DiTTO-point conditioned on M and trained to predict the final state of ρ.
Fig 2B: Five-year climate forecast (extrapolation). DiTTO is trained on the global temperature data from years 2013 to2015, validated from 2016 to 2017, and extrapolated from 2018 to 2022. The contour plots on the left and right columns
correspond to snapshots from five years of measured and forecasted global temperature distribution, respectively. The
four rows correspond to the global temperature profiles in Spring, Summer, Autumn, and Winter.
4

Diffusion-inspired Temporal Transformer Operator (DiTTO)

References[1] Aniruddha Bora, Khemraj Shukla, Shixuan Zhang, Bryce Harrop, Ruby Leung, and George Em Karniadakis.
Learning bias corrections for climate models using deep neural operators. arXiv preprint arXiv:2302.03173, 2023.
[2] Johannes Brandstetter, Daniel Worrall, and Max Welling. Message passing neural pde solvers. arXiv preprint
arXiv:2202.03376, 2022.
[3] Qianying Cao, Somdatta Goswami, and George Em Karniadakis. Lno: Laplace neural operator for solving
differential equations. arXiv preprint arXiv:2303.10528, 2023.
[4] Shuhao Cao. Choose a transformer: Fourier or galerkin. Advances in neural information processing systems,
34:24924–24940, 2021.
[5] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner,
Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words:
Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929, 2020.
[6] Ruchi Guo, Shuhao Cao, and Long Chen. Transformer meets boundary value inverse problems. arXiv preprint
arXiv:2209.14977, 2022.
[7] Jayesh K Gupta and Johannes Brandstetter. Towards multi-spatiotemporal-scale generalized pde modeling. arXiv
preprint arXiv:2209.15616, 2022.
[8] Jacques Hadamard. Sur les problèmes aux dérivées partielles et leur signification physique. Princeton university
bulletin, pages 49–52, 1902.
[9] Kai Han, Yunhe Wang, Hanting Chen, Xinghao Chen, Jianyuan Guo, Zhenhua Liu, Yehui Tang, An Xiao,
Chunjing Xu, Yixing Xu, et al. A survey on vision transformer. IEEE transactions on pattern analysis andmachine intelligence, 45(1):87–110, 2022.
[10] Zhongkai Hao, Zhengyi Wang, Hang Su, Chengyang Ying, Yinpeng Dong, Songming Liu, Ze Cheng, Jian Song,and Jun Zhu. Gnot: A general neural operator transformer for operator learning. In International Conference on
Machine Learning, pages 12556–12569. PMLR, 2023.
[11] Zhongkai Hao, Chengyang Ying, Zhengyi Wang, Hang Su, Yinpeng Dong, Songming Liu, Ze Cheng, Jun Zhu, andJian Song. Gnot: A general neural operator transformer for operator learning. arXiv preprint arXiv:2302.14376,
2023.
[12] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In
Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016.
[13] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08415, 2016.
[14] Geoffrey Hinton, Li Deng, Dong Yu, George E Dahl, Abdel-rahman Mohamed, Navdeep Jaitly, Andrew Senior,
Vincent Vanhoucke, Patrick Nguyen, Tara N Sainath, et al. Deep neural networks for acoustic modeling in speech
recognition: The shared views of four research groups. IEEE Signal Processing Magazine, 29(6):82–97, 2012.
[15] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural
Information Processing Systems, 33:6840–6851, 2020.
[16] James W Hurrell, Marika M Holland, Peter R Gent, Steven Ghan, Jennifer E Kay, Paul J Kushner, J-F Lamarque,
William G Large, D Lawrence, Keith Lindsay, et al. The community earth system model: a framework forcollaborative research. Bulletin of the American Meteorological Society, 94(9):1339–1360, 2013.
[17] Eugenia Kalnay, Masao Kanamitsu, Robert Kistler, William Collins, Dennis Deaven, Lev Gandin, Mark Iredell,
Suranjana Saha, Glenn White, John Woollen, et al. The ncep/ncar 40-year reanalysis project. In Renewable
Energy, pages Vol1_146–Vol1_194. Routledge, 2018.
[18] Georgios Kissas, Jacob H Seidman, Leonardo Ferreira Guilhoto, Victor M Preciado, George J Pappas, and Paris
Perdikaris. Learning operators with coupled attention. The Journal of Machine Learning Research, 23(1):9636–
9698, 2022.
[19] Nikola B. Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew M. Stuart,and Anima Anandkumar. Neural operator: Learning maps between function spaces. CoRR, abs/2108.08481, 2021.
[20] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neuralnetworks. Advances in Neural Information Processing Systems, 25, 2012.
[21] Varun Kumar, Leonard Gleyzer, Adar Kahana, Khemraj Shukla, and George Em Karniadakis. Crunchgpt: A
chatgpt assisted framework for scientific machine learning. arXiv preprint arXiv:2306.15551, 2023.
5

Diffusion-inspired Temporal Transformer Operator (DiTTO)[22] Remi Lam, Alvaro Sanchez-Gonzalez, Matthew Willson, Peter Wirnsberger, Meire Fortunato, Alexander Pritzel,
Suman Ravuri, Timo Ewalds, Ferran Alet, Zach Eaton-Rosen, et al. Graphcast: Learning skillful medium-range
global weather forecasting. arXiv preprint arXiv:2212.12794, 2022.
[23] Zijie Li, Kazem Meidani, and Amir Barati Farimani. Transformer for partial differential equations’ operator
learning. arXiv preprint arXiv:2205.13671, 2022.
[24] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart,and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint
arXiv:2010.08895, 2020.
[25] Xinliang Liu, Bo Xu, and Lei Zhang. Ht-net: Hierarchical transformer based operator learning model for
multiscale pdes. arXiv preprint arXiv:2210.10890, 2022.
[26] Yue Liu, Zhengwei Yang, Zhenyao Yu, Zitu Liu, Dahui Liu, Hailong Lin, Mingqing Li, Shuchang Ma, Maxim
Avdeev, and Siqi Shi. Generative artificial intelligence and its applications in materials science: Current situationand future perspectives. Journal of Materiomics, 2023.
[27] Zichao Long, Yiping Lu, Xianzhong Ma, and Bin Dong. Pde-net: Learning pdes from data. In International
conference on machine learning, pages 3208–3216. PMLR, 2018.
[28] Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear
operators via deeponet based on the universal approximation theorem of operators. Nature machine intelligence,3(3):218–229, 2021.
[29] Lu Lu, Xuhui Meng, Shengze Cai, Zhiping Mao, Somdatta Goswami, Zhongqiang Zhang, and George Em
Karniadakis. A comprehensive and fair comparison of two neural operators (with practical extensions) based onfair data. Computer Methods in Applied Mechanics and Engineering, 393:114778, 2022.
[30] Tung Nguyen, Johannes Brandstetter, Ashish Kapoor, Jayesh K Gupta, and Aditya Grover. Climax: A foundation
model for weather and climate. arXiv preprint arXiv:2301.10343, 2023.
[31] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pages 8162–8171. PMLR, 2021.
[32] Oded Ovadia, Adar Kahana, Panos Stinis, Eli Turkel, and George Em Karniadakis. Vito: Vision transformeroperator. arXiv preprint arXiv:2303.08891, 2023.
[33] Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani,
Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, et al. Fourcastnet: A global data-driven highresolution weather model using adaptive fourier neural operators. arXiv preprint arXiv:2202.11214, 2022.
[34] Ahmad Peyvan, Khemraj Shukla, Jesse Chan, and George Karniadakis. High-order methods for hypersonic flows
with strong shocks and real chemistry. Journal of Computational Physics, 490:112310, 2023.
[35] Md Ashiqur Rahman, Zachary E Ross, and Kamyar Azizzadenesheli. U-no: U-shaped neural operators. arXiv
preprint arXiv:2204.11127, 2022.
[36] Maziar Raissi, Paris Perdikaris, and George E Karniadakis. Physics-informed neural networks: A deep learning
framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal ofComputational physics, 378:686–707, 2019.
[37] Hendrik Ranocha, Michael Schlottke-Lakemper, Andrew Ross Winters, Erik Faulhaber, Jesse Chan, and Gregor
Gassner. Adaptive numerical simulations with Trixi.jl: A case study of Julia for scientific computing. Proceedingsof the JuliaCon Conferences, 1(1):77, 2022.
[38] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International
Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pages 234–241. Springer, 2015.
[39] Mehrnaz Rouhi Youssefi and Doyle Knight. Assessment of cfd capability for hypersonic shock wave laminarboundary layer interactions. Aerospace, 4(2):25, 2017.
[40] Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P Kingma. Pixelcnn++: Improving the pixelcnn withdiscretized logistic mixture likelihood and other modifications. arXiv preprint arXiv:1701.05517, 2017.
[41] Michael Schlottke-Lakemper, Gregor J Gassner, Hendrik Ranocha, Andrew R Winters, and Jesse Chan.
Trixi.jl: Adaptive high-order numerical simulations of hyperbolic PDEs in Julia. https://github.com/
trixi-framework/Trixi.jl, 09 2021.
[42] Michael Schlottke-Lakemper, Andrew R Winters, Hendrik Ranocha, and Gregor J Gassner. A purely hyperbolic
discontinuous Galerkin approach for self-gravitating gas dynamics. Journal of Computational Physics, 442:110467,
06 2021.
6

Diffusion-inspired Temporal Transformer Operator (DiTTO)[43] Dule Shu, Zijie Li, and Amir Barati Farimani. A physics-informed diffusion model for high-fidelity flow fieldreconstruction. Journal of Computational Physics, 478:111972, 2023.
[44] Chenyang Si, Ziqi Huang, Yuming Jiang, and Ziwei Liu. Freeu: Free lunch in diffusion u-net. arXiv preprint
arXiv:2309.11497, 2023.
[45] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using
nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2256–2265. PMLR,
2015.
[46] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint
arXiv:2010.02502, 2020.
[47] Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Daniel MacKinlay, Francesco Alesiani, Dirk Pflüger,and Mathias Niepert. Pdebench: An extensive benchmark for scientific machine learning. Advances in Neural
Information Processing Systems, 35:1596–1611, 2022.
[48] Tapas Tripura and Souvik Chakraborty. Wavelet neural operator for solving parametric partial differential equations
in computational mechanics problems. Computer Methods in Applied Mechanics and Engineering, 404:115783,
2023.
[49] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, andIllia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[50] Ting Wang, Petr Plechac, and Jaroslaw Knap. Generative diffusion learning for parametric partial differential
equations. arXiv preprint arXiv:2305.14703, 2023.
[51] Qingsong Wen, Tian Zhou, Chaoli Zhang, Weiqi Chen, Ziqing Ma, Junchi Yan, and Liang Sun. Transformers intime series: A survey. arXiv preprint arXiv:2202.07125, 2022.
[52] Hao Xu, Haibin Chang, and Dongxiao Zhang. Dl-pde: Deep-learning based data-driven discovery of partial
differential equations from discrete and noisy data. arXiv preprint arXiv:1908.04463, 2019.
[53] Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. arXiv preprint arXiv:1605.07146, 2016.

7

Diffusion-inspired Temporal Transformer Operator (DiTTO)

Supplementary Information
Appendix A:
A.1
A.1.1

learn mappings between function spaces using projections.
FNOs utilize the Fourier transform. They are effective
and easy to implement,. FNOs are accurate, especiallyfor smooth and periodic problems [29]. While the Fourier
kernel is continuous it is necessary to use discrete versions
for operator learning. Consequently, FNOs can be computationally costly when working with high-dimensional
problems requiring many Fourier modes.

Detailed methods

Background
Operator learning

The standard use of ML models for scientific computations
involves fitting a function to map numerical inputs to outputs. These inputs are ordinarily coordinates, materials,
boundary conditions, etc., and the outputs are usually solutions of forward PDEs. An example is physics-informed
neural networks (PINNs) [36], which use a deep neuralnetwork to solve PDEs by embedding elements from the
PDE problem into the loss function. So, the network trainson the given data while using prior information about the
problem it is solving. A major drawback is that for eachproblem, one needs to re-train the network, which is computationally expensive. This includes any changes to theparameters defining the problem.

A.1.2

Transformers and attention

First presented by Vaswani et al. [49], transformers havebeen widely used in the ML community. Transformers
introduce a new type of mechanism called the scaled dotproduct attention. The attention module attempts to gather
context from the given input. It operates on a discrete
embedding of the data composed of discrete tokens.
The original architecture was proposed for natural language processing purposes, where one encodes sentencesusing their enumerated locations in the vocabulary. Since
then, their usage has been extended to many other domains, and they outperform many different deep learning
architectures in a wide variety of tasks. These domains
include time series analysis [51] and computer vision [9].
For example, Vision Transformers (ViT) [5] split images
into small patches, tokenize them, and apply the attention
mechanism. In addition, they are computationally lighter
than other alternatives and can be easily parallelized.

Operator learning seeks to overcome this problem. Instead of fitting a function, one fits a mapping between
two families of functions. Consider a generic family ofd-dimensional time-dependent PDE problems of the form:Lu(x, t) = f (x, t), x ∈ D, t ∈ [0, tf inal ]Bu(x, t) = g(x, t), x ∈ ∂D, t ∈ [0, tf inal ] , (4)u(x, 0) = I(x),x∈D
The differential operator L and forcing term f defines thePDE, the boundary operator B while g defines the solution on the boundary, tf inal is the final physical time, I
is the initial condition, and D is a Euclidean domain in
Rd with boundary ∂D. We assume that the problem (4) iswell-posed [8], so a unique solution exists.

Transformers are becoming increasingly popular in the
SciML community. They have been used for operator learning in many different ways [23, 25, 11, 10]. These methods
show much promise by using the attention mechanism tofind connections between points in the physical domain tofunction values. Some methods emphasize the attention
mechanism itself [4, 6] and adapt it to PDE-related problems. Others utilize existing transformer models to solve
PDE problems more easily [21]. In this work, we employ
Let I be a function space containing initial conditions of elements from the original Transformer architecture as part(4). Then there exists another space U that contains their of the proposed neural network architecture.
respective solutions. We define an operator G : I −→ U :A.1.3

G(I)(x, t) = u(x, t),(5)where I ∈ I, x ∈ D, and t ∈ [0, tf inal ]. Each initialcondition I ∈ I is mapped into its corresponding solution
u ∈ U. The goal is to approximate the operator G using aneural network.

Diffusion models

A diffusion model is a generative deep learning model thatuses a Markov chain to produce samples that match a given
dataset [45]. These models aim to learn the underlying distribution of a given dataset. After learning this distribution,
they are used to generate new samples of similar properties
The first SciML operator learning method, called Deep- to those found in the training set.
ONet, was proposed by Lu et al. [28]. The main compoIn [15], Ho et al. introduced a new type of diffusion model
nents of a DeepONet are two neural networks: the branchcalled denoising diffusion probabilistic models (DDPM).
and the trunk. Each network can be a fully connected
It consists of a forward diffusion process and an inverse
neural network, convolutional, or any other architecture.
one. In the forward case, Gaussian noise is incrementally
Usually, the branch inputs are functions, and the trunk inadded to the original sample for a given number of iterputs are coordinates. DeepONets learn projections fromations. For a sufficiently large number of iterations, thethe functions to a vector space, so they can map input
noise completely destroys the original signal. Then, infunctions to output functions at specific points.
the reverse diffusion process, the goal is to reconstruct
Another operator learning approach is the Fourier neural the original signal by performing iterative denoising steps
operator (FNO) [24, 19]. FNOs, similarly to DeepONets, using a neural network. Diffusion models have been used8

Diffusion-inspired Temporal Transformer Operator (DiTTO)

for SciML purposes, especially for generative artificial in- A.2 DiTTO: Diffusion-inspired temporaltelligence purposes [50, 43, 26]. While we are not usingtransformer operator
their generative capabilities in this work, we briefly explain
their standard training procedure.
As described in Methodology, we propose using temporalconditioning, instead of conditioning on noise as in SecWe present a mathematical formulation based on the works
tion A.1.3. We use the initial condition of a differentialof Ho et al. [15] and Nichol et al. [31]. Given a data disequation and infer the entire temporal process. It opertribution x0 ∼ q(x0 ), we define a forward noising process
ates on various initial conditions, hence it is performingq which produces steps x1 , . . . , xT by adding Gaussian
operator learning according to Section A.1.1. In addition,
noise at time t with variance βt ∈ (0, 1) as follows:
after training, the inference is possible on any real temporal value. Hence, the inference is continuous in time.
T
As
shown in the results, not only interpolation in time isY
q(x1 , . . . , xT |x0 ) :=q(xt |xt−1 ),(6) possible, but also extrapolation.
t=1

To train the DiTTO network we gather data of multiple
time-dependent procedures, varying in the initial condition.
pThen,
we use the network architecture in Methodology, andq(xt |xt−1 ) := N (xt ; 1 − βt xt−1 , βt I).
(7)
select a training strategy (for example, temporal-bundling),
to fit the data. We elaborate on these steps in the following
Given a sufficiently large T and a well-behaved schedule subsections.
βt , the latent xT is nearly an isotropic Gaussian distribution. From (7), we see that xt is drawn from√ a conditional
Gaussian distribution with mean µt = 1 − βt xt−1 and A.2.1 Training dataset
variance σt2 = βt . In practice, this is done by randomly
sampling a noise level parameter ε ∼ N (0, I), and setting: To train a neural network using the formulation presented
in Methodology, we require a large set of initial conditions (inputs) and corresponding solutions (outputs). Letppxt = 1 − βt xt−1 + βt ε.
(8) {I m (x)}M
m=1 be a set of initial conditions with corresponding analytic solutions {um (x, t)}M
m=1 , where M is theThus, if we know the exact reverse distribution q(xt−1 |xt ), desired number of training samples. Each sample consists
we can sample xT ∼ N (0, I) and run the process in reverse of an initial condition and a PDE solution at the relevantmM
to get a sample from q(x0 ). However, since q(xt−1 |xt ) timesteps. In practice, {u (x, t)}m=1 are numerical apdepends on the entire data distribution, we approximate it proximations of the analytic solutions and not analytic
using a neural network with hyperparameters θ as follows: solutions which are often unavailable. Furthermore, the solutions are discretized in space using a grid that partitions
the domain D. We emphasize that for all m = 1, . . . , Mpθ (xt−1 |xt ) := N (xt−1 ; µθ (xt , t), Σθ (xt , t)). (9) and t = 0, . . . , T , um (x, t) is a matrix, and its dimensions
depend on the spatial discretization parameters, i.e., theThe neural network needs to learn the mean and variance number of nodes along each axis.
to complete the backward diffusion process. Importantly, We denote the forward process corresponding to the mT
using the formulation in (8), in each step, it is sufficient th initial condition and solution by {xmn }n=0 , wheremto know βt , xt , and ε to approximate xt−1 . Then, the net- xm
n := u (x, tn ). We define the following datasets:
work is used autoregressively to reconstruct x0 . Assuming
we know the schedules {βt }Tt=1 , we can view the neuralnetwork as the following mapping:
X = {(xm
0 , tn )| n = 1, . . . , T, m = 1, . . . , M }. (11)
Y = {xmn | n = 1, . . . , T, m = 1, . . . , M }(xt , ε) −→ xt−1 .
(10)
So, each solution of the PDE is transformed into T pairsof samples that correspond to the mapping described in
Equation (3).

In each step, the neural network performs a denoising operation, mapping xt to a slightly less noisy signal xt−1 .
Including the noise level parameter ε is essential for thedenoising operation. During training, various noise levels
are sampled. Knowing the specific noise level that distinguishes between consecutive states xt and xt−1 , is crucial
for effective denoising. Without this explicit knowledgeof the noise level, the denoising process would become
significantly more complicated, and the network may notconverge. Thus we have a conditional denoising operation,
conditioned on ε (or equivalently on the timestep with βt ).

A.2.2

The DiTTO neural network architecture andparameters

We use the architecture described in Methodology. We
mention that for the temporal input t, we use an embedding
mechanism based on the original Transformer positional
encoding [49]:9Diffusion-inspired Temporal Transformer Operator (DiTTO)
pos,
2i/demb
10000 pos(12)
P E(pos,2i+1) = cos,
10000(2i+1)/demb
where demb is the desired embedding dimension. Each
scalar t is mapped into a vector of size demb .
P E(pos,2i) = siniterating over the entire process, we iterate over a random
subsequence. Recall that for the m-th initial condition inT
the training set, the full process is {xmn }n=1 . Instead, wetake a set of random subsequences Sm ⊂ {0, 1, . . . , T },PMsuch that m=1 |Sm | = αM T for some α < 1. Choosing
α = 0.05 means we only use 5% of the given samples ineach epoch. The new DiTTO-s loss is given by:

We now describe the loss function used as a target for training the network. Let Oθ be the neural network describedM
in Methodology with hyperparameters θ. The goal of Oθ
m
1 X X ||Oθ (xm
0 , tn ) − xn ||2
is to learn the mapping described in (3), using the dataset, (14)
lossα :=
αM T m=1
ε + ||xmn ||2(11). We split this dataset into training, validation, andn∈Smtesting sets. We split them in a way that makes sure thatno initial conditions from the validation and testing sets
After each epoch, we randomly sample Sm again using aappear in the training set.
uniform distribution. That way, given a sufficiently large
Diffusion models are often trained with a probabilistic loss number of epochs, we expect to cover a significant portionfunction. However, since we learn a PDE operator, other of samples in the dataset.
loss functions commonly used for SciML applications are
more fitting. Consequently, we train the network with a
A.2.4 DiTTO-point
mean relative L2 loss:

loss :=T
M X
X||Oθ (xm , tn ) − xm ||21
M T m=1 n=10

ε + ||xm
n ||2

n

,

(13)

The architecture shown in Figure 1 can be used for problems in different dimensions. A way to accomplish that is
to use d-dimensional convolutions in all the convolutional
layers, where d ∈ {1, 2, 3} is the physical dimension of the
problem. Using high-dimensional convolutions requiresa large amount of memory. Furthermore, many physicalproblems are not defined on structured grids. Thus, to use aneural operator framework, it is often necessary to project
their solutions onto regular grids [29], which requires achange to the geometry of the problem.

where ε is a small number used to prevent a zero denominator and stabilize the loss. The inputs and outputs ofthe model are d-dimensional, so they are converted intoa one-dimensional array by column stacking (flattening)inside the loss function when needed. We describe theloss for the entire dataset for simplicity, but in practice, we To address these issues, we propose DiTTO-point, anotherdivide it into batches.
variant of DiTTO that solves high-dimensional problems
Iterating over the entire dataset (11) can be time- using exclusively 1-D convolutions. With DiTTO-point
consuming. For Mtrain initial conditions in the training we treat the domain as a set of points in space instead ofset, we have Mtrain · T samples. So, the number of train- a structured d-dimensional grid. A domain with N points
ing steps scales linearly with T . So the number of training is defined as a N × d matrix, where each row represents
a spatial coordinate. Similarly, we define the solution on
samples is very large for fine temporal discretizations.
this domain using a vector of size N , corresponding to the
A similar problem occurs in generative diffusion models. values of the solution at each point of the domain. Using
The original DDPM [15] requires hundreds of forward this formulation, regardless of the original dimensionality
passes to produce good results. Later works suggested of the problem, the input to DiTTO-point would always beways to improve the performance aspect of DDPMs. For of size N × d. This enables the use of 1-D convolutions
example, Song et al. [46] suggest using non-Markovian on this data, where d is the number of input channels.
processes to improve the computational cost. Nichol etal. [31] present a way to significantly reduce the number However, directly using the architecture shown in Figure 1of necessary steps by sub-sampling the original diffusion with 1-D convolution does not work, without any modprocess. Both methods focus primarily on the inference ifications, on high-dimensional problems. Importantly,
speed. However, in the case of DiTTO, inference is imme- switching from a structured grid to a set of coordinates
diate. In Methodology, we explained that we do not view results in the loss of spatial information. This is especiallyx0 , x1 , . . . , xT as an iterative process. Instead, we treat true when the order in which the coordinates appear is noteach sample individually, significantly increasing the infer- necessarily related to their physical distance. To solve thisence speed compared to generative models such as DDPM. issue, we use another layer of positional encoding (seeEquation (12)) based on the spatial coordinates of the grid.
Hence, we focus on speeding up the training process.
We apply this layer to the beginning of the overall architecture before we concatenate its output with the relevant
A.2.3 DiTTO-s
initial condition in the latent space. This enables DiTTOWe propose DiTTO-s, a faster variant of DiTTO that relies point to keep high-dimensional spatial information whileon a sub-sampling mechanism similar to [31]. Instead of using a one-dimensional architecture.
10Diffusion-inspired Temporal Transformer Operator (DiTTO)

A.2.5

DiTTO-gate

We propose another variant of the architecture shown in
Figure 1, called DiTTO-gate. Now, we modify the behavior of the skip connections of the U-Net decoder. Analysis
of diffusion models shows that the U-Net skip connections
introduce high-frequency features into the decoder [44].
For scenarios with fine details and sharp features, we put
extra emphasis on the skip connections. So, we introduce
a gate component, which operates directly on the skip
connections. This component is composed of a standard
convolutional block, as described in Methodology. So, in
DiTTO-gate, we add such a component to each level of thedecoder and use it on its incoming skip connection.

Appendix B: Hypersonic flow data Figure 3: 2D domain of the double cone problem. The
gray area is the domain and the highlighted lines show thegeneration and details

surface of the double cone, which is a slip wall boundary
condition. The small line at the bottom is a symmetry conWe first generate a training dataset by solving the 2D Euler
dition. The inflow and outflow boundaries are also shown.
equations on a fluid domain around a double-cone geometry. The governing equations are∂U ∂F ∂G
+
+
= 0,
∂t
∂x
∂y

Appendix C:

Extrapolation in time

(15)

Extrapolation is a challenging problem due to the inherent
nature of the data-driven surrogate networks to overfit thetraining distribution. We discuss a training strategy thatpartially alleviates the difficulties associated with extrapo  lation in time in using DiTTO.
ρuρvρ ρu 
 ρu2 + p  ρvu  During the training of DiTTO, the model is exposed onlyU =  ,F = ,G = 
.
ρv
ρuv 
ρv 2 + p  to the first 100 time steps. We train DiTTO with three
types of time-series modeling strategies demonstrated inρEu(ρE + p)u(ρE + p)(16) part a) of Figure 4 - i) autoregressive (look-forward winIn Equation (15), t is time, x, y are the spatial coordinates, dow lf = 1), ii) temporal-bundling (1 < lf < nt) [2],u and v are the x-direction and y-direction velocities, ρ is and iii) mapping (lf = nt).th We analyze their ability todensity, and p is the pressure. The total energy in Equa- extrapolate beyond the 100 time step. Specifically, we
train 6 different DiTTO models with lf =1, 5, 10, 20, 50,
tion (16) is illustrated as100(= nt). During the training, DiTTO learns a mappingu(x̄, t) → u(x̄, t + lf ) such that t + lf is less than thep1ρE =+ ρ(u2 + v 2 ), γ = 1, 4(17) number of time steps, nt = 100. Hence, each trajectory in
γ−1 2
the training dataset is split into nt − lf + 1 sub-trajectories.
During the inference stage, DiTTO leaps from û(x̄, t) to
We solve the system of equations using Trixi.jl numerical û(x̄, t + lf ). Because the ground truth is not available,
framework [41], which employs entropy stable discontin- except at t = 0, we consider u as the ground truth and ûuous Galerkin spectral element (ES-DGSEM) approach as a prediction from DiTTO.
[37, 42, 34] to solve hyperbolic and parabolic systems ofequations. ES-DGSEM features high accuracy and stability and employs adaptive mesh refinement to automatically C.1 Navier-Stokes Equations
adapt the mesh resolution to high gradient regions of the
As an example, we train DiTTO on the two-dimensional
flow field. The 2D physical domain is shown in Figure 3,
incompressible Navier-Stokes equation for a viscous, inwhere the boundary conditions are specified. Let M becompressible fluid in vorticity form is given by:
the free-stream Mach number. The domain is initializedusing constant uniform values for primitive variables as 
∂ ω + u · ∇ω = ν∆ω + f,
u = M , v = 0, p = 1.0, and ρ = 1.4. Each simulation  t(x, y) ∈ (0, 1)2 , t ∈ (0, tf inal ]is performed for a time t ∈ [0, 0.04],. The solution values ∇ · u(x, y, t) = 0,
ω(x, y, 0) = ω0 ,(x, y) ∈ (0, 1)2
are shown at 201 snapshots corresponding to equidistanceinstances of time.
(18)
where the vector of conservative variables, x-direction andy-direction flux vectors, are described as11

Diffusion-inspired Temporal Transformer Operator (DiTTO)

Figure 4: Temporal-bundling for efficient extrapolation. a) demonstrates 3 types of time-series modeling strategies
for extrapolating beyond the training interval. We consider a time series with 20 time steps. b) visualizes the error
accumulation for DiTTO models with different look-forward windows (lf ) trained in the time interval 0-100 and
extrapolated from 100-200. c) shows the relative percentage L2 errors at the 200th time step, corresponding to DiTTO
models with different lf. The results in b) and c) come from the test dataset with previously unseen 100 initial conditions.
The symbols in (b-c) denote the corresponding final time.
where ω is the vorticity, u is the velocity field, ν is theviscosity, and ∆ is the two-dimensional Laplacian operator. We consider periodic boundary conditions. The
source term f is given by f (x, y) = 0.1(sin(2π(x + y)) +cos(2π(x+y))), and the initial condition ω0 (x) is sampledfrom a Gaussian random field according to the distribution
N (0, 73/2 (−∆ + 49I)−5/2 ).
The dataset consists of 1000 trajectories with 200 time
steps. It was randomly split into training, validation, andtesting datasets in the ratio 80:10:10.
In Figure 4(b,c), we observe the minimum error with thelowest uncertainty occurs at lf = 20, suggesting that thetemporal-bundling technique [2] offers a sweet spot between the autoregressive and the mapping strategies for

extrapolating in time, with lower rates of error accumulation. Next, we investigate the extrapolation abilities of thetemporal-bundling strategy for the climate problem.
C.2

Climate Dataset

We use a publicly available climate dataset provided
by the Physical Sciences Laboratory meteorological
data: https://psl.noaa.gov/data/gridded/index.
html [17]. This data contains measurements of parameters over time, projected onto a spatial grid. We use thedaily average surface temperature data (at 1000 millibar
Pressure level) from January 1, 2013, to December 31,
2015, for training; January 1, 2016, to December 31, 2017,for validating; and January 1, 2018 to December 31, 202212

Diffusion-inspired Temporal Transformer Operator (DiTTO)

for testing (see Figure 2B). The data is projected to a spatial grid of dimensions 144 × 72, which corresponds toa resolution of 2.5◦ in both latitude and longitude. We
split the data to training, validation and test sets in theratio 30:20:50, for investigating the real-time long-term
temporal extrapolation capability of DiTTO.

a single trajectory consisting of 1095 global temperaturesnapshots arising from the initial condition on Jan. 1, 2013.
We utilize temporal-bundling with lf = 365 to split thistrajectory into 730 sub-trajectories, each comprising 365
time steps. We further reduce the computational costs associated with the training and inference by performing aproper orthogonal decomposition and training DiTTO to
The experiments shown in Figure 4 motivated the adoplearn the evolution of the eigen coefficients correspondingtion of the temporal-bundling technique with lf = 365 forto the five dominant eigenvectors estimated from the global
training DiTTO. The training dataset can be interpreted as
temperature profiles in the training dataset.

13