tl;dr: diffusion language models often maximize the ELBO, a lower bound on log-likelihood, rather than maximizing log-likelihood directly.

$$\operatorname{ELBO}(\theta;x)\leq \log p_\theta(x)$$

This writeup is a walkthrough of what it is, and mostly why does it solve an intractable problem in minimizing the loss of diffusion language models.

Disclaimer: this is not a very rigorous mathematical article (some details are intentionally left out for the sake of understanding) but hopefully not misleading.

What is a probabilistic model?

This primer might not be related to what ELBO is yet, but bear with me. Say you want to predict what college major a person is, a probabilistic model can be a probability distribution over the college majors. This would take the form of something like $p(major)$ e.g. $compsci = 0.3$, $\quad law = 0.22$, $\quad humanities = 0.12$ etc.

But this current model $p(major)$ is not very useful because it models everybody the same way. As in, it does not care what a person is like. A more useful model would be the one that also considers what the student is like, that is, $p(major | student)$. This way, the distribution of majors would change depending on the person. This is one definition of a probabilisitic model: a model that assigns a distribution over possible outcomes/categories, potentially conditioned also on an input (in this case, what the student is like)

Now what does that have to do with language models?

What language models do is: Assign probabilities to a given sequence of text $x$. Coming from our previous definition of a probabilisitic model, the probability tells us if $x$ is a plausible sentence or not. With this being said, sentences that make sense should be assigned higher than gibberish. For example,

Once when I was six years old I saw a magnificent picture in a book, called True Stories from Nature, about the primeval forest.

This small sentence (From The Little Prince, Antoine de Saint-Exupéry) should be assigned a higher probability, let’s say $0.3$ than

pqwo ususodoi ehrevvevvevvvv llplwokkeoke

(complete nonsense), which should be something like $0.000000002$. So essentially: a language model is a probabilistic model that assigns a probability to a given sentence, which tells you how likely this sentence would occur.

In Practice

Modern (at least autoregressive ones like ChatGPT) language models assign a probability distribution over a set of vocabulary tokens $\mathcal{V}$, given it’s previous tokens.

$$P(x_t = v \mid x_{<t})$$

Essentially, it gives you what token/word is likely to be next. So for example when we forward pass $x_{<t} = ``\text{the cat sat on the}"$ to the network, we have a probability distribution over $\mathcal{V}$. Likely words are assigned higher probabilities:

$$P(x_t = \text{``mat"} \mid x_{<t}) = 0.9$$

is higher compared to maybe $P(x_t = \text{``mitochondria"} \mid x_{<t}) = 0.002$.

At this point, you might be thinking, this is different from what we agreed with? because we agreed that language models assign a probability over whole texts, not predicting what token is next? We are actually still doing the same thing, it’s just that probability for whole texts/sentences is now factorized over next token probabilities.

$$P(x) = \underbrace{P(x_1, x_2, \dots, x_n)}_{\text{what we agreed with}} = \prod_{t=1}^{n} \underbrace{P(x_t \mid x_1, \dots, x_{t-1})}_{\text{what LMs do in practice}}$$

A central challenge in probabilistic models

There are 2 things you can do with probabilistic models

1. Given a data instance $x$, we can get how the probability of this data is to occur. How likely is $x$, essentially.

2. Generating new $x$ (sampling) from this model. In the case of language models, we are generating new sentences from the model.

In the case of our language models, $P(x)$ tells us exactly that. For sampling, we run $P(x_t \mid x_1, \dots, x_{t-1})$ and pick the most likely¹ token over and over again (this is how Autoregressive models work too).

Luckily, doing both is easy for language models. As has been discussed before, we can estimate a sentence $x$ by just forwarding the sentence through the network, or running the forward over and over again to sample/generate a new sentence from the model. I think, or rather i believe not all probabilistic models are easy in both².

Many Ways to Evaluate P($x$)

Circling back to our college major example with the distribution of $P_{major}$, evaluating the probability of college major, let’s say “medicine” can look like

$$P_{major}(\text{medicine}) = \frac{\#(\text{students in medicine})}{\text{total students}}$$

Which gives us the probability of, say picking a random person and them being in medicine. But of course, like many other things in life, language models are not that simple. How we estimate the probability of a text via a language model is to .forward() the neural network, which is a lot more complicated than just a lookup table of frequencies³. I think the reason why we moved on from only using frequencies is because we’re not gonna be able to assign an appropriate probability value to a piece of text just because it did not show up verbatim in the training data⁴ eventhough the text makes sense to a human reader. We want a model that is rich, in the sense that it can generalize to sentences it hasn’t seen during training.

Latent Variables

We mentioned earlier that it’s better to predict the college major of a person if we also include the properties of the student. Let’s say $\text{student} = \{\text{courses taken},\; \text{desk location},\; \text{supervisor}\}$. We’ll have a better model if we do

$$P(major | \{\text{courses taken},\; \text{desk location},\; \text{supervisor}\})$$

Now suppose we don’t observe the student’s attributes directly. These unobserved attributes are latent variables, and we have to consider their possible values. We can apply the same idea to language models. We call the $\text{student}$ information a latent variable that will make our college major estimation better. We can apply the same concept to language models. Let’s say we have the same sentence as before (shout out Antoine once again):

Once when I was six years old I saw a magnificent picture in a book, called True Stories from Nature, about the primeval forest.

This time, we have a latent variable for this sentence, we call

$$\text{sentence\_attribute} = \{\text{genre},\; \text{year},\; \text{audience},\; \text{country}\}$$

With $\text{sentence\_attribute}$ defined like this, we can have multiple instances of this set of attribute that are possible explanations of the text. One instance could be

$$\text{sentence\_attribute}_1 = \{\text{fantasy},\; \text{1990-2000},\; \text{children},\; \text{france}\}$$

But $\text{sentence\_attribute}_1$ is not the only plausible set of attributes that can explain the original text. Some can argue that the following $\text{sentence\_attribute}_2$ can also be a valid set of attribute that explains the text:

$$\text{sentence\_attribute}_2 = \{\text{classic},\; \text{1990-2000},\; \text{adults},\; \text{france}\}$$

With this example, we can see that, there are many possible latent vairables that can explain our data with variying levels of likeliness⁵

There are so many possible latent variables that can explain $x$, we have to consider all of them, even the unlikely ones like

$$\text{sentence\_attribute}_3 = \{\text{horror},\; \text{1990-2000},\; \text{babies},\; \text{france}\}$$

to get the true $P(x)$ by summing over all of them. Because however unlikely, sometimes they are still not technically impossible, with a really small probability value. In this example, $\text{sentence\_attribute}$ is constrained into only 4 attributes (genre, year, audience, country) and probably only several labels per category. But in practice, we can deal with even larger dimensions with even more labels per categories or even deal with non-discrete latent variables! When that happens, summing over every possible latent variables

$$P^\theta(x) = \sum_{z \in \mathcal{Z}} p^\theta(x \mid z)\,p(z)$$

or in the continuous case,

$$P^\theta(x) = \int p^\theta(x \mid z)\, p(z)\, dz$$

which are very impractical. In the case of diffusion language models, intermediate sequences act as latent variables and have the same length as the clean input. Under uniform token corruption, each position has $|\mathcal{V}|$ possible token values. That’s $|\mathcal{V}|^{\text{sequence\_length}}$ possible configurations for just one intermediate sequence! Computing the exact likelihood generally requires summing over entire trajectories of these intermediate states. This is definitely a problem. Because when we train diffusion language models, we want to compute

$$-\log(P^\theta(x))$$

Therefore, we have to find a workaround for this.

Workaround: Variational Posterior

Let us choose a tractable variational posterior $q(\text{sentence\_attribute} \mid x)$, which is the distribution of latent variables given the text. For the sake of generality, let’s say the latent variable $\text{(sentence\_attribute)}$ is just some $z$. Note that $\text{sentence\_attribute}$ is just an illustration. In practice, most latent variables are less obvious (it can be a vector of “meaningless”⁶ numbers, etc.) Ideally, $q(z \mid x)$ should assign $z$‘s that are likelier explanations of $x$‘s higher probabilities, like $\text{sentence\_attribute}_1$ for example.

Derivations and Why It Is Now “Computable”

The intractable problem we want to compute is

$$-\log(P^\theta(x)) = -\log\left(\sum_{z \in \mathcal{Z}} p^\theta(x \mid z)\,p(z)\right)$$

which we cannot do reasonably, since $\mathcal{Z}$ is too large. We introduce $q(z \mid x)$ into the expression

$$\begin{aligned} -\log(P^\theta(x)) &= -\log\left(\sum_{z \in \mathcal{Z}}\frac{q(z \mid x)}{q(z \mid x)} p^\theta(x \mid z)\,p(z)\right) \\ &= -\log\left(\sum_{z \in \mathcal{Z}}q(z \mid x) \frac{p^\theta(x \mid z)\,p(z)}{q(z \mid x)}\right) \\ &= - \log \left(\mathbb{E}_{q(z \mid x)}\left[\frac{p^\theta(x \mid z)\,p(z)}{q(z \mid x)}\right]\right) \end{aligned}$$

In case you (mostly I myself, lol) forgot, how it became an expectation is from

$$\mathbb{E}_{p(z)}[f(z)] = \int p(z)\, f(z)\, dz$$

At this point, since we have an expectation, we can estimate it using the Monte Carlo method⁷ therefore it is now tractable, given that sampling from $q(z \mid x)$ (where you got your latent variables from) is also affordable. However, we still have a problem.

Biasing an Unbiased Estimator

At the very core, we want to compute $-\log p^\theta(x)$. But we’ve established that evaluating $p^\theta(x)$ exactly is intractable, but we have rewritten it as an expectation:

$$p^\theta(x) = \mathbb E_{q(z\mid x)} \left[\frac{p^\theta(x\mid z)p(z)}{q(z\mid x)}\right].$$

We can approximate this expectation using a Monte Carlo sample average. Draw $K$ independent samples $z_k\sim q(z\mid x)$, then compute

$$\widehat p_K^\theta(x) = \frac1K\sum_{k=1}^{K} \frac{p^\theta(x\mid z_k)p(z_k)}{q(z_k\mid x)}.$$

Which makes it tractable. Now, Monte Carlo sample average here is an unbiased estimator. An unbiased estimator is an estimator for that if you repeated the sampling many times, your estimates would average out to the true value⁸. Now the problem is, there is a $\log$ term outside of that estimator

$$-\log(\widehat p_K^\theta(x))$$

Had we computed the exact $p^\theta(x)$, taking its negative log would give us the exact objective. But we instead have a random estimate $\widehat p_K^\theta(x)$. Applying the negative log generally introduces bias: on average, we overestimate the true negative log-likelihood, even though individual estimates may fall above or below it.

Lower Bounding

We mitigate this problem by opting to optimize something else that is not exactly the negative log likelihood, but can be estimated without bias. By Jensen’s Inequality, since $\log$ is concave,

$$\begin{aligned} \log \mathbb{E}[\cdot] &\geq \mathbb{E}[\log(\cdot)] \\ \\ \log \left(\mathbb{E}_{q(z \mid x)}\left[\frac{p^\theta(x \mid z)\,p(z)}{q(z \mid x)}\right]\right) &\geq \mathbb{E}_{q(z \mid x)}\left[\log\left( \frac{p^\theta(x \mid z)\,p(z)}{q(z \mid x)}\right)\right] \\ \\ -\log \left(\mathbb{E}_{q(z \mid x)}\left[\frac{p^\theta(x \mid z)\,p(z)}{q(z \mid x)}\right]\right) &\leq - \mathbb{E}_{q(z \mid x)}\left[\log\left( \frac{p^\theta(x \mid z)\,p(z)}{q(z \mid x)}\right)\right] \\ \end{aligned}$$

That final right-hand side is the negative ELBO (NELBO)! (ELBO itself is NELBO without the negative, the inequality right before it) We minimize it using yet another Monte Carlo estimates. When we do that, we’d get a true negative log likelihood that’s always at least as low as the negative ELBO.

Footnotes

1. This is called greedy decoding (although arguably it's no longer sampling bc it's deterministic), i believe there is a whole research direction in language model sampling ↩

2. i am pretty sure ↩

3. although early language models do, look up n-gram language models ↩

4. although there are ways to do so, but of course, they are limited ↩

5. For the sake of this example, we'll assume The Little Prince is more likely to be a children's book rather than for adults, so we can say $\text{sentence\_attribute}_1$ is more likely to explain $x$ than, say $\text{sentence\_attribute}_2$ ↩

6. not readily interpretable at face value ↩

7. which is basically just sampling it a buncha times and averaging instead of summing over all possible latents ↩

8. if this sentence is hard to read for you, i suggest learning Statistical Inference™ ↩