Tackling the intractable: why posterior inference matters

Diffusion models generate data by transforming noise through a series of steps. In real-world applications, however, we often need them to conform to certain constraints — think of generating an image with specific features, or producing text that adheres to a given context. This leads us to the problem of posterior inference, where we want to sample data that fits both the original prior distribution and the additional constraints. The catch? This posterior inference is typically intractable.

Relative Trajectory Balance

At its core, RTB is about being able to fine-tune a diffusion model prior into a posterior model, and sample according to the posterior distribution of interest. Ideally the posterior model should not drift too far from the prior — we want the samples to remain plausible under the original prior — but it should nonetheless change appropriately to capture the posterior of interest.

When tackling the posterior sampling problem, the goal is to sample from a posterior distribution $p_{\text{post}}(\mathbf{x}) \propto p_\theta(\mathbf{x})\, r(\mathbf{x})$, where $p_\theta(\mathbf{x})$ is our diffusion model prior and $r(\mathbf{x})$ is the constraint. Importantly, $r(\mathbf{x})$ does not need to be differentiable: it can be a measurement, a reward from another model, a human preference — anything, really.

Our RTB objective introduces a constraint that enforces the correct relationship between the prior and the posterior processes:

This ensures that the likelihood of the denoising trajectory under the posterior model $p_\phi^{post}$ stays close to the likelihood of the same trajectory under the prior $p_{\theta}$, while learning the scalar $Z_{\phi}$ (in log scale, for numerical stability) and sampling proportionally to $p_\theta(\mathbf{x})\,r(\mathbf{x})$. Thanks to its design, the advantage of RTB lies in its ability to achieve unbiased posterior sampling — proof in the paper.

Highlight results

To show the effectiveness of RTB we provide a diverse set of empirical results across multiple domains. Some highlights:

In vision, RTB applied to image generation achieves state-of-the-art results, improving on classifier guidance and RL-based baselines in both the diversity of the generated posterior samples and their closeness to the true posterior. Below is an example of fine-tuning an unconditional diffusion prior.

In stable diffusion, we can fine-tune a large text-to-image model to achieve high scores according to human preference, on par with state-of-the-art methods.

In language, we show we can fine-tune a posterior over a discrete language diffusion model prior to perform coherent posterior sampling for text infilling.

In reinforcement learning control tasks, we report state-of-the-art results when training a posterior behaviour policy from a diffusion behavioural prior, achieving comparable or better results than other recent methods.

Conclusion

Relative Trajectory Balance is a flexible objective that allows efficient and unbiased sampling from a posterior distribution, given a diffusion prior and a reward or likelihood. The empirical results across vision, language and control underscore the potential of RTB to set new benchmarks across AI applications. Whether you are diving into research or developing practical solutions, the advances presented here are poised to change how we generate and refine data.

To investigate the shortcut learning tendencies of deep models we introduce a new training setup: the Wisconsin Card Sorting Test for Machine Learners (WCST-ML).

The Wisconsin Card Sorting Test for Machine Learners

We envision a simple, yet powerful, setting. Imagine you were given this limited set of images and labels:

You are then asked to provide a label for the following image. You can choose between 0, 1 or 2, according to your understanding of the task:

The problem with this question lies in its ill-defined nature. The images and labels in the training matrix are fully correlated with respect to at least two explicable high-level cues — for example shape and colour. At training time the classification task can be solved perfectly by relying on either feature; at test time, an image showing an unseen feature combination requires generalisation.

If a human were to solve this task, several factors could play a role in the final decision. Prior knowledge about the world, the task, or even other related tasks might influence our beliefs about the underlying task to solve — yes, in a Bayesian sense. Ideally we would form a set of hypotheses (hopefully containing both shape and colour as alternatives), and make a choice given our prior and our observations, updating as we see more data. The important part is that the choice we make is informative of our priors, and of our biases. Absent a significantly different prior among participants, labelling the image 0 would display a bias toward shape, while labelling it 1 would showcase a colour-based task logic.

During training, deep learning models have their own biases, and have often been observed to collapse to easy-to-learn cues. Similarly to the degenerate classification case above, available datasets typically underspecify the task to be solved, and data is hardly ever comprehensive enough to fully constrain decision-making.

Models agree on which shortcuts to take

In our work we run extensive experiments to observe how different architectures react when put through a similarly underspecified visual discrimination task. Our first interesting finding is that, despite the different inductive biases induced by different architectures, we observe largely similar behaviour across models. Models reliably prefer certain cues — like colour — more than others. Colour is a particularly easy cue to fit in the simple example above, since we have some intuition about how it might be solved by something akin to channel selection in the input space.

Experiments on more complex datasets, such as UTKFace, show that cues such as ethnicity, gender and age also rank similarly across architectures, exposing potentially sensitive biases when making inference without appropriate measures or human-in-the-loop solutions.

Interestingly, we find these biases are explicable from the point of view of the loss landscape: the set of solutions biased toward the preferred cues occupies a far greater volume — and tends to be flatter — than those corresponding to the averted cues.

We call this fully correlated test WCST-ML, and propose a simple framework for exploring any such biases. As long as a dataset has multiple labels for each input, and enough elements to generate a fully correlated dataset, WCST-ML should be readily applicable. Much more in-depth findings are shared in the ICLR publication, which helps us understand what these shortcut biases look like during training — and potentially how to avoid them. These findings shed light on shortcut learning in deep models, and emphasise the importance of solutions such as active human intervention to remove model biases that may cause negative societal impacts.

Mitigating the bias with diffusion counterfactuals

In a follow-up work we propose a solution based on ensemble diversification. An effective approximation to the Bayesian ideal above is to let an ensemble of models entertain a range of diverse hypotheses while observing the degenerate training data. The diversity must be functional, and is therefore not trivial to enforce.

We propose a method based on synthetic counterfactuals sampled from an appropriately trained Diffusion Probabilistic Model (DPM). We find that DPMs cross stages of training in which they are more likely to produce samples with novel feature combinations — an originative stage.

Even when lacking fidelity, these samples can be leveraged directly for ensemble diversification through model disagreement, enforcing functional diversity across models.

In the table below you can see the fraction of models that attend to a given cue (out of an ensemble of 100) when trained without diversification — the baseline case above — and when trained with our diversification method.

This demonstrates that we can get models to attend to non-shortcut cues, even without the need for expensive additional data collection.

The full article is on PalermoLive.

Cambridge researchers are helping to design the healthcare robotics of the future by investigating a robot-assisted approach to training doctors in medical examinations. To read the full article, follow this link.

We have published part of this work in the leading robotics journal Soft Robotics. The research article is titled “Action Augmentation of Tactile Perception for Soft-Body Palpation”. In it we delve further into the complex realm of physical medical examination: we use a robotic platform and a mathematical framework both to understand and to employ complex palpation strategies, in order to achieve appropriate sensory perception for robotic medical diagnosis.

A robot that can autonomously learn to play piano by mimicking the dynamics of different human playing styles.

Our publication, in collaboration with Cheryn Ng, in the international journal PLOS ONE, is titled “Gaussian Process inference modelling of dynamic robot control for expressive piano playing”. In the paper we propose a Gaussian-Process-based mathematical framework for a robot to use sound feedback to autonomously learn how to play piano according to 10 different play styles.

Robotics and medicine: one of our publications in the journal Autonomous Robots, exploiting the application of machine learning and tactile sensing to aid in the autonomous detection of abnormal lumps in soft bodies.

In collaboration with Perla Maiolino and Fumiya Iida.

One of our articles, “Efficient Bayesian Exploration for Soft Morphology-Action Co-optimization”, with Perla Maiolino and Fumiya Iida, proposing a mathematical framework to perform the active exploration of robot control and morphology to aid in discrimination tasks.

Presentation