Extending Deep Learning Emulation Across Parameter Regimes to Assess Stochastically Driven Spontaneous Transition Events

Authors: ira j s shokar
peter h haynes
rich r kerswell
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-10 23:59:58

Abstract: Given the computational expense associated with simultaneous multi-task learn- ing, we leverage fine-tuning to generalise a transformer-based neural network emulating a stochastic dynamical system across a range of parameters. Fine-tuning a neural network with a dataset containing a set of parameter values yields a 40-fold reduction in required training size compared to training ab initio training for each new parameter. This facilitates rapid adaptation of the deep learning model, which can be used subsequently across a large range of the parameter space or tailored to a specific regime of study. We demonstrate the model’s ability to capture the relevant behaviour even at interpolated parameter values not seen during training. Applied to a well-researched zonal jet system, the speed-up provided by the deep learning model over numerical integration and the ability to sample from the probabilistic model makes uncertainty quantification in the form of statistical study of rare events in the physical system computationally feasible. Our code is available at https://github.com/Ira-Shokar/Stochastic-Transformer.

Main paper: 
 introduction: 
 Given the computational complexities inherent in turbulence modelling, investigations have delved into cost-effective alternatives to physics-based solvers. Machine learning has emerged as a promising direction, given its capacity to learn non-linear transformations and delegate a substantial portion of the requisite computation to offline training (Vinuesa & Brunton, 2022) .Most machine learning methods, however, face challenges in accurately generalising beyond their training distribution. An example would include emulation of a dynamical system, in which physical parameters governing the PDE (Partial differential equation) differ from those seen during training. One approach to overcome this is transfer learning (Pan & Yang, 2010) . In the case of PDEs, learned weights of a base trained network are frozen and additional layers are trained to generalise to the new tasks (Pellegrin et al., 2022; Goswami et al., 2022; Chen et al., 2021) . Alternatively, the transformer architecture (Vaswani et al., 2017) has been shown to perform extremely well in generalisation tasks without additional layers. This is due to the ability of the attention mechanism to conduct self-supervised learning during pre-training (Devlin et al., 2019) .In this work, we pre-train a transformer based network on a dataset comprising trajectories generated from a single parameter of a PDE, to learn a representative basis that captures the underlying dynamics of a stochastic dynamical system. This is followed by fine-tuning on a smaller dataset comprising trajectories from a range of parameter values, to generalise across parameter regimes. The neural network provides a resource-efficient emulator of the PDE. We show that autoregressive roll-outs remain indefinitely stable, while the neural network is capable of capturing the statistical properties across parameter regimes. 
methodology 2.1 physical system -quasi-geostrophic turbulence: 
 In the domain of earth system modeling, parameterisation schemes often integrate stochastic elements to approximate the impact of small-scale processes not explicitly resolved by the model, thus capturing their inherently chaotic nature (Berner et al., 2015) . In this study, we introduce a deep learning emulator spanning a parameter range of a well-studied idealised Geophysical Fluid Dynamics (GFD) system-beta-plane turbulence (Rhines, 1975; Galperin & Read, 2019) . This system serves as a useful simplest model for many aspects of flows in atmospheres and oceans. Modelled as a single-layer flow on a doubly-periodic domain, turbulence is generated via stochastic forcing, ξ(x, y, t), in the vorticity equation:∂ t ζ + u • ∇ζ + β∂ x ψ = ξ -µζ + ν n ∇ 2n ζ ∈ R Nx×Ny (1)with velocity field u(x, y, t) = (-∂ y ψ, ∂ x ψ) (note the sign convention used widely in GFD), stream function, ψ and the relative vorticityζ = ∂ x v -∂ y u.The parameter β represents the Rossby parameter, the rate of planetary rotation at a given latitude. Energy dissipation occurs through linear damping at a rate µ, while hyperviscosity, characterised by a coefficient ν and order n, removes vorticity at small scales. The forcing term ξ(x, y, t) takes the form of white-in-time noise injected into an annulus in wavenumber space. Simulations are run at a resolution N x = N y = 256 -see Appendix A.1 for further details.In turbulence modelling, a Reynolds decomposition is typically employed as an averaging procedure to separate coherent structures from fluctuations, here that takes the form of an eddy-mean decomposition in the zonal direction to separate the dynamics: u(x, y, t) = u(y, t) + u ′ (x, y, t) (where u(y, t) = 1 Lx Lx 0 u dx and L x = 2π is the domain size in x). This gives equations for the zonally averaged zonal velocity field U (y, t) = u(y, t) and the associated eddy fluctuation fields, ζ ′ (x, y, t):EQUATIONwhereEEN L = ∂ y (ζ ′ v ′ ) -∂ y (ζ ′ v ′ ) -∂ x (ζ ′ u ′ )denote the eddy-eddy nonlinear interaction terms. Here we see that the evolution of U depends on the nonlinear two-way interaction between the eddies and the mean flow, given by the Reynolds stress term (ζ ′ v ′ ). Forcing noise, ξ, is solely applied to the eddy fields, and therefore its effect on the mean flow is indirectly filtered through non-linear processes. The influence of β, akin to ξ, also manifests indirectly through the Reynolds stress term.As with most PDEs, varying the parameters of the governing equation gives rise to diverse phenomena. In this system, varying the parameter β gives rise to a range of range of dynamics, seen in the latitude-time plots in Figure 1 , in which not just the number of jets varies (maxima in U -yellow in colour), but also the dynamic tendencies, namely frequency of nucleation of a new jet, coalescence of two jets, and latitudinal jet translation.We train a deep learning emulator of equation 2 tailored to accommodate varying values of β. We employ the Stochastic Transformer (ST) (Shokar et al., 2023) , a transformer-based probabilistic model, that projects forward to predict an ensemble of plausible subsequent states { Ũt+1 }.Physical models are generally formulated as PDEs that are first-order in time, however given we are not operating with full physical states, but with the coarse grained variable U , in line with other work (Price et al., 2023) , we found it advantageous to condition on additional temporal context. The non-linear nature of equation 3 means that there is no closed form solution for ζ ′ and as such conditioning on additional temporal histories of U can be thought of as providing an approximate a closure for Reynolds stress term in equation 2. Consequently, the ST is structured to take time histories of U t-N :t , of length N , β and Gaussian noise ϵ ∼ N (0, 1) as inputs -in practice we found only conditioning on the last N = 2 sequence members was required. The ST is initially pre-trained for β = 0.9 using 200,000 snapshots, followed by fine-tuning for β = {0.3, 0.6, 1.2, 1.5, 1.8, 2.1, 2.4, 2.7} with 5,000 snapshots for each value. The ST produces an ensemble of trajectories originating from identical initial conditions, showcased here over a duration of 500 time units, for the previously unseen parameter values of β = {0.45, 0.75, 1.65, 2.55}. The ST produces qualitatively plausible trajectories, adeptly generalising to the different regimes.Generative models are often conditioned by prompt or by class information in which an encoder network is employed to provide an embedded representation (Ramesh et al., 2021) . However, as the additional conditioning in physical systems comes in the form of continuous variables (here ϵ and β), we can include them directly via cross-attention. The ST outlined in Shokar et al. (2023) implements cross-attention to learn an adaptive state-dependent weighted contribution for stochastic noise ϵ, given time histories U t-N :t , to model the conditional probability p(U t+1 |U t-N :t , ϵ). Sampling ϵ facilitates ensemble generation with different noise realisations, akin to sampling ξ in equation 1.In this work, we require the transformer to also account for the influence of the parameter, β, in modelling p(U t+1 |U t-N :t , ϵ; β). To achieve this, we introduce β as an additional feature preceding the W K ∈ R (D+1)×D (key) and W V ∈ R (D+1)×D (value) transformations, with D = N y -256. This differs from handling of ϵ which undergoes separate transformations via W Kϵ ∈ R D×D and W V ϵ ∈ R D×D before concatenation to facilitate a state-dependent adaptive weighing of the noise during attention. By incorporating β directly into the linear transformations, we parameterise the model, independent of the current state U t-N :t . Additionally, including temporal histories, U t-N :t , within the key and value vectors to facilitate learning of temporal correlations, results the following:Q = W Q U t-N :t ∈ R D (4) K = {W K {U t-N :t , β} d , W Kϵ ϵ} s ∈ R D (5) V = {W V {U t-N :t , β} d , W V ϵ ϵ} s ∈ R D (6)where each W denotes a learnable weight matrices, Q denotes the query vector and {.} s , {.} d here represent concatenations over the sequence dimension and the feature dimensions respectively. Attention is calculated as standard:Attention(Q, K, V ) = A(Q, K)V = softmax QK T D -1 2V , where softmax denotes the softmax operation to introduce sparsity in the attention matrix A.Following the approach outlined in Shokar et al. ( 2023), we introduce noise ϵ solely in the initial of M transformer blocks, to only induce a single forcing when making a forecast of U t+1 , aligning with the way stochastic PDEs are typically forced. Consistent with conventional practice in employing cross-attention, all subsequent attention blocks still include conditioning on β, givingK = W K {U t-N :t , β} d , V = W V {U t-N :t , β} d and Q = W Q U t-N :t insubsequent transformer blocks -see Appendices B.1 and B.2 for further details on the architecture and loss function.The ST is initially pre-trained with a dataset of 200,000 snapshots where β = 0.9, followed by fine-tuning with β = {0. 3, 0.6, 1.2, 1.5, 1.8, 2.1, 2.4, 2 .7} with 5,000 snapshots for each value. This range was chosen as above it dynamics display zonostrophy (persistent jets with minimal nucleation, coalescence, or jet translation), while below it, friction dominates, suppressing jet formation. The pre-trained weights are then used for initialisation with fine-tuning refining the model to generalise to the dynamics corresponding to various values of β. The initial learning rate used for the finetuning is a factor of 50 smaller than that for pre-training, as not to perturb the weights too far from learned solution during pre-training -see Appendix B.3 for further details. 
ensemble trajectories: 
 It is not possible to accurately predict a single trajectory due to the stochastic nature of the system, but it is possible to generate an ensemble of equally-likely trajectories. With the ST trained on a range of values of β, we compare model-generated trajectories with numerical simulations for various parameter values. To reduce training, the ST is only fine-tuned on a limited selection of values of β outlined above. We found this was sufficient as the ST is capable of reproducing the expected dynamics on unseen parameter values when interpolating between seen values of β.Figure 1 displays the zonally-averaged flow, for four values of β unseen during training, with each regime exhibiting distinct dynamics. The ST emulations closely align with reference numerical solutions, with accurate representation of jet nucleation and coalescence frequencies, as well as the latitudinal translation rate. Importantly we observe that the ST does not become physically inconsistent after a number of autoregressive steps.The comparison between ST emulations and reference trajectories not only confirms the model's adeptness in capturing the system's dynamics but also highlights its robustness across a range of variable parameter values. To provide quantitative assessment of the performance of the ST, we analyse the ability of the ST to reproduce key statistical properties of the stochastically driven jet system. These properties include spatial and temporal correlations, as well as spectral composition, with results shown in Appendix C. The ST's predictions align well with the expected statistical behaviour, providing confidence in its reliability for various scientific and engineering applications.In chaotic or stochastically driven systems predicting transition events, such as nucleation and coalescence events, are of interest, however large ensembles are required for accurate estimate of transition probabilities. As outlined in Shokar et al. (2023) , the ST provides a five-order-of-magnitude speedup over traditional numerical methods, facilitating the generation of such ensembles.Figure 2 .a displays a reference coalescence event, while Figure 2 .d displays a reference nucleation event, both generated via numerical integration with β = 1.35, a parameter value not seen during training. Each plot shows dashed lines prior to each event to display the initial conditions used to generate ensembles and measure the time interval between each initial condition and event of interest. In Figures 2.b,c,e,f the ST generates PDFs that show good agreement with the PDFs produced using numerical integration. This demonstrates the ST's capability of effectively capturing the dynamics of the physical system while generalising to unseen regimes. This allows for computationally efficient exploration of physical systems where previously prohibitively expensive.Sampling the network in this manner provides an estimate of aleatoric uncertainty, with the model naturally demonstrating reduced uncertainty as the forecast approaches the event occurrence. Moreover, it exhibits lower uncertainty when predicting coalescence events compared to nucleation events, consistent with observations from numerical studies in Cope (2021) . ) showing the the time to a coalescence event, with respect to the reference event in shown a, with the initial conditions 100 time units before the reference coalescence event. The green curve shows the PDF from numerical integration, comprising an ensemble of size 1000. The red curve shows the PDF from the ST, comprising an ensemble of 100,000. c, Shows the same information as b, but for initial conditions 25 time units before the reference coalescence event. d-f, Show the same information as a-c, but for a reference nucleation event. We observe agreement between the green and red curves indicating that the SLT is successfully capturing the expected distributions. 
conclusion: 
 The fine-tuned Stochastic Transformer is capable of emulating a stochastically driven fluid system, and is successfully shown to generalise to unseen parameter conditions, accurately capturing system dynamics across a spectrum of diverse parameter regimes. Demonstrating its ability to reproduce statistical properties of the system along side its computational efficiency when compared to traditional numerical integration, the ST facilitates economical exploration of the system parameter space, thereby enabling in-depth analyses of spontaneous transition events and other dynamic phenomena.Future work will look to extending the ST to include varying multiple parameters physical {β, µ, ν, k f , ε} to comprehensively characterise the state space of the system. This computationally efficient emulator enables cost-effective exploration of system dynamics, facilitating investigations previously impractical with numerical integration. Tasks include identifying the critical parameter β governing transitions between preferred jet numbers and examining how β affects the frequency of spontaneous transition event frequencies. Work will also look to expand beyond beta plane turbulence to emulate a diverse array of systems, both deterministic and stochastic, in order to thoroughly evaluate model robustness. Additionally, our aim is to extend the architecture to accommodate 2D and 3D inputs, enabling us to effectively model a broader spectrum of systems.A NUMERICAL SIMULATIONS OF QUASI-GEOSTROPHIC TURBULENCE Direct numerical simulation of vorticity in equation 1 is performed on a 2D doubly-periodic square domain (x, y) ∈ [0, 2π] 2 , using a pseudo-spectral method using the GeophysicalFlows (Constantinou et al., 2021) package in Julia (Bezanson et al., 2017) -and details are as in Shokar et al. (2023) .To promote the spontaneous emergence of zonal flows, rather than being directly forced, the fluid is stirred using a stochastic vorticity force ξ(x, y, t) with zero mean injected onto an annulus of wavevectors in Fourier space centred around a mean radial wavenumber k f with thickness δk = 1.Simulations are initialised from rest, ζ = 0, and parameters remain constant throughout the integration. The time for total kinetic energy to reach a statistically steady state during spin-up depends on the damping rate µ, typically achieved by µt = 2 -3 (here t = 50 -75 time units). Only data points beyond this point are used for training and testing the deep learning model.In all experiments values µ = 4 × 10 -2 , ν = 1, and an energy injection rate of ε = 10 -4 , with a principal forcing wavenumber of k f = 16 and a time step of δt = 4 × 10 -2 were chosen to be in the regimes in which coherent zonal jets form, with the jets exhibiting behaviour with a wide range of time variability. 
 
 Appendix: The neural network architecture we employ is the Stochastic Transformer (ST) (Shokar et al., 2023) . Figure 3 illustrates the ST architecture that comprises a multi-head (MHA) attention blocks and multi-layer perceptrons (MLP) -that consists of two linear layers and a GELU non-linearity between the two. The attention blocks allows the model to weight the importance of different parts of the input sequence and the conditioning inputs, here ϵ and β, when making predictions, while the MLPs allow it to learn complex non-linear relationships. The MHA blocks and MLPs each have skip connections around them to enhance gradient propagation towards the input during the backpropagation phase of training. The details of the how the stochastic and parametric conditioning are applied to the MHA are visualised in Figure 4 .To incorporate temporal information of the sequence of time histories, as is standard with Transformer architectures, a learnable temporal encoding vector (Kazemi et al., 2019) is added to the transformer inputs U t-N :t , before the concatenation of ϵ and β to indicate they are not part of the temporal sequence.Assuming spatial homogeneity for the forcing term ξ, the beta-plane system demonstrates latitudinal symmetry, denoted by ∆U (y, t) = U (∆y, t), where ∆ signifies a shift operator in the y dimension. This is result of the system's periodic boundary conditions and forcing symmetries. The transformer itself is not a translation invariant architecture. One approach is to replace each linear layer with a convolution operation. However, to avoid limiting the model's ability to only attend to local spatial connections, we chose to instead use the the method of slices (Budanur et al., 2015a; b) to remove the phase information of the data. This allows the model to learn weights in a phase-invariant manner, independent of any shift ∆y in the physical space, before reintroducing the phase information.The method of slices involves first taking a discrete Fourier transform, F, to transform the input data, U t into the spectral domain, Ût = F(U t ) ∈ C D , where D = N y . We obtain the phase information, ϕ, from the first Fourier mode using ϕ = arg Ût . The phase is extracted from U t only, ϕ(t), as also removing the phase from the previous time steps in the time history would prevent the transformer from also learning the latitudinal drifting dynamics that we observe in the system. We store this phase and construct a phase-aligned solution, by shifting each of the preceding timesteps with respect to ϕ(t), U ′ t-N :t = Z t-N :t e -ϕ(t)•K , such that the first Fourier mode of U ′ t is a pure cosine and K = [-N y /2 -1, ..., N y /2] ∈ Z D is the wavenumbers in spectral space. U ′ t-N :t is then input to the ST to obtain a prediction of U ′ t+1 , this way the ST handles inputs that have been translated to be identical manner, given any shift ∆y. We reintroduce the phase information to the output from the transformer U ′ t+1 via ϕ(t), U t+1 = U ′ t+1 e ϕ(t)•K , before an inverse Fourier 𝐒𝐭𝐨𝐜𝐡𝐚𝐬𝐭𝐢𝐜 𝐓𝐫𝐚𝐧𝐬𝐟𝐨𝐫𝐦𝐞𝐫 (𝐒𝐓) Arrows indicate the forward pass, where the ST receives an input including the value of β and a short temporal history of U t-N,t , with N denoting the sequence length. The model's weights are learned independently of the phase, ϕ, which represents the argument of the first mode of Ût ∈ C D . Noise vector, ϵ is appended to the space-time histories. The architecture comprises M transformer blocks, where 'MHA' stands for multi-headed attention using scaled dot-product attention -see Figure 4 for details. 'Norm' refers to layer normalisation, and 'MLP' is a positionwise multi-layer perceptron. The phase information is reintroduced to produce a forward trajectory of Ũt+1 .Discrete Fourier Transform 𝑈 2 !"#:! = 𝐹𝐹𝑇(𝑈 !"#:! ) Discrete Fourier Transform 𝑈 8 2 ′ !%& = 𝐹𝐹𝑇(𝑈 8 ′ !%& ) 𝑈 8 2 !%& = 𝑒 ' ) * 𝑈 8 2 ′ !%& Shift Phase Back 𝑈′ !"#:! 𝑀x Inverse Fourier Transform 𝑈′ !"#:! = 𝐼𝐹𝐹𝑇(𝑈′ C !"#:! ) Phase Align 𝑈′ C !"#:! = 𝑒 "' ) * 𝑈 2 !"#:! Calculate Phase 𝜙 = 𝑎𝑟𝑔 𝑈 2 ! (&) {𝑈′ !"#:! ,transform, to transform the outputs back to physical space, U t+1 = F -1 (U ′ t+1 ) ∈ R D . This is visualised in Figure 3 . As such the model acts in a translation equivalent manner. We found that including this symmetry was important for producing the correct PDFs shown in Figure 2 .The ST is pre-trained and fine-tuned using the following loss function, comprising of the Continuous Ranked Probability Score (CRPS) (Matheson & Winkler, 1976) , a proper scoring rule that generalises the Mean Absolute Error (MAE), and a spectral loss to ensure that energy is preserved at all scales:EQUATIONTheCRPS(U, Ũ ) = 1 m m i=1 U -Ũ (i) -1 2m 2 m i=1 m j=1 Ũ (i) -Ũ (j)is commonly used metric for probabilistic forecasts in meteorology and was demonstrated by (Pacchiardi et al., 2022) to an effective alternative to adversarial training methods. Here U is the truth trajectory from the training dataset, Ũ (i) is the i th member of the prediction ensemble, where the model produces an ensemble of size m for time step t + 1. The network is trained to minimise the difference between each predicted ensemble member and the truth trajectory U t+1 , while simultaneously maximising the dissimilarity between each individual ensemble member Ũ (i) t+1 given to the inclusion of the second term in the CRPS. While performing a forward pass over an an ensemble increases training costs, in practice an m = 2 is sufficient.Neural networks have been shown to exhibit a spectral bias (Rahaman et al., 2019) , inadequately capturing higher frequency modes due to their diminished contribution to overall energy compared to the larger-scales. In the context of modelling physical data, the preservation of not only large-scale properties but also energy across all scales is essential. For this reason a spectral loss is employed, this involves employing the discrete Fourier transform, denoted as A is the attention matrix -a mask is applied to the upper triangular elements to prevent forward looking, as is applied to the decoder transformer in Vaswani et al. (2017) , before being scaled by √ D and a softmax transformation is applied. This illustrates the case where the batch size = 1 and the number of heads = 1. The non-stochastic parametric self-attention in the following transformer blocks operates in the exact same manner, simply without the transformation and concatenation of ϵ, but still include concatenation with β. between these two quantifies. Through experimentation, we found λ = 1 to weight these terms achieved convergence during training.While a time step of δt = 4 × 10 -4 is required for stability of the numerical integration, sampling at such small intervals is unnecessary for the neural network to capture the dynamics. Therefore data input to the network is sampled at time intervals 2500 times larger than the numerical integration time step. Each time step of the system is represented as a one-dimensional vector U t (t) ∈ R D , with the target output being the subsequent time step U t+1 (y), also obtained from the numerical integration.Using a Bayesian hyperparameter sweep (Biewald, 2020) we determine the best-performing model. The metric used to minimise was the compose loss function outlined in equation 7. Through this sweep we found that the following hyperparameters were optimal: batch size: 256; optimiser: Adam; initial pre-training transformer learning rate: 5e-4; initial fine-tuning transformer learning Figure 5 : PDFs displaying joint probabilities of U , ∂ y U and ∂ t U . a, e, i, m, for visualisation of the 3D PDF we average over ∂ t U to obtain the density of values in the u-∂ y U space obtained via numerical integration for different values of β = {0.45, 0.75, 1.65, 2.55} unseen during training. b, f, j, n show the same u-∂ y U space from an evolution generated using the ST. c, g, k, o show the u-∂ t U space by averaging over ∂ y U for the numerical integration values and d, h, l, p show the u-∂ t U space for evolution generated using the ST. Here we can see that again the SLT shows good agreement with the numerical integration, covering the full space of dynamics, across the parameter range. rate: 1e-5, learning rate scheduler: exponential decay with decay rate: γ=0.99; pre-training epochs: 250, fine-tuning epochs: 250, CRPS ensemble size: 2; N (length of time history used to make next forecast): 2, number of transformer blocks: 4, number of transformer heads: 16.For a quantitative assessment of longer-term evolutions, we examine statistical properties of the flow. We generate Probability Density Functions (PDFs) for individual U values and their derivatives, ∂ y U and ∂ t U . These PDFs serve to measure the learned spatial and temporal correlation. We conduct this analysis over a time integration period of 5000 time units, producing two PDFs: p(U ) representing the PDF from numerical integration, U , and q( Ũ ) representing the PDF of forecasted values derived from the SLT, Ũ .In Figures 5 we plot As β increases, we do observe discrepancy between the PDFs generated via numerical integration and the ST for the joint distribution of U -∂ y U . We hypothesise that this is due to a reduced frequency of spontaneous events as zonostrophy increases, leading to a data imbalance. Future work shall explore this phenomena, however overall we see very good agreement between the numerical integrations and the ST emulations.The proportion of energy in each wavenumber, or the energy spectra, is an important physical property in fluid dynamics. For this we assess time-averaged power spectral density plotted for both the numerical integrations (in green) and the the ST (in red) over the same time integration period of 5000 time units as used above, plotted in Figure 6 . We observe a high level of agreement across all scales, paying particular attention on the higher wavenumbers where data-driven approaches often fail. This is all the more important in the case of 2D turbulence due to an inverse energy cascade. This means that correct representation of the small scales is vital due to their upstream influence on the large scale dynamics. These findings, in conjunction with the aforementioned evaluation metrics, allow one to place confidence in the ST's capacity to faithfully emulate the system dynamics, over the range of parameter values, even for those unseen during training. We believe this positions the ST as a valuable tool for investigating dynamical phenomena observed in the beta-plane system, facilitated by its computational speed-up. 


Decision: Authors proposes using a pre-trained transformer model fine-tuned on task-specific data to emulate complex stochastic dynamical systems. Paper shows improved computational efficiency and generalization capabilities. I recommend authors provide details on model training and improving the clarity of the paper for the camera-ready version. If possible, I encourage authors to expand the experiments beyond a single system and compare to prior work in the area.

Reviews: review: title: title: Octavi's review
main_review: main_review: I don't see the novelty of the paper. The authors pre-train a transformer and then fine-tune for task specific downstream problems. 
The novelty might be in the use case - stochastic dynamical systems - but there is a lot of good work in this topic that is not even mentioned here. Also, the writing could be improved, it is hard to understand.

For instance: 
Alexandre Cortiella, Kwang-Chun Park, Alireza Doostan,
Sparse identification of nonlinear dynamical systems via reweighted ℓ1-regularized least squares
Computer Methods in Applied Mechanics and Engineering
title: title: Review of Extending Deep Learning Emulation Across Parameter Regimes to Assess Stochastically Driven Spontaneous Transition Events
main_review: main_review: The paper discusses the extension of deep learning emulation for assessing stochastic events across different parameter regimes, focusing on spontaneous transition events driven by stochastic factors. It employs transformer-based networks to emulate complex dynamical systems, highlighting a significant computational cost reduction and the ability to generalize across diverse scenarios. This research presents a novel approach in combining deep learning with stochastic dynamics, aiming to enhance the understanding and prediction of unpredictable events in various scientific domains.

The manuscript presents an innovative integration of deep learning with stochastic dynamics, using the Stochastic Transformer (ST) to emulate complex systems across various parameter regimes. This approach not only demonstrates a significant computational efficiency but also extends the model's generalization capabilities to novel conditions. The use of pretraining for the SST across a broad spectrum of parameters showcases an interesting application of current PDE research with deep learning in stochastic modeling.

However, the application to a single system, though understandable given scope constraints, leaves open the question of the methodology's scalability and effectiveness across diverse domains. Future work could benefit from exploring a range of systems to validate the approach's universality, as also highlighted in the final section of the paper.

Minor comment:
- Second sentence in 4. Conclusions: typo for "it's" where it should be "its".

Overall, the paper is articulate and detailed, effectively communicating its contributions and findings. Given its novel application and the depth of analysis, it would undoubtedly spur valuable discussions at the workshop. I recommend acceptance of this work.