L EARNED C OARSE M ODELS FOR E FFICIENT T URBULENCE S IMULATION
A BSTRACT
Turbulence simulation with classical numerical solvers requires high-resolution grids to accurately resolve dynamics.
Here we train learned simulators at low spatial and temporal resolutions to capture turbulent dynamics generated at high resolution.
We show that our proposed model can simulate turbulent dynamics more accurately than classical numerical solvers at the comparably low resolutions across various scientiﬁcally relevant metrics.
Our model is trained end-to-end from data and is capable of learning a range of challenging chaotic and turbulent dynamics at low resolution, including trajectories generated by the state-of-the-art
Athena++ engine.
We show that our simpler, general-purpose architecture outperforms various more specialized, turbulence-speciﬁc architectures from the learned turbulence simulation literature.
In general, we see that learned simulators yield unstable trajectories; however, we show that tuning training noise and temporal downsampling solves this problem.
We also ﬁnd that while generalization beyond the training distribution is a challenge for learned models, training noise, added loss constraints, and dataset augmentation can help.
Broadly, we conclude that our learned simulator outperforms traditional solvers run on coarser grids, and emphasize that simple design choices can offer stability and robust generalization.
I NTRODUCTION
Turbulent ﬂuid dynamics are ubiquitous throughout the natural world, and many scientiﬁc and engineering disciplines depend on high quality turbulence simulations.
Examples range from design problems in aeronautics (Rhie, 1983) and medicine (Sallam & Hwang, 1984) to scientiﬁc problems of understanding molecular dynamics (Smith, 2015), weather (Beljaars, 2003), and vast galactic systems (Canuto & Christensen-Dalsgaard, 1998).
Predicting turbulent dynamics is challenging, and ﬁnding a general solution to the governing equations of ﬂuid dynamics, the Navier-Stokes Equation, is a famously open problem in Mathematical Physics (Jaffe, 2006).
The challenge is due in large part to the fact that turbulence is chaotic: small-scale changes in the initial conditions can lead to a large difference in the outcome.
Nonetheless, over the past several decades, high-quality numerical solvers have been engineered that integrate the governing partial differential equations (PDEs) and can maintain accuracy over long integration periods even for complex dynamics.
These solvers often must operate over high-resolution grids and therefore require substantial computational resources.
Otherwise, high-frequency information is lost due to hard-to-model “numerical viscosity”, and the simulated dynamics can further diverge from the underlying equations due to the chaotic dynamics.
Learned simulation represents a promising avenue for efﬁcient ﬂuid simulation because learned models can potentially adapt to capture the large-scale dynamics on coarse grids.
Learned simulators vary in the extent to which they incorporate components from classical solvers, like those that learn closures and subgrid discretization (Duraisamy et al., 2019; Freund et al., 2019; Pathak et al., 2020;
Um et al., 2021; Kochkov et al., 2021), versus adopting a pure Machine Learning (ML) approach (Kim et al., 2018; Lusch et al., 2018; Sanchez-Gonzalez et al., 2020;
Wang et al., 2020; Li et al., 2020c;a; Pfaff et al., 2021).
Advantages of the pure ML route are that the same model can be used without specialized knowledge of the domain and that they do not involve the challenge of interfacing with PDE solvers.
However, fully learned simulators often work well only when conditions are similar to the training distribution, which can limit their stability over time and their generalization capabilities.
Our primary contribution is to demonstrate that fully learned simulators can learn to predict turbulent dynamics on coarse grids more accurately than classical solvers.
This is true over a range of performance metrics, but the learned simulators particularly excel in their ability to preserve high frequency information.
While learned simulators can be prohibitively unstable, we show that the simple practice of tuning training noise and the temporal downsampling factor can reliably produce stable rollouts.
We introduce a simple model that combines neural networks architectures for grids (Dilated CNNs (Yu & Koltun, 2016),
ResNets (He et al., 2015)) with recent advancements on general purpose learned simulation on graphs (encode-process-decode models (Sanchez-Gonzalez et al., 2020; Pfaff et al., 2021) trained with noise (Sanchez-Gonzalez et al., 2018; 2020; Pfaff et al., 2021)).
We implement this as well as a set of learned models from the literature that have been proposed for modeling turbulence, and show that Dilated ResNet (Dil-ResNet) outperforms the often more complicated, more heavily parameterized models.
Rather than looking at a single problem, we evaluate our models over a variety of challenging turbulent and chaotic domains, showing that
Dil-ResNet is generally able to capture a variety of turbulent dynamics out-of-the-box.
This includes two environments generated by the new Athena++ solver(Stone et al., 2020), a state of the art (magneto-)hydrodynamics simulator used for astrophysics, where extremely high-resolution and large-scale simulations are often required.
Finally, we evaluate the ability of the learned models to generalize to initial conditions, rollout durations, and environment sizes outside of the training distribution.
Generalizing out of distribution is a known challenge for learned models, and we document where our models fail and what model and training choices can improve generalization.
R ELATED W ORK
Turbulence simulation is a crucial step in problems across disciplines.
These include forecasting problems, like predicting the spread of a wildﬁre (Leonardi et al., 2019) or the evolution of a hurricane (Zhang et al., 2020), engineering challenges, like designing more aerodynamic machinery (Rhie, 1983), and scientiﬁc questions, like understanding the physics of insect ﬂight (Dickinson et al., 1999) or dynamics of celestial objects (Fielding et al., 2018).
Across applications, a variety of statistical measures are relevant for quantifying physical variables under turbulence (Velocity Autocorrelation, variable histograms, Cooling, Energy spectra).
The primary challenge for simulating turbulence is accurately capturing small- and large-scale ﬂow dynamics, and their interplay, simultaneously.
Traditional approaches to computational ﬂuid dynamics integrate the governing PDEs describing the ﬂuid ﬂow, given initial and boundary conditions.
State-of-the-art turbulence simulators, such as Athena++ , are built upon efﬁcient algorithms that prioritize the preservation of high-order statistics.
Such simulators are used frequently in scientiﬁc applications to numerically compare model predictions with experimental observations, but can be extremely–often prohibitively–expensive to run.
Thus, there is signiﬁcant practical interest in developing more efﬁcient simulators.
Deep learning models for predicting physical dynamics have traditionally been used for modeling low-dimensional dynamical systems with graph-like (Battaglia et al., 2016; Sanchez-Gonzalez et al., 2018) and grid-like speciﬁcations (Raissi et al., 2017; Chen et al., 2018), but there is a growing body of work for ﬁne-grained modeling of complex dynamics, as in particle-based liquids and soft-body physics (Li et al., 2019; Ummenhofer et al., 2020; Sanchez-Gonzalez et al., 2020).
Learned models are often used in computer graphics to speed up (Wiewel et al., 2019; Um et al., 2018; Ladický et al., 2015; Holden et al., 2019) or super-resolve (Xie et al., 2018) simulations.
A number of papers predict the steady-state ﬂow around 2D airfoil wings, using CNN (Thuerey et al., 2020; Bhatnagar et al., 2019) or GCN (Belbute-Peres et al., 2020) architectures, or the transient dynamics of various low
Reynolds number ﬂows using mesh representations (Pfaff et al., 2021).
In recent years, a number of models have been introduced speciﬁcally for learning turbulent dynamics.
Many of these use neural networks to model some part of an otherwise classical, theoretically motivated numerical integration setup (Duraisamy et al., 2019; Freund et al., 2019; Pathak et al., 2020;
Um et al., 2021; Kochkov et al., 2021;
Tompson et al., 2016), PDE formulation (Raissi et al., 2017; Champion et al., 2019), or embedding space (Lusch et al., 2018).
Other models incorporate
Learned simulator ρ P v x v y v z 2 (cid:31) P r e d i c t i o n 1D KS Equation (KS-1D)
Time 2 (cid:31)
G r o u n d
T r u t h  3D Compressible Decaying Turbulence (CompDecay-3D) 3D
Mixing Layer (CoolMixLayer-3D)
Ground truth Prediction 2D Incompressible Decaying Turbulence (Incomp-2D)
b c d e a
Initial
Ground truth Prediction Initial
GT Pred Init
CNN enc dCNN×N ...
Encoder
Decoder Processor x y f e a t u r e s
CNN dec ... some information about turbulence simulation into the architecture, but like ours, are learned endto-end in a supervised way.
Wang et al.
(Wang et al., 2020) implemented a fully learned, neural network “Turbulent Flow Network” (TF-Net) that applies physically inspired decomposition with learned parameters and processes the components with U-Nets for ﬂuid simulation.
Li and colleagues introduced “Fourier Neural Operators” for turbulence simulation (Li et al., 2020a), which apply neural networks that convolution parameters in the Fourier domain.
This model is part of a broader family of work (Portwood et al., 2019; Li et al., 2020b) which combines basis transforms (Fourier,
Multipole, etc) with neural operators that process the transformed input.
We implement models from these papers as baselines to compare to our learned simulator.
M ODEL 3. L EARNED S IMULATION F RAMEWORK ON G RIDS
We train a simulation model s : X → X that maps X t to X t +∆ t , where X is a space of n + dimensional tensors which represent the state variables for each point on a grid with n spatial dimensions and one feature dimension at time t .
The true physical dynamics applied over K time steps yields a trajectory of states which we denote ( X t 0 , ..., X t K ) .
The “rollout” trajectory produced by iteratively applying s for K steps is denoted ( ˜ X t 0 , ˜ X t 1 , ..., ˜ X t K ) , where ˜ X t 0 =
X t 0 represents initial conditions given as input.
Our learnable simulators (Fig.
1a) implement s using neural networks denoted NN (with weights θ ), and predict the next state as the difference from the current state, ˜ X t +∆ t =
s
( ˜ X t ; θ ) =
˜ X t + NN ( ˜ X t ; θ ) .
3. M ODELS
In this work we study “fully learned“ models, which learn the physical dynamics end-to-end without any hard-coded solver components.
We implemented a novel Dilated ResNet model, as well as several baseline models.
Detailed model parameters can be found in the Appendix.
Dilated ResNet Encode-Process-Decode (Dil-ResNet)
We designed this architecture to combine the encode-process-decode paradigm (Sanchez-Gonzalez et al., 2018; 2020) with dilated convolutional networks (Fig. 1a).
The encoder and decoder each consist of single linear convolutional layer.
The processor consists of N = 4 dilated CNN blocks (dCNN n ) connected in series with residual connections (Yu & Koltun, 2016;
He et al., 2015).
Each block consists of 7 dilated CNN layers with dilation rates of (1, 2, 4, 8, 4, 2, 1).
A dilation rate of N, indicates each pixel is convolved with pixels that are multiples of N pixels away (N=1 reduces to regular convolution).
This model is not specialized for turbulence modeling per se , and its components are general-purpose tools for improving performance in CNNs.
Residual connections help avoid vanishing gradients, and dilations allow long-range communication while preserving local structure.
All individual CNNs use a kernel size of 3 and have 48 output channels (except the decoder, which has an output channel for each feature).
Each individual CNN layer in the processor is immediately followed by a rectiﬁed linear unit (ReLU) activation function.
The Encoder CNN and the Decoder CNNs do not use activations.
Encode-U-Net-Decode (U-Net)
We designed this architecture as a simpliﬁcation of the TF-Net baseline explained below.
It uses a simple linear CNN encoder and decoder (same as Dil-ResNet) and a standard U-Net block (Ronneberger et al., 2015).
As in (Ronneberger et al., 2015), each CNN stack is made of 3 CNNs, followed by activations.
Other hyperparameters of the U-Net also matched those used in (Ronneberger et al., 2015).
Turbulent Flow Net (TF-Net) (Wang et al., 2020) TF-Net uses a domain-speciﬁc variation of UNets to model turbulence, along with other architectural elements inspired by RANS-LES Hybrid
Coupling in the encoder.
Hyperparameters matched those used in (Ronneberger et al., 2015; Wang et al., 2020).
Constrained Turbulent Flow Net (Con-TF-Net) (Wang et al., 2020)
Wang et al.
also explored a constraint term in the loss which penalizes the divergence, which was constrained to be zero everywhere for the PDE in their incompressible turbulence experiments.
We implemented a similar model, adapting the constrained function to match constraints or preserved quantities in our datasets.
For KS-1D, the mean is constrained to be 0; for I NCOMP -2D, the divergence is constrained to be 0, and for C OMP D ECAY -3D the total energy is constrained to be the same as input’s total energy.
Constrained Dilated ResNet (Con-Dil-ResNet)
This model is Dil-ResNet (described above) with the additional constraint terms from (Con-TF-Net) added to the loss function.
This allows us to examine the effect of the loss term separately from the other aspects of the TF-Net architecture.
Fourier Neural Operator (FNO) (Li et al., 2020a)
This model uses the neural operator formalism introduced in earlier work (Li et al., 2020c;b), applied in the Fourier domain.
We implemented FNO for 1D and 2D spatial domains as in (Li et al., 2020a)
Hyperparameters are from (Li et al., 2020a).
CNN padding We use periodic padding for spatial axes to implement periodic boundary conditions.
For the C OOL M IX L AYER -3D, which has a ﬁxed boundary condition along the vertical axis, we mask out the value at the boundary when computing training loss and set it to the ground truth at each step of the rollout.
We also augment the input state with a binary feature to distinguish boundary pixels.
3. T RAINING Loss
All models 1 are trained to predict ∆ X .
We use a mean square error loss (cid:96) ( X t , X t +∆ t ) =
MSE( NN ( X t ; θ ) , ∆ X ) to optimize parameters θ .
Input X t and target ∆ X
= X t +∆ t − X t features are normalized to be zero-mean and unit variance, which seemed to improve training speed.
The additional mean square error loss components for constraint preservation in Con-Dil-ResNet and
Con-Dil-ResNet are added to the main loss with a relative weight of 1.
Training noise Optionally, we trained with Gaussian random noise with ﬁxed standard deviation σ added to the input X t and subtracted from the target ∆ X , as this has been shown to improve stability of rollouts and prevent error accumulation by training the model to correct for small errors (Sanchez-Gonzalez et al., 2018; 2020; Pfaff et al., 2021).
See Appendix for more details.
E XPERIMENTAL T URBULENT D OMAINS
We use four PDEs exhibiting chaotic or turbulent dynamics representative of a wide range of problems across science and engineering.
(Fig.
1b-e).
These environments range from the classic problems (KS Equations, Incompressible Turbulence) to more challenging environments that incorporate compressibility and non-uniform initial conditions.
The training data is generated at high resolution
Spatial and temporal downsampling factors and additional details about the environments are shown in the Appendix.
1D Kuramoto-Sivashinsky Equation (KS-1D): A PDE that generates unstable, chaotic dynamics in 1D, solved using Fourier spectral method (Kuramoto, 1978; Sivashinsky, 1977).
The state consists of a scalar velocity v per 1D grid point.
While KS-1D is not technically turbulent, it is a well-studied chaotic equation that is useful for assessing the ability of our learned models to capture dynamics that are highly unstable and nonlinear.
2D Incompressible Decaying Turbulence (I NCOMP -2D):
Fluid ﬂow modeled by Navier-Stokes in which small-scale eddies decay into large-scale structures due to the inverse energy cascade, solved using direct numerical simulation (Kochkov et al., 2021).
The state consists of a vector velocity v x , v y per 2D grid point.
These simulations are relevant to open questions in atmospheric dynamics (Boffetta & Ecke, 2012a).
3D Compressible Decaying Turbulence (C OMP D ECAY -3D): Decaying transonic turbulent ﬂow under Navier-Stokes, assuming adiabatic equation of state with adiabatic index γ = 5 / 3 .
The state consists of scalar density ρ , vector velocity v x , v y , v z , and scalar pressure P per 3D grid point.
Energy at each grid point is E = 1 / 2 ρv 2 +3 / 2 P .
Simulations were carried out with Athena++ (Stone et al., 2020), a state-of-the-art simulator used in astrophysics.
All real ﬂuids are somewhat compressible, and this is particularly non-negligible when modeling gases (Pope, 2000).
In astrophysics problems, for which the Athena++ was developed, understanding the properties of these ﬂows plays a crucial role in regulating planet, star, black hole, and galaxy formation (Brandenburg & Nordlund, 2011).
3D Compressible Turbulence with Radiative Cooling Mixing Layer (C OOL M IX L AYER -3D):
Turbulent mixing resulting from the Kelvin-Helmholtz instability, caused by velocity differences across the interface between ﬂuids of different densities, solved with Athena++ .
The state is speciﬁed in the same way as in C OMP D ECAY -3D.
Mixing involves strong cooling, leading to net ﬂow from the low-density phase into the mixing layer.
This environment was additionally challenging because of its open boundaries and non-uniform initial conditions.
This process is common in atmosheric ﬂows as well as many aspects of galaxy formation (Fielding et al., 2020).
R ESULTS
We evaluate the learned simulator models in terms of stability, performance, efﬁciency, and generalization.
Unlike numerical solvers, which are PDE speciﬁc, learned simulator models can be designed to be reusable across domains.
In Figure 1, we show that the same general-purpose architecture (Dil-ResNet) and loss, trained on each of the domains, learns to capture a range of qualitatively diverse turbulent dynamics across domains (Fig.
1b-e, videos: sites.google.com/view/learned-turbulencesimulators).
All results are evaluated on held-out test trajectories sampled from the same distribution used for training (except for generalization sections).
All units (except time) are normalized and therefore dimensionless.
Spatial coarsening Numerical solvers are known to lose information, particularly high-frequency information, when applied to coarse grids.
We compared our Dil-ResNet learned simulator to
Athena++ run on spatially coarsened grids (see Section 5 for comparison with learned models).
We chose the energy ﬁeld E = 12 ρv 2 + 32 P as the quantitative metric of 3D turbulence for the main text, as it summarizes performance on all state variables.
We show results for Athena++ run at resolutions of 32 3 (the same resolution as the learned models) and at 64 3 (higher resolution than the learned models, lower than the ground truth).
The “ground truth” data is from Athena++ run at 128 3 .
All grids are downsampled to 32 3 for comparison.
We compared models in terms of the RMSE
[Figure]
E n e r g y F i e l d R M S
[Figure]
L o g E n e r g y
S
p e c t r u m
R M S
[Figure]
n
[Figure]
r
[Figure]
y
S
p e c t r u m
R M
[Figure]
of the predicted Energy Field and the Log Energy Spectrum (See Appendix for deﬁnition).
We used the energy ﬁeld as the main variable as it is a combination of all ﬁve state variables, but see Videos and Appendix for per-variable results.
The learned simulator outperforms the comparable resolution Athena++ 32 3 across a variety of metrics, despite having no built-in specializations for turbulent dynamics.
In Figure 2a-b, we show
RMSE for the Energy Field for t < 1 , which corresponds to the initial phase of the turbulence decay seen during training (white window in Fig 2b).
Energy Field is given as E = 12 ρv
+ 32 P ,
and it implicitly summarizes performance on all state variables (other metrics are shown in the Appendix).
The learned simulator outperforms both the same- and higher-resolution Athena++ rollouts in terms of the Log Energy Spectrum (Fig. 2c,d), as the Athena++ simulators lose high frequency components that the learned simulators preserve (Figure 2e-i).
Log Energy Spectrum is computed by (1) taking the 3-D Fourier transform, (2) computing the amplitude of each component, (3) taking the histogram of the 3-D Fourier transform over frequency amplitude ( (cid:113) k 2 x
+ k 2 y + k 2 z ) to get a 1-D PSD, and (4) then taking the log.
We looked at a range of other physically relevant metrics and ﬁnd that the learned simulator outperforms the comparably coarse 32 3 Athena++ but not 64 3 Athena++ on predicting feature histograms, the phase histograms (pressure v. density and entropy v. pressure), and mean squared error for each feature.
The learned simulator outperformed both {+32 3 and 
+}
Athena++ simulators on higher order velocity autocorrelations, as well as spectrum error.
Since different scientiﬁc questions rely on different metrics, the tradeoffs of learned versus physics based simulators may vary across applications.
These metrics are deﬁned in the Appendix.
Stability and training noise While scientiﬁc simulators are typically designed to be stable over time, a common failure mode in learned models is that small errors can accumulate over rollouts and lead to a domain shift.
One reason for this is that, as the model is fed its most recent prediction back in as input for predicting future steps, its distribution of input states begins to deviate from
[Figure]
that experienced at training, where it fails to generalize and can make arbitrarily poor predictions.
We found that adding Gaussian noise σ = 0 .
01 to the inputs X t during training led to less accurate one-step predictions (Fig. 3a), but more stable trajectories (Fig. 3b). This is of particular importance for models that take a very large number of small steps.
This is presumably because the training distribution has broader support and the model is optimized to map deviant inputs back to the training distribution.
Temporal coarsening An advantage of learned simulators is that they can exploit a much larger step size than the numerical solver, as they can discover efﬁcient updates that capture relevant dynamics on larger timescale.
This allows for faster simulation.
See Videos and Appendix for qualitative examples of Dil-ResNet trained on a large range of exponentially increasing coarse timesteps in KS-1D, I NCOMP -2D and C OMP D ECAY -3D, which the model can adapt to (note that a separate model is trained for each ∆ t ).
Quantitatively, though the one-step error is at its lowest when using smaller time steps (Fig. 3c), the rollout error has an optimal time step at around 0.032 (Fig.
3d).
This demonstrates the tradeoff between large and small time steps.
Large time steps ( > 0 . 032 )
cause predicting the next state to become more challenging.
Small time steps ( < 0 . 008 ), which require more simulator steps for the same duration, often yield unstable models because they provide more opportunities for errors to accumulate (Fig. 3d)
(e.g. for some ∆ t , 1 s / ∆ t steps are required).
However, they can still be stabilized to some extent with training noise (Fig. 3f).
[Figure]
the various learned models had comparable performance, though the Dilated ResNets (Dil-ResNet, Con-Dil-ResNet) consistently have the lowest error.
In KS-1D (a), the noise harmed performance; in I NCOMP -2D (b) the noise particularly beneﬁts the FNO rollouts; in C OMP D ECAY -3D (c) the noise mainly stabilized rollouts.
Comparison across learned models In Figure 4 we compare the learned models in terms of the performance across the different primary turbulence datasets.
We ﬁnd that Dilated ResNets (Dil-ResNet, Con-Dil-ResNet) are the top performing models, numerically outperforming more specialized architectures.
However, we ﬁnd that all learned models (ours and the baselines) produce qualitatively good rollouts compared to Athena++ 32 3 , suggesting learned simulators are generally a good approach to turbulence simulation on coarse grids (see Videos).
We also ﬁnd training noise improves stability across learned models.
All models shown in Figure 4 are evaluated such that the ﬁrst predicted step is the 7th step of the ground truth trajectory.
This is to allow comparison with
TF-Net, which takes C = 6 context states as input.
We show in further detail the predictions of the different models in the Appendix.
Running time Learned simulators can accelerate simulations not only by coarsening in time and space, but also by running off-the-shelf on specialized GPU hardware, unlike e.g., Athena++ which is currently CPU speciﬁc.
By downsampling in time ( ∆ t = 0 .
→ 0 . 032 ) and space ( 128 3 → 32 3 ), the Dil-ResNet model trained on C OMP D ECAY -3D running on a single GPU speeds up the wall-clock simulation time by 1000x compared to Athena++ running on an 8-core CPU (further details in Appendix).
GPU solvers can be developed for many applications and can greatly accelerate simulation; however, these can require a lot of expertise and time to develop, and will still require high spatial and temporal resolution grids.
Constraint satisfaction and preserved quantities Traditional solvers often implement constraints to preserve conserved quantities.
Learned simulators do not necessarily learn such constraints and may fail to capture these (see Appendix).
We found that training with noise prevents this in some cases: it helps Dil-ResNet keep the mean value of the KS-1D velocity, and
the I NCOMP -2D
divergence bounded near zero.
However, in C OMP D ECAY -3D using noise does not seem to help preventing the model from predicting drifts in total mass, vector momentum and energy, even though they are ﬁxed across the whole dataset.
We speculate this is because these preserved quantities only account fordegrees of freedom out of a 3D state that is made of 163840 ( = 5 · 32 3 ) output variables.
Generalization to longer rollouts Dil-ResNet can stably unroll KS-1D for much longer than the trajectories seen in training.
This is expected, as the distribution of states is stationary in time (see
Appendix).
By contrast, the distribution of states in C OMP D ECAY -3D is not stationary in time, as the ﬂow decays progressively into smaller and smaller structures.
We can use this to construct a more challenging generalization setting by training Dil-ResNet on early stages of the decay only ( t < 1 ), and testing on longer durations in which the turbulence develops further than observed during training.
( 1 < t < 4 in Fig.
2b, d).
We ﬁnd Dil-ResNet remains competitive with 32 3 Athena++ for turbulence that has decayed for up to twice longer t < 2 than observed during training.
Generalization to different initial conditions We varied the ratio of solenoidal to compressive components in the initial velocity ﬁeld in C OMP D ECAY -3D
(Fig.
5a)
in order to test the models’ ability to generalize to novel initial conditions.
Compared to Athena++ at 32 3 and 64 3 , we found that the models generalize well to more solenoidal but not to more compressive initial conditions, possibly due to faster turbulence decay under compressive conditions 4 .
We ﬁnd that Dil-ResNet (with or without noise) has particular difﬁculty generalizing to the more compressive components.
This is somewhat ameliorated with the added constraint to the loss function (Con-Dilated-ResNet).
Results for other models are similar (see Appendix).
Generalization to larger boxes We tested the generalization capability of Dil-ResNet to
C OOL M IX L AYER -3D boxes with different size L in the x and y direction (the length along the vertical z axis perpendicular to the mixing layer remained constant).
We quantiﬁed the predicted cooling velocity (small laminar ﬂow in the low-density phase perpendicular to the interface, see Appendix for more details) which is of scientiﬁc relevance and known to depend on the box size.
Speciﬁcally, for
C OOL M IX L AYER -3D, the dynamics undergo a transition when the box width increases relative to the cooling length,which is difﬁcult to study because simulations for these widths essentially prohibitive.
We found that unless Dil-ResNet is trained on a range of sizes, the predicted cooling velocity does not follow the right trend, and even when trained on a range of sizes, generalization beyond the training
[Figure]
range is not reliable across seeds (Fig. 5b).
We speculate that achieving this form of generalization would require either stronger inductive biases, more sophisticated dataset engineering, or both.
C ONCLUSIONS
Learned simulation techniques have been advancing steadily in recent years.
Our results demonstrate that learned simulators can outperform comparably coarse solvers in challenging turbulence settings on a range of scientiﬁcally relevant statistical measures, suggesting that state-of-the-art learned simulators can be useful for efﬁcient, accurate simulation even in chaotic domains.
Learned simulators can perform particularly well in terms of preserving high frequency information.
The Dilated ResNet model had lower log Energy Spectrum RMSE than both 32 3 Athena++ and the higher resolutions 64 3 Athena++ , and the RMSE remained lower even beyond the duration seen during training.
In our experiments, we found that using training noise and temporal downsampling improved the stability and accuracy of extended rollouts.
Out-of-distribution generalization is a known challenge for learned simulation (Duraisamy et al., 2019) and for deep learning methods generally, and we ﬁnd that the coarse Athena++ models outperform the learned models on more extreme generalization conditions.
However, certain training choices can make learned models more robust.
In particular, physically-informed regularization (e.g., Con-Dil-ResNet and Con-TF-Net) improved generalization to different initial conditions on
C OMP D ECAY -3D, and for C OOL M IX L AYER -3D, augmenting the training data with a larger range of box sizes improved generalization performance on an even larger range.
Future work should explore methods from other areas of deep learning for improving adversarial robustness and domain generalization, such as feature denoising (Xie et al., 2019) and discriminative objectives (Tzeng et al., 2017; Lample et al., 2017).
More broadly, we conclude that a key potential role for learned simulators in the near term is distilling expensive, high-resolution engineered simulators into low-resolution algorithms to strike better performance-versus-efﬁciency tradeoffs.
Our approach can be applied effectively to diverse and challenging environments, which suggests that even generic CNN-based models may apply well to low-resolution simulation on a wide range of physical domains which can be represented as grids.
A CKNOWLEDGMENTS
We would like to thank Meire Fortunato, Matthew Grimes and Lyuba Chumakova for helpful discussions and comments on the manuscript.
A A PPENDIX
[Table]
[Table] [Table]
Spatial downsample
[Table]
[Table]
(64, 128, 256, 512, (64x3, 128x3, 256, 128, 64) 256x3, 512x 256, 128, 64)
[Table] [Table]
[Table]
B A DDITIONAL M ODEL D ETAILS
Encode-U-Net-Decode (U-Net)
We designed this architecture as a simpliﬁcation of the TF-Net baseline explained below.
It uses a simple linear CNN encoder and decoder (same as Dil-ResNet) and a standard U-Net block (Ronneberger et al., 2015).
The standard U-Net block consists of 7 layers.
The ﬁrst 3 layers each consist of a CNN stack followed by a factor-of-2 downsampling via striding the input.
The fourth layer of the U-Net is another CNN stack, which at this point operates over a grid that is 8 times smaller than the input along each spatial axis.
The last 3 layers each consist of a factor-of-2 upsampling using nearest neighbors, concatenation with the output of one of the ﬁrst three layers (in reverse), and another CNN stack.
This downsampling/upsampling process allows long-range communication similar to the Dilated ResNet, while also preserving local structure.
As in (Ronneberger et al., 2015), each CNN stack is made of 3 CNNs, followed by activations.
Other hyperparameters of the U-Net also matched those used in (Ronneberger et al., 2015).
Turbulent Flow Net (TF-Net) (Wang et al., 2020) TF-Net uses a domain-speciﬁc variation of
U-Nets to model turbulence, along with other architectural domain-speciﬁc elements inspired by
RANS-LES Hybrid Coupling in the encoder.
The model implements a custom encoder by taking a sequence of the most recent C = 6 states as inputs, and uses this to build three separate input arrays, by applying speciﬁc combinations of learned spatial ﬁlters, learned temporal ﬁlters, and differences of those.
Each of these components are processed by a separate U-Net encoder (ﬁrst 4 layers of the
U-Net), then the latent outputs of each layer are concatenated and fed into a shared U-Net decoder (ﬁnal 3 layers).
Hyperparameters also matched those used in (Ronneberger et al., 2015; Wang et al., 2020).
Constrained Turbulent Flow Net (Con-TF-Net) (Wang et al., 2020)
Wang et al.
also explored a constraint term in the loss which penalizes the divergence, which was constrained to be zero everywhere for the PDE in their incompressible turbulence experiments.
We implemented a similar model, adapting the constrained function to match either constraints or preserved quantities in our datasets.
For KS-1D, the mean is constrained to be 0; for I NCOMP -2D, the divergence is constrained to be 0, and for C OMP D ECAY -3D the total energy is constrained to be the same as the total energy of the input.
Constrained Dilated ResNet (Con-Dil-ResNet)
This model is Dil-ResNet (described above) with the additional constraint terms from (Con-TF-Net) added to the loss function.
This allows us to examine the effect of the loss term separately from the other aspects of the TF-Net architecture.
Fourier Neural Operator (FNO) (Li et al., 2020a)
This model uses the neural operator formalism introduced in earlier work (Li et al., 2020c;b), applied in the Fourier domain.
This model applies matrix multiplications in the Fourier-transformed grid with complex weights learned independently for each component, as well as linear updates The model combines pixel-wise learned linear embeddings in the spatial domain The model combines learned linear layers, applied to each point of the spatial grid separately, with matrix multiplication on the Fourier-transformed grid using learned weights.
A unique matrix of weights is learned for point in the Fourier grid.
Low-pass ﬁltering is optionally implemented by truncating all but K modes along each dimension in the Fourier-transformed grid.
The authors explored using a version of the model that takes a stack of past states as input.
Since the state is fully-Markov, we opted for the version of the model that takes a single state as input as it is more comparable to the majority of the other models.
Unlike previous models, which rely on spatially local convolutions and depth for long-range communication, this model applies long-range communication at every step in the Fourier domain.
We implemented this model for 1D and 2D spatial domains as in (Li et al., 2020a).
We found difﬁculties running this model on the 3D datasets due to compute and memory constraints.
We speculate this is due to time and space complexity of the Fourier transform (when taking into account the backwards pass) which would probably require a custom gradient re-materialization strategy.
This may be speciﬁc to TensorFlow which required us to reimplement the non-differentiable tf.signal.rfft3d as a sequence of tf.signal.rfft2d and tf.signal.fft .
B. A DDITIONAL DATASET DETAILS B.1. 1D K URAMOTO -S IVASHINSKY (KS) E QUATION
This is a well-studied 1D PDE that generates unstable, chaotic dynamics in 1 dimension (Kuramoto, 1978; Sivashinsky, 1977) with periodic boundaries.
The ground truth simulations are computed using the Fourier Spectral Method.
Initial condition was set to cos ( w 1 x + φ 1 )(1+sin ( w 2 x + φ 2 )) , where x ranges from 0 to 2 π , φ i are sampled uniformly from [0 , 2 π ) , and w 1 are integers sampled uniformly from [1 , 12) .
We perform a warmup from this initial condition for 75 simulation time units.
B.1. 2D I NCOMPRESSIBLE D ECAYING T URBULENCE
This models ﬂuid ﬂow described by the Navier-Stokes equations in which small-scale eddies decay into the large-scale structures due to the inverse energy cascade (Boffetta & Ecke, 2012b).
The underlying simulations were performed by solving the incompressible Navier-Stokes equations using a direct numerical simulation (DNS) ﬁnite-volume solver.
Boundaries along both dimensions are periodic, and the initial conditions consist of random velocity ﬁelds with small-scale variation.
Initial conditions for different trajectories were obtained by sampling a high-resolution velocity ﬁeld from a log-normal distribution of amplitude 1 and wavenumber k = 10 .
We perform a warmup forsimulation time units, this is done to discard the transient ﬂow that is heavily inﬂuenced by the underlying numerical scheme.
[Table]
Incompressible
Compressible
Radiative Cooling
Decaying Decaying
Mixing Layer
[Table]
DNS
Athena++ Athena++
[Table]
[Table]
v v x , v y ρ, v x , v y , v z , P ρ, v x , v y , v z , P
[Table]
n
[Table]
a
[Table] [Table]
[Table]
up refers to the initial transient from initial conditions which is inﬂuenced by the underlying numerical scheme and discarded from evaluation and training.
Figures are dimensionless.
[Table]
B.1. 3D C OMPRESSIBLE D ECAYING T URBULENCE
This models decaying transonic Navier-Stokes turbulent ﬂow in a 3D cubic box with periodic boundary conditions (e.g., Orszag & Patterson, 1972).
These simulations adopt an adiabatic equation of state with a constant adiabatic index γ = 5 / 3 .
Simulations were carried out with Athena++ (Stone et al., 2020).
The initial turbulence is driven on scales ≥ L , the size of the box.
The turbulent driving pattern in the initial condition is split into its compressive and solenoidal components using a Helmholtz decomposition.
The relative strength of these two components is varied from purely compressive to purely solenoidal.
The initial driving pattern results in a root-mean-squared velocity of √ 2 c s , where c 2 s =
( P/ρ ) is the sound speed of the ﬂuid.
The initial conditions are varied across trajectories by randomizing the phase of the spectral components, leading to different pattern in real space.
We perform a warmup for 0 .
05 simulation time units.
Compressible turbulence is ubiquitous across problems in science: all ﬂuids are somewhat compressible, and this is usually nonnegligible when modeling gases (Pope, 2000).
In astrophysics problems, for which the Athena++ was speciﬁcally developed, understanding the dynamics and properties of these ﬂows on small and large scales plays a crucial role in regulating planet, star, black hole, and galaxy formation (Brandenburg & Nordlund, 2011).
B.1. 3D C OMPRESSIBLE R ADIATIVE C OOLING M IXING L AYER D YNAMICS
These simulations model the interplay of radiative cooling and mixing that results from turbulence driven by the Kelvin-Helmholtz instability, which can arise when there is velocity difference across the interface between two ﬂuids of different densities.
The simulations are set up as a boundary problem initialized with a low-density ﬂuid ( ρ = 0 .
on the top half of the domain ( z > 0 ) moving in the positive x direction ( v x = 2 . 04 , v y = v z = 0 ), and a high-density ﬂuid ( ρ = 1 . ) moving in the negative x direction in the bottom half of the domain ( v x = − 2 .
04 , v y
= v z = 0 ).
The initial pressure is set to 1.
The boundary conditions are periodic in both x and y , and ﬁxed in z .
To break the symmetry, small perturbations to v z are added along the boundary.
These perturbations are changed across trajectories by randomizing the phase of the spectral components, leading to different patterns in real space.
We perform a warmup for 2 .
32 simulation time units.
Simulations were carried out with Athena++ (Stone et al., 2020).
This data presents a few unique challenges compared to the others.
First, it was the only domain with ﬁxed boundary conditions.
Second, because turbulent dynamics are limited to the vicinity of the mixing layer, and because the behavior of the ﬂuid above, at, and below the mixing layer is markedly different, there is less data representative of each type of ﬂuid dynamic behavior.
B. A DDITIONAL MODEL DETAILS
The learned simulator consisted of a CNN encoder, a dilated CNN processor with residual skip connections, and a CNN decoder.
The parameters of the model are listed in Table A.1.
CNN parameters All individual CNNs use a kernel size of 3 and have 48 output channels (except the decoder, which has an output channel for each feature).
Each individual CNN layer in the processor is immediately followed by a rectiﬁed linear unit (ReLU) activation function.
The Encoder
CNN and the Decoder CNNs do not use activations.
Dilated connections The dilation rate in the CNN ﬁlter introduces space between each element in the ﬁlter.
Whereas a standard CNN ﬁlter with kernel 3 would operate over 3 adjacent pixels, a dilated CNN ﬁlter with kernel 3 and dilation 2 will operate over 3 pixels spaced at intervals of 2 in each dimension, spanning a ﬁeld 5 pixels wide.
This increases the scale of the CNN kernels while preserving the number of parameters and the resolution.
CNN padding For each dimension in X with periodic boundary conditions, we implemented periodic padding.
For dimensions with a ﬁxed boundary condition, we forced the boundary to a value rather than letting the model predict it, and masked the loss so the model was not trained to predict boundary conditions.
The tensor was padded with repetitions of the boundary value.
We also augmented the input state with a feature that indicated ﬁxed-boundary versus non-boundary states using a one-hot vector, so the model could distinguish them.
B. A DDITIONAL TRAINING DETAILS
Training noise In some cases, we trained with Gaussian random noise with ﬁxed variance σ added to the input X t of the loss function.
Note that this not only affects the input to the neural network, but also slightly modiﬁes the target ∆ X = X t +∆ t − X t (and also impacts the variance used to normalize targets).
Loss At training time we sample pairs of input-output states (separated by the model time step) from the trajectories, and perform gradient updates based on a single step of the model.
We do not rollout the model during training.
Optimization We optimized the loss using an Adam optimizer.
We trained the models for up to 10M steps, with exponential learning rate decay annealed from 1 e − 4 to 1 e − 7 in the ﬁrst 6M steps.
Models usually reached convergence at around 5M steps.
Training took up to a week on an NVIDIA
V100 GPU.
Hyper-parameter optimization We did not perform exhaustive hyper-parameter optimization, except on the scale of the noise (which we scanned for each domain) and some informal tuning of the depth of the dilated blocks and latent size (to make sure the network had enough capacity for the most challenging domain).
Note that achieving optimal performance on test trajectories from the training distribution was not part of the scope of this work, and there are likely hyperparameters that would improve performance over our results.
Validation All of the research was performed by looking at performance on the validation sets.
The test set was completely held-out until ﬁnal evaluation prior to writing the paper.
Early stopping For the Mixing Layer Turbulence (C OOL M IX L AYER -3D), which was more prone to overﬁtting, we used early stopping based on the validation performance.
For all other models we simply evaluated the model as it was at the end of training.
B. A DDITIONAL RESULTS
Three instances of each model were trained 3 times using 3 different initialization seeds.
Bar plots indicate median performance and bars indicate min-max performance.
We chose the energy ﬁeld
E = 12 ρv 2 + 32 P as the quantitative metric of 3D turbulence for the main text, as it summarizes performance on all state variables.
Figures and videos for 3D environments (C OOL M IX L AYER -3D and C OMP D ECAY -3D) show a single slice of the 3D grid at y = 0 , displaying the x and z coordinates in the horizontal and vertical axis, respectively.
Downsampling in time Fig.
B.6 as well as Videos show models for KS-1D, I NCOMP -2D, and
C OMP D ECAY -3D models trained and working well on a wide range of time step sizes.
Note that for a fair comparison in Fig.
3d,f performance is averaged across only time steps that are predicted for all models (e.g. multiples of the largest time step, 256).
Downsampling in space All comparisons across spatial resolutions are always obtained by ﬁrst downsampling the data into a common 32 3 grid, using an approach that preserves mass, momentum, and energy (Same approach used by Athena++ ).
Fig.
B.2 (top row) shows the downsampling comparisons equivalent to Fig.
2g-i for each of the state variables independently.
Running time For 1 simulation time unit of C OMP D ECAY -3D, the CPU runtime for Athena++ with 8 CPU processors is ∼ 4s, ∼ 60s, and ∼ 1000s for 32 3 , 64 3 , and 128 3 resolutions, respectively (quartic scaling, as the time step is also scaled with the resolution to compensate for numerical viscosity).
In comparison, the learned model’s runtime is 1s on an NVIDIA V100 GPU, and 20-30s on a 8-core CPU.
Note that the learned model runs at reduced spatial and temporal resolution, but preserves the dynamics of the high resolution 128 3 Athena++ simulation.
Simulations in
Athena++ may be faster if implemented for GPUs.
However, because scientiﬁc simulators like
Athena++ are specialized, each simulator’s GPU implementation requires specialized engineering effort, whereas learned models can take advantage of methods designed more generally for deep learning.
Cooling velocity Having reliable models that generalize to larger boxes increases the applicability of learned models to scientiﬁc domains by enabling experiments in regions of parameter space that would otherwise be prohibitively expensive to simulate.
For example, for C OOL M IX L AYER3D, the dynamics undergo a transition when the box width increases relative to the cooling length ( v turb t cool ), which is difﬁcult to study because simulations for these widths essentially prohibitive.
Thus, we want to understand what affects the learned simulator’s ability to generalize to a range of box widths (the length L x and width L y of the input tensor X ) not previously seen in the training data.
For C OOL M IX L AYER -3D, we can look at the cooling velocity, the average inﬂowing velocity at the low density ﬂuid boundary that develops as a means to resupply the energy that has been radiated away in the mixing layer.
Cooling velocity is a useful metric because it is a scalar quantity that depends on the box width.
Furthermore, understanding how turbulent dynamics at a mixing layer pull heat from the surroundings is scientiﬁcally relevant for questions in astrophysics (Fielding et al., 2020).
We evaluate generalization to different box sizes ( L x = L y ∈ [0 . 25 , 0 .
(Fig. 5b).
We ﬁnd that the model trained on box length
L x = L y = 0 .
75 (orange) is not able to produce trajectories with the correct cooling velocity for other box sizes.
However, we also ﬁnd that augmenting the dataset with data from a range of box sizes ( L x = L y ∈ [0 . 5 , 0 . 75 , 1 . 0 , 1 . 25] ) improves accuracy of the cooling velocity estimate for the unseen box lengths, although the generalization outside the training domain remains imperfect and unreliable across seeds.
We speculate that achieving this form of generalization would require stronger inductive biases and/or more sophisticated dataset engineering.
Other Statistical Metrics
The choice of metrics can impact the utility of learned simulators for scientiﬁc applications.
Here we describe a few physically motivated metrics we use to evaluate the learned models.
We use the power spectral density, which we quantify for the total energy in Figure 2, to quantify how well the different models are preserving information at different resolutions.
Simulators at coarse resolution are known to lose high-frequency information, and we are speciﬁcally interested understanding how learned and physics-based simulators compare in terms of preserving detailed, high-frequency components.
To quantify the power spectral density, we n -D Fourier transform, take the magnitude at each point of the Fourier transformed signal to get the n -D PSD, and convert this to a 1D spectrum by taking the mean magnitude over points in the n -D frequency grid with the same frequency vector length to get a 1D power spectral density PSD ( k ) over a scalar frequency k = (cid:112)(cid:80) k 2 x i for spatial dimensions x i .
Frequencies with a vector of length > 1 (the length of the environment along 1 dimension) are cutoff, leaving only frequencies that fall within the sphere of radius 1.
We plot log PSD and compute spectrum error over these frequencies.
The spectra for different state variables are shown in Figure B.1 for the coarsened Athena++ models and DilResNet, and in Figure B.3 for the learned models.
The error over time for the spectra is shown in
Figures B.2 and B.2.
The probability distribution functions (PDFs) of the ﬂuid quantities provides an important basis for understanding turbulent ﬂows.
Due to the inherently chaotic nature of turbulence, statistical measures are often the most robust measures of their properties.
The shape and widths of PDFs are common metrics for characterizing turbulent ﬂows(Federrath & Klessen, 2012).
Moreover, joint PDFs of ﬂuid quantities (e.g., velocity-temperature, or pressure-entropy) are used to understand more complex ﬂow behavior and to characterize which ﬂow components are responsible for observed behavior (Fielding et al., 2020).
Joint distribution functions are often especially useful when certain physical processes only take place in speciﬁc region of phase space (e.g., combustion or radiative cooling).
PDFs for different state variables and joint distributions (entropy v. pressure, density v. pressure) are shown in
Figure B.1 for the coarsened Athena++ models and Dil-ResNet, and in Figure B.3 for the learned models.
The error over time for these statistics is shown in Figures B.2 and B.2.
Another measure of interest are spatial autocorrelation measures, particularly velocity autocorrelation.
Like the histograms, these capture statistical structure and provide a metric that is more invariant to the chaotic behavior of turbulence.
The p th order autocorrelation is closely related to the velocity structure function, S p ( (cid:96) )
=
E r , u , (cid:107) u (cid:107) =
1 [ (cid:107) v r + (cid:96) u − v r (cid:107) p ] , for velocity v , deﬁned at a position vector r (e.g., v r ), with unit vector u multiplied by a scalar (cid:96) .
Seminal self-similarity arguments that underpin turbulent theory lead to the prediction that S p ( (cid:96) )
∝ (cid:96) p/ 3 (Kolmogorov, 1941; Frisch, 1995).
Structure and autocorrelation functions, therefore, provide important measures on the ability of learned simulators to capture subtle, but scientiﬁcally essential properties of turbulent ﬂows.
First- and second-order autocorrelation functions for different state variables are shown in Figure B.1 for the coarsened Athena++ models and the Dil-ResNet model and Figure B.3 for all models trained on
C OMP D ECAY -3D.
The error for each of these functions over time are shown in Figures B.2 and B.4.
Model comparison Fig.
B.7 shows predicted rollouts for KS-1D of different learned models presented in Figure 4a.
Fig.
B.3 shows statistics of the ﬁnal frame of the predicted rollouts for different learned models presented in Figure 4.
The error of each of these statistics over time is shown in Fig.
B.4.
Videos show the rolled out predictions for different models for I NCOMP -2D and
C
OMP
D
ECAY -
3D.
KS-1D generalization Figs.
B.8 and B.9 show generalization to larger domains and longer trajectories.
While the learned simulations do not perfectly capture the ground truth, we see that qualitative features of turbulence are preserved across the rollout.
C OMP D ECAY -3D generalization Figs.
B.8 and B.9 show quantitative generalization performance of the Dil-ResNet model to larger domains and longer trajectories on KS-1D. While the learned simulations do not perfectly capture the ground truth, we see that qualitative features of turbulence are preserved across the rollout.
Figures B.2 and B.4 show different metrics for different variables across rollout (pink background indicates the region not seen during training).
Videos show the predicted rollouts when generalizing to different initial conditions and to longer rollouts.
[Figure] [Figure]
A u t o c o rr ( k
[Figure]
[Figure] [Figure]
[Figure]
u t
o c o rr ( k = 2 )
c o rr
[Figure]
i
[Figure]
[Figure]
[Figure]
π
[Figure]
T
[Figure]
h
Example
[Figure]
Example Example
Example 2 π P r e d i c t i o n
D o w n s a m p l e 2 π P r e d i c t i o n D o w n s a m p l e 2 π P r e d i c t i o n
D o w n s a m p l e 2 π P r e d i c t i o n D o w n s a m p l e 2 π P r e d i c t i o n
D
[Figure]
n s a m
[Figure]
l
[Figure] [Figure]
π P r e d i c t
i
[Figure]
n
[Figure]
i
[Figure]
i
[Figure] [Figure]
4 π
G r o un d
T r u t h
S p a t i a l c oo r d i n a t e
Example Example
Example Example
[Figure]
2 π 4 π
P r e d i t i o n S p a t i a l c oo r d i n a t e
[Figure]
[Figure]
π
E x a m p l e 0 G r o un d T r u t h 2 π
P r e d .
2 π
E x a m p l e 1 G r o un d T r u t h 2 π
P r e d .
2 π
E x a m p l e 2 G r o un d T r u t h 2 π
P r e d .
2 π
E x a m p l e 3 G r o un d T r u t h
[Figure] [Figure]
[Figure]
Added noise during training improves stability and therefore leads to lower error on longer rollouts.
Cutoff point for TF-Net with no noise has RMSE 2.99, and cutoff bar for Dil-ResNet goes to Inf (numerical error).
As shown in Figure 5, we see that adding the constraint to the loss function makes the model generalize more reliably to the more compressive regime, and that training noise promotes stability for longer rollouts..
Note that these results are numerically different from Figure 5 because rollouts were started from the sixth timestep to allow comparison with TF-Net, which takes as input a sequence of length C = 6 .