Naive Bayes is all you need?

Naive Bayes is all you need ? 苏剑林 · 2023-06-08 · Machine Learning language model attention LLM bayesian

Sorry for the clickbait-ish title. After writing NBCE: Extending the Context Length of LLMs Using Naive Bayes], I began to feel that Naive Bayes and the attention mechanism share a lot of common features. After working through the derivation a bit further, I found that attention can in fact be viewed as a generalized, parameterized version of Naive Bayes. If that's the case, doesn't "Attention is All You Need]" also imply "Naive Bayes is all you need"? That's where the title of this post comes from.

Below I'll walk through my train of thought and analyze how the attention mechanism can be understood from the perspective of Naive Bayes.

Naive Bayes

This post mainly considers language models, which model $p(x_t|x_1,\cdots,x_{t-1})$. By Bayes' rule, we have

\begin{equation}p(x_t|x_1,\cdots,x_{t-1}) = \frac{p(x_1,\cdots,x_{t-1}|x_t)p(x_t)}{p(x_1,\cdots,x_{t-1})}\propto p(x_1,\cdots,x_{t-1}|x_t)p(x_t)\end{equation}more

Using the independence assumption $p(x_1,\cdots,x_{t-1}|x_t) = \prod\limits_{j=1}^{t-1} p(x_j|x_t)$, we get

\begin{equation}p(x_t|x_1,\cdots,x_{t-1}) \propto \prod_{j=1}^{t-1} p(x_j|x_t)p(x_t)\end{equation}

Applying Bayes' rule once again $p(x_j|x_t)=\frac{p(x_t|x_j)p(x_j)}{p(x_t)}\propto\frac{p(x_t|x_j)}{p(x_t)}$, we obtain

\begin{equation}p(x_t|x_1,\cdots,x_{t-1}) \propto \frac{1}{[p(x_t)]^{t-2}}\prod_{j=1}^{t-1} p(x_t|x_j)\end{equation}

Taking the logarithm of both sides gives

\begin{equation}\log p(x_t|x_1,\cdots,x_{t-1}) = \sum_{j=1}^{t-1}\log p(x_t|x_j) - (t - 2) \log p(x_t) + \text{const}\end{equation}

A generalized result

We carried out the same derivation in NBCE: Extending the Context Length of LLMs Using Naive Bayes], and as in that post, we generalize the above expression to:

\begin{equation}\log p(x_t|x_1,\cdots,x_{t-1}) = (1 + \beta)\mathcal{P}[\log p(x_t|x_j)] - \beta \log p(x_t) + \text{const}\end{equation}

Here $\beta$ is treated as a hyperparameter to be tuned, and $\mathcal{P}$ is some form of pooling. Let's focus on the case where $\beta=0$, with pooling implemented as a weighted average, i.e.

\begin{equation}\log p(x_t|x_1,\cdots,x_{t-1}) = \sum_j a_{t,j} \log p(x_t|x_j) + \text{const}\label{eq:nb-core}\end{equation}

Here $a_{t,j}$ is a function of $x_{t-1}$ and $x_j$.

Some readers might ask: can this generalized expression still be called Naive Bayes? I think it can be viewed as a generalized form of Naive Bayes, since ordinary Naive Bayes can be seen as an equal-weighted average over the various $\log p(x_t|x_j)$ terms, whereas here we've replaced that with a more general weighted average. That said, by choosing $a_{t,j}$ to be a function of $x_{t-1}$ and $x_j$, we highlight the role of $x_{t-1}$, which mitigates one of the drawbacks of Naive Bayes — its disregard for order. So more precisely, equation $\eqref{eq:nb-core}$ is a combination of a 2-gram language model and Naive Bayes.

Attention emerges

Next, if we further parameterize $\log p(x_t|x_j)$, we arrive at something resembling attention. It's not hard to see that $p(x_t|x_j)$ is essentially the old Skip-Gram model from Word2Vec, whose conventional modeling approach is "embedding + inner product + softmax", i.e.

\begin{equation}p(x_t|x_j) = \frac{e^{v(x_j)\cdot w(x_t)}}{Z(x_j)},\quad Z(x_j) = \sum_{x_t\in Vocab}e^{v(x_j)\cdot w(x_t)}\end{equation}

So we simply take

\begin{equation}\log p(x_t|x_j) = v(x_j)\cdot w(x_t) + \text{const}\end{equation}

Substituting this into equation $\eqref{eq:nb-core}$, we get

\begin{equation}\log p(x_t|x_1,\cdots,x_{t-1}) = \left(\sum_j a_{t,j} v(x_j)\right)\cdot w(x_t) + \text{const}\label{eq:nb-core-2}\end{equation}

If we pull out the expression in the parentheses and treat it as a general-purpose feature-fusion operation, it's actually just ordinary attention. In other words, using a single layer of attention to build a language model is, in effect, generalized Naive Bayes.

Of course, we haven't yet pinned down $a_{t,j}$. In the previous section we said that $a_{t,j}$ is a function of $x_{t-1}$ and $x_j$, and it must also be normalized (as a weighted average), so a natural choice — just like in Skip-Gram — is "embedding + inner product + softmax":

\begin{equation}a_{t,j} = \frac{e^{q(x_{t-1})\cdot k(x_j)}}{Z_t},\quad Z_t = \sum_{j=1}^{t-1} e^{q(x_{t-1})\cdot k(x_j)}\end{equation}

Substituting this into equation $\eqref{eq:nb-core-2}$ gives us the now-standard dot-product attention. Of course, this isn't the only way to do it — there's also additive attention and so on — but the main reason dot-product attention is preferred is that it can be parallelized while being relatively memory-efficient.

Stacking and residuals

No matter how it's parameterized, a single layer of Naive Bayes always has limited capacity, so we need to further increase the model's complexity. From a neural network perspective, the main way to increase model complexity is to add depth, i.e. to stack layers on top of each other. So how should we understand this stacking from the perspective of probability distributions? The answer is: latent variable models.

A latent variable model, roughly speaking, introduces a latent variable $z_1,z_2,\cdots,z_{t-1}$ such that

\begin{equation}p(x_t|x_1,\cdots,x_{t-1}) = \int p(x_t|z_1,\cdots,z_{t-1})p(z_1,\cdots,z_{t-1}|x_1,\cdots,x_{t-1})dz_1 \cdots dz_{t-1}\end{equation}

In plain terms, this fits a more complex distribution by superposing simpler ones, in the same spirit as a GMM (Gaussian Mixture Model). Following the earlier discussion, we again model $p(x_t|z_1,\cdots,z_{t-1})$ with Naive Bayes, which, at the feature level, corresponds to a single layer of attention. As for $p(z_1,\cdots,z_{t-1}|x_1,\cdots,x_{t-1})$, following the characteristic structure of autoregressive models, we decompose it as

\begin{equation}p(z_1,\cdots,z_{t-1}|x_1,\cdots,x_{t-1}) = \prod_{k=1}^{t-1} p(z_k|x_1,\cdots,x_k)\end{equation}

This way, each $p(z_k|x_1,\cdots,x_k)$ takes exactly the same form as $p(x_t|z_1,\cdots,z_{t-1})$, and so it too can be modeled with Naive Bayes. For simplicity, we can define $z_k$ as a continuous variable, and define $p(z_k|x_1,\cdots,x_k)$ as a Dirac distribution], so the integral can be computed directly, and the result turns out to be a stack of two attention layers.

Finally, there's one more key component in the Transformer: residual connections. In effect, these generalize equation $\eqref{eq:nb-core}$ to

\begin{equation}\log p(x_t|x_1,\cdots,x_{t-1}) = \log p(x_t|x_{t-1}) + \sum_j a_{t,j} \log p(x_t|x_j) + \text{const}\end{equation}

which can be understood as a form of pooling that emphasizes the role of 2-grams — a kind of prior. As for the remaining components, such as the FeedForward layers and LayerNorm layers, these don't involve interaction between tokens, and can be understood as a more elaborately parameterized form of Naive Bayes.

Admittedly, this kind of broad-brush explanation may feel a bit forced. But my original intent wasn't to give a precise account of Transformers or attention — rather, I was hoping that the Naive Bayes perspective might yield some new ideas about length generalization. Unfortunately, I haven't yet arrived at the results I was hoping for. Still, even though this may look like a bit of self-indulgent theorizing, I do believe the Naive Bayes and latent-variable-model perspective outlined above still has room to be developed further — for instance, it seems plausible that we could use the Naive Bayes viewpoint to explain why in-context learning works in attention-based language models.

Summary

This post has laid out the connection between Naive Bayes and the attention mechanism, showing that attention can be viewed as a form of generalized Naive Bayes. From this perspective, we can also gain further insight into aspects of attention such as layer stacking and residual connections.