christopher beckham

reflecting on past work – offline model-based optimisation

2025-07-01T00:00:00+00:00

reflecting on past work -- offline model-based optimisation

1. intro - what is model-based optimisation?
2. last year's work
- 2.1. ranking validation metrics
3. 🪵🔥 reflection, and future work
- 3.1. can we combine online and offline mbo?
- 3.2. links, and open source
4. References

In this post, I introduce offline model-based optimisation and share some highlights from research I published at the tail end of my PhD [1]. The focus is on issues of validation and extrapolation – two topics that remain frustratingly underexplored in the field. Along the way, I offer an honest critique of my own work, reflect on broader challenges in offline MBO, and point to a few promising directions forward.

While I’m no longer working in academia, this topic still matters a lot to me. If any of it resonates with you – whether you’re a researcher, engineer, or just curious, I’d love to chat, albeit in a bandwidth-conscious capacity.

1. intro - what is model-based optimisation?

In model-based optimization (MBO), the goal is to design inputs—often tangible, real-world objects—that maximise a reward function reflecting some desired property. Crucially, this reward function exists in the real world, meaning that evaluating it requires physically testing the input and observing its outcome. For example, in protein design, the input might be a molecular specification, and the reward could be the protein’s ability to bind to a specific kind of receptor in the body. In this setting, the true objective $f(x) \in \mathbb{R}^{+}$ measures binding affinity, and querying it involves running expensive lab experiments. We will refer to this "true" objective function as the reward oracle.

In theory, we could optimise the reward oracle $f$ directly by combining it with a search algorithm, learning from past evaluations to guide the search for better inputs. This involves learning a "surrogate" model $\ft$, which is exploited by the search algorithm and is also refined iteratively based on feedback provided by the reward oracle. This is called online model-based optimisation ("online MBO"). In practice however, querying $f$ is often prohibitively expensive and slow, and we need many samples if we want the surrogate model to learn enough about the underlying problem in order to guide the search towards promising candidates.

But if we already have an existing dataset of input-reward pairs (here $y = f(x)$ is the "reward"), we can use it to pre-train the surrogate model. The idea is that a well-trained surrogate may then require far fewer queries during the online phase to guide the search effectively. This pre-training stage is known as offline model-based optimisation ("offline MBO").

Since $\ft$ is an approximation however, it is particularly brittle. It can behave unpredictably on regions of input space which are not represented in the training data (what we call "out-of-distribution"), and this is especially problematic since we assume the best input candidates are not already yet in the training set. Such models are also vulernable to "adversarial examples", where by naively optimising for the highest-scoring candidate with respect to the surrogate often produces one which is no longer plausible according to real world constraints. This makes it an interesting yet challenging area of research.

1.1. formalising offline mbo

In the offline setting, we assume access to a labelled dataset of input-output pairs, denoted $\mathcal{D} = \{(x_i,y_i)\}_{i=1}^{N}$, where $y_i \in \mathbb{R}^{+}$ and $y_i = f(x_i)$. The goal is to learn a model of the reward – typically a surrogate function $\ft$ – without querying the true function during training or model selection. While much of the offline literature focuses on a surrogate $\ft$, it's arguably more useful to think in terms of a generative model that defines a joint distribution $\pt(x,y)$. Via Bayes' rule, it can factorise in one of two ways.

The first is $\pt(y|x)p(x)$, and we can think of $\pt(y|x)$ as the probabilistic form of the surrogate, i.e. $\ft(x)$ parameterises the mean of the distribution. We can interpret this as: first we sample $x$ from some prior (e.g. the training distribution), then we predict its reward. While this is technically a generative model, it is not particularly useful for sampling high-scoring candidates as the only sampling we do is from the training data (and not a model we have learned).

In practice, a common baseline takes a hybrid approach which doesn't quite correspond cleanly to this. This involves sampling $x$ from the training data $x \sim \ptrain(x)$, which is then iteratively updated by ascending the gradient of $\ft(x)$ (which is typically the mean of $\pt(y|x)$). While this produces inputs with higher predicted reward, it abandons the semantics of the above factorisation and tends to produce poor inputs when scored against the reward oracle.¹

The second factorisation is $\pt(x|y)p(y)$, which we can think of as saying: first choose the desired reward $y$, then find an input which has that reward. Since $\pt(x|y)$ is a conditional generative model, not only can we target high reward regions, we can also avoid generating implausible inputs since it is a mechanism built into the model. (While generative models are not totally invulnerable to generating implausible inputs, they still do a lot better than discriminative models as plausibility is built into the model by design, i.e. likelihood.)

For the remainder of this work, we will define our joint generative model $\pt(x,y)$ as the second factorisation:

\begin{align} \pt(x,y) = \pt(x|y)\ptrain(y), \end{align}

where $\ptrain(y)$ is the empirical distribution over the rewards in the training set.

1.2. reward-based extrapolation

The key idea which seperates MBO from regular generative modelling is that we don't just want to generate any kind of sample from the model. We would like to generate samples whose real reward $y$ is as large as possible, as these have the most real world utility. The difficulty lies in the fact that these (extremely) high scoring samples do not exist in the training set, otherwise MBO would be a much simpler task where we only need to generate things that plausibly look like what is already in the training set. This means MBO has to extrapolate – it has to learn what constitutes low and medium-scoring samples, and infer what a high-scoring sample may look like.

This also implies that the behaviour of the generative model needs to somehow be "tweaked" at generation time. For instance, we have defined a generative model $\pt(x,y)$ to be the following:

\begin{align} \pt(x,y) = \pt(x|y)\ptrain(y), \end{align}

where $\ptrain$ is the empirical distribution of $y$'s observed in training. If we simply sample according to this strategy, we will only sample conditioned on the kinds of reward seen in the training set. To rectify this, we could switch out the prior for another distribution $\widehat{p}(y)$, one which reflects a larger distribution of rewards. For instance, if $\ptrain(y)$ reflects a range of values from $[0,100)$, perhaps the new prior reflects those from $[100,200]$. From this, we can define the "extrapolated" model:

\begin{align} \widehat{\pt}(x,y) = \pt(x|y)\widehat{p}(y). \end{align}

(I am using the widehat notation '$\widehat{\pt}$' to symbolise 'higher', a version of $\pt$ which is biased towards high scoring samples, rather than something implying a statistical approximation.)

Ideally we would like to find an "extrapolated" model $\widehat{\pt}(x,y)$ such that it maximises the average reward coming from the reward oracle, which we will simply call the "test reward":

\begin{align} m_{\text{test-reward}}(\tilde{p}) = \mathbb{E}_{x \sim \tilde{p}(x,y)} f(x), \tag{1} \end{align}

and therefore we wish to maximise $m_{\text{test-reward}}(\ptnew)$. In other words, we want to find a $\pt(x|y)$ and $\widehat{p}(y)$ such that samples produced by the former have as large of a reward as possible, according to the reward oracle. Since this equation involves $f$ which is too expensive to compute during training or model selection, it is only intended to be executed at the very end of the machine learning pipeline. But this does not help us during training or model selection.

To rectify this, we could simply replace $f$ with the surrogate model $\ft$. However, $\ft$ has also only been trained on the same empirical distribution of rewards, and we cannot expect it to score inputs conditioned on e.g. $[100,200]$ reliably, as this is clearly out-of-distribution. One approach is to split the dataset into low-to-moderate scoring examples and high-scoring examples. For instance, if our original dataset only represented samples with reward in $[0,100]$, then we could for instance split it into $[0,50]$ (low-to-moderate) and $[50,100]$ for high scoring (see Fig. 1). The low-to-moderate split is used to train the generative model, while the latter forms a validation set.

Both data splits (which is just the full dataset) can actually used to train a validation proxy, $\fvalid$. It makes sense to evaluate $\ptnew$ against this because it has been trained on the full distribution of rewards coming from the dataset. Since $\fvalid$ has "seen" samples in $[50,100]$, even if the generative model hasn't, it can still produce inputs conditioned on this range and we can use the validation proxy to validate it. Therefore, this setup allows us to measure not just generalization, but generalization specifically in the context of reward extrapolation.

Figure 1:If we imagine our dataset (left) being split into low-to-medium (train) and high-scoring samples (validation), we first train a generative model on just the training set (the low-to-medium scoring samples). From this, we instantiate the extrapolated model, generate high scoring samples, then validate these with the validation proxy. The validation proxy in turn is trained on the full dataset, which makes it an appropriate "psuedo gold standard" the generative model can be scored against.

1.3. ‼️ why evaluation is difficult (and misunderstood)

With the rapid progress in generative modeling over the past few years, our approach to evaluation has evolved. In earlier eras of machine learning, it was common to assess models based on likelihood over a test or validation set – a natural outcome of maximum likelihood estimation, where the goal is to find parameters $\theta$ that maximise the probability of the observed data.

Because of the extremely rapid advances in generative modelling in the past few years, the way we have performed evaluation has changed. In the olden days of machine learning, it was more common to evaluate machine learning models by way of likelihood on a test or validation set. This is a natural consequence of maximum likelihood estimation, which states that we wish to find a model which best "explains" the data, i.e. find parameters $\theta$ such that the parameterised model assigns the highest average likelihood across all samples. However, likelihood is only concerned with how plausible pre-collected samples are, rather than whether samples generated from the model itself satisfy a useful notion of preference. (Also, likelihood isn't a particularly accurate measure of sample quality. [2], [3], [4]) Such preferences can be encoded with a reward function $f$, but this is typically expensive to compute as it reflects a real world process (i.e. $y = f(x)$ is like asking a human rater to evaluate $x$).

As mentioned in Sec. 1.2, a principled strategy is to approximate $f$ with $\fvalid$ and continue forward. Even if $\fvalid$ is an approximation, it actually serves as a useful anchor for the generative model. This is because even though it is only trained on low-to-moderate scoring inputs, we can measure its ability to generate high-scoring inputs against the validation proxy which has technically seen high scoring inputs during training. Compared to other MBO literature, I make a very explicit distinction between validation and testing which does not seem to be well-respected, and I partly suspect it's because there is a conflation between "real world" MBO and "academic" MBO. (These are terms I created, and the latter is not meant to be read in a disparaging sense.)

By "academic MBO" I simply mean doing MBO in the context of academic research, i.e. publishing papers. In this situation it may not be practically feasible to evaluate the reward oracle $f$, for instance in the case where the benchmark data involves an extremely expensive human evaluation (e.g. protein synthesis). To rectify this, some MBO datasets are actually based on simulation environments, and the same simulation provides a reward oracle which can be used to score the data.

Since the simulator is just a function that can be freely executed in silico with negligible monetary cost, researchers can (intentionally or not) "violate the spirit" of offline MBO by abusing the simulator and constantly querying it during training or model seleection.² This is especially enticing in academia because there is an overwhelming bias towards pushing things that "beat SOTA" or are "novel". Conversely, in "real world" MBO there is already a safeguard against abusing the ground truth and that is time and money. Therefore, in order to respect the economic burden associated with MBO, a validation set needs to set aside as this is ultimately what we will use in the real world before sending off samples to be tested.

Apart from simulation environments, most MBO datasets are really just finite collections of data from a real world problem. Since the reward oracle is infeasible to compute for academic research, a "test proxy" $\ftest$ is trained on the entire dataset and used as an approximation to the reward oracle. Like with the simulator, this can be easily abused, and necessitates the use of a seperate validation proxy $\fvalid$.

Due to the different types of rewards oracles already mentioned, below is a table explaining what they are for:

name	what is it
reward oracle	$f(x)$: real world reward model, extremely expensive to compute. This may also refer to a simulation environment's reward model.
"proxy" oracle	- $\ft$: a regression model trained on the training set. While it is a discriminative model, it can be "hacked" to act as a generative model. In this article, I prefer to use generative modelling terminology, in which case $\pt(y \vert x)$ is used instead. However, in this post I prefer to use the conditional density $\pt(x \vert y)$ will usually be referred to, instead of the other terms.
	- $\fvalid$ : typically not defined in literature, but this is specifically a proxy oracle intended for model selection and hyperparameter tuning. It is trained on the combined training and validation set. Here we will call it the validation proxy.
	- $\ftest$ : proxy oracle trained on train + valid + test set (all of the data). This typically exists in "academic MBO" where the ground truth is also too impractical to compute at test time. Here we will call it the test proxy.

1.4. training, validation, and testing

As discussed in Sec. 1.2, we need to measure not just generalisation, but extrapolation. If our validation set follows the proposed setup in Fig. (1), then we can just approximate Eqn. (1) by introducing some approximate reward model $\tilde{f}$:

\begin{align} m_{\text{reward}}(\tilde{p}; \tilde{f}) &= \mathbb{E}_{x \sim \tilde{p}(x,y)} \tilde{f}(x). \tag{2} \end{align}

From this, the function $m_{\text{reward}}(\widehat{p_{\theta}}; \fvalid)$ now constitutes our first validation metric. By "validation metric" we simply mean some function which measures the ability of the model to extrapolate. More generally, it may not only be a function of an approximate oracle $\tilde{f}$, but also other things such as the validation set itself. (We will discuss some other ones later.)

Note that while Eqn. (2) is a principled and reasonable approach to determining how well $\pt(x|y)$ extrapolates, this is just one possible validation metric of many. On one hand, it is quite interpretable: assuming a fixed $\fvalid$, Eqn. (2) is maximised when samples produced from $x \sim \pt(x|y), y \sim \pvalid(y)$ produce the largest average reward. On the other hand, $\fvalid$ is an approximate model and shares the same vulnerabilities to adversarial examples and overconfidence as many other regression model. Therefore, validation metrics go beyond Eqn. (2), and may involve measuring other aspects of the generative model or data.

So far we have discussed the need to measure extrapolation (Sec 1.2), as well as the lack of a validation set which is crucial to measuring it (Sec 1.3). From this we can motivate a very principled and reasonable train-validate-test recipe, which is the following:

Inputs: Split total dataset $\mathcal{D}$ into: $\dtrain$, $\dvalid$, and $\dtest$. Ensure that the valid and test sets contain higher reward inputs, as per Sec. 1.2.
Training: Train the generative model $\pt(x|y)$ on $\dtrain$. Also, if the validation metric necessitates it, train a validation proxy $\fvalid$ on $\dtrain \cup \dvalid$.
Validation: Use $\dvalid$ and/or $\fvalid$ for model selection / hyperparameter tuning.
Final evaluation: assuming we already have a recipe for generating high scoring samples from the model, score those samples with either the reward oracle (if we operate in "real world MBO"), or test proxy (if we operate in "academic MBO").
- If we need the test proxy, train a $\ftest$ on $\mathcal{D}$, and measure the average reward via $m_{\text{reward}}(\widehat{p_{\theta}}; \ftest)$.

Finally: note that for the "real world MBO" step in "final evaluation", since we'll be sending off samples to the real world, it is much more data efficient to first re-train the best model on the entire dataset $\mathcal{D}$ using the same hyperparameter configuration, and then use that to generate samples.

In the absence of a reward oracle which can be judiciously evaluated, we need to turn to cheap-to-compute validation metrics. We already saw one in Eqn. (2), and there are many others which can be conceived of. Given a list of these metrics a-priori, how can we figure out which one performs the best for our task?

2. last year's work

Let us begin with a summary of everything so far:

(1) In offline model-based optimisation we wish to learn a reward-conditioned generative model from an offline dataset of input-reward pairs. The rewards are originally obtained from a ground truth reward "oracle", which is assumed to be too expensive to query during training or validation of the generative model.
(2) Evaluating samples from a generative model is a difficult task. Firstly, likelihood-based evaluation is not sufficient to evaluate the quality of outputs. Secondly, samples ideally need to be evaluated by human feedback (which is perfectly encapsulated by the notion of a reward oracle). Lastly, models trained need to extrapolate beyond the rewards they were trained on, as the better they can extrapolate, the more impactful they will be in the real-world.
(3) Evaluation is difficult, often neglecting a validation set. This may be related to the confusion between "real world" and "academic" MBO. In "academic MBO", the reward oracle is replaced with a test proxy or simulator. While these are technically useful and cheap-to-compute, non-sparing use of these fundamentally violate the spirit of offline MBO, whose emphasis is on trying to extract as much value as possible from the available data without resorting to expensive reward oracle queries.
(4) (Repeating last section's paragraph) In the absence of a reward oracle which can be judiciously evaluated, we need to turn to cheap-to-compute validation metrics. We already saw one in Eqn. (2), and there are many others which can be conceived of. Given a list of these metrics a-priori, how can we figure out which one performs the best for our task?

The work I published last year addresses these points.

To implement the train-valid-test protocol described, some technical considerations were needed. Experiments were implemented with Design Bench, a popular MBO benchmarking framework [5]. Design Bench imposes a reward threshold $\gamma$ which dictates which samples are assigned to the training set. For example, any samples whose $y \leq \gamma$ are assigned to the training set, and the rest is obscured from the user (in an API-like sense). Because of this, all of the remaining samples $\gt \gamma$ are not assigned to a validation set – in fact, the library does not prescribe one at all. Two possible solutions are:

(1) Simply hold out some small part of the training set as the validation set. This respects the intended design of the library, but effectively reduces the size of the training set and therefore handicap model performance compared to other Design Bench-based models which use the full training set. (In Fig. 1 left, $\dtrain$ is shown here, so imagine cutting out some portion of this as the validation set.)
(2) Define that all samples whose $y \gt \gamma$ belong to the validation set (Fig. 1, right). Since the validation proxy $\fvalid$ is always trained on the combined train+valid split, this means it is trained on the full dataset. This technique does not respect the intended design of the library, even if its motivation is quite principled.

I chose (2), which is illustrated in Fig. (3)-right. However, this requires some nuance when it comes to interpreting the relationship between the validation proxy $\fvalid$ and the test oracle $\ftest$. If the dataset is based on a simulator, then we already have "$f$" and we can train $\fvalid$ on the full dataset as the simulator can be treated as the ground truth from which the dataset's samples were drawn from.

Figure 3. Left: Design Bench's API exposes a training set which is all samples which fall below the threshold gamma (shown as the red line). Right: By considering samples which exceed gamma, we can define a validation set consistent with Sec. 1.3. The validation proxy is then trained on both the validation and training splits, which effectively is the full dataset, if and only if the reward oracle is either from the real world or a simulator.

Otherwise, if $\ftest$ is actually an approximate test oracle, then by Design Bench's definition it has been trained on the full dataset. This means training a validation proxy would involve training on all of the data and therefore be equivalent to a test oracle. But the latter needs to have seen more data to be a useful tool to measure generalisation once training and validation is completed. Therefore, in this situation, we let the test surrogate remain as the "gold standard" which has been trained on all of the data, and we only allow the validation proxy to be trained on a subset of the full dataset. Concretely, this would be the training set, plus an X% subsample of any examples whose reward exceeds $\gamma$. This is shown below in Fig. (4), and in this illustration X% is 50%.

Figure 4.If the true reward oracle cannot be evaluated at test time, Design Bench defines the test proxy as being trained on the full dataset. This however contradicts our definition of a validation proxy, which is trained in the same way (left). To resolve this, we can define the validation set as being a random subsample of samples greater than gamma (in the figure, this is 50%). That way, the test oracle still gets to be trained and defined with respect to the full dataset and serve its purpose as a gold standard to measure generalisation.

Therefore, for Design Bench, if we are dealing with a task for which no simulator environment exists, then we have to use a test proxy. That means invoking Fig. (4) for determining the precise train, valid, and test splits. Otherwise, if a simulator already exists, then we invoke Fig. (3).

2.1. ranking validation metrics

Now, all that is left is a validation metric. This metric is a function of the generative model $\pt$, and may also be a function of the validation proxy $\fvalid$ and validation set $\dvalid$. (To keep notation light, we will assume that the metric $m$ can take arbitrary number of arguments, even though so far we see it as a function of the first two.)

We already saw one of these metrics, which is simply Eqn. (2) but with $\fvalid$ substituted for $\tilde{f}$, which is just $m_{\text{reward}}(\tilde{p}; \fvalid)$:

\begin{align} m_{\text{reward}}(\ptnew, \fvalid) & = \mathbb{E}_{x \sim \ptnew(x,y)} \fvalid(x) \\ & = \mathbb{E}_{x \sim \pt(x|y)\pvalid(y)} \fvalid(x). \tag{3} \end{align}

This metric doesn't particularly care about how "calibrated" the model is. For instance, if we condition on $y = 50$ and get an example whose reward according to $\fvalid$ is $1000$, the model doesn't get penalised for it. The only thing that matters is that the samples from $\ptnew$ score as high as possible on average. Otherwise, is this is concerning, another validation metric is the "agreement" [6], which measures the extent to which the validation proxy agrees with the supposed label of the input generated by the model:

\[m_{\text{agreement}}(\tilde{p}; \tilde{f}) = \mathbb{E}_{x \sim \tilde{p}(x,y)} (y - \tilde{f}(x))^2. \tag{4}\]

In our case, if we substitute in $\tilde{f} = \fvalid$ and $\tilde{p} = \ptnew$ we get:

\[m_{\text{agreement}}(\ptnew; \fvalid) = \mathbb{E}_{x \sim \ptnew(x,y)} (y - \fvalid(x))^2. \tag{4}\]

For example, if we sample $y=50$ to generate an example and this is what $\fvalid$ also agrees with it and predicts the same value, then the resulting loss will be zero. More generally, this metric selects for generative models which can correctly produce samples in the extrapolated regime, according to the validation proxy.

In principle, these metrics could be combined together as a sum (or a weighted sum), but this slightly complicates the analysis as we also have to determine suitable scaling factors for each term.

Other validation metrics I defined were:

$\mathcal{M}_{\text{FD}}(\tilde{p}; \tilde{f}, \dvalid)$: Frechet Distance (FD) ([4], [7]) between the distribution of samples coming from $\ptnew$ and the validation set. Note that this is not the same as Frechet Inception Distance (FID), which uses the ImageNet-pretrained Inception network as a feature extractor. Here, we define the feature extractor as being some suitable bottleneck in $\fvalid$, as we want to leverage features which are specific to the domain at hand.
$\mathcal{M}_{\text{DC}}(\tilde{p}; \tilde{f}, \dvalid)$: The "density and coverage" metric proposed in [8], which is an improved version of the precision and recall metric originally proposed in [9]. This metric was originally motivated to tease out two important factors which determine how close two distributions are: sample quality and mode coverage, which can be thought of as precision and recall, respectively. While these terms can be individually computed, here I simply sum both terms, simply treating it as an alternative metric which can be compared to FD.
$\mathcal{M}_{\text{C-DSM}}(\tilde{p}; \dvalid)$: The conditional denoising diffusion loss [10] but evaluated on the validation set. Essentially, we are asking how well the model can denoise high scoring samples that it has never seen before. Since DDPMs are likelihood-based models, this is also a likelihood-based loss and therefore may not correlate well with sample quality. However, it is trivial to incorporate as a validation metric since it is already defined as a training loss.

Going back to the purpose of this work, we ask: what validation metrics work best, and how do we measure it? Ideally, evaluating validation metrics requires access to the reward oracle, as we need to measure them up against some gold standard. That’s where simulation environments become interesting: they give us access to a something which very closely mimics a real world oracle, letting us test how well different validation metrics correlate with the actual ground truth. The idea is to use this setup to run a large-scale comparison of metrics across many simulated datasets, so we can better understand which validation metrics are most trustworthy when we don’t have access to the ground truth. Ideally, this gives us actionable guidance for real-world MBO deployments.

To evaluate the effectiveness of a validation metric, we conduct a large-scale empirical study. Specifically, we train a wide variety of model configurations, log the value of each validation metric, and assess how well these values correlate with the corresponding true test reward, as computed by Equation (1). For any given validation metric, this yields a scatter plot where the x-axis denotes the metric's value and the y-axis represents the true reward under the "extrapolated" model $\ptnew(x, y)$. This also makes it possible to compute the Pearson correlation, i.e. how does the test reward (y-axis) behave in relation to the validation metric?

We perform this study using denoising diffusion probabilistic models (DDPMs) [10], chosen for their flexibility and strong performance in generative modeling. Holding the DDPM backbone architecture fixed, we vary several hyperparameters—including network width, reward dropout probability³, and reward guidance strength. Each unique combination of these hyperparameters defines a distinct configuration.

The results are illustrated below for several continuously-valued datasets from Design-Bench⁴. In particular, the Ant, Kitty, and Hopper environments provide simulation reward oracles, making them especially well-suited for this type of analysis. For completeness, we also include the Superconductor dataset, which uses a test proxy but still serves as a valuable point of comparison.

Figure 5. Each dataset is a subfigure, and each column is a validation metric. Each metric is plotted against M_test_reward , and points are colour-coded according to a diffusion-model specific generation parameter called label guidance. The Pearson correlation ρ is shown above each plot. Since each validation metric in this figure is designed to be minimised (negative signs are added to metrics which typically are maximised), we are interested in metrics which are most negatively correlated with the test reward.

Some metrics are plotted as their negatives (e.g., $-\mathcal{M}{\text{DC}}$ and $-\mathcal{M}{\text{reward}}$) to maintain consistency across all plots. Although these metrics are originally defined to be maximised, we negate them so that all metrics are presented as quantities to be minimised. Since this also applies to the test reward $\mathcal{M}_{\text{test-reward}}$, the best validation metric is one which is the most strongly negatively correlated with it.

Since the above plots are a lot of information to process, we can just jump straight to the figure which barplots the Pearson correlation for each of these experiments:

Figure 6.

The above figure differs a little from the one before it, as we actually have three additional groups of experiments on the right corresponding to "c.g." in parentheses. These correspond to the "classifier guidance" variant of diffusion [11]. I won't go into details here, but you can think of this variant as really defining a special joint distribution $\pt(x,y) \propto p_{\beta}(y|x)^{w}\pt(x)$ where $p_{\beta}(y|x)$ is a regression model also trained on the training data. (So it's like a probabilistic form of "$\ft$", only that here we use subscript $\beta$ as $\theta$ is already assigned to the generative model.) Conversely, "c.f.g." [12] can simply be thought of as just $\pt(x,y) = \pt(x|y)p(y)$ but with some algebra applied to Bayes' rule such that we condition on an "implicit" classifier $\pt(y|x)$.

Overall, if we count which validation metric was most negatively correlated with the test reward for each dataset-guidance configuration, agreement is the most performant, followed by Frechet Distance.

Lastly, the results we obtained are shown below in Fig. (7).

Figure 7. The c.f.g. and c.g. variants of our diffusion are shown in the bottom two rows. The numbers denote the normalised mean performance of samples with respect to the full dataset, which means values greater than 1 indicate that (on average) a reward higher than anything in the full dataset was obtained. Compared to the many other variants commonly tested against in Design Bench, we obtain competitive performance.

3. 🪵🔥 reflection, and future work

In the name of transparency and introspection, I will discuss what I think could have been done better.

The fundamental question we are trying to answer is: given a list of validation metrics a-priori, which are most useful as substitutes for the reward oracle? We exploited the fact that simulation environments exist – which grant access to the reward oracle – and then measure how correlated they are on four datasets. While this work isn’t about setting new benchmark records, the results are very encouraging (Fig. (7)). That said, it’s worth noting that this correlation is measured on the same data used to select the best metric, so there’s an inherent optimism bias. The ideal thing to do would have been to demonstrate that these metrics perform well on other downstream tasks. (The irony of this is not lost on me, but by the time I realised its significance I was very burned out from the project.)

The paper [1] was also tough to write, and honestly, so was this blog post. Maybe it’s because there were just too many ideas bouncing around at once: the importance of using a validation set, how to design that set to test extrapolation, thinking about MBO through the lens of generative modeling, and on top of that, proposing the use diffusion models – which, at the time, hadn’t really been explored in offline MBO. It would’ve been much simpler to just stick to the message of "validation sets, but for extrapolation". But at the time, that felt almost too obvious, like writing a paper just to say that validation sets are useful (which we already take for granted). But things that are supposedly "obvious" sometimes get published and go on to accrue hundreds of citations at the least, so maybe my barometer for that is miscalibrated.

I'm not sure how I would have approached the project if I did it again. When I reflect on past and ongoing industry work in similar problems, I notice a recurring pattern: we lean heavily on proxy metrics for validation. They’re cheap, measurable, and give us a sense of progress. But they’re also brittle and riddled with edge cases. It’s hard not to conclude that, sooner or later, all roads lead to human feedback models. It’s the only thing which consistently captures what we actually care about, even if it's noisy and expensive. So what does this mean then for this type of research? Maybe we ought to focus on best approximating the reward oracle and incorporate the best of both worlds in a single model: reliable, domain-specific priors, but also cheap and heterogenous sources of human feedback. The cost of obtaining labels is still an issue, but based on current trends it seems like the ever-increasing scale of foundation models will progressively allow for fewer-shot fine-tuning or prompting on top of them, which is a label efficient solution.

3.1. can we combine online and offline mbo?

Lastly, online and offline MBO feel somewhat siloed, when really one leads to another. Ideally, we want to build a good inductive prior in the offline setting and then segue into online to refine the model with real interactions. But in practice, offline MBO is only concerned with models which produce high-scoring samples "out of the box" with respect to the reward oracle, not whether that same model can be effectively used by an online learning algorithm to efficiently query it.

A real-world MBO workflow might appear as the following:

(1) We start with offline data, e.g. past experiments, human preferences, etc.
(2) Train a generative model on that data. It could be a conditional model $\pt(x|y)$ or unconditional $\pt(x)$. We may even decide to train a training proxy $\ft(x)$ (which I will just "lump" into the generative model category).
(3) Use generative model + some search algorithm to propose a small batch of high-scoring candidates and query those candidates with the reward oracle to obtain their labels.
(4) Add the newly-obtained (input, label) pairs to the dataset.
(5) Retrain or fine-tune the generative model on updated dataset.
(6) Repeat steps (3)-(5) until budget is exhausted.

For (3), examples of "use generative model" may include:

Using the generative model as a prior, e.g. if $\pt(x,y)=\pt(x|y)p(y)$, then the search algorithm can initialise its starting point via a sample from $\pt(x|y)$.
The search algorithm uses either $\pt(x)$ or $\pt(x|y)$ to evaluate the density (i.e., the plausibility) of any input it has proposed. Evaluating the density is possible with certain models such as time-continuous diffusion [13] and normalising flows.

Here is a more concrete sketch of the algorithm with an added twist. First, to avoid any bias in assessing generalisation performance due to optimism, we use validation proxy $\fvalid$ even during the online mode, and save the final evaluation with $f$ until the very end. We also define a budget $T_{\text{max}}$, which is how many evaluations we can perform in online mode:

(1) Assume offline dataset $\mathcal{D}$, split into $\dtrain$, $\dvalid$, and $\dtest$. $\fvalid$ is also trained on $\dtrain \cup \dvalid$.
(2) Train $\pt$ on $\dtrain$. (Here, $\pt$ can refer to any density deemed useful, for instance $\pt(x)$, or $\pt(x,y)$.)
(3) Online mode. For timestep $t = 1, \dots, T$:
- (3a) Use $\pt$ with search algorithm to sample a batch of high-scoring candidates.
- (3b) Obtain labels of candidates with $\fvalid$, compute mean reward $r_t$ over the batch, save this value.
- (3c) Update $\tilde{\mathcal{D}}$ with previously obtained labels and fine-tune $\pt$ with it.
(4) Compute discounted sum of reward: $G_T = \sum_{t=1}^{T} \gamma^{t-1}r_t$ for discount rate $\gamma$.

The twist is an idea I took from RL, which is the discounted sum of rewards. This sum is meant to encode the notion that rewards obtained earlier carry larger weight than later, as each successive evaluation progressively increases the overall cost in querying the reward oracle. By choosing this as the validation metric, we favour generative models and search algorithms which produce high-scoring candidates as cheaply as possible.

(Hats off to my co-author Alex, who really instilled a sense of MBO needing to be cost-effective. I think this idea really hits at the heart of that.)

3.2. links, and open source

Here are some things you may find useful:

🛠️ [validation-metrics-offline-mbo]: the original code for my paper. This uses the DDPM style of diffusion model from Ho et al.
🛠️ [offline-mbo-edm]: this is a bit more minimalistic and has a more up-to-date diffusion model which is EDM. Not only is this more performant, it generalises existing diffusion models which grants a lot of flexibility when it comes to deciding how to sample.

Design Bench can take some time to setup, so whichever repo you look at I highly recommend you consult the installation readme I wrote here here. As of time of writing, the mainline branch for Design Bench has broken urls for its datasets, so you should switch to my branch:

git clone https://github.com/brandontrabucco/design-bench
git checkout chris/fixes-v2
cd design-bench
pip install . -e

4. References

[1]

C. Beckham, A. Piché, D. Vazquez, and C. Pal, “Exploring validation metrics for offline model-based optimisation with diffusion models,” Transactions on machine learning research, 2024.

[2]

F. Huszár, “How (not) to train your generative model: Scheduled sampling, likelihood, adversary?,” Arxiv preprint arxiv:1511.05101, 2015.

[3]

L. Theis, A. v. d. Oord, and M. Bethge, “A note on the evaluation of generative models,” Arxiv preprint arxiv:1511.01844, 2015.

[4]

M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, G. Klambauer, and S. Hochreiter, “GANs trained by a two time-scale update rule converge to a Nash equilibrium,” Corr, vol. abs/1706.08500, 2017, Available: http://arxiv.org/abs/1706.08500

[5]

B. Trabucco, X. Geng, A. Kumar, and S. Levine, “Design-bench: Benchmarks for data-driven offline model-based optimization,” in International conference on machine learning, PMLR, 2022, pp. 21658–21676.

[6]

A. Kumar and S. Levine, “Model inversion networks for model-based optimization,” Advances in neural information processing systems, vol. 33, pp. 5126–5137, 2020.

[7]

D. Dowson and B. Landau, “The Fréchet distance between multivariate normal distributions,” Journal of multivariate analysis, vol. 12, no. 3, pp. 450–455, 1982.

[8]

T. Kynkäänniemi, T. Karras, S. Laine, J. Lehtinen, and T. Aila, “Improved precision and recall metric for assessing generative models,” Advances in neural information processing systems, vol. 32, 2019.

[9]

M. S. Sajjadi, O. Bachem, M. Lucic, O. Bousquet, and S. Gelly, “Assessing generative models via precision and recall,” Advances in neural information processing systems, vol. 31, 2018.

[10]

J. Ho, A. Jain, and P. Abbeel, “Denoising diffusion probabilistic models,” Advances in neural information processing systems, vol. 33, pp. 6840–6851, 2020.

[11]

P. Dhariwal and A. Nichol, “Diffusion models beat gans on image synthesis,” Advances in neural information processing systems, vol. 34, pp. 8780–8794, 2021.

[12]

J. Ho and T. Salimans, “Classifier-free diffusion guidance,” Arxiv preprint arxiv:2207.12598, 2022.

[13]

Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole, “Score-based generative modeling through stochastic differential equations,” Arxiv preprint arxiv:2011.13456, 2020.

Footnotes:

While online MBO also does a sort of hill climbing on the surrogate, the difference is that the resulting input is validated against the reward oracle, and this data is used to update the model.)

This should not be interpreted as discouraging "sim2real" experiments, where simulators are used to pre-train a model which is then adapted to a real world task. The difference is that if you treat the simulator as a training scaffold, then you need an external reward function to measure real performance. Otherwise, you're just evaluating on the same thing you are training on.

Conditional diffusion models are typically trained with dropout on the conditioning variable $y$ (in our case, the reward). This makes them act as both unconditional and conditional models.

⁴

Continuous datasets were used as DDPM operates on continuous values. While discrete variants do exist, I did not explore these. The simplest way to extend this work to discrete datasets is to use a discrete VAE to encode samples into a continuous latent space and perform diffusion there instead.

paper thoughts – questionable research practices (QRPs) in machine learning

2024-10-07T00:00:00+00:00

paper thoughts: questionable research practices (QRPs) in machine learning

I enjoyed reading this paper (Leech, Gavin and Vazquez, Juan J and Yagudin, Misha and Kupper, Niclas and Aitchison, Laurence, 2024) and wish I wrote it (or was a part of writing it!). Each section in this post can be cross-referenced with their Table 1:

Stable Diffusion-generated "depressed scientist with a knife in his mouth". I think that was the prompt I used for it.

1. Tuning hyperparameters further after test

Another common way to leak information is to tune on the test set: training a model, evaluating it on the test set, and then doing further hyperparameter search or testing again with a different evaluation metric. … The resulting models are in some sense being implicitly fitted to the test set (since we use the test score as a signal to build the next model).

(n.b I originally wrote this here on hackernews but I develop the text a bit further here)

This is very true, and I would argue there is a very prevalent misunderstanding (or just ignorance) towards the distinction between a validation and test set. When the distinction is actually made between the two, the idea is that one performs model selection on the validation set, i.e. find the best HPs such that you minimise (or maximise) some metric that operates on that subset of the data. Once you've found your most performant model according to that metric, you then evaluate that same metric on the test set. Why? Because that becomes your unbiased measure of generalisation error. Note that in a production setting you'll want to get an even better model by re-training on all the data available (train + valid + test) under those ideal HPs but that's completely fine: if somebody asks you what the generalisation error of the re-trained model is, you simply point them to the test set metric computed on the model you trained beforehand, the one where you followed the train-valid-test pipeline.

This distinction goes against the publish or perish mentality of academia. Since reviewers (and by association, researchers) are obsessed with "SOTA", "novelty", and bold numbers, a table of results purely composed of metrics computed on the test set is not easily controllable from the point of view of actually "passing" the peer review process (if you want to be ethical about it). Conversely, what's easily controllable is a table full of those same metrics computed on the validation set: just perform extremely aggressive model selection until your best model gets higher numbers than all the baselines in the table. However, rather than report separate tables for validation and test set, the common QRP is to just treat them as one and the same.

Admittedly, it is very anxiety-inducing to leave your fate up to a held-out test set whose values you can't optimise for, especially when your career is at stake. Interestingly, if your validation set numbers were great only for the test set, it would indicate you were "overfitting" via model validation. That would suggest either making the model search less aggressive or going for a simpler model class. The latter approach is called Occam's razor, but does our field really encourage simplicity? (See "Superfluous cogs" at Sec. 3.3.1 of (Leech, Gavin and Vazquez, Juan J and Yagudin, Misha and Kupper, Niclas and Aitchison, Laurence, 2024))

To distinguish this from classic contamination (training on test data), Hosseini et al. [2020] call this ‘over-hyping’ and note that it biases results even if every iteration of the cycle uses cross-validation properly.

It goes back even further than that, see (Domingos, Pedro, 2012) (back in the olden days before we had deep learning):

Cross-validation can help to combat overfitting, for example by using it to choose the best size of [model] to learn. But it is no panacea, since if we use it to make too many parameter choices it can itself start to overfit.

Even with cross-validation, we have to mitigate against this. The easiest solution is to simply hold out a test set which is independent from the cross-validation procedure. One can even have each fold of cross-validation serve as a test set (in the context of an "inner" cross-validation which handles training and model selection).

2. Over/underclaiming

To be done. There is a piece I'd like to write about the weirdness of evaluation metrics in generative models.

EDM diffusion models - a Jupyter implementation, and how they are implemented in practice

2024-05-31T00:00:00+00:00

EDM diffusion models - a Jupyter implementation, and how they are implemented in practice

1. Introduction

I wrote a self-contained implementation of NVIDIA's EDM in a Jupyter notebook. You may find it useful for the following reasons:

You want to explore diffusion models on toy datasets (why? Because toy datasets train fast and require little compute);
You want to understand diffusion models from a more fundamental perspective (toy datasets are great for that) (also see Ali Rahimi's NeurIPS '17 talk on why simple toy experiments are great);
You want to try out new schedulers which fit under the particular framework proposed by EDM (e.g. defining novel parameterisations of $\sigma(t)$ and $s(t)$).
You want a "full" implementation of the general algorithms, i.e. being able to arbitrarily specify $\sigma(t)$ and $s(t)$. (The original repo hardcodes those parameterisations.)
You want to generate fun mp4's showing the diffusion trajectory of a 1D dataset (like below).

You can find the code here.

EDM trained on a bimodal Gaussian dataset. Here we see the diffusion trajectory from t=100 (prior distribution) to t=0 (data distribution).

The notebook itself is self-contained and makes no reference to this blog post, however the following information may be useful to you if you want to know more about how EDM's algorithms are implemented for two major open source diffusion projects (diffusers and ComfyUI, the latter of which uses k-diffusion).

2. Schedulers

A particular class of scheduler algorithms implemented by k-diffusion (and by extension, ComfyUI and HuggingFace' diffusers) have rather opaque names because they don't clearly reflect their origins, nor are they complete implementations of the original algorithms from which they were derived. These algorithms are actually based on the "Algorithms 1 & 2" proposed in NVIDIA's EDM (Karras, Tero and Aittala, Miika and Aila, Timo and Laine, Samuli, 2022). Essentially, they are the deterministic and stochastic variants (respectively) of an ODE, one which was designed to encapsulate all of the different diffusion model variants used at the time. That ODE has the following form:

\begin{align} \label{orgd307d46} \dd x & = \Big[ \frac{\dot{s}(t)}{s(t)}x - s(t)^{2} \dot{\sigma}(t) \sigma(t) \nabla_{x} \log p\big(x/s(t); \sigma(t)\big) \Big] \dd t, \end{align}

where: $s(t)$ is some time-dependent scaling function of the input $x$; $\sigma(t)$ is the time dependent noise variable; and $\dot{s}(t) = \frac{\partial s(t)}{\partial t}$ and $\dot{\sigma}(t) = \frac{\partial \sigma(t)}{\partial t}$ . Along with other hyperparameters (such as how precisely the timesteps are discretised), this ODE is able to generalise the deterministic components of the sampling algorithms found in other papers.

Relationship between Algorithms 1 & 2 and how they are implemented in k-diffusion and diffusers.

Both algorithms are "somewhat" implemented in the following libraries:

k-diffusion, which takes the name sample_heun.
diffusers, which takes the name HeunDiscreteScheduler;
Both have more computationally efficient variants (i.e. without the second order steps) called sample_euler and EulerDiscreteScheduler, respectively;
but there are differences to the original paper and sometimes the implementations are not complete (in the case of diffusers).

The way Algorithm 1 is presented in the EDM paper (Karras, Tero and Aittala, Miika and Aila, Timo and Laine, Samuli, 2022) is "general" since it assumes no particular $\sigma(t)$ and $s(t)$ (see Figure 2). In terms of how they are implemented in practice:

In k-diffusion it's called sample_heun, but only if we pass $\gamma = 0$ into the function so that the stochastic part (the noise injection component proposed by Alg. 2) disappears. Also it's assumed $\sigma(t)=t$ and $s(t)=1$, so it's not the general form of the algorithm as shown in Fig. 2.
In diffusers it's called HeunDiscreteScheduler. Also it's the same parameterisation as k-diffusion, but $\gamma = 0$ is only supported (as of 21/06/2024) which means that it literally does turn into Algorithm 1. (The reason for only supporting $\gamma = 0$ seems to stem from the awkwardness of the fact that the the U-Net backbones in diffusers require discrete indices to be passed to the forward method instead of continuous values, which means one would have to back-calculate the appropriate "timestep" from $\hat{\sigma}$.)

Algorithm 2 in (Karras, Tero and Aittala, Miika and Aila, Timo and Laine, Samuli, 2022) is basically a stochastic variant of Algorithm 1, but the paper does not present the general form of the algorithm. Rather, it assumes $\sigma(t) = t$ and $s(t) = 1$ (see Figure 2). In terms of code:

k-diffusion implements it with that specific choice of $\sigma(t)=t$ and $s(t)=t$.
For diffusers, HeunDiscreteScheduler does not support $\gamma > 0$ yet and so there is no support for Algorithm 2 per se. However, EulerDiscreteScheduler does (confusingly).
While the specific choices of $\sigma(t)$ and $s(t)$ are well justified (they perform the best empirically), having the more general forms of the algorithms would open them up to exploring different forms of the general ODE.

Quite frankly, I'm not the biggest fan of these scheduler names because they don't reflect the fact they are EDM-specific algorithms (even if the attribution is there in the docstrings). Why can't we simply just implement one mega-algorithm called edm_sampler and allow the option for a use_second_order flag as well as gamma so that it encapsulates everything? Or at least use class names like edm_deterministic_sampler, edm_stochastic_sampler, etc. I suppose the reason whh they are named so "generically" (for lack of a better term) is that the general-form ODE proposed by EDM really does encompass (as of time of writing) "more or less" all of the diffusion variants commonly used. Therefore, to just give it a name like "{Euler,Heun}DiscreteScheduker" is not unreasonable.

Lastly, there is one additional algorithm which shares the same naming convention as the others but really has nothing to do with the EDM paper. This is the "ancestral sampling algorithm" based off Jonathon Ho's DDPM paper (Ho, Jonathan and Jain, Ajay and Abbeel, Pieter, 2020). In diffusers it's called EulerAncestralDiscreteSampler (see here) and in k-diffusers it's called sample_euler_ancestral (see here). More info on that is in Sec. 2.1.

Algorithm 1 from EDM. Here the sigma and scaling functions remain arbitrary.

Algorithm 2 from EDM. Here we can see a specific parameterisation for the sigma and scaling functions. Otherwise, the general form of this algorithm should match Algorithm 1 on top of the stochastic component.

2.1. Ancestral sampler

Both k-diffusion and diffusers have a version of the Euler-based version of Algorithm 2. To make matters even more confusing, the sample_euler_ancestral algorithm is basically the Euler variant of Algorithm 2 but with the EDM-specific noise injection mechanism cut out in favour of ancestral sampling. Ancestral sampling is detailed in Appendix F of the continuous-time score matching paper from Song et al. (Song, Yang and Sohl-Dickstein, Jascha and Kingma, Diederik P and Kumar, Abhishek and Ermon, Stefano and Poole, Ben, 2020). The update rule for this is:

\begin{align} x_{i+1} & = x_i + (\sigma_{i+1}^2 - \sigma_{i}^2) s_{\theta}(x_i, \sigma_i) + \sqrt{\frac{\sigma_{i+1}^2(\sigma_{i}^{2}-\sigma_{i+1}^2)}{\sigma_i^2}}z_i \\ & = x_i + (\sigma_{i+1}^2 - \sigma_{i}^2) s_{\theta}(x_i, \sigma_i) + \frac{\sigma_{i+1}}{\sigma_i} \sqrt{\sigma_i^2 - \sigma_{i+1}^2} z_i, \end{align}

where $s_{\theta}(x, \sigma) = (x - D(x; \sigma) / \sigma^2$ and $z_i \sim \mathcal{N}(0, \mathbf{I})$. (Unlike in (Song, Yang and Sohl-Dickstein, Jascha and Kingma, Diederik P and Kumar, Abhishek and Ermon, Stefano and Poole, Ben, 2020), I am being consistent with the rest of this post by denoting $\sigma_0$ as the highest noise scale as $\sigma_{T-1}$ as the smallest.)

This equation can basically be seen as doing the ODE step (the first two terms on the RHS) but then injecting noise $\sim \mathcal{N}(0, \sigmaup)$. For reasons not clear to me yet, this is not the exact same as what's implemented in k-diffusion (see here and here), which implements something seemingly a lot more complicated:

\begin{align} x_{i+1} = x_i + (\sigmadown - \sigma_{i}) s_{\theta}(x_i, \sigma_i) + \underbrace{\text{min}\Big(\frac{\sigma_{i+1}}{\sigma_i} \sqrt{\sigma_i^2 - \sigma_{i+1}^2}, \sigma_{i+1}\Big)}_{\sigmaup} z_i, \end{align}

and $\sigmadown = \sqrt{\sigma_{i+1}^2 - \sigmaup^2}$. (I've also redefined $\sigmaup$ here to also include the min.) (If anyone knows more about this, please reach out so I can update this post.)

3. Conclusion

In conclusion, I have shared a Jupyter implementation of EDM on toy datasets, as well as elucidate (pun intended) some of the opaque naming conventions used in the practical implementations which implement EDM's algorithms.

A deep dive into conditional variational autoencoders

2023-04-27T00:00:00+00:00

A deep dive into conditional variational autoencoders

Updates:

(06/10/2023) Some fixes to the text, and an overhaul of the discussion section to connect it better with the rest of the post.

1. Introduction
- 1.1. Contributions
2. The generative and inference process
3. Experiments
- 3.1. When z and y are independent
  - 3.1.1. Controllable generation
- 3.2. When z and y are dependent
4. Discussion
- 4.1. cVAEs and conditional Gaussian diffusion models
5. Conclusion
6. Appendix
7. References

1. Introduction

This is yet another excerpt from my upcoming PhD thesis. I actually wanted to write this several years ago after some really painful experiences I had with getting conditional VAEs to work on a generative modelling project I was working on. To the best of my knowledge, I haven't seen any paper that talks in depth about these difficulties and so I am quite happy to finally share them with everyone. VAEs – despite their conceptual simplicity – can be difficult to understand and even more so for its conditional variants. In this post I will dive into the theory of conditional VAEs, derive interesting equations which elucidate their behaviour, and corroborate those insights on a simple toy dataset.

1.1. Contributions

This post focuses almost exclusively on conditional VAEs, but it also equally applies to unconditional ones. The contributions of this post are as follows:

In Section 2 we present cVAEs through an unconventional but rather enlightening perspective, inspired by Esmaeili et al. (2018). This involves thinking about the VAE as parameterising two separate pathways (the generative and inference process), and the evidence lower bound can be derived as the KL divergence between these two. While their work was derived assuming unconditional VAEs, we consider the conditional case as well.
We discuss two parameterisations of a cVAE: one where the conditioning variable $\yy$ and latent variable $\zz$ are assumed to be either independent (Section 2.1) or dependent (Section 2.2). The former is a useful parameterisation if one is interested in performing controllable generation.
Through the lens of mutual information estimation, we elucidate the difficulties involved in training such a class of models (Section 2.4). In particular, we show such that optimising VAEs is a careful balance between ensuring that sample quality and diversity are adequate for both the generative and inference processes.
We present experiments corroborating our theoretical analyses on a toy 2D dataset consisting of two Gaussian clusters, where a cVAE must be trained to sample from either of the two clusters correctly (Section 3).
We discuss how VAE training issues can be avoided by considering adversarial learning or hybrid adversarial / VAE-style models.

2. The generative and inference process

VAEs are typically derived by starting off by assuming a latent variable model of the form $\pt(\xx,\zz)$, and noting that integrating this expression over $\zz$ to obtain $\pt(\xx)$ is intractable. There is a further assumption that $\pt(\xx,\zz) = \pt(\xx|\zz)p(\zz)$, but we also don't know what inputs $\xx$ correspond to what $\zz$, and since deriving $\pt(\zz|\xx)$ is also intractable we need to introduce a separate network $\qp(\zz|\xx)$ to do the job for us.

Inspired by esmaeili2018structured, we can actually derive the ELBO for a VAE by framing the training objective as a minimisation over the KL divergence between two pathways, each encoded by their own joint distribution. (We already saw the first of these joint distributions previously, which is the generative distribution denoted $\pt$.) Since we're also talking about conditional VAEs, we will be dealing with an additional latent variable $\yy$, but unlike $\zz$ we have labels for this. We assume that $\yy$ encodes some semantic meaningful label of interest, for instance the class of a digit, or the identity of an object.

The first pathway is the inference process, denoted $\qp(\xx,\zz,\yy)$. It factorises into $\qp(\zz|\xx,\yy)q(\xx,\yy)$, and to obtain a sample $(\zz,\xx,\yy)$ from this joint we simply perform the following:

\begin{align} \label{eq:inference} \xx, \yy & \sim q(\xx, \yy) \ \ \text{(ground truth)} \tag{2a} \\ \zz & \sim \qp(\zz|\xx, \yy) \tag{2b} \end{align}

where $q(\xx,\yy)$ is the ground truth data distribution, and $\qzgivenx$ is our learnable variational posterior, subscripted with $\phi$. The inference process is concerned with extracting latent representations from actual samples from the data distribution. This is to be contrasted with the generative process, in which samples are generated as the following:

\begin{align} \label{eq:generative} \zz, \yy & \sim p(\zz,\yy) \tag{3a} \ \ \text{(prior)} \\ \xx &\sim \pt(\xx|\zz,\yy) \tag{3b}, \end{align}

where $p(\zz,\yy)$ is prescribed beforehand. (We will talk a little more about this shortly.)

Since joint distribution for both processes are $\ptgreen(\xx,\zz,\yy)$ and $\qp(\xx,\zz,\yy)$ and we can derive their KL distribution as follows:

\begin{align} \label{eq:case1} \argmax_{\color{green}{\theta}, \color{purple}{\phi}} & -\kldiv \Big[ \qp(\XX,\ZZ,\YY) \ \| \ \ptgreen(\XX,\ZZ,\YY) \Big] \\ & = \mathbb{E}_{\qp(\xx,\zz,\yy)}\big[ \log \frac{\pt(\xx,\zz,\yy)}{\qp(\xx,\zz,\yy)} \big] \tag{4a} \\ & = \mathbb{E}_{\qp(\zz|\xx,\yy)}\big[ \log \frac{\pt(\xx | \yy, \zz)p(\yy,\zz)}{\qp(\zz|\xx,\yy)} \big] - \mathbb{E}_{q(\xx,\yy)} \log q(\xx, \yy) \tag{4b} \\ & = \mathbb{E}_{\qp(\xx,\zz,\yy)}\big[ \log \frac{\pt(\xx | \yy, \zz)p(\yy, \zz)}{\qp(\zz|\xx,\yy)} \big] - \text{const.} \tag{4c} \\ & = \mathbb{E}_{\qp(\xx,\zz,\yy)} \big[ \log \pt(\xx|\yy,\zz) \big] + \mathbb{E}_{\qp(\zz|\xx,\yy)} \big[ \log \frac{p(\yy, \zz)}{\qp(\zz|\xx,\yy)} \big] - \text{const.} \tag{4d} \\ & = \mathbb{E}_{\qp(\zz,\xx,\yy)}\big[ \log \pt(\xx|\yy,\zz) \big] - \kldiv\Big[ \qp(\ZZ|\XX, \YY) \| p(\ZZ,\YY)\Big], \tag{4e} \end{align}

which gives us the typical formulation of the ELBO which we see in most VAE papers.

2.1. When z and y are independent

At this point, we have to specify what $p(\zz,\yy)$ is, and we have two options. The first is to assume that $p(\zz,\yy) = p(\zz)p(\yy)$, i.e. they are independent. This means that the joint distribution of the generative process factorises into:

\begin{align} \pt(\xx,\zz,\yy) = \pt(\xx|\zz,\yy)p(\zz)p(\yy) \tag{5} \end{align}

which leads us to the following ELBO:

\begin{align} & -\kldiv \Big[ \qp(\XX,\ZZ,\YY) \ \| \ \ptgreen(\XX,\ZZ,\YY) \Big] \tag{6a} \\ & \myeq{\text{if ind.}} \mathbb{E}_{\qp(\zz,\xx,\yy)}\big[ \log \pt(\xx|\yy,\zz) \big] + \mathbb{E}_{\qp(\zz,\xx,\yy)}\big[ \log \frac{\pgreen(\zz)}{\qp(\zz|\xx,\yy)} \big] + \log \pgreen(\yy) \tag{6b} \\ & = \text{likelihood} - \kldiv\Big[ \qp(\ZZ|\XX,\YY) \| p(\ZZ) \Big] + \text{constants}. \tag{6c} \end{align}

Here, $p(\yy)$ is some prior for $\yy$ but it falls out of the KL term since it is a constant, so we need not worry about it. All that is left is to define a prior for $p(\zz)$, and in practice this is most often an isotropic Gaussian distribution. The graphical model for the $\color{green}{\text{generative process}}$ is also shown in Figure 1.

Figure 1: Graphical model of the generative process, corresponding to the case where p(z,y) = p(z)p(y).

Such a factorisation may be useful to encode if we are seeking to learn disentangled representations. For instance, if we were learning a conditional VAE over SVHN digits (where $y$ encodes the identity of the digit), perhaps we would like for our VAE to learn a $\zz$ that encodes \emph{everything else} in the image apart from the digit itself (for instance background details and font style). This would make for a very controllable generative process where we could arbitrarily mix and match style and content variables from different examples to create new ones.

2.2. When z and y are dependent

Otherwise, $\pgreen(\zz,\yy) = \pgreen(\zz|\yy)\pgreen(\yy)$ and $\pgreen(\zz|\yy)$ is the conditional prior. This means that the joint distribution factorises into:

\begin{align} \pt(\xx,\zz,\yy) = \pt(\xx|\zz,\yy)p(\zz|\yy)p(\yy) \tag{7} \end{align}

The conditional prior can either be fixed (i.e. each possible value of $\yy$ gets mapped to a Gaussian), or it can be learned, in which case we denote it as $\pt(\zz|\yy)$. In this case the ELBO can be written as:

\begin{align} & -\kldiv \Big[ \qp(\XX,\ZZ,\YY) \ \| \ \ptgreen(\XX,\ZZ,\YY) \Big] \tag{8a} \\ & = \mathbb{E}_{\qp(\zz,\xx,\yy)}\big[ \log \pt(\xx|\yy,\zz) \big] + \mathbb{E}_{\qp(\zz,\xx,\yy)}\big[ \log \frac{p(\zz|\yy)}{\qp(\zz|\xx,\yy)} \big] + \log p(\yy) \tag{8b} \\ & = \text{likelihood} - \kldiv\Big[ \qp(\ZZ|\XX,\YY) \ \| \ p(\ZZ|\YY) \Big] + \text{constants}. \tag{8c} \end{align}

Consequently, the graphical model for the $\color{green}{\text{generative process}}$ is shown in Figure 2.

Figure 2: Graphical model of the generative process, corresponding to the case where p(z,y) = p(z|y)p(y).

2.3. The role of the beta term

Let us look at both versions of the ELBO, equations 6(c) and 8(c), and write them as minimisations over $\thetagr, \phip$:

\begin{align} \text{dep.} \rightarrow & \min_{\thetagr, \phip} -\mathbb{E}_{\qp(\zz,\xx,\yy)}\big[ \log \pt(\xx|\yy,\zz) \big] + \beta\kldiv\Big[ \qp(\ZZ|\XX,\YY) \ \| \ p(\ZZ|\YY) \Big] \tag{9a} \\ \text{indep.} \rightarrow & \min_{\thetagr, \phip} -\mathbb{E}_{\qp(\zz,\xx,\yy)}\big[ \log \pt(\xx|\yy,\zz) \big] + \beta\kldiv\Big[ \qp(\ZZ|\XX,\YY) \ \| \ p(\ZZ) \Big] \tag{9b}, \end{align}

where 'dep' and 'indep' are shorthand for 'dependent' and 'independent'. Also note that since the independent case is assuming $p(\zz,\yy) = p(\zz)p(\yy)$ we could also define $\qp(\zz|\xx,\yy) = \qp(\zz|\xx)$ to remove the dependence on $\yy$, but to keep notation consistent we will leave it in for the remainder of this post.

What makes VAE training difficult to get right is the interplay between the two terms in each equation. The first equation is maximising the likelihood of the data with respect to samples from the inference network. In order for this to happen, $\zz$ should encode as much information about $\xx$ as possible through the variational posterior $\qp$, which is our learned encoder. At the same time however, the second term is working against the first, because it is enforcing that each per example variational posterior must be close to the prior distribution¹. Since the prior is not a function of $\XX$ it implies that some information about $\XX$ in the encoding pathway has to be lost. Essentially, we are trading off between sample quality with respect to:

the inference pathway, which is $\qp(\zz,\xx,\yy) = \qp(\zz|\xx,\yy)q(\xx,\yy)$, where $q(\xx,\yy)$ is the ground truth joint distribution;
and the generative pathway, which is $\pt(\zz,\xx,\yy) = p(\zz,\yy)\pt(\xx|\zz,\yy)$,

and hence why it is useful to know that the evidence lower bound in Eqn. (9) is a direct result of minimising the KL divergence between those two distributions.

In practice, what one observes with a VAE as a function of $\beta$ is the following:

if $\beta$ is too small then samples from the prior distribution $\zz \sim p(\zz)$ will not look as good as samples from the variational encoder $\zz \sim \qp(\zz|\xx,\yy)$;
if $\beta$ is too large then sample quality with respect to both will be degraded, and hence the search for $\beta$ is a careful balance between the two extremes;
and if $\beta$ is 'just right', sample quality with respect to both should be 'ok'.

In Figure 5 we show images from an unconditional VAE illustrating this trade-off for MNIST.

Figure 5a: β=0, so no KL regularisation. Samples from the generative process look bad and ones from the generative process look good.

Figure 5a: β=1000, too much regularisation, so both pathways look bad. Since the encoder q(z|x) has essentially collapsed into p(z), all samples look the same.

Figure 5a: β=0.05. Samples from both processes look decent.

2.4. A mutual information perspective for the KL term

This aforementioned loss of information due to $\kldiv\big[ \qp(\ZZ|\XX,\YY) \ \| \ p(\ZZ, \YY) \big]$ can be theoretically shown, by re-writing the KL term to be the sum of a mutual information term and another KL divergence term.

For the dependent case:

\begin{align} \text{dep.} & \rightarrow \kldiv \Big[ \qp(\ZZ|\XX,\YY) \| p(\ZZ|\YY) \Big] \\ & = \mathbb{E}_{\qp(\zz,\xx,\yy)} \log \frac{\qp(\zz|\xx,\yy)}{p(\zz|\yy)} \tag{10a} \\ & = \mathbb{E}_{\qp(\zz,\xx,\yy)} \log \Big[ \frac{\qp(\zz|\xx,\yy)}{p(\zz,\yy)} \cdot \frac{\qp(\zz)}{\qp(\zz)} \Big] \tag{10b} \\ & = \mathbb{E}_{\qp(\zz,\xx,\yy)} \log \Big[ \frac{\qp(\zz|\xx,\yy)}{\qp(\zz)} \cdot \frac{\qp(\zz)}{p(\zz,\yy)} \Big] \tag{10c} \\ & = \mathbb{E}_{\qp(\zz,\xx,\yy)} \log \frac{\qp(\zz|\xx,\yy)}{\qp(\zz)} + \mathbb{E}_{\qp(\zz,\yy)} \frac{\qp(\zz)}{p(\zz,\yy)} \tag{10d} \\ & = I_{\phip}(\ZZ; \XX, \YY) + \kldiv[ \qp(\ZZ) \| p(\ZZ|\YY) ] - \underbrace{\mathbb{E}_{\qp(\yy)} \log p(\yy)}_{\text{const}} \tag{10e} \end{align}

Similarly, for the independent case we obtain:

\begin{align} \text{indep.} & \rightarrow \kldiv \Big[ \qp(\ZZ|\XX,\YY) \| p(\ZZ) \Big] \nonumber \\ & = \kldiv \Big[ \qp(\ZZ|\XX) \| p(\ZZ) \Big] \nonumber \\ & = I_{\phip}(\ZZ; \XX, \YY) + \kldiv[ \qp(\ZZ) \| p(\ZZ) ] - \text{const}. \tag{10f} \end{align}

In either of the two cases, the minimisation of their respective KL terms implies minimising the mutual information between $\XX$ and the pair $(\ZZ,\YY)$, denoted as $I_{\phip}(\ZZ; \XX, \YY)$. Therefore, when we increase $\beta$ we are inevitably reducing the information $\ZZ$ stores about $\XX$ with respect to the encoder $\qp$.

2.5. A mutual information perspective between Z and Y

In the previous section we showed how minimising the KL term in the ELBO involves also minimising the mutual information between $\ZZ$ and $\XX,\YY$ through its decomposition in Eqn. (10e) and (10f), and that it is a consequence of trying to match the generative and inference distributions. Furthermore, the extent to which we try to minimise this equation affects the relative difference in sample quality between $\zz$'s which are sampled from the prior distribution versus ones generated with the variational distribution.

Minimising the mutual information between $\ZZ$ and $\YY$ for $\ZZ,\YY$ independent VAEs is also important since we want the two variables to encode completely separate concepts. For instance, it is common in image datasets for $\YY$ to encode something semantically desirable about $\XX$, for instance the identity of the object in the foreground or what category it belongs to. If our dataset is labelled such that $\YY$ is assigned such semantic meaning, then we would like $\ZZ$ to encode everything else that is not related to $\YY$.

From Sec. 2.4 we showed that minimising the per-example KL means also minimising $I_{\phip}(\XX; \ZZ)$. In actuality it would be nice to instead minimise the mutual information between $\ZZ$ and $\YY$ (even though this term is not present in the equation), but the issue is that $\XX$ \emph{also encodes} information about $\YY$, and so trying to drive down $I_{\phip}(\ZZ; \YY)$ would inevitably mean we need to drive down $I_{\phip}(\ZZ; \XX)$, but this degrades sample quality². In the absence of extra supervisory signal³ that could potentially encourage the network to only encode the `non-label' parts of $\XX$ in $\ZZ$, we are stuck with a very difficult optimisation problem.

2.5.1. Practical considerations

Figure 6: In practice, if too much information about Y is encoded in Z via the inference network, then the conditioned Y for the decoder may have little to no influence on the output (the corresponding edge is shown as a dotted red line).

In practice, if the KL term is not large enough (Eqn. (9b)) then the decoder $\pt(\xx|\zz,\yy)$ will ignore the $\YY$ variable. This is presumably because $\ZZ$ will contain too much information about $\YY$ which in turn renders it irrelevant with respect to the decoder (Figure 6). This is an issue because it prevents us from performing controllable generation. Essentially, given some input $\xx$ if we can encode it into its (independent) factors of variation $\zz, \yy$ then we could easily swap out $\yy$ with a new label $\yy'$ and decode to produce a different kind of output (see Sec. 6.4 for an example):

\begin{align} (\xx, \yy) & \sim \mathcal{D} \tag{12a} \\ \yy' & \sim p(\yy) \tag{12b} \\ \zz & \sim \qp(\zz|\xx,\yy) \tag{12c} \\ \xx' & \sim \pt(\xx|\zz,\yy') \tag{12d} \end{align}

If the KL term is not weighted high enough however then $\yy'$ won't make any difference whatsoever. Unfortunately, it is difficult to tell whether this is happening through monitoring the ELBO. Basically, one will need to figure out via cross-examination what the 'largest' value for the KL term can be before $\yy$ gets ignored by the decoder.

3. Experiments

We now present some experiments on a toy 2D dataset for both variants of cVAE. The dataset consists of two Gaussians, and the ground truth is:

\begin{align} p(\xx) = \sum_{i \in \{0,1\} }p(\xx,\yy_i) = \sum_{i \in \{0,1\}} p(\xx|\yy_i)p(\yy_i), \end{align}

where $p(\xx|\yy=0) = \mathcal{N}(\xx; [-2.5, 1]^{T}, 2\mathbf{I})$, $p(\xx|\yy=1) = \mathcal{N}(\xx; [6,-2]^{T}, 2 + \mathbf{I})$, and $p(\yy=0) = p(\yy=1) = \frac{1}{2}$. Samples from this distribution are visualised below in Figure 3.

Figure 3: Illustration of the toy 2D dataset used. The dataset comprises of two Gaussians, each corresponding to one of two binary labels (y=0 or y=1).

For the following experiments, we train a single hidden layer MLP for both the encoder and decoder. The encoder is a mapping $\mathbb{R}^{2} \rightarrow \mathbb{R}^{h} \rightarrow \mathbb{R}^{2}$ which means the latent variable is also two-dimensional, for interpretability sake. Likewise, the decoder is of a similar mapping.

For the following experiments, we wish to illustrate the behaviour of a conditional VAE with respect to the following attributes: (1) whether $\ZZ$ and $\YY$ are dependent or not; (2) as a function of increasing the KL regularisation coefficient $\beta$. Furthermore, we wish to illustrate both behaviours in input space $\mathcal{X}$ as well as latent space $\mathcal{Z}$. For convenience, both the input and latent spaces are two-dimensional, and subsequent figures will make it clear which space is being visualised.

Concretely, the encoder $\qp(\zz|\xx)$ is an MLP $\mathbb{R}^{p=2} \rightarrow \mathbb{R}^{h} \rightarrow \mathbb{R}^{p=2}$ for $h$ hidden units. Likewise, the decoder takes on a similar structure.

3.1. When z and y are independent

First we show $\beta = 0$, illustrated in Figure 3. Samples from the inference process are shown in $\color{purple}{\text{purple}}$ and those from the generation process in $\color{green}{\text{green}}$, similar to the notation that we have been using so far in the equations. For instance if we consider the inference process: for a given $(\xx, \yy)$ from the data distribution, we sample $\zz \sim \qp(\zz|\xx,\yy)$ and then we reconstruct by sampling $\tilde{\xx} \sim \pt(\xx|\zz,\yy)$. The corresponding reconstruction error is shown in the title (the squared L2 norm between the original points and their reconstructions), and we can see that the error is small enough we can essentially consider it to be zero. However, things don't look so good for the generative process: for a given $\zz \sim p(\zz)$, we can either choose to decode with $\pt(\xx|\zz,\yy=0)$ or $\pt(\xx|\zz,\yy=1)$, and these more or less fall in the same region. This indicates that choosing $\yy$ does not make a difference to the generated samples (recall Fig. 6 in Sec. 2.5.1). What we would like to see is the samples from the prior falling into their respective clusters.

Figure 3a: β = 0. Here, there is no weight on the KL term, so reconstructions are good and there so is the inference process. However, samples from p(z|y=0) or p(z|y=1) (when decoded) fall in the same region. Overall, with respect to the generative process, sample quality and sample diversity are bad.

We can also visualise samples in latent space as well as the distributions for $p(\zz)$ as well as the conditional inference distributions $\qp(\zz|\yy_i)$, and this is shown below in Fig. (3b). (Note that $\qp(\zz)$ the inference marginal itself is also just the weighted sum of both of these distributions, weighted by their prior probability $q(y=i)$.)

Figure 3b: β = 0, showing samples in z space, which is also two-dimensional. The prior distribution p(z) is shown as the green sphere. We can see that there significant mutual information between Z and Y here, and this is because it is easy to tell apart the two clusters.

In Figure 4a, if we choose $\beta = 0.01$, it looks as though some of the green points have been pulled to their respective cluster but there is still some overlap between the two categories and we don't see any clear pattern of separation. At the very least, sample diversity is superior to that in Figure 1 because at least the green points are sufficiently spread out to cover the two clusters of the data. The reconstruction error for the inference process has only taken a minor hit, increasing from roughly zero to $\approx 0.02$. In Figure 4b, we can see that the marginal $\qp(\zz)$ is a little closer to the prior, but it's still easy to make out the two separate clusters belonging to the different $\yy$'s, so $I_{\phip}(\ZZ; \XX, \YY)$ is still reasonably large.

(Figure 4a, top) Reconstructions are decent and there so is the inference process. Samples from the generative process still do not appear to respect their clusters but unlike Figure 1 we see an acceptable level of sample diversity here, since those samples are covering more regions of the data distribution. Overall, with respect to the generative process, sample quality is bad but sample diversity is good.

(Figure 4b, bottom) Samples from q(z) are somewhat close to the prior p(z). We can see that there is significant mutual information between Z and Y here, and this is because it is easy to tell apart the two clusters.

Finally, in Figure 5 for $\beta = 1$ we finally see that the green points get matched to their respective clusters. Unfortunately, the inference process has degraded and reconstruction error has significantly increased as as result ($\approx 1.61$). We can also see this qualitatively for the rightmost cluster, where reconstructions lie on a very narrow subspace instead of being more evenly distributed across the cluster. Therefore, we can say that with respect to both processes, sample quality is \emph{very good} but sample diversity has \emph{degraded}. Lastly, note that in Figure 5b the two condtionals $\qp(\zz|\yy=0)$ and $\qp(\zz|\yy=1)$ are more or less the same, which indicates roughly zero mutual information between $\ZZ$ and $\YY$. Because of this, the autoencoder will now be `incentivised' to make use of $\yy$ since it will obviously be a useful variable to leverage use when maximising the log likelihood of the data (assuming $\beta$ is not too large, since it controls the degree to which the optimisation focuses on the likelihood term).

(Figure 5a, top): Sample diversity has suffered with respect to both inference and generative distributions, and this can be seen in the right-most cluster (points lie on a narrow line). Quantitatively, the encoder q_Φ(z|x,y) is no longer very accurate, with a reconstruction error of ~1.61.

(Figure 5b, bottom) q_Φ(z) looks more or less the same as p(z). Here, the distributions q_Φ(z|y=0) and q_Φ(z|y=1) appear roughly the same, so we can say that I_Φ(Z; Y) is small.

3.1.1. Controllable generation

One benefit of training a $\ZZ,\YY$ independent VAE is that we can perform controllable generation more easily (or at least hope to) compared to the dependent variant. For instance, if $\ZZ$ and $\YY$ encode the non-semantic and semantic parts of the input, we could generate a novel example by combining the semantic content of one input with the non-semantic content of another. In this case, $\YY$ is a binary random variable indicating the cluster:

\begin{align} (\xx,\yy) & \sim \mathcal{D} \tag{13a} \\ \zz & \sim \qp(\zz|\xx,\yy) \tag{13b} \\ \xx' & \sim \pt(\xx|\zz,1-\yy) \tag{13c} \end{align}

Similar to Sec. 3.1 we illustrate this with increasing values of $\beta$ starting from zero. See Figures 7(a,b,c) and their associated captions.

Figure 7a: β = 0. Label swapping doesn't seem to do anything (pink points don't switch cluster).

Figure 7b: β = 0.01. Label swapping has a marginal effect but label-swapped samples in pink are spread out between both clusters.

Figure 7c: β = 1.0. Label swapping looks like it works now, albeit at the cost of sample diversity for the right-most cluster.

As we can see, when $\beta$ is large enough we see the label swapping experiments properly take effect.

3.2. When z and y are dependent

When $\zz$ and $\yy$ are dependent then $p(\zz,\yy) = p(\zz|\yy)p(\yy)$. Either we fix the conditional prior $p(\zz|\yy)$ a-priori and manually define both $p(\zz|\yy=0)$ and $p(\zz|\yy=1)$, or we learn the conditional prior instead, in which case we can substitute the term with $\pt(\zz|\yy)$ instead. Learning the conditional prior simply means including four extra parameters in $\theta$ that comprise the mean and variance of the Gaussians corresponding to $\yy=0$ and $\yy=1$.

In Figures 8(a,b,c) we produce similar plots to that of Sec. 3.1.

Figure 8a: β = 0 with the learned conditional prior. Reconstruction error shown in the title.

Figure 8b: β = 0.01 with the learned conditional prior. Reconstruction error shown in the title.

Figure 8c: β = 1.0 with the learned conditional prior. Reconstruction error shown in the title.

We also show an additional set of plots showing what the samples look like in latent space, as well as where the learned conditional priors $\pt(\zz|\yy=0)$ and $\pt(\zz|\yy=1)$ are located. These are shown below in Figure 9.

Figure 9a: β = 0 with the learned conditional priors, shown in green.

Figure 9b: β = 0.01 with the learned conditional priors, shown in green.

Figure 9c: β = 1.0 with the learned conditional priors, shown in green.

Here, we observe something interesting: each posterior $\qp(\zz|\yy_i)$ has been matched to its respective conditional prior $\pt(\zz|\yy_i)$, and we can explicitly show this by rewriting the KL loss to remove the $\XX$ in the conditioning part of $\kldiv\big[ \qp(\ZZ|\XX,\YY) \ \| \ \pt(\ZZ |\YY) \big]$:

\begin{align} & \min_{\phip, \thetagr} \kldiv\Big[ \qp(\ZZ|\XX,\YY) \ \| \ \pt(\ZZ | \YY) \Big] \tag{14a} \\ & = \min_{\phip, \thetagr} \mathbb{E}_{\qp(\xx,\zz,\yy)} \Big[ \log \frac{\qp(\zz|\xx,\yy)}{\pt(\zz|\yy)} \Big] \tag{14b} \\ & = \min_{\phip, \thetagr} \mathbb{E}_{\qp(\xx,\zz,\yy)} \Big[ \log \frac{\qp(\zz|\xx,\yy)}{\qp(\zz)} \cdot \frac{\qp(\zz|\yy)}{\pt(\zz|\yy)} \cdot \frac{\qp(\zz)}{\qp(\zz|\yy)} \Big] \tag{14c} \\ & = \min_{\phip, \thetagr} \mathbb{E}_{\qp(\xx,\zz,\yy)} \Big[ \log \frac{\qp(\zz|\xx,\yy)}{\qp(\zz)} \Big] + \mathbb{E}_{\qp} \Big[ \log \frac{\qp(\zz|\yy)}{\pt(\zz|\yy)} \Big] + \mathbb{E}_{\qp} \Big[ \log \frac{\qp(\zz)}{\qp(\zz|\yy)} \Big] \tag{14d} \\ & = \min_{\phip, \thetagr} \mathbb{E}_{\qp(\xx,\zz,\yy)} \Big[ \log \frac{\qp(\zz|\xx,\yy)}{\qp(\zz)} \Big] + \mathbb{E}_{\qp} \Big[ \log \frac{\qp(\zz|\yy)}{\pt(\zz|\yy)} \Big] - \mathbb{E}_{\qp} \Big[ \log \frac{\qp(\zz|\yy)}{\qp(\zz)} \Big] \tag{14e} \\ & = \min_{\phip, \thetagr} I_{\phip}(\ZZ; \XX, \YY) + \underbrace{\kldiv\Big[ \qp(\ZZ|\YY) \| \pt(\ZZ|\YY) \Big]}_{\text{match these two!}} - I_{\phip}(\ZZ; \YY). \tag{14f} \end{align}

We emphasise the second term, which is the KL divergence between the variational posterior marginalised over $\XX$ and conditional prior.⁴

4. Discussion

So far we have seen that the ability for either conditional VAE to be able to decode samples from the prior is heavily dependent on the value of $\beta$ that is chosen. From Section 2.4 we showed that this inevitably comes at a cost, which is reducing the mutual information between $\XX$ and $\ZZ$ with respect to the encoder $\qp$. This means that sample quality becomes degraded. Based on what we have seen so far we can say the following about $\beta$:

(1) For any type of VAE (conditional or unconditional), it is crucial to tune $\beta$ (the `per-example' KL) in order to balance the trade-off between sample quality and sample diversity with respect to both the inference and generative distributions. The effect of increasing $\beta$ however means reducing the mutual information between $\ZZ$ and $\XX$, and this degrades sample quality with respect to either process (Section 2.4).
(2) For $\ZZ,\YY$ dependent cVAEs, increasing $\beta$ increases the strength of the KL term which matches the conditional priors to their respective variational posteriors (which we showed in Eqn. (14f)). As per (1) however this also means sample quality degrades.
(3) For $\ZZ,\YY$ independent cVAEs, we do not want $\ZZ$ to contain any information about $\YY$. While we do not explicitly have a term for $I_{\phip}(\ZZ;\YY)$, an increase in $\beta$ implicitly decreases it since it corresponds to decreasing $I_{\phip}(\ZZ;\XX)$. Again, as per (1) this corresponds to a decrease in sample quality.

While there is a vast literature proposing improved variants of the VAE, arguably its core design is too restrictive, and that there is always going to be a trade-off between the quality of the inference and generative distributions. We can also highlight this difficult with only a few lines of derivations. To keep things simple, let us assume an unconditional VAE and therefore a KL between the following joints:

\begin{align} & \min_{\phip,\thetagr} \kldiv \Big[ \qp(\XX,\ZZ) \ \| \ \ptgreen(\XX,\ZZ) \Big] \tag{15a} \\ & = \mathbb{E}_{\qp(\zz,\xx)} \log \frac{\qp(\zz,\xx)}{\pt(\xx,\zz)} \tag{15b} \\ & = \mathbb{E}_{\qp(\zz,\xx)} \log \frac{\qp(\zz|\xx)q(\xx)}{\pt(\xx|\zz)p(\zz)} \tag{15c} \\ & = \mathbb{E}_{\qp(\zz,\xx)} \log \Big[ \frac{\qp(\zz|\xx)q(\xx)}{\pt(\xx|\zz)p(\zz)} \cdot \frac{\qp(\zz)}{\qp(\zz)} \Big] \tag{15d} \\ & = \underbrace{\mathbb{E}_{\qp(\zz,\xx)} \log \frac{\qp(\zz|\xx)}{\qp(\zz)}}_{I_{\phip}(\XX; \ZZ)} + \mathbb{E}_{\qp} \log \frac{\qp(\zz)}{p(\zz)} + \mathbb{E}_{\qp} \log \frac{q(\xx)}{\pt(\xx|\zz)}, \tag{15e} \end{align}

where we see that the first term is a minimisation of the mutual information between $\ZZ$ and $\XX$ with respect the inference network.

This begs the question as to what could be done to make it easier to train cVAEs while minimising the loss of sample quality. Arguably the most difficult variant to get `right' is the $\ZZ,\YY$ independent VAE, because we have the added constraint that $\ZZ$ should not contain any information about $\YY$, but in order to reduce $I_{\phi}(\ZZ;\YY)$ we also need to inevitably reduce $I_{\phi}(\ZZ; \XX)$ as well (Section 2.5). While one could `hack' the ELBO by decreasing $\beta$ while also adding a term which is intended to maximise $I_{\phip}(\ZZ;\YY)$, from personal experience such attempts have not worked at all. This is most likely because the likelihood term is simply contradicting everything else: recall that it is maximising the log likelihood of the data $\xx$ given latent code $\zz$ from the encoder, and the mutual information between $\XX$ and $\ZZ$ must be large in order to do that.

From personal experience, $\ZZ,\YY$ independent generative models are trivial to get working with GANs. However, they they sit on the opposite spectrum of the sample quality and diversity trade-off: VAEs suffer in terms of the former while GANs suffer in terms of the latter. In order to combine the best of both worlds, many works have been proposed to either combine VAEs and GANs (makhzani2015adversarial, larsen2016autoencoding, mescheder2017adversarial) or propose GANs which can also perform inference (chen2016infogan, dumoulin2016adversarially, donahue2016adversarial). For instance, one of the simplest additions to do this is `InfoGAN' chen2016infogan, which simply proposes that one adds an extra output branch to the discriminator to predict any of the latent codes passed into the generator (i.e. $\zz$ and $\yy$). Then the final loss is the usual two-player minimax game but both generator and discriminator optimise their parameters to minimise this prediction loss. While the original motivation of this paper was to mitigate loss of sample diversity (as is common with GANs), another benefit is that the discriminator $D$ can act as an inference network.

Another possible solution is to simply forego the idea of trying to optimise the two distributions (generative and inference) to be close to each other since it results in contradicting losses (for instance, likelihood vs per-example KL). Instead, we could simply train a deterministic autoencoder – whose only purpose is inference – but in parallel train a sampler network (e.g. a GAN) to learn its own distribution $\ptgreen(\zz)$ to match $\qp(\zz)$ makhzani2015adversarial. In this setup, the sampler network and the autoencoder are designed such that they should complement rather than contradict each other, and we can think of $\ptgreen(\zz)$ as actually learning what would be the prior $p(\zz)$ for a regular VAE. (It is worth noting that one popular VAE variant – the `vector-quantised' autoencoder van2017neural – actually learns the prior in a post-hoc fashion like we have proposed, but they learn this as an autoregressive model.)

Concretely, we could learn a kind of VAE where the per-example KL term is instead replaced with an `adversarial' divergence $\fdiv$ ⁵ between $\qp(\ZZ)$ and a learnable prior $\ptgreen(\ZZ)$ (Figure X):

\begin{align} \label{eq:cvae:adv_ae} \min_{\thetagr, \phip} \ -\mathbb{E}_{\qp(\zz,\xx)} \log \pt(\xx|\zz,\yy) + \lambda \fdiv \Big[ \qp(\ZZ) \| \pt(\ZZ) \Big], \tag{16} \end{align}

where samples $\zz \sim \ptgreen(\zz)$ are computed via $\zz = G_{\theta}(\eta)$ for some simple prior $p(\eta)$. While equation \ref{eq:cvae:adv_ae} is no longer an ELBO, one can think of there existing an actual ELBO which is just the likelihood term plus the per-example KL between $\qp(\XX|\ZZ,\YY)$ and $\ptgreen(\ZZ)$, which in our case is not computable in closed form since $\ptgreen(\zz)$ is implicitly represented by the generator.

Proposed adversarial autoencoder, where the prior p(z) is now replaced with a learnable prior p_θ(z) which is implemented with a GAN generator G_θ. The generator is trained to map samples from a simple prior p(η) to those in z space. The corresponding loss in Equation \ref{eq:cvae:adv_ae} is similar to the ELBO, but rather than a per-example KL term we compute an adversarial loss between the marginals q_Φ(z) and p_θ(z).

While this is an interesting start, one issue is that samples from the learned prior may not necessarily decode into plausible looking samples. For instance, if $\qp(\zz)$ is not sufficiently smooth but rather 'spiky', then samples from $\ptgreen(\zz)$ which don't fall into one of those spikes may not decode into a plausible image. In such a case, we may need to add extra regularisation in the form of an \emph{additional} adversarial loss which ensures that decoded samples from $\ptgreen(\xx)$ should be indistinguable from those from the real data distribution $q(\xx)$. In that case, we should really be training a GAN to match the generative and inference pathways:

\begin{align} \label{eq:cvae:adv_ae_joint} \min_{\thetagr, \phip} \ -\mathbb{E}_{\qp(\zz,\xx)} \log \ptgreen(\xx|\zz,\yy) + \fdiv \Big[ \qp(\ZZ, \XX) \| \ptgreen(\ZZ, \XX) \Big]. \tag{16} \end{align}

It turns out that Eqn. (17) without the likelihood term is equivalent to the bidirectional GAN (dumoulin2016adversarially, donahue2016adversarial). In particular, these models propose the following loss:

\begin{align} \min_{\thetagr, \phip} \fdiv \Big[ \ptgreen(\XX,\ZZ) \| \qp(\XX,\ZZ) \Big]. \tag{17} \end{align}

However, one downside is the learned inference encoder $\qp(\zz|\xx)$ is inaccurate and cannot be used to faithfully reconstruct examples, due to the fact that there is no explicit reconstruction (likelihood) loss. As mentioned in dumoulin2016adversarially however one could learn a separate encoder in a post-hoc fashion to address this issue. This issue is also addressed in li2020decomposed.

4.1. cVAEs and conditional Gaussian diffusion models

Diffusion models can be seen as multi-latent generalisations of VAEs ho2020diffusion, and are theoretically very closely related to score-based generative models (see weng2021diffusion for derivations showing their equivalence for the case where the distributions are Gaussian). Instead of just a single latent variable $\zz$, we have many noisy versions of $\xx$ which we denote $\xx_1, \dots, \xx_T$ for $T$ denoising diffusion timesteps (but we can think of this collection of as variables as just $\zz$ for convenience). Apart from this, the main differences are:

There is no inference network $\qp$, instead $q$ is fixed and we have a joint distribution which is the forward process $q(\xx_0, \dots, \xx_T)$ where larger $t$ corresponds to progressively noisier data;
all $\xx_t$ for $t \in \{1, \dots, T\}$ are the same dimension as $\xx_0$;
and $q(\xx_T|\xx_{t-1}) \approx q(\xx_T)$ for sufficiently large total number of timesteps $T$, and we denote the prior $p(\zz) = q(\xx_T)$.

As for conditional diffusion models, some commonly used variants of diffusion are not derived from the conditional ELBO. They're usually modifications done to the reverse conditional to also condition on $\yy$, to give $\pt(\xx_{t-1}|\xx_t, \yy)$. If we denote the collection of noisy random variables $\xx_1, \dots, \xx_T$ as just $\zz$, we can think of that sort of model's decoder as $\pt(\xx|\zz,\yy)$ instead of $\pt(\xx_0|\xx_1, \dots, \xx_T, \yy)$. Therefore, these formulations can be seen as fancier $\ZZ,\YY$ dependent VAEs. To the best of my knowledge, I have not seen a formulation analogous to the $\ZZ,\YY$ independent case.

Figure 10: left: flow graph for an unconditional VAE; right: flow graph for an unconditional diffusion model. For both we illustrate the inference pathway and generative pathway. To be consistent with VAE notation, we have used x instead of x₀ and z instead of x_T.

5. Conclusion

In conclusion, we have:

Derived conditional VAEs through the lens of minimising the KL divergence between two distributions: the inference and generative distributions, which comprise the two halves of a variational autoencoder.
Introduced two conditional variants, corresponding to whether $\ZZ$ and $\YY$ are independent and dependent. For the independent case, we highlighted its usefulness in controllable generation.
Discussed the need to carefully balance the weight of the KL term, which balances the trade-off between sample quality and coverage with respect to the inference and generative distributions. We also derived a mutual information based interpretation of the KL term in order to elucidate its effect on training.
Presented experiments on toy 2D datasets which corroborate our theoretical observations.
Discussed how one can avoid the difficulty of optimising a VAE by instead training an adversarial autoencoder and learning a prior distribution to match the inference marginal. From this discussion we arrive at the bidirectional GAN, which has a very close relationship with the VAE in the sense that the former minimises an arbitrary $f$ divergence between the generative and inference pathways while the latter uses the forward KL divergence. This brings us full circle!

6. Appendix

6.1. Derivation of Esmaeli's joint KL

Here we derive the main equation presented in esmaeili2018structured. This corresponds to the unconditional VAE, without $\yy$ conditioning.

\begin{align} \color{green}{\theta}, \color{purple}{\phi} & = \argmax_{\color{green}{\theta}, \color{purple}{\phi}} -\mathcal{D}_{\text{KL}}\Big[ \qp(\ZZ,\XX) || \pgreen(\ZZ, \XX) \Big] \tag{10a} \\ & = \mathbb{E}_{\qzx} \Big[ \log \frac{\pzx}{\qzgivenx q(\xx)} \Big] \tag{10b} \\ & = \mathbb{E}_{\qzx} \Big[ \log \frac{\pxgivenz p(\zz)}{\qzgivenx q(\xx)} \Big] \tag{10c} \\ & = \mathbb{E}_{\qzx} \Big[ \log \frac{\pxgivenz p(\zz)}{\qzgivenx q(\xx)} \cdot \frac{\ptgreen(\xx)}{\ptgreen(\xx)} \cdot \frac{\qp(\zz)}{\qp(\zz)} \Big] \tag{10d} \\ & = \mathbb{E}_{\qzx} \Big[ \log \frac{\pxgivenz}{\ptgreen(\xx)} + \log \frac{\qp(\zz)}{\qzgivenx} + \log \frac{\ptgreen(\xx)}{q(\xx)} + \log \frac{p(\zz)}{\qp(\zz)} \Big] \tag{10e} \\ & = \mathbb{E}_{\qzx} \Big[ \log \frac{\pxgivenz}{\ptgreen(\xx)} + \log \frac{\qp(\zz)}{\qzgivenx} \Big] + \mathbb{E}_{q(\xx)} \Big[ \log \frac{\ptgreen(\xx)}{q(\xx)} \Big] + \\ & \ \ \ \ \mathbb{E}_{\qp(\zz)} \Big[ \log \frac{p(\zz)}{\qp(\zz)} \Big] \tag{10f} \\ & = \mathbb{E}_{\qzx} \Big[ \underbrace{\log \frac{\pxgivenz}{\ptgreen(\xx)}}_{\circleone} - \underbrace{\log \frac{\qzgivenx}{\qp(\zz)}}_{\circletwo} \Big] - \underbrace{\kldiv\Big[ q(\XX) \| \ptgreen(\XX) \Big]}_{\circlethree} - \\ & \ \ \ \ \ \underbrace{\kldiv\Big[ \qp(\ZZ) \| p(\ZZ)}_{\circlefour} \Big], \tag{10g} \end{align}

where:

$\ptgreen(\xx) = \int_{\zz} \ptgreen(\xx|\zz)p(\zz) d \zz$, the marginal distribution of the data with respect to the generative process. This is also called the marginal likelihood.
$\qp(\zz) = \int_{\xx} \qp(\zz|\xx)q(\xx) d\xx$ , the marginal distribution over the latent code with respect to the inference process. This is also called the inference marginal.

6.2. Conditional case

We can derive the conditional case by adding $\yy$ wherever it is necessary. Starting from Eqn. (10f), we derive the following:

We can subsequently refine this equation depending on the factorisation of $p(\zz,\yy)$, which we do below.

6.2.1. z and y are independent

For the sake of space, I will simply use $\qp$ to refer to the full joint distribution $\qp(\zz,\xx, \yy)$. For $p(\zz,\yy) = p(\zz)p(\yy)$, we get:

\begin{align} & \mathbb{E}_{\qp} \Big[ \log \frac{\ptgreen(\xx|\zz,\yy)}{\ptgreen(\xx)} - \log \frac{\qp(\zz|\xx,\yy)}{\qp(\zz)} \Big] + \mathbb{E}_{q(\xx)} \Big[ \log \frac{\ptgreen(\xx)}{q(\xx)} \Big] + \\ & \ \ \ \ \ \mathbb{E}_{\qp(\zz,\yy)} \Big[ \log \frac{p(\zz)}{\qp(\zz)} + \log p(\yy) \Big] \tag{12a} \\ & = \mathbb{E}_{\qp} \Big[ \log \frac{\ptgreen(\xx|\zz,\yy)}{\ptgreen(\xx)} - \log \frac{\qp(\zz|\xx,\yy)}{\qp(\zz)} \Big] + \mathbb{E}_{q(\xx)} \Big[ \log \frac{\ptgreen(\xx)}{q(\xx)} \Big] + \\ & \ \ \ \ \ \mathbb{E}_{\qp(\zz)} \Big[ \log \frac{p(\zz)}{\qp(\zz)} \Big] + \mathbb{E}_{q(\yy)} \log p(\yy) \tag{12b} \\ & = \mathbb{E}_{\qp} \Big[ \underbrace{\log \frac{\ptgreen(\xx|\zz,\yy)}{\ptgreen(\xx)}}_{\circleone} - \underbrace{\log \frac{\qp(\zz|\xx,\yy)}{\qp(\zz)}}_{\circletwo} \Big] - \underbrace{\kldiv\Big[ q(\XX) \| \ptgreen(\XX) \Big]}_{\circlethree} \\ & \ \ \ \ \ - \underbrace{\kldiv\Big[ \qp(\ZZ) \| p(\ZZ)}_{\circlefour} \Big] + \text{const.} \tag{12c} \end{align}

Here, $p(\yy)$ can fall out of the optimisation since it's just a constant. However, since it's a prior we can set it to whatever it is we want it to be, either the actual empirical distribution of $\yy$ for our dataset or another distribution.

6.2.2. z and y are dependent

Again, starting from Eqn. (10f), if we assume that $p(\zz,\yy) = p(\zz|\yy)p(\yy)$ then:

\begin{align} & \mathbb{E}_{\qp(\zz,\xx,\yy)} \Big[ \log \frac{\ptgreen(\xx|\zz,\yy)}{\ptgreen(\xx)} - \log \frac{\qp(\zz|\xx,\yy)}{\qp(\zz)} \Big] + \mathbb{E}_{q(\xx)} \Big[ \log \frac{\ptgreen(\xx)}{q(\xx)} \Big] + \\ & \ \ \ \ \mathbb{E}_{\qp(\zz,\yy)} \Big[ \log \frac{p(\zz|\yy)p(\yy)}{\qp(\zz)} \Big] \tag{13a} \\ & \mathbb{E}_{\qp(\zz,\xx,\yy)} \Big[ \log \frac{\ptgreen(\xx|\zz,\yy)}{\ptgreen(\xx)} - \log \frac{\qp(\zz|\xx,\yy)}{\qp(\zz)} \Big] + \mathbb{E}_{q(\xx)} \Big[ \log \frac{\ptgreen(\xx)}{q(\xx)} \Big] + \\ & \ \ \ \ \mathbb{E}_{\qp(\zz,\yy)} \Big[ \log \frac{p(\zz|\yy)}{\qp(\zz)} + \log p(\yy) \Big]. \tag{13b} \\ & \mathbb{E}_{\qp(\zz,\xx,\yy)} \Big[ \log \frac{\ptgreen(\xx|\zz,\yy)}{\ptgreen(\xx)} - \log \frac{\qp(\zz|\xx,\yy)}{\qp(\zz)} \Big] - \kldiv\Big[ q(\XX) \| \ptgreen(\XX) \Big] + \\ & \ \ \ \ -\kldiv \Big[ \qp(\ZZ) \| p(\ZZ|\YY) \Big] + \mathbb{E}_{q(\yy)} \log p(\yy). \tag{13c} \\ \end{align}

Here, we need to choose what $p(\zz|\yy)$ is. Either it can be a fixed distribution (i.e. a distribution is pre-assigned for each possible value of $\yy$), or it could also be learned, in which case we can denote it as $\ptgreen(\zz|\yy)$.

6.3. Per-example KL versus marginal KL

Thanks to Eqn. (10e) we can just re-arrange its terms to express their relationship as the following:

\begin{align} \kldiv \Big[ \qp(\ZZ|\XX,\YY) \| p(\ZZ|\YY) \Big] - I_{\phip}(\ZZ; \XX, \YY) = \kldiv[ \qp(\ZZ) \| p(\ZZ|\YY) ] + \text{const.} \tag{10e} \end{align}

Therefore, minimising the marginal KL on the RHS of this equation means:

(1) Making $I$ larger, for a fixed per-example KL (first term on the LHS);
(2) or making per-example KL smaller, for fixed $I$.

(1) seems beneficial because increasing $I$ means $\ZZ$ loses less information about $\XX$, but this only makes sense in the context of a $\ZZ,\YY$ dependent VAE.

We also note the RHS of this equation was proposed in kumar2017variational, but for unconditional VAEs.

6.4. Z-Y independent cVAE on SVHN

Here is an artifact from an old research project I did involving controllable generation. We were trying to do style/content swaps for images from SVHN – here, one can think of the content as being $\yy$, the identity of the SVHN digit. For each row:

x1 is $\xx_1$, x2 is $\xx_2$. Their corresponding labels are the digits, e.g. $\yy_1$ will be 18. $\yy_2$ depends on what column we are looking at.
recon is the reconstruction of $\xx_1$, as per the inference process.
x1_c, x2_s says: take the content of $\xx_1$ and the style from $\xx_2$. This means, we sample $\xx \sim \ptgreen(\xx|\yy_1,\zz_2)$, where $\yy_1$ is the identity of $\xx_1$, and $\zz_2 \sim \qp(\zz|\yy_2,\xx_2)$.
x2_c, x1_s says the opposite: take the content of $\xx_2$ and the style from $\xx_1$. This means, we sample $\xx \sim \ptgreen(\xx|\yy_2,\zz_1)$, where $\yy_2$ is the identity of $\xx_2$, and $\zz_1 \sim \qp(\zz|\yy_1,\xx_1)$.

Figure 7: An example of a Z,Y independent conditional VAE trained on a modified version of SVHN.

7. References

beckham2023thesis Beckham, C. (2023). PhD thesis dissertation. (Work in progress.)
kingma2013auto Kingma, D. P., Welling, M., & others, (2019). An introduction to variational autoencoders. Foundations and Trends in Machine Learning, 12(4), 307–392.
kingma2019introduction Kingma, D. P., Welling, M., & others, (2019). An introduction to variational autoencoders. Foundations and Trends in Machine Learning, 12(4), 307–392.
esmaeili2018structured Esmaeili, B., Wu, H., Jain, S., Bozkurt, A., Siddharth, N., Paige, B., Brooks, D. H., … (2018). Structured disentangled representations. arXiv preprint arXiv:1804.02086, (), .
burgess2018understanding Burgess, C. P., Higgins, I., Pal, A., Matthey, L., Watters, N., Desjardins, G., & Lerchner, A. (2018). Understanding disentangling in beta-VAE. arXiv preprint arXiv:1804.03599, (), .
child2020very Child, R. (2020). Very deep VAEs generalize autoregressive models and can outperform them on images. International Conference on Learning Representations, (), .
ho2020diffusion Ho, J., Jain, A., & Abbeel, P. (2020). Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33(), 6840–6851.
dhariwal2021diffusion Dhariwal, P., & Nichol, A. (2021). Diffusion models beat GANs on image synthesis. Advances in Neural Information Processing Systems, 34(), 8780–8794.
kumar2017variational Kumar, A., Sattigeri, P., & Balakrishnan, A. (2017). Variational inference of disentangled latent concepts from unlabeled observations. arXiv preprint arXiv:1711.00848, (), .
dumoulin2016adversarially Dumoulin, V., Belghazi, I., Poole, B., Lamb, A., Arjovsky, M., Mastropietro, O., & Courville, A. (2016). Adversarially Learned Inference. In , International Conference on Learning Representations (pp. ). : .
donahue2016adversarial Donahue, J., Kr\"ahenb\"uhl, Philipp, & Darrell, T. (2016). Adversarial feature learning. arXiv preprint arXiv:1605.09782, (), .
nowozin2016f Nowozin, S., Cseke, B., & Tomioka, R. (2016). F-gan: training generative neural samplers using variational divergence minimization. Advances in neural information processing systems, 29(), .
zhang2019variational Zhang, M., Bird, T., Habib, R., Xu, T., & Barber, D. (2019). Variational f-divergence minimization. arXiv preprint arXiv:1907.11891, (), .
makhzani2015adversarial Makhzani, A., Shlens, J., Jaitly, N., Goodfellow, I., & Frey, B. (2015). Adversarial autoencoders. arXiv preprint arXiv:1511.05644, (), .
larsen2016autoencoding Larsen, A. B. L., S\onderby, S\oren Kaae, Larochelle, H., & Winther, O. (2016). Autoencoding beyond pixels using a learned similarity metric. In , International conference on machine learning (pp. 1558–1566). : .
mescheder2017adversarial Mescheder, L., Nowozin, S., & Geiger, A. (2017). Adversarial variational bayes: unifying variational autoencoders and generative adversarial networks. In , International conference on machine learning (pp. 2391–2400). : .
chen2016infogan Chen, X., Duan, Y., Houthooft, R., Schulman, J., Sutskever, I., & Abbeel, P. (2016). InfoGAN: interpretable representation learning by information maximizing generative adversarial nets. Advances in neural information processing systems, 29(), .
van2017neural Van Den Oord, A., Vinyals, O., & others, (2017). Neural discrete representation learning. Advances in neural information processing systems, 30(), .
goodfellow2020generative Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., … (2020). Generative adversarial networks. Communications of the ACM, 63(11), 139–144.

Footnotes:

One may wonder whether it is more appropriate to instead modify the KL term to be less 'strict' and match $\qp(\ZZ|\YY)$ with $p(\ZZ)$ instead, and we discuss this in Sec. 6.3.

While it is possible in principle to derive an additional loss term which specifically penalises $I(Z; Y)$ (e.g. with Monte Carlo approximation or with adversarial learning), from personal experience it came with very little success. I suspect it is because such a term only works if the likelihood term is sufficiently downweighted, but this causes sample quality to suffer and we just end up with the same problem as we do with the original KL term.

If one had a highly supervised dataset of 'paired' examples $(\xx^{(i)}_1, \xx^{(i)}_2)$ where $\xx_1$ and $\xx_2$ only dithered by $\YY$ (i.e. all other factors of variation remained the same) then it would perhaps be much easier to learn this style of VAE, but such datasets are usually not reflective of the real world.

⁴

Interestingly, another mutual information term falls out of the derivation and it is negative. Since Eqn. (14f) is framed as a minimisation, minimising the negative of this is really maximising it, so $\phip$ is also being updated to maximise the mutual information between $\ZZ$ and $\YY$ with respect to the encoder $\qp$.

⁵

The tilde emphasises that GANs are approximating a particular $f$-divergence, see goodfellow2020generative.

Techniques for label conditioning in Gaussian denoising diffusion models

2023-01-27T00:00:00+00:00

Techniques for label conditioning in Gaussian denoising diffusion models

Updates

(12/02/2023) Clean up some derivations, re-label some equations.
(02/02/2023) Expanded classifier-free guidance section, talking about the relationship between w and the dropout probability. Also explaining that classifier-free may be beneficial in a semi-supervised scenario.

1. Preliminaries
2. Classifier-based guidance
3. Classifier-free guidance
- 3.1. Sources of confusion
- 3.2. Benefits
4. Conditional ELBO
5. Conclusion
6. References

In this very short blog post, I will be presenting my derivations of two widely used forms of label conditioning for denoising diffusion probabilistic models (DDPMs) ho2020denoising. I found that other sources of information I consulted didn't quite get the derivations right or were confusing, so I'm presenting my own reference here that I hope will serve myself and others well.

1. Preliminaries

I go over some preliminaries here. If you're new to DDPMs, it may be best to read this first!

DDPMs can be derived by first starting off with the evidence lower bound, which can be expressed as:

\begin{align} \label{eq:elbo} \log p(\xx) & \geq \text{ELBO}(\xx) \\ & = \mathbb{E}_{q(\xx_0, \dots, \xx_T)} \Big[ \underbrace{-\log \frac{p(\xx_T)}{q(\xx_T|\xx_0)}}_{L_T} - \sum_{t > 1} \underbrace{\log \frac{\pt(\xx_{t-1}|\xx_t)}{q(\xx_{t-1}|\xx_t, \xx_0)}}_{L_t} - \underbrace{\log \pt(\xx_0|\xx_1)}_{L_0} \Big. \tag{0} \end{align}

Using typical DDPM notation, $\xx_0 \sim q(\xx_0)$ is the real data, and $q(\xx_t|\xx_{t+1})$ for $t \in \{1, \dots, T\}$ defines progressively noisier distributions (dictated by some noising schedule $\beta_t$), and $\pt(\xx_{t-1}|\xx_t)$ parameterises a neural net which is trained to reverse this process. In practice, $\pt$ is re-parameterised such that it in turn is a function of a noise predictor $\epst(\xx_t, t)$ which is trained to predict only the noise in the image that is generated via $\xx_t \sim q(\xx_t|\xx_0)$:

\begin{align} \pt(\xx_{t-1}|\xx_t) = \mathcal{N}(\xx_{t-1}; \frac{1}{\sqrt{\alpha_t}}\Big( \xx_t - \frac{1-\alpha_t}{\sqrt{1-\alphabar_t}} \epst(\xx_t, t)\Big), \sigma(\xx_t, t)). \end{align}

As a further simplification, each of the $T$ KL terms in the ELBO can be simplified to the following noise prediction task:

\begin{align} \mathcal{L}(t) = \mathbb{E}_{\xx_0, \xx_t, \epsilon_t} \big[ \frac{\beta_t^2}{2\sigma_t^2 \alpha_t(1-\alphabar_t)} \| \epsilon_t - \epsilon_{\theta}(\xx_t, t)\|^{2} \big]. \end{align}

In practice, a biased version of the loss is used which removes the weighting term inside the square brackets. This has the effect of upweighting the loss in favour of noiser images (i.e. $\xx_t$ for large $t$):

\begin{align} \mathcal{L}_{\text{simple}}(t) = \mathbb{E}_{\xx_0, \xx_t, \epsilon_t} \big[ \| \epsilon_t - \epsilon_{\theta}(\xx_t, t)\|^{2} \big]. \end{align}

The following derivations hinge on one important equation that relates diffusion models to score matching song2020score. I take the following from Lilian Weng's blog weng2021diffusion:

\begin{align} \st(\xx_t, t) \approx \nabla_{\xx_t} \log q(\xx_t) = -\frac{\epst(\xx_t, t)}{\sqrt{1 - \alphabar_{t}}}. \tag{1} \end{align}

The question we would like to answer in the following sections is: given an unconditional diffusion model $\pt(\xx)$, how can we easily derive a conditional variant $\pt(\xx|\yy)$?

Figure 1: Algorithm for sampling. Taken from ho2020denoising.

2. Classifier-based guidance

Through Bayes' rule we know that:

\begin{align} q(\xx_t|y) = \frac{q(\xx_t, y)}{q(y)} = \frac{q(y|\xx_t)q(\xx_t)}{q(y)}. \end{align}

Taking the score $\nabla_{\xx_t} \log q(\xx_t|y)$, we get:

\begin{align} \nabla_{\xx_t} \log q(\xx_t|y) & = \nabla_{\xx_t} \log q(y|\xx_t) + \nabla_{\xx_t} \log q(\xx_t) - \underbrace{\nabla_{\xx_t} \log q(\yy)}_{= 0} \\ & \approx \nabla_{\xx_t} \log q(\yy|\xx_t) - \frac{\epst(\xx_t, t)}{\sqrt{1-\alphabar_t}}, \ \ \text{(using eqn. (1))} \tag{2a} \end{align}

where in the last line we make clear the connection between the score function and the noise predictor $\epst$ weng2021diffusion. We could also use Equation (1) to do the same thing to the LHS of Equation (2a):

\begin{align} \nabla_{\xx_t} \log q(\xx_t|y) & \approx \nabla_{\xx_t} \log q(\yy|\xx_t) - \frac{\epst(\xx_t, t)}{\sqrt{1-\alphabar_t}} \\ \implies -\frac{\epst(\xx_t, \yy, t)}{\sqrt{1-\alphabar_t}} & \approx \nabla_{\xx_t} \log q(\yy|\xx_t) - \frac{\epst(\xx_t, t)}{\sqrt{1-\alphabar_t}} \tag{2b} \end{align}

If we re-arrange for $\epst(\xx_t, \yy, t)$ in Equation (2b) we finally get:

\begin{align} \epst(\xx_t, y, t) & \approx \epst(\xx_t, t) - \sqrt{1-\alphabar_t} \nabla_{\xx_t} \log q(\yy|\xx_t) \tag{2c} \end{align}

The only thing left to do is to approximate the ground truth classifier $q(\yy|\xx_t)$ with our own classifier $\pphi(\yy|\xx_t; t)$. This will be defined as:

\begin{align} \epstt(\xx_t, y, t) := \epst(\xx_t, t) - \sqrt{1-\alphabar_t} \nabla_{\xx_t} \log \pphi(\yy|\xx_t; t). \tag{2d} \end{align}

This classifier should be trained on the same distribution of images from the forward process $q(\xx_0, \dots, \xx_T)$ so that appropriate gradients are obtained during sampling. Note that I have also conditioned on $t$ as well, which I suspect probably is a good extra supervisory signal so that the classifier knows what timestep $\xx_t$ is coming from.

In practice, we can also define the weighted version as follows, which allows us to balance between (conditional) sample quality and sample diversity:

\begin{align} \label{eq:cg_supp} \underbrace{\bar{\epstt}(\xx_t, t, y; w) := \epst(\xx_t, t) -\sqrt{1-\bar{\alpha}_t} w \nabla_{\xx_t} \log \pphi(y|\xx_t; t)}_{\text{classifier-free guidance, plug this into Fig. 1}}. \tag{2e} \end{align}

Even though this form of label conditioning requires an external classifier, it is quite a simple and principled derivation, and therefore I like it. Essentially, from an unconditional diffusion model $\pt(\xx)$ we are inducing a conditional variant at generation time. One can think of sampling from this particular noise predictor as sampling an $\xx \sim p_{\theta,\phi}(\xx|\yy) \propto \pphi(\yy|\xx)\pt(\xx)$ for a given $\yy$.

I think one interesting aspect of this formulation is that, since the induced conditional model is a function of both the unconditional model $\pt(\xx)$ and the classifier $\pphi(\yy|\xx)$, the entire generative model could be improved by switching out either component in isolation with an updated version. This could be useful if:

it is too expensive to re-train the diffusion model at regular intervals. Since classifiers are a bit faster to train, one strategy could be to update (retrain) the classifier at more frequent intervals than the diffusion model.
One wishes to leverage a pre-trained + frozen unconditional diffusion model for transfer learning with their own prescribed classifier.

3. Classifier-free guidance

The idea behind classifier-free guidance is that one could simply instead condition on $\yy$ in the reverse process, i.e. use $\pt(\xx_{t-1}|\xx_{t}, y)$ instead of $\pt(\xx_{t-1}|\xx_t)$. In our case, this would be conditioning on $\yy$ for the noise predictor $\epst(\xx_t, y, t)$. However, the authors also propose learning the unconditional version at the same time for the same model, which means that during training $\yy$ random gets dropped with some probability $\puncond$. When the label does get dropped, it simply gets replaced with some null token, so we can think of $\epst(\xx_t, t) = \epst(\xx_t, y = \emptyset, t)$. (In practice, dhariwal2021diffusion found that a $\puncond$ of 0.1 or 0.2 works well.)

The reason for this algorithm is so that a variant of Equation (2c) can be derived without depending on an external classifier. From Bayes' rule, we know that:

\begin{align} \pt(\yy|\xx_t) = \frac{\pt(\yy,\xx_t)}{\pt(\xx_t)} = \frac{\pt(\xx_t|y)p(\yy)}{\pt(\xx_t)}, \end{align}

and that therefore the score $\nabla_{\xx_t} \log \pt(\yy|\xx_t)$ is:

\begin{align} \nabla_{\xx_t} \log \pt(y|\xx_t)= \nabla_{\xx_t} \log \pt(\xx_t|y) + \underbrace{\nabla_{\xx_t} \log p(\yy)}_{= 0} - \nabla_{\xx_t} \log \pt(\xx_t). \end{align}

We simply plug this into Equation (2c) (as well as re-introduce $w$) to remove the dependence on $q(y|\xx_t)$:

\begin{align} \bar{\epst}(\xx_t, y, t; w) & := \epst(\xx_t, t) -\sqrt{1-\bar{\alpha}_t} w \nabla_{\xx_t} \log \pt(y|\xx_t) \\ & = \epst(\xx_t, t) -\sqrt{1-\bar{\alpha}_t} w \Big[ \nabla_{\xx_t} \log \pt(\xx_t|y) - \nabla_{\xx_t} \log \pt(\xx_t) \Big] \\ & = \epst(\xx_t, t) -\sqrt{1-\bar{\alpha}_t} w \Big[ \frac{-1}{\sqrt{1-\bar{\alpha}_t}} \epst(\xx_t, y, t) - \frac{-1}{\sqrt{1-\bar{\alpha}_t}} \epst(\xx_t, t) \Big] \\ & = \epst(\xx_t, t) + w \epst(\xx_t, y, t) - w \epst(\xx_t, t) \\ & = \underbrace{\epst(\xx_t, t)}_{\approx \nabla_{\xx_t} \log p(\xx)} + w \Big( \underbrace{\epst(\xx_t, y, t) - \epst(\xx_t, t)}_{\approx \nabla_{\xx_t} \log p(\yy|\xx)} \Big). \tag{3a} \end{align}

From Equation (3a) we can see that the term being multiplied by $w$ is (roughly) the score induced by the implicit classifier that defined by the diffusion model itself. Note that Equation (3a) could also be re-written as:

\begin{align} \underbrace{\bar{\epst}(\xx_t, y, t; w) := (1-w)\epst(\xx_t, t) + w \epst(\xx_t, y, t)}_{\text{classifier-free guidance, plug this into Fig. 1}}, \tag{3b} \end{align}

3.1. Sources of confusion

Equation (3b) appears to be almost the same as Equation 6 of dhariwal2021diffusion, though in their paper all the signs appear to be flipped and $(1+w)\epst(\xx_t,t) - w\epst(\xx_t, y, t)$ is used instead. I'm not sure if this is an oversight or something wrong in my own derivations, but we can just think of it as another way to formulate Equation (3b); essentially, if you substitute in $-w$ instead of $w$ for the weighting, you would get:

\begin{align} \underbrace{\bar{\epst}(\xx_t, y, t; w)|_{w = -w} = (w+1)\epst(\xx_t, t) - w \epst(\xx_t, y, t)}_{\text{classifier-free guidance, alternate formulation, plug this into Fig. 1}}. \tag{3c} \end{align}

A minor confusion I had with this paper stemmed from the fact that there are two parameters which are used to create a modified score estimator: $\puncond$ is used at training time to weight the unconditional score estimator $\epst(\xx_t, t)$, and $w$ is used at generation time to weight the conditional score estimator $\epst(\xx_t, y, t)$ without using $\puncond$. Since we use dropout on $\yy$ at training time with probability $\puncond$, we can actually think of the predicted score as being a Bernoulli random variable of the form:

\begin{equation} \epst(\xx_t, y, t; w)\big|_{w=1-\puncond} =\begin{cases} \epst(\xx_t, y=\emptyset, t) & \text{with probability $\puncond$}.\\ \epst(\xx_t, y, t) & \text{otherwise}, \end{cases} \end{equation}

and therefore the expected value of this variable would be the following (as per the definition of a Bernoulli random variable):

\begin{align} \bar{\epst}(\xx_t, y, t; w)\big|_{w=1-\puncond} & = \puncond \epst(\xx_t, t) + (1-\puncond) \epst(\xx_t, y, t). \tag{3d} \end{align}

Here, we can see that the relationship between $w$ and $\puncond$ is through $w = 1 - \puncond$, but we actually don't want to stick with this definition at test time since it also assumes $w \in [0,1]$. This means that Equation (3b) is only ever going to be a convex combination between the unconditional and conditional scores. Conversely, letting $w \in \mathbb{R}^{+}$ lets us be as aggressive as we need to be with guiding the diffusion model.

3.2. Benefits

One potential benefit from the classifier-free formulation is that the implicit classifier and unconditional model share the same set of weights $\theta$. If we assume that the knowledge about the unconditional model in $\theta$ can 'transfer' over to the conditional part (and vice versa), then this formulation would make a lot of sense in a semi-supervised scenario where one may have significantly more unlabelled examples than labelled ones. The unlabelled ones can be trained with the unconditional score estimator, and hopefully improve the performance of the conditional variant.

4. Conditional ELBO

The previous two methods involve turning an unconditional diffusion model into a conditional one by either leveraging an explicit classifier (classifier guidance) or deriving an implicit one (classifier-free guidance). For the classifier-guided variant, the new conditional model can be written as:

\begin{align} p_{\theta,\phi}(\xx|\yy; w) & \propto \underbrace{\pphi(\yy|\xx)^{w}}_{\text{explicit}} \pt(\xx). \end{align}

For classifier-free, this classifier is implicit, and the balance between the two following terms isn't just via $w$ at generation time but also through the training hyperparameter $\puncond$:

\begin{align} \pt(\xx|\yy; w) & \propto \underbrace{\pt(\yy|\xx)^{w}}_{\text{implicit}} \pt(\xx). \end{align}

When we compare both formulations in this manner, we might also ask ourselves, what's stopping us from just training a conditional model $\pt(\xx|\yy)$ directly, rather than through the product of a classifier and an unconditional model? This is certainly possible, via the conditional ELBO. This would correspond to taking Equation (0) and adding $\yy$ to each conditional distribution, as well as converting the prior $p(\xx_T)$ to a learned conditional prior $\pt(\xx_T|\yy)$:

\begin{align} \log p(\xx|\yy) & \geq \text{ELBO}(\xx, \yy) \\ & = \mathbb{E}_{q(\xx_0, \dots, \xx_T, \yy)} \Big[ \underbrace{-\log \frac{\pt(\xx_T|\yy)}{q(\xx_T|\xx_0,\yy)}}_{L_T} - \sum_{t > 1} \underbrace{\log \frac{\pt(\xx_{t-1}|\xx_t,\yy)}{q(\xx_{t-1}|\xx_t, \xx_0, \yy)}}_{L_t} \\ & - \underbrace{\log \pt(\xx_0|\xx_1, \yy)}_{L_0} \Big]. \tag{4} \end{align}

To me, this is the most theoretically rigorous way to derive a conditional diffusion model. (In fact, this has already been used in lu2022conditional for speech diffusion!) Oddly enough, this doesn't appear to be the way that labelling is done in practice. Ironically, in the variational autoencoder literature this is how almost all conditional variants are derived, and diffusion models are just multi-latent generalisations of VAEs which learn $T$ latent codes instead (with the added constraint that the dimensionality of those codes are the same as the input dimensionality). I suspect this is probably because, unlike in the case of VAEs, one has to think carefully about how $\yy$ can be conditioned on in the forward process, especially if $\yy$ is not the same dimension as $\xx$.

For more details about this kind of model, I highly recommend you read my other post where I talk about lu2022conditional and implement a proof-of-concept that also works on discrete labels (through MNIST). I also show that one of the hyperparameters used in the training of this model also acts like a sort of knob that allows one to control between sample quality and diversity.

5. Conclusion

I will summarise everything with some key bullet points:

Classifier-based / classifer-free guidance allow us to imbue unconditional diffusion models with the ability to condition on a label.
Classifier-based guidance requires an external classifier, but decomposing the model into two modules may be beneficial from the point of view of retraining or fine-tuning on new data.
Classifier-free guidance does not require an external classifier, but requires an extra hyperparameter $\puncond$ during training. Since the same weights are used to parameterise both the implicit classifier and unconditional score estimator, it may be useful in a semi-supervised learning scenario.
A more theoretically direct approach to conditioning on labels is to derive a Gaussian DDPM via the conditional ELBO (Equation (4)), but would require some extra derivations and model assumptions to be made. A conditional ELBO-based approach is used in lu2022conditional, and I speak about it here.
All three variants allow for weighting trading off between sample quality and diversity.

6. References

ho2020denoising Ho, J., Jain, A., & Abbeel, P. (2020). Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33(), 6840–6851.
song2020score Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., & Poole, B. (2020). Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, (), .
classifierfree Ho, J., & Salimans, T. (2022). Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598, (), .
dhariwal2021diffusion Dhariwal, P., & Nichol, A. (2021). Diffusion models beat GANs on image synthesis. Advances in Neural Information Processing Systems, 34(), 8780–8794.
lu2022conditional Lu, Y., Wang, Z., Watanabe, S., Richard, A., Yu, C., & Tsao, Y. (2022). Conditional diffusion probabilistic model for speech enhancement. In , ICASSP 2022-2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) (pp. 7402–7406).
weng2021diffusion Weng, L. (2021). What are diffusion models? lilianweng.github.io, (), .
sohn2015learning Sohn, K., Lee, H., & Yan, X. (2015). Learning structured output representation using deep conditional generative models. Advances in neural information processing systems, 28(), .

Learning the conditional prior over classes for image diffusion

2022-09-24T00:00:00+00:00

$\newcommand{\xx}{\boldsymbol{x}}$ $\newcommand{\yy}{\boldsymbol{y}}$ $\newcommand{\pt}{p_{\theta}}$ $\newcommand{\QQ}{\boldsymbol{Q}}$ $\newcommand{\mm}{\boldsymbol{m}}$ $\newcommand{\alphabar}{\bar{\alpha}}$ $\newcommand{\mt}{\mu_{\theta}}$ $\newcommand{\epst}{\epsilon_{\theta}}$ $\newcommand{\betatilde}{\tilde{\beta}}$ $\newcommand{\deltatilde}{\tilde{\delta}}$ $\newcommand{\linspace}{\text{linspace}}$ $\newcommand{\embed}{\text{embed}_{\theta}}$

Learning the conditional prior over classes for image diffusion

If you found this useful and wish to cite it, you can use this corresponding Bibtex entry:

@misc{beckham2022_condprior,
  author = {Beckham, Christopher},
  title = {Tech report: Learning the conditional prior over classes for image diffusion},
  year = {2022},
  howpublished = {\url{https://beckham.nz/2022/09/24/cond-diffusion.html}}
}

The Github repository for this code can be found here.

Unconditional diffusion
The conditional variant
Reproduction
References

Unconditional diffusion

Let us start with a quick refresher for a typical (unconditional) diffusion model (see [1] for more details). The $t$-step forward process can conveniently be derived to give us the following:

\begin{align} \label{eq:uncond_fwd_t_step} q(\xx_t|\xx_0) = \mathcal{N}(\xx_t; \sqrt{\alphabar_{t}}\xx_0, (1-\alphabar_{t}) \mathbf{I}) \end{align}

The reverse process can be formulated as:

\begin{align} \label{eq:uncond_reverse} \pt(\xx_{t-1}|\xx_t) = \mathcal{N}(\xx_{t-1}; \mt(\xx_t, t), \betatilde_t \mathbf{I} ) \end{align}

where:

$\betatilde_t = \frac{1-\alphabar_{t-1}}{1 - \alphabar_{t}} \beta_t$;
$\mt(\xx_t, t) = \frac{1}{\sqrt{\alpha_t}} \xx_t - \frac{\beta_t}{\sqrt{\alpha_t} \sqrt{1-\alphabar_{t}}}\epst(\xx_{t}, t)$; that is, we parameterise a neural network that tries to predict the noise.

The conditional variant

In this report I consider the paper proposed by [2], which proposes a conditional variant of diffusion for speech synthesis. It’s presented as a very principled derivation, in the sense that their formulation can be seen as a generalisation of unconditional diffusion. The formulation is a bit math heavy, and I will simply defer you to their work if you want more details. Essentially, in their formulation $\yy$ is actually ‘noisy’ version of $\xx$ rather than a discrete label, so it presents itself as being a conditional diffusion model that can map between two domains (clean vs noisy speech). However, soon I will explore modifying this formulation so that one can instead condition on a discrete label.

Instead of deriving the one-step (conditional) forward first $q(\xx_t | \xx_{t-1}, \yy)$ and then subsequently deriving $q(\xx_t | \xx_{0}, \yy)$, the authors take the opposite approach and derive first $q(\xx_t | \xx_{0}, \yy)$:

\begin{align} q(\xx_t | \xx_0, \yy) = \mathcal{N}(\xx_t; (1-\mm_t) \sqrt{\alphabar_{t}} \xx_0 + \mm_t \sqrt{\alphabar_{t}} \yy, \delta_{t} \mathbf{I}) \nonumber \end{align}

where the mean is characterised by a convex combination between the two terms, with $\mm_t \in [0,1]$ being the interpolation coefficient. Essentially, as we run the forward diffusion, the mean progressively moves towards a scaled $\yy$ under some variance schedule $\delta_t$. Because of the interpolation, we can see that $\yy$ must be of the same dimensionality as $\xx$ as well. The prior distribution for some $\yy$ can be expressed as:

\begin{align} \pt(\xx_T |\yy) = \mathcal{N}(\xx_T; \sqrt{\alphabar_T} \yy, \delta_{T} \mathbf{I}) \nonumber \end{align}

so that $\yy$ essentially parameterises the mean of the conditional prior $\pt(\xx_T|\yy)$ up to the scaling factor $\sqrt{\alphabar_T}$.
From which they show that under a particular derivation for $\delta_t$, marginalising $q(\xx_t|\xx_0,\yy)$ over $\yy$ recovers the unconditional forward process $q(\xx_t|\xx_0)$:

\begin{align} \delta_t = (1 - \alphabar_{t}) - m_t^2 \alphabar_t \nonumber \end{align}

hence this is a conditional generalisation of forward diffusion. $\delta_t$ here is analogous to the $\beta_t$ in the unconditional version. After some crazy long derivations, the authors derive the reverse process. This takes in a similarly convenient form, which is basically:

\begin{align} \label{eq:cond_reverse} \pt(\xx_{t-1}|\xx_{t}, \yy) = \mathcal{N}(\xx_{t-1}; \mt(\xx_t, \yy, t), \deltatilde_{t} \mathbf{I} ) \nonumber \end{align}

where:

$\deltatilde_t = \frac{\delta_{t|t-1} \cdot \delta_t}{\delta_{t-1}}$ (analogous to $\betatilde_t$ in the unconditional version of the learned reverse process)
$\delta_{t|t-1} = \delta_{t} \Big( \frac{1 - m_t}{1 - m_{t-1}} \Big)^{2} \alpha_{t} \delta_{t-1}$ (this is also the variance term of $q(\xx_t | \xx_{t-1}, \yy)$, so its analogue in the unconditional $q(\xx_t | \xx_{t-1})$ would be $\beta$)

Reproduction

I had issues with trying to get this to work on images. The paper didn’t go into details about how the ‘noisy’ $\yy$ is constructed, apart from some comments that this is produced by interpolating different audio signals to corrupt the original $\xx$. There were two strategies I conceived:

(1) generate a noisy $\yy$ via some interpolation $\yy = \lambda \xx + (1-\lambda)\boldsymbol{\epsilon}$ where $\lambda$ is sampled from some distribution whose support is in $[0,1]$ and $\epsilon \sim \mathcal{N}(0, \mathbf{I})$. For instance, we could define $\lambda \sim \text{Uniform}(a,b)$ for some hyperparameters $a$ and $b$ (this is also called mixup [3]).
(2) learn an embedding for $\yy$. Essentially, let $y'$ denote the \emph{actual label} of the image, and let $\yy = \text{embed}_{\theta}(y')$ be a learnable embedding layer that maps a label to a tensor of the same dimensions as the original image, whose parameters will also be updated in unison with those from the noise predictor $\epst$ (hence the subscript $\theta$). This formulation therefore requires a modification to Algorithms 1 and 2 in the original paper.

I initially tried (1) but found results vary quite a bit depending on $\yy$. If it is too close to the original image then you recover it, if it’s too high then you either get bad looking images or you get the same images every time (as if the noise does not matter). For (2) I just ended up getting mode dropping in general.

For this I decided to go with (2), since (1) is more akin to a domain translation type of problem and I want to do conditional generation in the more general sense (conditioned on a discrete label or token). What I found solved it was to carefully consider the schedule for $\mm$. The original paper set $\mm$ to:

\begin{align} \mm_t = \sqrt{ (1-\alphabar_{t}) / \sqrt{\alphabar_t}} \nonumber \end{align}

and remark that it gives $\mm_0 = 0$ and $\mm_T \approx 1$, but this only holds for a specific parameterisation of $\beta_t$. in particular the two choices of $\beta$ that were chosen for the experiments in [2]:

Base CDiffuSE: $\beta_t \in [1 \times 10^{-4}, 0.035]$
Large CDiffuSE: $\beta_t \in [1 \times 10^{-4}, 0.0095]$

Of course, one could enforce $\mm$ to be within this range by setting $\mm$ to be something like $\linspace(0, 1, T)$, but I found that changing this did not make a difference when it came to getting my experiments working. In fact, setting $\mm = \text{linspace}(0, \tau, T)$ seemed to resolve this, for some ‘reasonable’ values of $\tau$. Note that $\mm_{T} = \tau$. In Figure 1 we can see randomly generated images for different values of $\tau$, after 10 epochs of training with a $\beta_{t} \in [0.0001, 0.02]$ (these are just defaults I inherited from the original difusion code). It appears that if $\tau$ is too ‘large’ (close to one) then samples exhibit low diversity, yet when $\tau$ is too small (close to zero) samples are not as likely to be as class consistent (i.e. reflect the label it was conditioned on). In general:

\begin{align} \pt(\xx_T | \yy) = \mathcal{N}(\xx_T; \tau \sqrt{\alphabar_T} \yy, \delta_{T} \mathbf{I}), \nonumber \tag{1} , \end{align}

so it seems like when $\tau$ is too small then the magnitude of the mean in $\xx_T \sim \pt(\xx_T|\yy)$ (which is a function of $\yy$) becomes diminished. When $\tau$ is too large, I hypothesis that maybe this is because the variance schedule becomes too large (see Figure 1).

Examining different values of $\tau$

$$\tau=0.3$$	$$\tau=0.5$$	$$\tau=0.7$$

$$\tau= 0.8$$	$$\tau= 0.9$$	$$\tau = 1.0$$

Since $\tilde{\delta_t}$ is the conditional analogue to $\tilde{\beta}$ (and that $\tilde{\delta_t}$ is also a function of $\mm$), we plot these two variables to see how the variance changes as a function of time. This is shown in Figure 2.

Figure 1: plotting the reverse diffusion variances for both the unconditional (left) and conditional variants (right).

Reverse diffusion

Using $\tau = 0.7$, we can visualise the reverse difffusion for ten randomly sampled images per class, for a grid of 100 generated examples in total.

Figure 2: visualising the reverse diffusion for each class in MNIST.

For the sake of time I have simply chosen to visualise what qualitatively looks to be the best $\tau$; however in practice we should select the one which gives us the best FID or likelihood on the data.

Visualising the learned embeddings

What is interesting is that we can visualise the learned y’s by the network; these are what we start off with when we run the learned reverse diffusion as described in Equation (1). In Figure 3 we visualise the learned embeddings $\yy = \embed(y')$ for $y' \in \{0, 1, \dots, 9\}$. We can see that the network has, in a sense, learned ‘pseudo examples’ that parameterise the mean of $p(\xx_T|\yy)$ for each of the classes.

Figure 3: Visualising the learned embeddings for each class in MNIST.

If I can shamelessly plug my own work here, in [4] we proposed a way to ‘incrementally’ fine-tune a generative model simply by adding the new classes’ indices to the lookup table that is the embedding layer. Surely, a similar thing could be done here (assuming that our diffusion model is ‘expressive’ enough to facilitate generating from these new classes).

References

[1] Ho, J., Jain, A., & Abbeel, P. (2020). Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33, 6840-6851.
[2] Lu, Y. J., Wang, Z. Q., Watanabe, S., Richard, A., Yu, C., & Tsao, Y. (2022, May). Conditional diffusion probabilistic model for speech enhancement. In ICASSP 2022-2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) (pp. 7402-7406). IEEE.
[3] Zhang, H., Cisse, M., Dauphin, Y. N., & Lopez-Paz, D. (2017). mixup: Beyond empirical risk minimization. arXiv preprint arXiv:1710.09412.
[4] Beckham, C., Laradji, I., Rodriguez, P., Vazquez, D., Nowrouzezahrai, D., & Pal, C. (2022). Overcoming challenges in leveraging GANs for few-shot data augmentation. arXiv preprint arXiv:2203.16662.

The obsession with SOTA needs to stop

2022-08-02T00:00:00+00:00

The obsession with SOTA needs to stop

Source: I generated this. I should be doing better things with my time.

(This is a rough draft. I may do some heavy edits to this over time. Stay tuned.)

One of my biggest frustrations with machine learning (I guess more specifically deep learning) has been the absolutely abysmal quality of reviews that has come out over the past few years. A very common trope is the rejection of a paper because it fails to ‘beat SOTA’ (state-of-the-art), amongst many other critiques (‘lacks novelty’, ‘too theoretical’, etc.). I want to specifically vent about the obsession of SOTA and why I think it is not only stupid, but incentivises bad research.

To start off, I want to get off my chest that I don’t have a problem with SOTA-chasing per se as a research contribution. If you come from an engineering mindset where you want to maximise the performance of some model on a real-world setting then sure, you would place lots of emphasis on obtaining state-of-the-art performance. We want to have the best performing algorithms deployed in the real world, else self-driving cars might drive off of cliffs and kill their owners. Obviously, there are many other research contributions one could pursue: optimising for other useful metrics (e.g. memory footprint, inference speed), comparing algorithms, performing ablations, testing hypotheses, new perspectives, and so forth. My main issue lies with reviewers (and authors) placing exhorbitant emphasis on “SOTA”. There are two main reasons why I have absolute disdain for ths:

(1) There is more to research than beating SOTA (duh, see the paragraph I wrote above on all the other things you could work on);
(2) comparisons between algorithms in the literature are generally flaky and often times not statistically significant.

In my experience reading papers in machine learning, many authors appear to lack a basic understanding for how to properly evaluate a machine learning model – or, maybe they do but want to appease silly reviewers who don’t know how to do it. After all, the ultimate unit of currency in academia is the (published) paper and this often incentivises sloppiness. Here are all of the things I have seen people do (or not do) when it comes to papers proposing SOTA contributions:

(a) There is no distinction between a validation and test set. Ergo, the validation set is the test set and authors basically just tune hyperparameters on the test set, leading to optimistic estimates of generalisation performance.
(b) Strawmanning the baseline and steelmanning the proposed algorithm. Because there is so much emphasis on ~novelty~ in the field, one is strongly incentivised to spend all of their hyperparameter tuning budget (GPUs) on squeezing every little bit of performance out of their proposed algorithm, when the same may not have been done for the baseline.
(c) No estimates of uncertainty are computed for the (supposedly better) proposed algorithm. So what if your method does 95.09% and the previously SOTA method was 95.01%? It could very well go in the opposite direction over multiple runs of the same experiment (more on that soon).
(d) The issue of confounding variables. Often times papers will simply just quote numbers from other papers. While it is certainly pragmatic to do in a very fast-paced field, there will often be many confounding variables because the experimental setup of the paper you are comparing against is completely different. Maybe they are using a different framework, maybe they have preprocessed the data differently, maybe a different optimiser was used, and so forth. I am not discouraging against this, but rather saying that if you decide to go with that approach everything needs to be taken with a grain of salt. This means that as an author, you need to be extremely modest with your claims of SOTA, considering how many confounding variables you did not control for. As a reviewer, this means that if the proposed algorithm gets a few percentage points lower than SOTA then you need to ask yourself: what exactly are the paper’s claims? If the paper is claims to beat SOTA by 5% and that is literally the only proposed contribution, then sure, you’d probably be justified in rejecting it either on the grounds of statistical flakiness (i.e. cofounders). Either way, there usually there are other contributions in a paper and these need to be evaluated against the authors’ claims and done so in a holistic manner. If the paper supposedly is worse than SOTA by 5% but is not central to the paper and/or the authors don’t claim to beat it, why are you arguing for its rejection based on the fact that they didn’t beat it? Shouldn’t the positive and negative points of a paper be weighted in proportion to the claims associated with those points and the results associated with those claims?

Both authors and reviewers are often ignorant of the above points, though my criticism lies more with reviewers because they are meant to be the gatekeepers of the literature (and I mean gatekeeping in a good way, not the malevolent elitist way). When reviewers are ignorant of it and are obsessed with SOTA, they end up accepting papers that claim SOTA but whose claims are barely supported statistically. This just adds noise to the literature and makes life harder for the honest researcher who is actually trying to beat SOTA in a a principled manner. Maybe that honest researcher finds that at the end of the day none of the algorithms compared really do any better than the other, but reviewers won’t care about their paper because ~novelty and SOTA wins above everything else~. On the other hand, reviewers can also end up rejecting papers on the basis of not beating SOTA because they value it disproportionally at the expense of all of the other interesting contributions and results that the paper may have proposed (see bullet point (e)).

Extra rambles

Furthermore, on the topic of measurements of uncerainty: what is the uncertainty being computed over (when confounding variables are controlled)? Computing uncertainty (variance) over random initialisation seeds is completely different to say, random subsampling or cross-validation over your training set. In the former case, one is measuring the behaviour of the algorithm when subjected to random initialisations: if algorithm A gets 85% +/- 5% accuracy and algorithm B gets 92% +/- 10% accuracy, then this would indicate that B is less stable and would probably need more repeated training runs so that we can select the model which performs best on the validation set. If we are randomly subsampling our data, then we are essentially measuring the stability of the algorithm with respect to what might happen if one were to collect the data in practice. For instance, if we performed cross-validation and algorithm A obtained 85% +/- 20% and algorithm B obtained 85% +/- 5%, then implementing and running algorithm A on our own dataset is a whole lot riskier since it may only give us an accuracy of 60% simply by chance (i.e. one standard deviation below the mean). I am bringing this specific example up because as an author, you may propose an algorithm which performs better with respect to dataset uncertainty than seed uncertainty, but in today’s day and age I don’t even think enough reviewers are nuanced enough (or simply care) to take this into account.

Conclusion

That is all for now.

Towards a more sane Mac OS user experience, and I am late to the party

2022-07-24T00:00:00+00:00

Towards a more sane Mac OS user experience, and I am late to the party

Despite making the transition to Mac OS from Windows as early as 2014, I don’t feel like I have made a particularly strong effort to become a ‘power user’ in the past eight years. By ‘power user’, I mean the kind of person that tweaks every little setting and configuration file and hotkey so that they have an optimised workflow that works for them. While I certainly regret not delving deeper into this (and I have my reasons), I figured that this is the year that I try and at least make an effort.

Mac OS is pretty decent out of the box. And I certainly would never go back to Windows for my work life and development-related things. However, over time the user experience gradually started to irritate me, and it’s hard to describe that in precise detail because I feel as if it’s a culmination of many different things of a relatively minor nature, rather than one or two blaring issues that really stick out. For instance, I recall Windows generally having a ‘snappier’ and more responsive UI, and I found it weird that the OS didn’t even ship with an option to let you snap windows to the edges of the screen; that has to be done with a third party app. Sometimes when I installed a new application, Spotlight had a hard time trying to find it, so instead of simply typing cmd+space <application name> I had to instead open Finder, navigate to /Applications and actually launch the application from there. Oh, speaking of Finder, I absolutely hate its search with a passion. Many times I had to spin up a terminal window and use ls and grep instead because it was much more reliable. Ughhhhh. When Apple introduced their split screen mode, it only let you split the screen with two applications and only on the horizontal axis. Why so feature incomplete??? Eventually I migrated to a decent tiling window manager called Amethyst. Tiling window managers automatically re-arrange your windows, though they can also be sometimes annoying. I’ll get to that at the end of this post.

Some months ago I went through the mental crusade of trying to identify what would make my user experience a whole lot better. I concluded that the thing I wanted the most for this OS was to have something like Spotlight where you can search all open applications so that you could easily switch to them by invoking a hotkey and typing some text. As it stands, Spotlight doesn’t have this feature (though it does let you search through other boatloads of crap), and the only ‘power user’ way to flick through apps is to use mission control with a swipe gesture or use cmd+tab (which is basically alt+tab on Windows). This feature is so nice that I also have it installed in my web browser, and it’s an absolute time-saver when you have so many tabs open:

Figure 1: Showing off the Vim Vixen plugin for Firefox. This adds some nice Vim keybindings to let you power through your tabs. Here, I simply press 'b' and then use 'tab' to cycle through all of my open tabs, or even search by typing in a string.

I knew that Linux had such apps available (dmenu, rofi) and so I was torn between staying on Mac or transitioning to Linux. When I was experimenting with that transition, I decided to do it using a virtual machine (I recommend Parallels) just so that I would have the best of both worlds in front of me; after all, if the virtual machine isn’t sluggish then I really see no reason why you would want to go through the hassle of dual booting. For a few weeks I played with an Ubuntu distro with i3 as my window manager and it felt pretty decent, apart from how absolutely painful it was to have to trudge through Stackoverflow for every little confusion I had with every little config file I had to modify. It’s basically config files galore, each one having its own esoteric syntax you have to learn and with no GUIs in sight to ease the transition. I also had to deal with the fact that I had to be careful that any keybindings I set up for the guest OS (Linux) did not interfere with the host OS (Mac). Ultimately, I didn’t feel super comfortable making the transition to my virtual Linux OS because the font rendering was so much better on Mac, not to mention the amazing trackpad experience which I didn’t want to give up. Also, if I did go the dual boot route, I didn’t want to risk succumbing to this sort of fate (context: a related thread on Hackernews):

Well, if you’re very opinionated regarding your setup, trying to force macOS into your ways won’t work, macOS is great and very easy to use and gives you zero problems but you have to adapt to it. I’ve also moved after many years of linux and I could not be happier. I like easy and I like to focus on getting my actual work done, I got tired of spending weeks personalizing stuff, dealing with drivers issues, tuning the trackpad, adjusting applications to work with different dpi screens, etc, etc. For me it was a never ending war and a lot of time wasted.

For now, I feel like I have some made some headway with finding something that works for me on Mac. The first is a very awesome tiling window manager called Yabai. Yabai requires some hacks in order to leverage all of its features, but it offers really awesome power-user stuff if you combine it with a hotkey daemon like skhd. Some of its features include:

being able to instantly switch workspaces without the annoying swipe animation (with hotkeys);
the ability to switch focus on windows in a workspace;
being able to swap the positions of windows in a workspace;
being able to resize windows with hotkeys;
adding rules to control which windows should be tiled and what should remain floating;
the ability to set window opacity;
…and loads more.

It’s not all roses, however. Some of these features require you to disable System Integrity Protection. Right now I am on a work laptop, and that is definitely something I would not be allowed to do (or even achieve since I don’t have admin rights). That ends up reducing the appeal of Yabai, and in that case it may not have much more to offer than any other tiling window manager that lets you customise hotkeys.

At least for my personal laptop, skhd has allowed me to assign hotkeys to other yabai functions. Here is me switching workspaces with cmd+alt+<number> (the number corresponds to the workspace number in mission control):

Figure 2: Cycling through workspaces with my hotkeys. Note that we don't have to endure a swiping animation either, so it's snappy.

These are the commands I have in my shkdrc file to make that work:

# focus on a specific workspace
cmd + alt - 1 : yabai -m space --focus 1
cmd + alt - 2 : yabai -m space --focus 2
cmd + alt - 3 : yabai -m space --focus 3

If you install choosem and combine it with skhd, you can basically get the rofi-style application search that I mentioned earlier. For instance, in my ~/skhdrc file I have the following:

shift + cmd - space : /Users/beckhamc/miniconda3/bin/choosem yabai focus

(The command choosem yabai focus actually makes use of yabai since its API includes a command that lets you get metadata on all of the open windows. That command in particular is yabai -m query --windows)

I have this functionality bound to cmd+shift+space. (cmd+space is still bound to Spotlight, which can still be useful for launching applications.)

Figure 3: This is like Spotlight, but I can search open applications and focus to them. It is absolutely crazy that Mac does not have this built into the OS.

Tiling window managers are great, but it can be annoying when they try and tile applications with small windows. Here are some applications I blacklisted in my Yabai config:

# use this: yabai -m query --windows
# to help with filtering window types
yabai -m rule --add title=' Preferences$' manage=off
yabai -m rule --add app='Finder' manage=off
yabai -m rule --add app='System Information'  manage=off
yabai -m rule --add app='TV'  manage=off
yabai -m rule --add app='choose' manage=off # float choosem window
yabai -m rule --add app='Emacs' manage=off # emacs has a quirky ux on mac

That is all for now.

My notes on discrete denoising diffusion models (D3PMs)

2022-07-11T00:00:00+00:00

My notes on discrete denoising diffusion models (D3PMs)

1. Updates
2. Introduction
3. Characterisation of the reverse process
- 3.1. Deriving the main equation
4. Parameterisation of the reverse process
5. Conclusion
6. References

1. Updates

(24/04/2023) Thanks to Yiwei Kwok who (in the comments section) pointed out an error in my derivation of the learned reverse process.
(22/04/2023) Thanks to James Ye who worked with me on deriving the equations. He also contributed some useful questions which in turn helped me write a better explanation.
(14/07/2022) Thanks to Alex Piché who spotted a potential error with my derivation in the original version of this blog post. I think the derivation I have is correct now.

2. Introduction

Here are some of my thoughts on a (semi-)recent paper that came out by Austin et al [1]. It proposes a variant of the (continuous) diffusion model for discrete data, which is typically an awkward modality to deal with since almost everything we do with deep neural networks is in continuous space. Many types of data can be represented as discrete, for instance text, molecules, graphs, and even images if we don't dequantise their pixel values (which are typically ordinal variables taking on values (0-255)). Representing data as discrete variables can also be a reasonable form of compression.

In this short blog post I won't be going into all of the details of the paper, but I will be mainly presenting some math I did to help me better understand the proposed algorithm.

Figure 1: illustration of the forward and reverse process in D3PM. Attribution: Austin et al.

3. Characterisation of the reverse process

Probably the most confusing aspect of the paper for me was Section 3.3, which explains how the reverse process $\pt$ is formulated. For some confusing reason (which looks to be empirically justified) instead of just having our neural network $\pt$ predict $\xx_{t-1}$ from $\xx_t$ (i.e. parameterise $\pt(\xx_{t-1} | \xx_{t})$), our neural network learns $\pt(\xx_0 | \xx_t)$, which is learning how to jump from $\xx_t$ directly to $\xx_0$. However, then the question is, how do we actually do proper diffusion, which is denoising iteratively one step at a time?

It took me some time to figure this out, but the key is actually in Equation (3) of [1]. Equation 3 in [1] basically says that for the forward process $q$, it is possible (that is, tractable) to reverse it if we condition on $\xx_0$, which is shown to be the following:

\begin{align} q(\xx_{t-1}|\xx_t, \xx_0) & = \frac{q(\xx_t | \xx_{t-1}, \xx_0) q(\xx_{t-1}|\xx_0) }{q(\xx_t | \xx_0)} \tag{1} \\ & = \text{Cat}\Big( \xx_{t-1}; \boldsymbol{p} = \underbrace{\frac{\xx_t \QQ_t^{T} \odot \xx_0 \bar{\QQ}_{t-1}}{\xx_0 \bar{\QQ}_t \xx_{t}^{T}}}_{\text{is this correct?}} \Big) \tag{1b} \end{align}

3.1. Deriving the main equation

In the original version of this blog post, I just stated Eqn. (1b) without really deriving it myself. However, there were a few people who were asking about how it was derived. After a long back and forth with James Ye (who helped me with the derivations) it looks like the equation is correct, but it took a lot of trial and error to get there.

We start off with some preliminaries:

$\xx_t \in \{0,1\}^{k}$ (for any $t$) is a discrete label represented as a one-hot vector.
$\tilde{\xx}_{t} \in [0,1]^{k}$ is the predicted variable (for any $t$), but since it is a probability distribution over the $k$ elements it is not one-hot.
$Q_t \in [0,1]^{k \times k}$ is a stochastic matrix whose rows sum to 1.

We have to be very careful here with our derivations because some of the expressions computed evaluate to vectors (probability distributions over the $k$ classes), and some evaluate to scalars, since they are individual probabilities computed by indexing into a probability distribution. I also find that it is crucial to make it clear when we are referring to the random variables themselves compared to their realisations, i.e. when they have been observed or conditioned on. Otherwise things get really confusing. First, let me re-write Eqn. (1) such that $X_t$ now denotes the random variable and $X_t = \xx_t$ its realisation:

\begin{align} q(X_{t-1} | X_t = \xx_t, X_0 = \xx_0) & = \frac{q(X_t = \xx_t|X_{t-1}, X_0 = \xx_0)q(X_{t-1} | X_0 = \xx_0)}{q(X_t = \xx_{t} | X_0 = \xx_0)} \tag{2} \\ & = \frac{q(X_t = \xx_t|X_{t-1})q(X_{t-1} | X_0 = \xx_0)}{q(X_t = \xx_{t} | X_0 = \xx_0)} \tag{2b} \\ \end{align}

where in expression (A) in (2b) we make use of the Markov property, since;

\begin{align} q(\xx_t|\xx_{t-1}, \xx_0) = q(\xx_t|\xx_{t-1}). \end{align}

We now define the individual terms in Eqn. (2b), starting from left to right in the numerator and then the denominator.

\begin{align} & \underbrace{q(X_t = \xx_t|X_{t-1} = \xx_{t-1})}_{1 \times 1} = \big[ \underbrace{\xx_{t-1}}_{1 \times k} \underbrace{Q_t}_{k \times k} \big] \underbrace{\xx_{t}^{T}}_{k \times 1} \tag{3} \end{align}

Note that this assumes we have observed $X_{t-1}$ as well, but we clearly haven't since we're trying to compute the conditional distribution over $\xx_{t-1}$ to begin with! Therefore, we ought to write Eqn. (3) by enumerating all the possible values $\xx_{t-1}$ could take, which is just the identity matrix $\mathbf{I}_{k}$. However, that means that $\xx_{t-1}$ disappears from Eqn. (3) and we get the following:

\begin{align} & \underbrace{q(X_t = \xx_t|X_{t-1})}_{k \times 1} = \big[ \underbrace{\mathbf{I}_k}_{k \times k} \underbrace{Q_t}_{k \times k} \big] \underbrace{\xx_{t}^{T}}_{k \times 1} = Q_t \xx_{t}^{T} \tag{4} \end{align}

(If the idea of having a non-observed variable on the conditioning side of the expression seems weird, I elaborate on this on Sec. 3.1.1.) Eqn. (4) is now a column vector however, and we want to keep things consistent by representing probability distributions or examples as row vectors. So let us abuse notation by redefining Eqn. (4) so that it's a row vector. That just means transposing the RHS expression of Eqn. (4) to be:

\begin{align} & \underbrace{q(X_t = \xx_t|X_{t-1})}_{1 \times k} := [Q_t \xx_{t}^{T}]^{T} = \xx_{t} Q_{t}^{T} \tag{4b} \end{align}

For the second term:

\begin{align} \underbrace{q(X_{t-1} | X_0 = \xx_0)}_{1 \times k} = \underbrace{\xx_0}_{1 \times k} \underbrace{\bar{Q}_{t-1}}_{k \times k} \tag{5} \end{align}

where $\bar{Q}_{t-1} = Q_{1}Q_{2} \dots Q_{t-1}$.

For the last (denominator) term:

\begin{align} \underbrace{q(X_t = \xx_{t} | X_0 = \xx_0)}_{1 \times 1} = \big[ \underbrace{\xx_0}_{1 \times k} \underbrace{\bar{Q}_{t}}_{k \times k} \big] \underbrace{\xx_t^{T}}_{k \times 1} \tag{6} \end{align}

Putting it all together, we now get:

\begin{align} \therefore q(X_{t-1} | X_t = \xx_t, X_0 = \xx_0) = \text{Cat}\Big(\xx_{t-1}; \frac{ \underbrace{\xx_{t}Q_t^{T}}_{\text{vector}} \odot \underbrace{\xx_0 \bar{Q}_{t-1}}_{\text{vector}} }{ \underbrace{\xx_0 \bar{Q}_t \xx_t^T}_{\text{scalar}} } \Big). \ \ \ \square \end{align}

In conclusion, yes, the result is consistent with the original equation in the D3PM paper. But its derivation should have been in the appendix.

3.1.1. Non-observed conditioning variables

In Eqn. (4) we saw an interesting kind of expression, one where the probability of a particular $X_t$ was being conditioned on a non-observed $X_{t-1}$. Before we elaborate on this, perhaps it is useful to consider all the different possible realisations of the expression $q(X_t | X_{t-1})$:

$q(X_{t} | X_{t-1}) \in [0,1]^{k \times k}$, what is the probability distribution over the different values $X_t$ can taken on, for some unspecified $X_{t-1}$?
$q(X_t|X_{t-1} = \xx_{t-1}) \in [0,1]^{1 \times k}$, what is the probability distribution over the different values of $X_t$ given that I have observed $X_{t-1}$ to be $\xx_{t-1}$?
$q(X_t = \xx_t | X_{t-1} = \xx_{t-1}) \in [0,1]$: what is the probability of observing $X_{t} = \xx_t$, given that I have observed $X_{t-1}$ to be $\xx_{t-1}$?
And lastly $q(X_t = \xx_{t} | X_{t-1}) \in [0,1]^{k \times 1}$: what is the probability of observing $X_t = \xx_t$, given that… well, nothing has been observed, so what does this mean?

We know via Eqn. (4) that $q(X_t = \xx_{t} | X_{t-1})$ is a column vector (i.e. a $k \times 1$ matrix) and that its entries encode the following:

\begin{align} q(X_t = \xx_{t} | X_{t-1}) = \begin{bmatrix} q(X_t = \xx_{t} | X_{t-1} = [1, 0, \dots, 0]) \\ q(X_t = \xx_{t} | X_{t-1} = [0, 1, \dots, 0 ]) \\ \dots \\ q(X_t = \xx_{t} | X_{t-1} = [0, 0, \dots, 1 ]) \\ \end{bmatrix} \end{align}

So its interpretation is also simple: as a distribution over all possible observed $X_{t-1}$'s, what is the probability of observing $X_t = \xx_t$? Or, simply consider the $j$'th element of the vector instead: what is the probability of observing $X_t = \xx_t$, if I did condition on the observation that $X_{t-1}$ was $j$?

I thank James Ye for asking this question, since it also had me confused. Hopefully my explanation suffices.

4. Parameterisation of the reverse process

While we know that $q(\xx_{t-1}|\xx_t, \xx_0) = q(\xx_{t-1}|\xx_t)$ due to the Markov property, when we derive reverse of the forward process we need to actually keep it in. In fact, rather than just doing away with $\xx_0$ completely we will instead marginalise it out:

\begin{align} q(\xx_{t-1}|\xx_{t}) & = \frac{\sum_{\xx_0} q(\xx_{t-1}, \xx_t, \xx_0)}{q(\xx_t)} \tag{7} \\ & = \frac{\sum_{\xx_0} q(\xx_{t-1} | \xx_t, \xx_0) q(\xx_0 | \xx_t) q(\xx_t) }{q(\xx_t)} \tag{7b} \\ & = \sum_{\xx_0} q(\xx_{t-1} | \xx_t, \xx_0) q(\xx_0 | \xx_t) \tag{7c} \\ & = \mathbb{E}_{q(\xx_0|\xx_t)} \ q(\xx_{t-1} | \xx_t, \xx_0) \tag{7d} \end{align}

Note that the expection is over $q(\xx_0|\xx_t)$, which we don't have! What we do have however is our learned reverse process, so we can just approximate this term with $\pt(\xx_0|\xx_t)$. I'm going to abuse notation here and call this $q_{\theta}$ since this is an 'amalgamation' of the forward process and our learned reverse process:

\begin{align} q(\xx_{t-1}|\xx_{t}) \approx \mathbb{E}_{\xx_0 \sim \pt(\xx_0|\xx_t)} \ q(\xx_{t-1} | \xx_{t}, \xx_0) = q_{\theta}(\xx_{t-1}|\xx_{t}). \tag{6} \end{align}

Taking the expectation on both sides of Equation (3) in [1], we can derive the following:

\begin{align} \mathbb{E}_{\pt(\xx_0|\xx_t)} \ q(\xx_{t-1}|\xx_t, \xx_0) & = q_{\theta}(\xx_{t-1}|\xx_t) \tag{8} \\ & = \mathbb{E}_{\pt(\xx_0|\xx_t)} \ \frac{q(\xx_t | \xx_{t-1}, \xx_0) q(\xx_{t-1}|\xx_0) }{q(\xx_t | \xx_0)} \tag{8b} \\ & = \mathbb{E}_{\pt(\xx_0|\xx_t)} \ \frac{q(\xx_t | \xx_{t-1}) q(\xx_{t-1}|\xx_0) }{q(\xx_t | \xx_0)} \tag{8c} \\ & = q(\xx_t | \xx_{t-1}) \ \mathbb{E}_{\pt(\xx_0|\xx_t)} \ \frac{q(\xx_{t-1}|\xx_0) }{q(\xx_t | \xx_0)} \tag{8d}. \end{align}

If the expectation is approximated by a single sample $\xx_0 \sim \pt(\xx_0|\xx_t)$ then it disappears and we get the following:

\begin{align} & \approx q(\xx_t|\xx_{t-1}) \frac{q(\xx_{t-1}|\xx_0)}{q(\xx_t|\xx_0)}. \tag{8e} \end{align}

Let's run through this line by line:

From (8b) to (8c), $q(\xx_t|\xx_{t-1}, \xx_0) = q(\xx_t|\xx_{t-1})$ due to the Markov property.
From (8c) to (8d) we can move $q(\xx_t|\xx_{t-1})$ outside the expectation since it does not depend on $\xx_0$.

I thank Yiwei Kwok for pointing out an error in the initial derivation.

Equation (8e) is implemented here in code, when x_start_logits=True. To be consistent with what's in the code, let us call fact1 (short for 'factor') the $q(\xx_t|\xx_{t-1})$ term and fact2 the term $q(\xx_{t-1}|\xx_0)$. The denominator isn't computed since the implementing method is only considering the logits, but this can easily be normalised at any time by taking the softmax.

fact1 = self._at(self.transpose_q_onestep_mats, t, x_t). This function call is implementing $\xx_{t} \QQ_{t}^{T}$.
fact2 = self._at_onehot(self.q_mats, t-1, jax.nn.softmax(x_start, axis-1). This function call is implementing $\xx_0 \bar{\QQ}_{t-1}$. x_start here is actually the predicted logits $\tilde{\pt}(\xx_0|\xx_t)$, which subsequently gets normalised with jax.nn.softmax(x_start).
Note that the multiplication of both factors is done in log space, so we add the terms, i.e. log(fact1*fact2) = log(fact1) + log(fact2).

5. Conclusion

I thank the original paper author Jacob Austin for addressing a confusion of mine in the code.

That is it for now! If you have any questions or spot errors in my equations, please reach out to me on Twitter or via email.

6. References

[1] Austin, J., Johnson, D. D., Ho, J., Tarlow, D., & van den Berg, R. (2021). Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34, 17981-17993.
[2] Ho, J., Jain, A., & Abbeel, P. (2020). Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33, 6840-6851.

Training GANs the right way

2021-06-28T00:00:00+00:00

Training GANs the right way

Welcome to my first blog post! I am years overdue in doing this. To preface things a bit:

I expect to update this over time, both in response to feedback as well as when I find new tricks or nuggets of wisdom that are worth sharing.
This post assumes you have some familiarity with GANs and how they work, and that you have implemented a GAN at least once. If not, you may find some sections confusing.
Also, because I don’t typically dabble with high-resolution GAN models (i.e. BigGAN, PGAN), you shouldn’t expect to see any tricks on how to get those models to work. My day-to-day research often involves training GANs of a ‘modest’ resolution (32px, 64px), so the tricks I explain here are mostly applicable to those models.

Updates:

(05/07/2020) Added a minimal SN-GAN implementation of MNIST on Colab here.
(04/07/2020) Added extra tip in ‘common pitfalls’ about image sizes.
(01/07/2020) Added section in how to deal with more than one noise variable as input; extra text on FID and Inception metrics.

If you found this useful and wish to cite it, you can use this corresponding Bibtex entry:

@misc{beckhamc_traininggans,
  author = {Beckham, Christopher},
  title = {Training {GAN}s the right way},
  year = {2021},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://beckham.nz/2021/06/28/training-gans.html}}
}

The key ingredient to stabilising GAN training
Structuring the code cleanly
Easy or common pitfalls to make
Plot your damn learning curves!
Because of Lipschitz, you don’t need to cripple your discriminator
Optimisers
- Pesky hyperparameter: ADAM’s epsilon
Evaluation metrics
- Which is the best?
Conditioning on more than one noise variable
Resources
Conclusion
References
Appendix
- The relationship between PyTorch and TF Inception scores

The key ingredient to stabilising GAN training

GAN training used to involve a lot of heuristics in order to minimise their various degenerate behaviours, like mode collapse or mode dropping. Commonly this meant setting up the generator and discriminator in such a way that the latter did not perform ‘too well’, because that would ultimately mean risking vanishing or exploding gradients from the discriminator. The worst degenerate behaviour is when the generator exhibits mode collapse, in which the model outputs the same image no matter what input its given. The other behaviour, mode dropping, also isn’t ideal because it means there are certain factors of variation in the data distribution that fail to get modelled. Despite the fact that we try to do our best to minimise this behaviour, I would argue here that GANs exhibit mode dropping behaviour by design. This is because GANs optimise something that isn’t identical to maximum likelihood, which is inherently mode covering and is the mechanism by which autoencoders are trained. (See [#ref:colin].)

In 2017 the Wasserstein-GAN (WGAN) [#ref:wgan] was published and made some very interesting theoretical and empirical contributions to GANs and adversarial training, most notably in stabilising them and making them significantly less of a hassle to train. To summarise the paper – and I hope I don’t butcher this explanation – the gist of the paper is that:

There exist generator distributions which do not converge to the data distribution under non-Wasserstein divergences (JS, KL, etc.);
ones that do converge under non-Wasserstein divergences converge for Wasserstein;
the discriminator (called the critic in the paper) under a WGAN gives non-saturating (clean) gradients everywhere, provided that the discriminator is K-Lipschitz with respect to its parameters (for small $K$).

Therefore, the key thing is to somehow ensure that the discriminator is $K$-Lipschitz (for some small $K$). In the original paper this was done via (1) weight clipping technique, and in later literature this was followed by (2) gradient penalty [#ref:wgan_gp] and then (3) spectral normalisation [#ref:sn]. A quick summary of these is that (1) is overly excessive in its regularisation, (2) is rather expensive to compute and only penalises certain parts of the function, and (3) regularises the entire function and is rather cheap to compute. In particular, spectral norm[#ref:sn] ensures that the Lipschitz constant of the network is upper bounded by one (1-Lipschitz), by constraining the largest singular value for each weight matrix in the discriminator network to be equal to one. We’ll talk about that later, but for now let’s go back to basics.

To briefly summarise WGAN [#ref:wgan], via the Kantorovich-Rubenstein (KR) duality, the Wasserstein distance between two distributions can be expressed as the following:

\[W(P_r, P_{\theta}) = \sup_{||D||_{L} \leq 1} \mathbb{E}_{x \sim P_r} D(x) - \mathbb{E}_{x \sim P_{\theta}} D(x),\]

where $P_r$ refers to the real distribution and $P_{\theta}$ the generator’s distribution, and $D(x) \in \mathbb{R}$. Unlike the original paper, I am using $D(\cdot)$ to denote the discriminator (critic) instead of $f(\cdot)$. The supremum is basically saying that the Wasserstein distance corresponds to the function $D$ that makes the following term as large as possible, i.e. make $\mathbb{E}_{x \sim P_r} D(x)$ very large and $\mathbb{E}_{x \sim P_{\theta}} D(x)$ very small. Equivalently, modifying the supremum so that it instead is over $K$-Lipschitz functions simply amounts to multiplying the distance by the scaling factor $K$, i.e:

\[K \cdot W(P_r, P_{\theta}) = \sup_{||D||_{L} \leq K} \mathbb{E}_{x \sim P_r} D(x) - \mathbb{E}_{x \sim P_{\theta}} D(x).\]

While we obviously cannot compute such a supremum, the basic idea is to just enforce that Lipschitz constraint on $D$ and treat it as a network that simply approximates the true Wasserstein distance. Quite simply, $W$ is the loss that the discriminator tries to maximise, while the generator tries to minimise it. In other words, $D$ tries to maximise the (approximated Wasserstein) distance between the real and generated distributions, while $G$ tries to make them as similar as possible. Therefore, we can write $D$ and $G$’s losses as:

\[\begin{align} \mathcal{L}_{D} & = \max_{D} \ \ \ \ \ \mathbb{E}_{x \sim P_r} D(x) - \mathbb{E}_{x \sim P_{\theta}} D(x) \\ & = \min_{D} \underbrace{\mathbb{E}_{x \sim P_{\theta}} D(x)}_{\text{make as small as possible}} - \underbrace{\mathbb{E}_{x \sim P_r} D(x)}_{\text{make as large as possible}} \\ \mathcal{L}_{G} & = \min_{G} - \underbrace{\mathbb{E}_{x \sim P_{\theta}} D(x)}_{\text{make as large as possible}} \end{align}\]

Note that for $D$ I am writing its loss both as a maximisation and a minimisation (which is just the negative of the maximisation). The former is convenient if you want to think about the two networks’ losses in terms of a minimax game, but the latter is more convenient for an actual implementation since we tend to minimise loss functions in code. Let us briefly compare this to JS-GAN’s formulation, which is sort of similar but involving logs and a sigmoid nonlinearity on $D$:

\[\begin{align} \mathcal{L}_{D} & = \max_{D} \mathbb{E}_{x \sim P_r} \log \sigma(D(x)) + \mathbb{E}_{x \sim P_{\theta}} \log(1 - \sigma(D(x))) \ \ \ \text{(1)} \\ & = \max_{D} \mathbb{E}_{x \sim P_r} \log D_r(x) - \mathbb{E}_{x \sim P_{\theta}} \log(D_f(x)) \ \ \ \text{(2)}, \\ \end{align}\]

where in (1) $\sigma(\cdot)$ denotes the sigmoid nonlinearity. In the following line I have decided to define two outputs on $D$ instead ($D_r$ for probability of real and $D_f$ for probability of fake, as if this was a two-class softmax) so that we can remove the sigmoid term and turn the summation term in (1) into a subtraction in (2) to be more reminiscent of WGAN’s loss. Finally, the log function can just be thought of as a non-linearity as well.

The reason for this little comparison is because it segways into an interesting detail that Martin Arjovsky wrote in the WGAN appendix. He says that, ultimately, a JS-GAN is quite similar to WGAN if you also apply similar Lipschitz constraints:

Figure 1: Interesting appendix detail about JSGAN, found in the supplementary material from WGAN.

I find this interesting, though what certainly muddies the waters a bit is that there evidence to suggest that in practice a Lipschitz’d JS-GAN performs better than a Lipschitz’d WGAN:

it seems to be the case in the spectral norm paper (though the comparisons between JS and WGAN actually used gradient penalty, not spectral norm);
some peeps here and here posted issues on getting SN to work with WGAN, and even the author of the spectral norm paper had trouble himself;
I vaguely recall achieving lower Inception scores on an old project that used a WGAN-SN.

Maybe it’s a mystery for now, but if it’s any consolation, theoretically both the regularised JSGAN and WGAN formulations are very similar and in practice you should just use JS-GAN, or even better, the hinge loss proposed in [#ref:sn]. Spectral normalisation has been in PyTorch for quite some time now, and it’s as simply as wrapping each linear and conv layer in your discriminator with torch.nn.utils.spectral_norm:

import torch
from torch import nn

# spec norm is built into pytorch, and it's very
# plug and play
from torch.nn.utils import spectral_norm as spec_norm

n_in = 10
n_out = 20

layer_l = spec_norm(nn.Linear(n_in, n_out))

layer_c = spec_norm(nn.Conv2d(n_in, n_out, kernel_size=3))

Prior to this work, these were many tricks that I saw on how to stabilise GAN training [#ref:gan_hacks]. While I don’t want to say that none of these have any utility anymore whatsoever, I have not had to use any of these. I feel however that any newcomers may stumble across such things and not realise that there are more straightforward ways to train GANs now.

Figure 2: Some GAN training hacks I stumbled across on this Github repo. These tricks date back five years and probably many of these tricks are not needed anymore.

Structuring the code cleanly

Clean code makes it easier to find bugs and reason about your code. Good abstractions are part of that. I like to have a train_on_batch(x) method, which performs a single gradient step over the batch x. Since GANs are a two-player game between $G$ and $D$ however, how do we structure the code? Which network should we train first?

Personally, I perform the gradient steps for the generator first, i.e., the $G$ step. But it doesn’t matter which way you do it, as long as you’re careful as to how you’re handling the gradients. I’m going to be writing PyTorch-inspired pseudo-code (it’s using Python and PyTorch syntax but it’s not meant to be run per se, not unless you define a bunch of variables and implement the undefined methods I have written).

####################################################################
# This code is not meant to be executed -- it's simply pseudocode. #
####################################################################

REAL = 1
FAKE = 0

def train_on_batch(x_real):
    
    opt_g.zero_grad()
    opt_d.zero_grad()
    
    # --------------
    # First, train G
    # --------------
    
    z = sample_z(x_real.size(0))
    x_fake = G(z)
    g_loss = gan_loss( D(x_fake), REAL)
    # This backpropagates from the output of D, all the
    # way back into G.
    g_loss.backward()
    # G's gradient buffers are filled, we can perform
    # an optimisation step.
    opt_g.step()
    
    # ------------
    # Now, train D
    # ------------
    
    # IMPORTANT: D's grad buffers are filled because
    # of what we did above.
    opt_d.zero_grad()
    
    # x_fake.detach() not necessary here but it stops
    # grads from backpropagating into G. Even if that
    # did happen howwever, the start of `train_on_batch` 
    # zeros both gradient buffers anyway, so it doesn't 
    # matter.
    d_loss = gan_loss( D(x_fake.detach()), FAKE) + \
             gan_loss( D(x_real), REAL )
    d_loss.backward()
    opt_d.step()
    
    return g_loss.detach(), d_loss.detach()

What if you want to do it the other way around? Easy, though I find it’s a bit more confusing:

####################################################################
# This code is not meant to be executed -- it's simply pseudocode. #
####################################################################

REAL = 1
FAKE = 0

def train_on_batch(x_real):
    
    opt_g.zero_grad()
    opt_d.zero_grad()
    
    # ------------
    # Now, train D
    # ------------

    z = sample_z(x_real.size(0))
    x_fake = G(z)
    
    # Call `x_fake.detach()` because we don't want to
    # backpropagate gradients into G. If you didn't
    # use detach(), then when you call `g_loss.backward()`
    # later on you'd have to supply `retain_graph=True`.
    # Let's not do such confusing things...
    d_loss = gan_loss( D(x_fake.detach()), FAKE) + \
             gan_loss( D(x_real), REAL )
    d_loss.backward()
    opt_d.step()
    
    # ------------
    # Now, train G
    # ------------
    
    opt_d.zero_grad()
    
    g_loss = gan_loss( D(x_fake), REAL)
    # This backpropagates from the output of D, all the
    # way back into G.
    g_loss.backward()
    # G's gradient buffers are filled, we can perform
    # an optimisation step.
    opt_g.step()
    
    return g_loss.detach(), d_loss.detach()

You may see in some GAN implementations the use of backward(retain_graph=True) in order to be able to backprop through a particular graph more than once. For instance, in the code block directly before this you would have to use retain_graph=True if I didn’t detach() the tensor x_fake for the block that trains D. I have avoided such use of that here since I personally find it more confusing to think about.

Why have I defined a function gan_loss here? Because for certain GAN formulations (like JS-GAN), it helps a lot with readability. Rather than doing this, which can be relatively less readable and more error prone:

d_loss = -log(D(x_real)) - log(1-D(x_fake))
g_loss = -log(D(x_fake))

You could instead just write (at the expense of a few extra lines):

bce = nn.BCELoss()

# 1 = real, 0 = fake
ones = torch.ones((x_real.size(0), 1)).float()
zeros = torch.zeros_like(ones)

d_loss = bce( D(x_real), ones ) + bce( D(x_fake), zeros )
g_loss = bce( D(x_fake), ones )

Easy or common pitfalls to make

Make sure that your generated samples are in the same range as your real data. For instance, if your real data is always in [-1, 1] but your fake data is in [0, 1], that is something that the discriminator can pick up on to easily distinguish real from fake, and could result in degenerate training. My own rule of thumb is to always do non-fancy preprocessing on the real inputs: simply put it in the range [-1, 1] by performing (x-0.5)/0.5, and make the output of your generator function tanh. When you want to visualise those images in matplotlib, simply denormalise by computing x*0.5 + 0.5.
If your generated images are not the same size (spatial dimension) as your real images, the discriminator can easily pick up on this and your generator loss will climb extremely high. This may seem like a no-brainer but I have made this mistake a few times. An example of this is MNIST: the default size of the images in torchvision.datasets.MNIST is 28x28 but your generator architecture may generate 32x32 images, so in this case make sure you add a torchvision.transform that resizes the MNIST digit to 32x32.
Batch norm can sometimes be unpredictable and result in wildly different generated images at training or test time (where in training time the batch statistics are computed over the minibatch, and at test time the moving averages are used). Usually I just use instance norm in place, of it. I’ve been bitten by batch norm’s intracacies too many times.

Plot your damn learning curves!

Most of the time when people come to me with GAN problems, they never come with graphs included and I usually have to ask for it. Sometimes I get no immediate response, and I wonder if it’s because they have to write code to do it, which certainly implies that their workflow isn’t very plotting-centric. Pessimistically it makes me wonder if the way most people debug things in deep learning is to just look at metrics flying down the terminal screen like this:

EPOCH: 001, TRAIN_D_LOSS: 0.4326436, TRAIN_G_LOSS: 1.352356, TIME: 435.353 SEC ...
EPOCH: 002, TRAIN_D_LOSS: 0.4521224, TRAIN_G_LOSS: 1.325623, TIME: 425.353 SEC ...
EPOCH: 003, TRAIN_D_LOSS: 0.4575744, TRAIN_G_LOSS: 1.657234, TIME: 422.533 SEC ...
EPOCH: 004, TRAIN_D_LOSS: 0.4025356, TRAIN_G_LOSS: 1.124543, TIME: 411.632 SEC ...
EPOCH: 005, TRAIN_D_LOSS: 0.4235636, TRAIN_G_LOSS: 1.457234, TIME: 450.353 SEC ...
...
...
...
(continue until your eyes get sore)

While our field does seem like the polar opposite of classical statistics (i.e. simple linear models, interpretability, etc.), we do have something in common with classical statistics, and that is in exploratory analysis. Just like classical statistics, deep learning should involve lots of exploratory analysis. Sure, that exploration isn’t going to be on 10,000-dimensional data (we can only see in three dimensions), but it will be on variables that your model is either optimising or measuring, which is extremely important to monitor if you want to know your network is training correctly. In our case, the losses are d_loss and g_loss, as defined in earlier pseudocode.

In Figure 3 I have illustrated these curves, left for a regular JS-GAN and on the right a spectrally normalised variant, SN-GAN.

Figure 3. Left: Unregularised JS-GAN; right: JS-GAN regularised with spectral normalisation ('SN-GAN'). Note that the discriminator loss is unnormalised and should be divided by 2 here, since it constitutes two terms (the real and the fake terms). In that case, for the right-most figure the generator and discriminator losses for SN-GAN are roughly the same value, though typically over time they will diverge very slowly. It does seem that this phenomena completely contradicts the 'Nash equilibrium' theory of GANs, which says that -- theoretically -- both networks should converge to the same loss, implying that neither network is able to reach a lower loss by fooling the other. This seems to generally not happen in practice. For more details on why, you can consult [[#ref:nash]](#ref_nash).

In my experience training SN-GAN (not JS-GAN), you want to be in a regime where the D loss is lower than the G loss but not by ‘too much’. For instance, if one loss is much larger or smaller than the other, this may be indicative of something problematic. For instance, if d_loss << g_loss (where << = significantly lower), this could mean your generator does not have enough capacity to model the data well and cannot compete well with the discriminator. Alternatively, if g_loss << d_loss, the discriminator may not be powerful enough, either. While you could use some sort of ‘balancing’ heuristic like roughly make both networks have the same # of learnable parameters, this is actually really crude. For instance, it’s expected $G$ have way more parameters than $D$ since it’s actually trying to model the data distribution (rather than simply distinguish between the two), so it doesn’t exactly set off alarm bells if your discriminator is only 5M parameters and your generator is 50M. Also, a 50M network with heavy weight decay does not have the same modeling capacity as one without it. Lastly, the issue certainly may not be in model complexity but rather the training dynamics. For example, if either network is really deep, are you making use of residual skip connections? If you’re concerned with vanishing gradients, you could also plot the average gradient norms here).

This is one of these things you will get a feel for after training lots of GANs. My best advice is to look at existing implementations of GANs on Github (ones that achieve good performance) and use their architectures as a base for your own work.

Because of Lipschitz, you don’t need to cripple your discriminator

Based on what I said about K-Lipschitz, it should certainly help (and not degrade gradients in any way) to give your discriminator a head start by training it for relatively more iterations than the generator. In other words, the better the discriminator is at distinguishing between real and fake, the better the signal that the generator can leverage from it. Note that as I mentioned earlier in this article, this logic did not make sense in the ‘pre-WGAN’ days because making the discriminator too good was detrimental to training.

In my own code, I simply keep track of the iteration number so that I can compute something like the following:

def train_on_batch(x, iter_, n_gen=5):
    # Generator
    ...
    ...
    if iter_ % n_gen == 0:
        g_loss.backward()
        opt_g.step()
        
    # Disc
    ...
    ...
    d_loss.backward()
    d_loss.step()

Where iter_ is the current gradient step iteration, and n_gen defines the interval between generator updates. In this case, since it’s 5, we can think of this as meaning that the discriminator is updated 5x as much as the generator.

Other people may do something like the following, preferring to leverage the data loader to perform such a thing:

def train(N):
    
    # number of discriminator iters per
    # generator iter
    n_dis = 5

    for iteration in range(N):

        for x_real in data_loader:
            # Update G
            train_on_batch_g(x_real)

        for _ in range(n_dis):
            for x_real in data_loader:
                # Update D
                train_on_batch_d(x_real)

Optimisers

The only optimiser I’ve really used for GANs is ADAM, and it seems like everyone else does as well. I don’t know why it works so well, but maybe it’s because all of our models have evolved over time to perform well on ADAM [#ref:adam_evolve], even if it may not necessarily be the right optimiser to always use.

[#ref:sn] presents a neat paper on various ADAM hps that they tried out, and which ones gave the best Inception scores:

Figure 4: barplot illustrating ideal ADAM hyperparameters over different GAN parameterisations. (Taken from spectral norm paper)

As one can see here, for SN-GAN, option (C) performs best for CIFAR10 and option (B) for STL-10. Unfortunately however, as that graph shows, you can get wildly different optimal hps depending on what kind of GAN variant you are working with. In my own experience, I have found betas=(0, 0.9), lr=2e-4 to be a reasonable starting point.

Pesky hyperparameter: ADAM’s epsilon

ADAM’s default epsilon parameter in PyTorch is 1e-8, which may cause issues after a long period of training, such as your loss periodically exploding or increasing. See the StackOverflow post here as well as the Reddit comments here.

There may be some alternatives that alleviate this issue, such as Adamax. But I have not tried using it yet with GAN training.

Evaluation metrics

The two most commonly used metrics are Inception and FID, though there have been a whole host of other ones proposed as well. Unfortunately, these metrics can often be rather frustrating to use, and I will detail as many reasons as I can for this below.

Here are some things you should keep in mind about FID:

Depending on how FID is calculated, it may only be appropriate to use it as a measure of underfitting, not overfitting. Looking at the code of pytorch-fid, which in turn is based on test, the pre-computed FID statistics for a particular dataset may be based on either the training set, the validation set, or both sets. For CIFAR10, it looks like it’s the training set. This means that if your generative model ‘cheated’ and simply memorised the training set, you’d get an FID of zero. I don’t know what % of GAN papers compute FID on train/valid/test but this seems like yet another confounding factor that makes reproducibility hard.
Like Inception (see above), these scores are different depending on what pre-trained weights are used for the Inception network, which typically differ depending on whether the implementation is in PyTorch or TensorFlow. It seems like pytorch-fid now uses the same weights as the Inception network in TF, but I recall a few years back having to use a different set of weights for PyTorch with this code. I managed to dig up some old numbers where I evaluated various hyperparameters of a GAN I trained using both TF and PyTorch implementations, and got wildly different results (see Appendix).

There are other problems too. See ‘A note on the inception score’ [5] for a comprehensive critique of these scores.

Either way, you should be monitoring these metrics during training, as opposed to simply training the GAN for a fixed number of epochs and evaluating them afterwards. Think about this as being no different from training a classifier and monitoring its validation set accuracy during training: you want to know if image quality is improving over time, and you want to know if you’re able to stop training early when the FID/Inception plateaus or gets worse. Since these metrics can be expensive to compute, set up your code so that the metrics are computed every few epochs, so that you don’t completely bottleneck your training time.

Which is the best?

Ultimately, I feel like whatever evaluation metric you use really depends on what your downstream task is. If you just want to generate pretty images then sure, use something like FID or Inception. If however you are using GANs for something like data augmentation, then an appropriate metric would be training a classifier on that augmented data and seeing how well it performs on a held-out set.

While GANs are great in many aspects, one of their biggest downsides is not having a theoretically straightforward way of evaluating likelihoods. VAEs more or less give you this, by allowing you to compute a lower bound on $p(x)$.

Conditioning on more than one noise variable

Sometimes you may want to train a more flexible generator, one that is able to be conditioned on more than one noise variable. For instance, rather than $G(z)$ you may want $G(z, c)$, where $c \sim p(c)$ comes from some other prior distribution, e.g. a Categorical distribution. If you train the GAN, you will most likely find that it completely ignores one noise variable and uses the other.

In order to resolve this, you should turn to InfoGAN [#ref:infogan]. Put simply, the discriminator should have two additional output branches, one to predict both latent codes fed into the generator, and you should train the discriminator to be able to predict these codes from the generated image. In other words, we want to maximise the mutual information between the original codes input $(z, c)$ and the generated image $G(z, c)$. One caveat however, we actually want to minimise such a loss with respect to both generator and discriminator; so this loss isn’t adversarial per se, but more like a ‘cooperative’ loss that both networks have to minimise. Here is some pseudo-code, assuming $p(c)$ here is a Categorical distribution:

def train_on_batch(x):

  opt_d.zero_grad()
  opt_g.zero_grad()
  # or opt_all.zero_grad()

  c = sample_c(x.size(0))
  z = sample_z(x.size(0))

  # generator loss
  x_fake = G(z, c)
  ...
  ...
  # discriminator loss
  d_realfake, _, _ = D(x_fake)
  ...
  ...

  # `opt_all` wraps both the generator
  # and discriminator parameters, since we
  # we want to minimise the infogan loss
  # wrt to both networks.
  opt_all.zero_grad()
  # imagine D has three output layers: one
  # to determine real/fake, and the other two
  # for our noise variables.
  _, d_out_c, d_out_z = D(x_fake)
  z_loss = torch.mean((d_out_z-z)**2)
  # you could also use mean squared error here,
  # but i'm using the x-entropy to be more correct
  # about it, since p(c) is a multinomial distribution.
  # this means that `d_out_c` has been transformed with
  # the softmax nonlinearity.
  c_loss = categorical_crossentropy(d_out_c, c).mean()
  infogan_loss = z_loss + c_loss
  infogan_loss.backward()
  opt_all.step()
  

Resources

[a]: PyTorch generative model collections. Figure 3’s JS-GAN was trained with this repo, spectral normalisation was added to the code by wrapping each linear and conv definition in GAN.py.
[b]: Christian Cosgrove’s minimal HingeGAN + spectral norm implementation on CIFAR10. It was very useful to me when I started getting serious about GANs a few years ago.

Conclusion

If you think I’ve missed something or you’ve spotted an error, please let me know in the comments section at the bottom of this page or reach out to me on Twitter!

References

[1]: https://colinraffel.com/blog/gans-and-divergence-minimization.html
[2]: Arjovsky, M., Chintala, S., & Bottou, L. (2017, July). Wasserstein generative adversarial networks. In International conference on Machine Learning (pp. 214-223). PMLR.
[3]: Gulrajani, I., Ahmed, F., Arjovsky, M., Dumoulin, V., & Courville, A. (2017). Improved training of Wasserstein GANs. arXiv preprint arXiv:1704.00028.
[4]: https://github.com/soumith/ganhacks
[5]: Barratt, S., & Sharma, R. (2018). A note on the inception score. arXiv preprint arXiv:1801.01973.
[6]: Farnia, F., & Ozdaglar, A. (2020, November). Do GANs always have Nash equilibria?. In International Conference on Machine Learning (pp. 3029-3039). PMLR.
[7]: https://parameterfree.com/2020/12/06/neural-network-maybe-evolved-to-make-adam-the-best-optimizer/
[8]: Miyato, T., Kataoka, T., Koyama, M., & Yoshida, Y. (2018). Spectral normalization for generative adversarial networks. arXiv preprint arXiv:1802.05957.
[9]: Chen, X., Duan, Y., Houthooft, R., Schulman, J., Sutskever, I., & Abbeel, P. (2016, December). InfoGAN: Interpretable representation learning by information maximizing generative adversarial nets. In Proceedings of the 30th International Conference on Neural Information Processing Systems (pp. 2180-2188).

Appendix

The relationship between PyTorch and TF Inception scores

import pandas as pd
import matplotlib.pyplot as plt
%matplotlib inline

df = pd.read_csv("./incep_tf_vs_pt.csv", names=["pytorch", "tensorflow"])
df

	pytorch	tensorflow
0	6.78	7.68
1	6.82	7.63
2	6.75	7.68
3	6.73	7.40
4	6.63	7.48
5	6.85	7.81
6	6.29	7.24
7	6.92	7.66
8	7.03	7.63
9	6.98	7.62
10	7.18	7.81
11	6.80	7.60
12	6.79	7.63
13	7.01	7.76
14	7.03	7.84
15	6.25	7.02
16	6.82	7.62

from sklearn import linear_model
import numpy as np
lm = linear_model.LinearRegression()

X = df['pytorch'].values.reshape(-1, 1)
y = df['tensorflow'].values.reshape(-1, 1)

lm.fit(X, y)

plt.scatter(X, y)
plt.ylabel('tensorflow IS')
plt.xlabel('pytorch IS')

xs = np.linspace(6, 7.5, num=100).reshape(-1, 1)
ys = lm.predict(xs)
plt.plot(xs, ys)
plt.grid()
lm_str = "y = %.2fx + %.2f" % (lm.coef_[0][0], lm.intercept_[0])
plt.title(lm_str)

name	what is it
reward oracle	\(f(x)\): real world reward model, extremely expensive to compute. This may also refer to a simulation environment's reward model.
"proxy" oracle	- \(\ft\): a regression model trained on the training set. While it is a discriminative model, it can be "hacked" to act as a generative model. In this article, I prefer to use generative modelling terminology, in which case \(\pt(y \vert x)\) is used instead. However, in this post I prefer to use the conditional density \(\pt(x \vert y)\) will usually be referred to, instead of the other terms.
	- \(\fvalid\) : typically not defined in literature, but this is specifically a proxy oracle intended for model selection and hyperparameter tuning. It is trained on the combined training and validation set. Here we will call it the validation proxy.
	- \(\ftest\) : proxy oracle trained on train + valid + test set (all of the data). This typically exists in "academic MBO" where the ground truth is also too impractical to compute at test time. Here we will call it the test proxy.

christopher beckham

reflecting on past work – offline model-based optimisation

reflecting on past work -- offline model-based optimisation

Table of Contents

1. intro - what is model-based optimisation?

1.1. formalising offline mbo

1.2. reward-based extrapolation

1.3. ‼️ why evaluation is difficult (and misunderstood)

1.4. training, validation, and testing

2. last year's work

2.1. ranking validation metrics

3. 🪵🔥 reflection, and future work

3.1. can we combine online and offline mbo?

3.2. links, and open source

4. References

Footnotes:

paper thoughts – questionable research practices (QRPs) in machine learning

paper thoughts: questionable research practices (QRPs) in machine learning

1. Tuning hyperparameters further after test

2. Over/underclaiming

EDM diffusion models - a Jupyter implementation, and how they are implemented in practice

EDM diffusion models - a Jupyter implementation, and how they are implemented in practice

1. Introduction

2. Schedulers

2.1. Ancestral sampler

3. Conclusion

A deep dive into conditional variational autoencoders

A deep dive into conditional variational autoencoders

Table of Contents

1. Introduction

1.1. Contributions

2. The generative and inference process

2.1. When z and y are independent

2.2. When z and y are dependent

2.3. The role of the beta term

2.4. A mutual information perspective for the KL term

2.5. A mutual information perspective between Z and Y

2.5.1. Practical considerations

3. Experiments

3.1. When z and y are independent

3.1.1. Controllable generation

3.2. When z and y are dependent

4. Discussion

4.1. cVAEs and conditional Gaussian diffusion models

5. Conclusion

6. Appendix

6.1. Derivation of Esmaeli's joint KL

6.2. Conditional case

6.2.1. z and y are independent

6.2.2. z and y are dependent

6.3. Per-example KL versus marginal KL

6.4. Z-Y independent cVAE on SVHN

7. References

Footnotes:

Techniques for label conditioning in Gaussian denoising diffusion models

Techniques for label conditioning in Gaussian denoising diffusion models

Updates

Table of Contents

1. Preliminaries

2. Classifier-based guidance

3. Classifier-free guidance

3.1. Sources of confusion

3.2. Benefits

4. Conditional ELBO

5. Conclusion

6. References

Learning the conditional prior over classes for image diffusion

Learning the conditional prior over classes for image diffusion

Table of contents

Unconditional diffusion

The conditional variant

Reproduction

Examining different values of \(\tau\)

Reverse diffusion

Visualising the learned embeddings

References

The obsession with SOTA needs to stop

The obsession with SOTA needs to stop

Extra rambles

Conclusion