Category 5 hurricanes, financial crashes, and global pandemics are just a few examples of rare events whose high risks necessitate understanding and mitigation. Developments in artificial intelligence (AI) could go a long way towards improving our ability to model and mitigate the impacts of such extreme events, but current training methods are often unable to deal effectively with outliers in data – which is exactly what extreme events are. If outliers are present in training data, they skew the AI’s expectations, but if they’re omitted entirely, models will wrongly assume they never occur. To address this shortcoming, the Photrek team, led by Dr Kenric Nelson, has developed a new training technique to design more robust AI systems that can cope with rare, extreme events.

Taking AI to the Next Level by Mimicking the Brain

While many of us are only just getting used to the age of artificial intelligence – still fumbling around for effective ways to incorporate ChatGPT into our workflows – researchers are already working on the next evolution: artificial general intelligence (AGI). Achieving AGI would mark a significant improvement in the reasoning capabilities of artificial systems. Current AI models are capable of replicating human performance in some areas, but not all. Seven-fingered hands and impossibly warped perspectives are common occurrences in AI-generated images. On the other hand, AGI represents the point where machines possess human-level intelligence that can be adapted and applied to a much broader array of tasks.

Experts disagree on how close we are to achieving this milestone, but it’s clear there’s work to be done, which raises the question of how we can improve upon the current generation of AI models. One approach is to try and mimic the behaviour of the human brain. The “predictive coding” hypothesis suggests that brains effectively operate as layers of prediction machines. Our senses collect data from the external world which is then passed up through the layers in the brain. Each of them predicts what data it expects to receive from the ones beneath it and compares that prediction to the data it actually receives. The brain’s aim is to minimise the difference between its own predictions and the real-world data as measured by a quantity called “free energy”. If the free energy is large, the brain is “surprised” and updates its modelling so that it can better predict how the world behaves in future.

Teaching Machines How to Reason Probabilistically

The corresponding process in machine learning is known as variational inference. During its training, an AI model will attempt to calculate a mathematical function called the posterior. It’s a probability distribution; it records the probabilities of different explanations for or hypotheses about the data. This is similar to how the higher levels of our brains have their own internal models of the world, which they use to generate forecasts and predictions. But while our brains learn continuously, current AI models are trained all in one go. From a given dataset, they estimate the posterior, trying to minimise error. This posterior then serves as the AI’s “predictions” in the future. By working with probability distributions, variational inference is a particularly suitable training method for AI models that will be working with complex systems. Fields like meteorology, finance, and public health involve lots of uncertainty and unknown variables. The best we can do is infer likelihoods from known data, such as the probability of a natural disaster or a global pandemic occurring. Variational inference means that AIs learn to think in terms of probabilities, to be aware of the uncertainties in their analysis.

For complex systems, the true posterior is impossible to calculate, so scientists restrict the AI to working with a family of simpler probability distributions – like the familiar bell curve. With each new data point, the model estimates which of the simple distributions is closest to the real posterior by calculating and then minimising free energy – the same way our brains work!

But one thing variational inference isn’t so good at is dealing with outliers. This is what Dr Kenric Nelson and his colleagues want to fix. As the founder of the technology company Photrek, Dr Nelson is particularly interested in the application of AI to complex systems like financial markets and the environment. Predicting extreme weather events, such as category 5 hurricanes or floods, could help to mitigate their devastating impacts on people’s lives and improve our abilities to predict them with sufficient time to prepare. But, to an AI, events such as these are outliers; typical training methods result in AI models downplaying the chances of such extreme events ever occurring. On the flip side, if outliers are included in the AI training data, then models can erroneously believe that they’re part of the norm.

Using Curved Information Spaces for More Robust Models

To address this problem, Dr Nelson and his team tried working with a different set of probability distributions aside from the usual bell curves. The problem with standard bell curves is that they decay rapidly at extreme values – we say they have very short tails – and this leads AI models to think that these extreme values can be ignored. Instead, Dr Nelson considered another family of distributions called the coupled exponentials. These distributions may look quite similar to bell curves, but can have longer tails – meaning that there’s a higher probability of obtaining extreme values. This enables AI models to find a suitable middle ground—not ignoring outliers while also not letting them skew the training process too much.

One of the novel features Dr Nelson found when considering this family of distributions is that they live in a non-trivial geometric space. Imagine we have a grid of points in space, then each point in the grid represents one of the specific distributions from the coupled exponential family. The same thing can be done for models working with bell curves, but in that case, you end up with an Euclidean space. In other words, you can treat the points as being spread out across a flat space where distances between points are calculated as usual. The concept of distance in this space is important, as it tells you which distributions are similar to one another and which ones are very different. As the machine learns, it adjusts its posterior distribution to something suitable that’s nearby. You don’t want the model jumping all over the place – that would be like someone completely changing their opinion whenever they read something new. 

The special thing about coupled exponentials is that their space is not flat, but curved. This works the same way that curved spacetime does in physics. Out in deep space, the fabric of spacetime is flat and matter particles will travel through it along straight lines. On the other hand, if particles are in the vicinity of a heavy object like the sun, then the fabric of spacetime is curved and the matter particles travel along elliptical pathways instead. The same happens for the AI models Dr Nelson considered. The innovation Dr Nelson and the Photrek team have developed is that they can learn over a geometry with very high curvature, while actually completing the measurements over a lower-curvature geometry guaranteed to have a finite measure of variance.

Improved Performance using Celebrity Faces

To test their new framework, Dr Nelson and his team trained the model on the Celebrity-A data set. CelebA is a set of 202,599 images of celebrity faces that have a high degree of variability, making it a particularly complex data set for models to parse. This made it an ideal choice to test the new techniques. They found that their methods led to a 10–25% improvement in perception metrics over using standard bell curves.

Dr Nelson’s work is an important generalisation of existing AI training methods to a wider class of probability distributions. Using these methods, outliers can be incorporated into training data so that models are aware of their existence, but in a way that doesn’t spoil their overall performance, potentially generating AI models that are more robust. This has clear applications to a number of fields that deal with complex situations and environments including finance, medicine, and weather. In particular, AI models trained via this approach can appreciate that natural disasters and extreme weather events should be considered even if they are rare, improving our ability to predict such occurrences and lowering the risks of catastrophe.