What is a distributional shift?

A “distributional shift” is the difference between the data distribution an AI is trained on and the data distribution it is deployed in. This word can also be used to refer to the effect of this shift on the AI.

When an AI is trained on a specific set of data, for example, learning to place strawberries into a bucket, the data could consist of multiple different scenarios of buckets in different positions. This set of scenarios is commonly referred to as the “training distribution”.

Now after training, we deploy our cleverly trained AI into the wild to hurl strawberries into buckets and regardless of where the bucket is, the AI supposedly succeeds! The set of scenarios after training, where our AI deals with new data, is called the “deployment distribution”.

But why is distributional shift such an essential part of the issues with AI Safety, such that it gets a name of its own?

We often assume that the training distribution is fairly similar to the deployment distribution. This is often fine - if the deployment distribution is (statistically) similar to the training distribution, models will continue to perform well, like they did in training.

The concept of a “distributional shift” is explained neatly by Rob Miles in this video, and if you would like to follow the analogy of the strawberry AI to grasp the same, keep reading.

Consider the situation where the buckets were all red in the training distribution but our deployed AI encounters a blue bucket. Remember, an untrained AI has no idea of what a “bucket” is, or a “strawberry”, or what it even means to “throw”. It learns these definitions through whatever it sees in its training scenarios. So the model might define a bucket to be a "red vaguely cylindrical object, even though in the real world, “red” has no relevance to what a bucket can be. (In other words, it has unintentionally added the color “red” as a must-have property of a bucket alongside “cylindrical” and “hollow”, thereby marking anything that doesn’t satisfy that property as “not a bucket”.)

Simply speaking, if there’s a bright red fire hydrant next to the blue bucket, the former has a larger chance of getting a strawberry chucked at it.

An AI good at its job would need to learn from a training distribution that features buckets of various colors, heights, radiuses and so on. But even then, we may still not be prepared for everything, especially when it comes to more complex tasks and more possible scenarios in the real world.

What would the AI do if it saw a bucket upside down?

If we were to refine the definition a bit more, the significant difference in data and how the AI performs in the deployed distribution and the training distribution is the “distributional shift”.

We generally attempt to tackle “distributional shifts” through various different methods. There is often a validation/evaluation/test distribution which we use right after training. This distribution features scenarios not present in the training distribution, to test how well the AI generalizes (how it responds to new unexpected scenarios). Other methods considered to handle the distributional shift are:

• Adversarial training, the approach of adding malicious examples to the training distribution to minimize the AIs accuracy
• Having highly varied data in the training distribution
• Monitoring the decrease in the performance of AI over the course of training and identifying the reason
• Bringing in other AI that’s learned from an entirely different training distribution for the same task.

The safety risk posed by “distributional shifts” lands on the shoulders of advanced AIs with goals far beyond placing strawberries into buckets. When we have AIs embedded in military infrastructure, educational systems, healthcare institutions and the like, with immediate impactful consequences, it becomes highly vital that the AI encounters a minimal amount of new data that is costly in performance.

(A worthy aspect to consider is whether we want the training distribution to perfectly resemble the real world, or come pretty close. Since that’s effectively equal to training the AI in the real world, incorrect behaviour that the AI may still be possible of performing and be trained away from might only emerge in the “actual” real world. However, in a training distribution that has a few more unrealistic features, these behaviours might emerge quicker, due to perhaps being more acceptable, and thus we can train them away earlier. Thus, there is a possibility that “distributional shift” in some scenarios may be desired and inevitable.)