Computers were born to simulate the world. The very first job of the first modern computer, the 1945 ENIAC, was to solve differential equations for a hydrogen bomb. More recently, equations have been replaced by data: machine learning methods absorb unfathomable amounts of information and learn to predict its patterns, without knowing the exact mechanics.
Flux is a machine learning library which bridges those two approaches. Its one weird trick is allowing almost any Julia code, including complex simulations, to run as part of a learning process. That in turn enables a hybrid approach to modelling: write down the rules that you do know, and use data to fill in the gaps. When there’s too much uncertainty for a pure simulation, but pure learning might not respect physical constraints, Flux gives you best of both worlds.
Flux doesn’t compare to TensorFlow or PyTorch for language models or modern AI training, but it has found an excellent niche among researchers, who have used it to find new insights about the world rather than blindly matching patterns. What follows is a small smattering of research projects that have used Flux, and which illustrate its modest but (I think!) meaningful impact on the world.
When a drive is interrupted by roadworks, you’ll spend some time thinking about alternative routes, before setting off with a rough plan. How does the brain know when planning is worthwhile, and when to stop? Research in Nature Neuroscience from Cambridge shows that an answer emerges from a simple computational model. They created an agent based on a recurrent neural network, but augmented with the ability to plan out its actions at roughly threefold faster than real time. After training it to solve mazes, they found that the model’s decisions about when to hesitate and when to explore had an uncanny match to human behaviour. They then examined brain recordings of rats and found that mysterious bursts of activity known as “hippocampal replay sequences” behaved just like the AI’s planning stage, showing that they are mental simulations rather than just memories, and bridging the gap between artificial intelligence and the old-fashioned kind.
Aircraft navigation can’t always rely on GPS, which is vulnerable to jamming. Earth’s magnetic field offers a solution, but measuring it is difficult due to interference from the vehicle itself. MIT researchers use machine learning to generate more reliable estimates of the true field. They then go further by incorporating a physical model for such estimates, developed during WWII, into the neural network, making its parameters tunable. Over several test flights the position error was kept below 45m, comparable to GPS errors of about 10m, showing that this can be a reliable and practical backup option.
Children with congenital heart defects often differ greatly from adult physiology and from each other, which makes one-size-fits-all clinical guidance impossible. This paper from Stanford shows that it’s feasible to make treatment predictions based on a “digital twin”, reconstructing a seven-year-old patient’s heart from CT scans, and running simulations to see the effect of interventions and spot vulnerabilities that are not visible from ECG results alone. Because the mechanistic simulation is complex and too slow for bedside use, they use the full model to generate data and train a high-fidelity surrogate neural network. That model can be run on a laptop, and allows clinicians to run what-if scenarios and create personalised treatment plans in minutes.
The city of Pohang in South Korea is not known for earthquakes, but in November 2017 residents were surprised by the second-largest ever one in the country, which injured dozens and left nearly two thousand homeless. The magnitude 5.5 quake turned out to be caused by pressure mismanagement in a nearby geothermal energy project. Managing pressure in underground reservoirs is complex, requiring high-fidelity fluid simulators that are impractical to run in real time. Work in Nature Scientific Reports from Los Alamos demonstrates how instead a differentiable physics model can be embedded in a neural network that learns to predict pressure changes based on the underground permeability map – and can do so in about 14ns per sample, or 40,000 times faster than the original finite-volume solver. Work like this is crucial for the climate, making both geothermal energy and CO2 sequestration more feasible.
These ideas are useful in the home, too. Heat pumps, batteries and solar cells are a promising way to reduce residential energy usage – about 30% of the world’s total – but only if they are intelligently coordinated, which normally requires accurate forecasts and physical models for each home. Work from DTU shows that reinforcement learning can achieve close to optimal energy self-sufficiency without special knowledge.
This paper explores similar ideas applied to national gas pipelines, embedding a physical simulation of the whole network within the optimiser to get high-fidelity, real-time control.
3D printing usually involves creating objects layer-by-layer, and slowly. Tomographic Volumetric Additive Manufacturing (TVAM) instead uses light beams to freeze the object within a photosensitive resin all at once. This paper increases the theoretical resolution (and thus practicality) of TVAM by switching from a ray-tracing-like approach to a full wave-optics model of light; differentiating through the physics simulation allows finding the optimal pattern of light beams, and custom derivative tricks make differentiation feasible, increasing the scale of objects that can be produced this way.
In Indonesia it’s challenging to predict wildfires, because they are sensitive not just to weather but to subterranean water levels and peat distribution. Work in Nature Scientific Reports combines surface weather models with a neural network for a more accurate forecast, and they found it able to predict peat fires that other models missed.
This is the world’s first image of a black hole, produced in 2019 from the Event Horizon Telescope (EHT). If you remember the glowing accretion disc from Interstellar, this is the real deal, at the centre of the massive M87 galaxy. Members of the EHT collaboration explored how an object the size of six billion suns can test general relativity, through images that appear elliptical rather than perfectly round. Synthesising faint flux signals from a global array of telescopes is a challenging enough inference task in itself; the project uses Flux to filter generated images.
Elsewhere in the universe, researchers investigated galactic winds in the starburst galaxy M82. They modelled the known physics, in a differential equation, with a small neural net taking the place of an unknown factor known as the “mass-loading rate”. Despite fitting to limited data – X-ray observations without velocity measurements – the combined approach revealed unknown structure in the physics.
QuEra and Harvard ran the “largest quantum machine learning experiment to date”, and used Flux as a classical baseline.
That’s all for now! If you hear of other cool work using Flux, please do let me know.