Simulation-Based Parallel Training
Anonymous Author(s)
AffiliationAddress email
Abstract
Numerical simulations are ubiquitous in science and engineering.
Machine learningfor science investigates how artificial neural architectures can learn from thesesimulations to speed up scientific discovery and engineering processes.
Most ofthese architectures are trained in a supervised manner.
They require tremendousamounts of data from simulations that are slow to generate and memory greedy.
In this article, we present our ongoing work to design a training framework thatalleviates those bottlenecks.
It generates data in parallel with the training process.
Such simultaneity induces a bias in the data available during the training.
Wepresent a strategy to mitigate this bias with a memory buffer.
We test our frameworkon the multi-parametric Lorenz’s attractor.
We show the benefit of our frameworkcompared to offline training and the success of our data bias mitigation strategy tocapture the complex chaotic dynamics of the system.
Introduction
Multiparametric dynamical simulations are crucial to scientific discovery and engineering applicationsfrom computational fluid dynamics [30] to biochemistry [1].
Let’s denote g such a simulation, X itsinput parameters and Y tX the output for a given time step t .
[Equation]
Recently, scientists have considered deep learning as a means to accelerate science by speeding upthose simulations and learning from them [37].
Most works train an artificial neural network f θ ina supervised manner, from data generated by the simulation process g .
The exact combination ofinputs and outputs of the training process depends on the application.
For instance, for a networktrained to replace the original simulator g with direct time predictions, as done in the seminal workof Raissi et al.
[33], we have f θ ( X, t )
≈ Y tX .
However, for the auto-regressive models proposed in[32, 4], we have f θ ( X, Y tX )
≈ Y t +1 X .
Classical training consists in generating a set of simulation data X train , Y train .
Where X train would berepresentative of a parameter set on which we want to infer with the trained neural network.
Efficienttraining f θ , and capturing all the complexity of the simulated system, generally requires tremendous amounts of data [6].
The dimension of Y train can thus quickly explode and lead to a massive memory footprint.
Additionally, the sequential generation of Y train can be prohibitively slow.
To overcome these limitations, one may restrict X train , or subsample Y train by reducing the number of time steps or observed dimensions.
To reduce X train does not necessarily hinder the quality of the training
Techniques ranging from adaptive experimental design [12] to importance sampling [17] help toidentify the parameters that will bring the most relevant information for the training of the network.
These techniques, however, require knowledge about the network response to different inputs X ,which is not feasible if X train is set beforehand.
We propose a framework that generates the simulation data simultaneously with the training process.
Data generation is highly parallelized and streamed to training to avoid storing any data in files.
Theframework also allows the steering of the training by selecting new parameters X based on somefeedback from the network on previously presented parameters.
This article presents our ongoingwork covering the first point.
Compared to classical training from files, the simultaneity of thetraining and the data generation presents specificities that must be addressed.
At any time τ of the training process, only the already generated data Y τ can be presented to thenetwork.
There is no guarantee that Y τ is representative of Y train .
Because of the limited number ofconcurrent simulations g and their iterative production of data, the two sets can have very distinctdistributions.
This setting commonly arises in online deep learning [8].
If we stream the data, usingthe last received time step Y τ directly for training, this is known to lead to poor performances relatedto catastrophic forgetting [18].
Therefore, to mitigate the bias in the data due to the discrepanciesbetween Y train and Y τ , we introduce a memory buffer of limited size.
It is associated with a datamanagement policy in charge of the selection of elements to construct batches and their eviction whenthe buffer is full.
Moreover, we rely on a simple sampling strategy of X train to improve the diversityof Y τ .
We validate our framework and the strategy to mitigate data bias on a model that captures thechaotic dynamics of Lorenz’s attractor [26].
First, we show that our data bias mitigation strategy issuccessful in producing more diverse batches than pure streaming learning.
Second, we show thatour framework can prevent the overfitting offline training suffers, by exposing the network to morediverse trajectories.
This overfitting typically occurs when the simulation data are too big to be stored,which imposes to limit the size of the dataset and consequently reduces its representativeness.
Third,we show that the framework can match the performances of offline training with a comprehensivedataset.
This illustrates there is no performance decrease using the online framework.
The main contributions of our ongoing work are:
• a deep learning framework orchestrating parallel data simulations with the online training ofneural architectures on multi-parametric dynamical systems;
• a strategy to mitigate the bias in data induced by the online nature of the framework.
Related Work
2.1 Deep Learning and Numerical Simulations
In recent years, deep learning has permeated the numerical simulations of multi-parametric dynamicalsystems.
Scientists proposed algorithms to substitute the traditional partial differential equations(PDE) solvers [33, 23, 4].
Instead of replacing entirely the original solver, others proposed toaugment the resolution of coarse solutions [39, 20].
Similarly, Erichson et al. presented a network toreconstruct the dynamics of a fluid from a few data points only [10].
Others have relied on numericalsimulations to study how machine learning interfaces with traditional dynamical system theories, likethe Koopman operator [28], or help to rediscover the governing equations of simulated systems [7].
All these approaches rely on supervised training, for which data mostly come from simulators.
Evenapproaches that could theoretically be data-free, like physics-informed neural networks [33], haverequired simulated data to work successfully on complex physics [27].
Most of the published work is only presented with fairly simple simulations, either in terms of thecomplexity of the governing equations or by the scale of the problem at stake.
Even though someworks propose to capture intrinsic properties of the studied physical systems, it is, to our knowledge, not yet clear how deep learning architectures trained on simple problems can scale to larger andmore complex ones, common to the industry and today’s scientific challenges.
For instance, usinggraph neural networks, Pfaff et al.
could infer dynamics with higher resolution than the one usedfor training [32].
Equivariance is another example to learn fundamental properties of a system, likesymmetries, and thus reduce the amount of data required for training [35].
2.2 Online Training
By generating the simulations along the training process the full dataset is not accessible anymore atall times.
In this setting referred to as online learning [8], data presented to the network may be biased.
It leads to poor generalization performances, like catastrophic forgetting [18].
Several techniqueshave been proposed to mitigate this catastrophic forgetting.
Either by tuning the loss for the optimalweights of the network not to change dramatically while being trained on new data [19, 21].
Anotherapproach consists in repeating the samples previously shown to the model [14].
Such a strategy hasalso been successfully applied in the context of deep reinforcement learning and coined as replaybuffer by the community [36, 34].
2.3 Parallel Frameworks for Deep Reinforcement Learning
Training frameworks presenting several degrees of parallelism have already been proposed by thedeep reinforcement learning community [31, 29, 16].
Especially, to achieve the acting and thelearning in parallel.
Similarly, in our framework, the data generation and the training are parallelized.
Liang et al.
propose, within the Ray project, a task-based formalism to abstract those different degreesof parallelism for deep reinforcement learning [24].
If some reinforcement learning experimentslead to massive parallel resource usage due to the very large number of actors to run, the simulationitself is not massively parallelized [2].
Indeed, these deep reinforcement learning framework supportsimulations that are either sequential or multi-threaded, but not scientific PDEs solvers requiringdistributed memory parallelization with MPI to address problems at a relevant scale.
Despite the similarity in the parallelism of the data generation and the training, the comparison withdeep reinforcement learning has limits.
We are not learning to interact with an environment but ratherlearning the environment itself.
Therefore there is no policy, and thus no concept of off-policy.
Allsamples are relevant, although some may bear more useful information for the training of the network.
It is also worth mentioning the work of Brace et al., which introduces a parallel framework to selectbiomolecular simulations to run, but not to train generically architectures on numerical simulations[3].
Training Framework
3.1 Architecture Overview
Our framework 1 is based on an adaptation of the open source Melissa architecture developed tomanage large-scale ensemble runs for sensibility analysis [38] and data assimilation [13].
Ourgoal is to enable the online (no intermediate storage in files) training of a neural model from datagenerated through different simulation executions.
These simulations are executed with differentinput parameters (parameter sweep), making the different members of the ensemble.
We target executions on supercomputers where simulations can be large parallel solver codes executedon several nodes and the training parallelized on several GPUs (Figure 1).
We give below an overviewof the framework, a detailed description being beyond the scope of this paper.
The framework relieson three main components:
[Figure] • The launcher orchestrates the execution of the simulation instances and the training.
The
[Figure] launcher interacts with the machine batch scheduler (e.g. Slurm or OAR) to request resources [Figure]
for the other component executions.
[Figure]
• Each client executes one simulation instance.
When starting, it dynamically connects to the
[Figure]
server.
As soon as it produces new data (a new time step), the data is sent to the server.
[Figure]
• The server runs on several GPUs to train the model in parallel.
On reception of new data
[Figure]
from one client, the data is first pushed into the local memory buffer.
When a new batch
is required for training, it is extracted from this buffer.
The management of this buffer is
detailed in the following section 3.2.
This framework thus alleviates the need for files to store intermediate results.
We target large-scaleexecutions generating large amounts of data that cannot be reasonably stored.
As an example, Melissahas been reported performing a sensibility analysis run executing 8000 simulation instances thatgenerated 48 TB of data processed online [38].
The secondary benefit of coupling data generation andtraining is that the server can control which simulation instances to run next, opening the possibilityfor adaptive experimental design.
The framework combines several levels of parallelism (at the ensemble, client, server, and communication levels) to enable efficient and scalable executions.
The simulation can be a parallel codewritten in C/C++, Fortran or Python, and combining the usage of MPI, OpenMP and Cuda.
Theserver relies on Pytorch Data Distributed Parallelism [22] for enabling multi-GPU training.
3.2 Data Bias Mitigation Strategy
The online data generation inherently carries 3 types of data bias (Figure 2):
• Intra-simulation biais.
Simulators produce dynamics as time series by discretizing time
and progressing iteratively.
After T time steps, only the data { Y tX , 0 ≤ t ≤ T } for the
simulation instance X are available for training.
• Inter-simulation bias.
Computational resources are limited and so often not all simulation
instances can be executed concurrently.
Let’s assume that only c simulations can execute
concurrently at any time.
At the end of the execution of the c first simulations training only
has access to the data { Y tX i , 0 < i ≤ c } .
• Memory bias.
Not only online training cannot have access to the not already generated data ,
but we also cannot keep all already generated data due to memory constraints.
We can, at
most keep a rolling sample of data that fits into a given memory budget
Therefore, batches cannot be drawn uniformly from the full dataset as performed with traditionalepoch-based training.
This bias has a detrimental effect on the training, partly due to catastrophicforgetting [18].
Some physical systems exhibit different behaviors depending on the parameter of the simulation.
Forinstance, Reynolds number for fluid dynamics [11].
The inter-simulation bias creates discrepanciesbetween Y
τ and Y train .
The sampling strategy consists in selecting parameter X so that Y τ is asrepresentative as possible of all the behaviors in Y train .
In the following, we adopt a simple strategythat randomly samples X among X train .
To mix the data further and address the intra-simulation and memory biases, we introduce a memorybuffer .
It acts like a cache of the data previously generated and sent to each trainer.
Without thememory buffer, data would be immediately streamed from the simulators to the trainers, a situationprone to catastrophic forgetting.
The memory buffer has a set capacity, s .
It is associated with adata management policy that specifies the rule to evict data when the buffer is full, and select data toconstruct the batches.
The current policy selects samples randomly from the buffer and erases themupon selection.
This avoids having samples repeated too many times in batches while the trainer iswaiting for new data from the simulators.
Due to the intra-simulation bias, the first simulation timesteps arrive first in the memory buffer.
To avoid yielding first batches with overly represented firsttime steps, we set a threshold s ready that signals when the buffer has enough elements to be sampled.
Because the training loop and the simulation are executed at different speeds on different processes,the memory buffer must also amortize the data production and consumption between the clients andthe trainer.
We set a second threshold s min that ensures there is always a minimal amount of samplesin the buffer, for the batch construction to be indeed random.
In summary, the key idea is that the simultaneous generation of simulation data creates a bias in thedata.
We rely on adequate sampling of the simulation parameter and a memory buffer to mix the datapresented to the network and guarantee the training quality.
Groups of concurrent simulations are executed sequentially.
As soon as a simulation time step Y t X i is computed, it is passed to the memory buffer.
Simultaneously, this cache of limited capacity yields batches from previously stored time steps.
The content of the buffer and the batches constructed from it are represented at two different times τ i and τ j .
At τ i the buffer is not yet full.
At τ j the buffer is full and contains some time steps from previous simulations.
Experiments
To prove the usefulness of our approach we examine the case of a neural network trained to capturethe underlying physics from the observation of complex dynamics.
For this purpose, we consider
Lorenz’s attractor [26].
Its dynamics are described by the system 1.
The attractor is chaotic.
Itmeans that, despite being deterministic, two points initially close have very distinct trajectories.
Thedynamics thus present a complex and rich behavior.
Nonetheless, the underlying non-linear equationis fairly simple.
The possibility to recover the dynamics of Lorenz using machine learning techniqueshas already been studied in [5, 9], even with partial observation [25].
Here, we train a neural networkto reconstruct the trajectory with an integration scheme of Euler, i.e. we predict f θ ( X, Y t )
≈ Y t +1 − Y t ∆ t ,where X = ρ and Y t =
( x t , y t , z t ) the coordinates of the system at time t given by Equation 1.
Themodel is auto-regressive as the trajectory is reconstructed iteratively at inference.
[Equation]
We consider the multi-parametric case for which the initial position X 0 and the parameter ρ vary.σ and β are respectively fixed to 10 and 8 / 3 .
The goal is to train a neural network that generalizesfor any initial position and ρ belonging to the interval [0 , 100] .
In the following experiments, initialpositions X 0 are always taken from the normal distribution N (15 , 30) .
Data are standardized withstatistics computed beforehand.
The model architecture also remains the same across the experiments.
It consists of 3 layers of 512 features.
The first layer has 4 input features for the position and theparameter ρ .
The output has 3 features for the velocity.
We use the SiLU activation function betweeneach intermediate layer [15].
For the memory buffer, we arbitrarily set s , s ready , and s min equivalentto 10, 5 and 3 simulations respectively.
In our experiments, we consider different datasets and training settings.
We always generate adifferent number of trajectories for values of ρ taken in [0 , 20 , 40 , 60 , 80 , 100] .
First, we have a fulloffline dataset consisting of 100 trajectories of 2000 time steps ∆ t of 10 − 2 second.
A second offlinedataset, we call restricted , consists of only 10 trajectories. A third offline dataset, referred to assubsampled , corresponds to 10000 trajectories sampled every 100 time steps. The restricted andsubsampled offline datasets are artificially reduced to illustrate the case of more complex simulationsfor which it might be difficult to obtain and store a representative dataset. As we compare differentdataset sizes and different training strategies (offline or online), for a fair comparison, we ensure thatthe quantity of information seen by the network is always the same.
This translates into a constanttotal number of batches. For online training, there is no notion of epoch anymore.
Hence, the trainingof the full, restricted, and subsampled datasets is respectively achieved on 100, 1000, and 100 epochs.
For the online dataset, we generate 10000 trajectories.
In the streaming setting, the trajectories aregenerated from simulations with increasing values of ρ . Data are directly gathered into batches andpassed to the trainer. A second online setting, noted as sampling , randomly samples ρ prior to the datageneration. The online setting, noted as sampling + memory buffer , combines the random samplingof the parameters and a memory buffer with random data selection policy.
For all training strategies,the batch size is fixed to 1024 time steps.
The validation dataset consists of 10 trajectories, 2 for eachvalue of ρ taken in the same interval as for the training set.
4.1 Memory Buffer Mitigates the Bias Induced by Online Training
To illustrate the efficiency of the two levels of control, namely the sampling of parameters and thememory buffer, in mitigating the data bias, we consider the batch statistics in the different training settings.
Figure 3 shows the mean and standard deviation of the batches for the different strategies.
The streaming strategy appears biased, exposing the training process to catastrophic forgetting thatcan explain the gap between train and validation losses in Figure 4.
We see that the coupling of thesampling strategy and the memory buffer makes the distribution closer to the one observed in theoffline setting, showing its capability to reduce bias. [Figure] y
[Figure] M ea n a nd s t a nd [Figure] Sampling and memory buffer correct the prejudicial bias observed in streaming batches.
4.2 Parameter Space Exploration with Online Training
Training and validation losses for the different training settings are displayed in Figure 4.
We observethat for the restricted and the subsampled offline datasets the validation losses are almost an orderhigher than the training losses, which indicates overfitting.
This shows that, even for simple equationslike Lorenz’s system (Equation 1), a neural network can experience difficulties in recovering thedynamics with only a few trajectories.
It stresses the importance of having a representative dataset.
For the same number of batches, the online training, with a random sampling strategy and the memorybuffer, presents training and validation losses that decrease similarly with the number of batches.
It isalso for this setting that the validation loss is minimal.
By presenting to the network more diversetrajectories, this strategy overcomes the overfitting observed with limited offline datasets.
4.3 Recovery of Offline Performances
Figure 4 shows the training dynamics with the full offline dataset and the online setting (sampling +buffer) are equivalent.
Both configurations correctly recover the dynamics of the attractor on a testtrajectory generated with ρ equal to 28 (Figure 5).
The predicted trajectories are stable as shown in Figure 5a.
They match the reference for a limited time, as seen in Figure 5b. The lack of coherenceafter some time is expected due to the chaotic nature of Lorenz’s attractor.
Compared to offlinetraining with a comprehensive dataset, which can be difficult to obtain for more complex systems, theonline strategy does not induce any degradation of the trained network performances. [Figure]
[Figure]
Conclusion
[Figure]
In this article, we have proposed a framework to train artificial neural networks on synthetic simulations generated along the training process.
We have shown this setting induces a bias in thedata that hinders the quality of the training.
We have introduced simple strategies to mitigate thebias: a random sampling of the simulation parameters and a memory buffer that allows mixingfurther the data presented to the network.
We have applied these strategies to train a network on Lorenz’s attractor.
We have shown that leveraging the online setting to generate more data, but stillkeeping the same number of batches as for an offline strategy, can lead to a lower validation loss. Wehave also shown the strategies were successful to match the performance of offline training with acomprehensive dataset in capturing the underlying dynamics of Lorenz’s atttractor, without incurringthe same memory footprint.
Such results need further testing with different use cases to be confirmed.
Advanced implementations of the memory buffer will also be further investigated like reservoir orimportance sampling for the selection of elements in the memory buffer [40, 17].