NBCE: Extending LLM Context Length with Naive Bayes

NBCE：使用朴素贝叶斯扩展LLM的Context处理长度 苏剑林 · 2023-05-23 · Machine Learning language model extrapolation LLM bayesian

Still playing with Naive Bayes in the age of LLMs? This is probably the first thought many readers have upon seeing the title. And yes, that's exactly what this is about — when the ancient Naive Bayes meets cutting-edge LLMs, something surprising happens: we can directly extend the context-processing length of existing LLMs without fine-tuning the model, without depending on the model architecture, with linear efficiency, and with results that look pretty good. This is the NBCE (Naive Bayes-based Context Extension) method proposed in this article.

Feeling for stones to cross the river #

Suppose $T$ is the token sequence to be generated, and $S_1,S_2,\cdots,S_n$ are several relatively independent Context sets that have been given (say, $n$ different paragraphs, at least not the kind where a single sentence gets split into two fragments). Suppose their combined length already exceeds the training length, but any single $S_k$ plus $T$ still fits within the training length. We need to generate $T$ conditioned on $S_1,S_2,\cdots,S_n$, i.e., estimate $p(T|S_1, S_2,\cdots,S_n)$.

more

Simply put, Naive Bayes is "Bayes' rule + independence assumption." By Bayes' rule:

\begin{equation}p(T|S_1, S_2,\cdots,S_n) \propto p(S_1, S_2,\cdots,S_n|T)p(T)\end{equation}

Here $\propto$ omits a constant factor unrelated to $T$. By the (conditional) independence assumption:

\begin{equation}p(S_1, S_2,\cdots,S_n|T) = \prod_{k=1}^n p(S_k|T)\end{equation}

so we have

\begin{equation}p(T|S_1, S_2,\cdots,S_n) \propto p(T)\prod_{k=1}^n p(S_k|T)\end{equation}

Applying Bayes' rule again, $p(S_k|T) \propto \frac{p(T|S_k)}{p(T)}$, gives

\begin{equation}p(T|S_1, S_2,\cdots,S_n) \propto \frac{1}{p^{n-1}(T)}\prod_{k=1}^n p(T|S_k)\end{equation}

or

\begin{equation}\log p(T|S_1, S_2,\cdots,S_n) = \color{red}{\sum_{k=1}^n \log p(T|S_k)} - \color{green}{(n-1)\log p(T)} + \color{skyblue}{\text{const}}\label{eq:nbce-1}\end{equation}

Here, both $\color{red}{p(T|S_k)}$ and $\color{green}{p(T)}$ can be computed directly with an existing LLM — in fact, any language model will do, since this is independent of architecture and requires no fine-tuning on long text. Among these, $\color{red}{p(T|S_k)}$ is the probability predicted from a single Context, while $\color{green}{p(T)}$ is the probability with no Context (or an empty Context). Multiple Contexts can be placed in the same batch for parallel computation, and the computational cost grows linearly with the number of Contexts.

Peeling back the layers #

Of course, Naive Bayes relies on the independence assumption, which limits its practical effectiveness. To make it "surpass its origins," let's further refine Eq. $\eqref{eq:nbce-1}$ — "peel back the layers" and "keep the essence" — in order to achieve better results.

First, let us denote $\log p(T|S) = [\log p(T|S_1),\cdots,\log p(T|S_n)]$, and

\begin{equation}\overline{\log p(T|S)} = \frac{1}{n}\sum_{k=1}^n \log p(T|S_k)\end{equation}

and set $\beta = n - 1$, so that Eq. $\eqref{eq:nbce-1}$ can be rewritten as

\begin{equation}\log p(T|S_1, S_2,\cdots,S_n) = \color{red}{(\beta + 1)\overline{\log p(T|S)}} - \color{green}{\beta\log p(T)} + \color{skyblue}{\text{const}}\label{eq:nbce-2}\end{equation}

Once rewritten in this form, two questions naturally arise: 1. If we treat $\beta$ as a hyperparameter to tune, might we get better results? 2. $\overline{\log p(T|S)}$ is just Average Pooling over $\log p(T|S)$ — would switching to a different pooling method (denoted $\mathcal{P}$) work better? That is,

\begin{equation}\log p(T|S_1, S_2,\cdots,S_n) = \color{red}{(\beta + 1)\mathcal{P}[\log p(T|S)]} - \color{green}{\beta\log p(T)} + \color{skyblue}{\text{const}}\label{eq:nbce-3}\end{equation}

So I experimented with these two questions on a 7B model, and the preliminary conclusion is: in reading-comprehension scenarios, Max Pooling combined with $\beta=0.25$, using Greedy Search, performs reasonably well overall — however, results obtained via Random Sample are essentially unreadable.

Final scheme #

Why does Greedy Search do well while Random Sample does poorly? We know that Random Sample means "sampling according to the distribution," and its poor performance implies that the result of Max Pooling isn't a well-formed distribution. Greedy Search, on the other hand, only cares about the token with the highest probability, not about whether the distribution itself is sensible — and the fact that it performs well tells us that the token with the highest probability is usually correct. The higher the probability, the lower the uncertainty. So, to improve the results of Random Sample, we change the pooling method to directly output the distribution with the lowest uncertainty:

\begin{equation}\begin{aligned} &\mathcal{P}[\log p(T|S)] = \log p(T|S_{\color{red}{k}}) \\[5pt] &\color{red}{k} = \mathop{\text{argmin}} \big\{H_1,H_2,\cdots,H_n\big\} \\[5pt] &H_i = -\sum_T p(T|S_i)\log p(T|S_i) \end{aligned}\end{equation}

Substituting this into Eq. $\eqref{eq:nbce-3}$ gives the final form of NBCE (Naive Bayes-based Context Extension). It's worth pointing out that although our starting point was Naive Bayes, the generalized Eq. $\eqref{eq:nbce-3}$ has already gone beyond the conventional scope of Naive Bayes, while still retaining its interpretability. It's easy to see that the form of Eq. $\eqref{eq:nbce-3}$ is quite intuitive: 1. The predictions from different Contexts are aggregated (or "voted on") via method $\mathcal{P}$ (with weight $\beta+1$), and the prediction with no Context is subtracted (with weight $\beta$); 2. The reason for subtracting the no-Context prediction is to make the model lean more toward incorporating the Context rather than answering purely from its own stored knowledge (note: a paper appearing on Arxiv three days later, Trusting Your Evidence: Hallucinate Less with Context-aware Decoding, proposed the same trick to reduce hallucination); 3. Different scenarios can use different values of $\beta$ — for example, tasks that require combining Context for reading comprehension might use a larger $\beta$, while tasks leaning toward free-form creation might use a smaller $\beta$. In my view, all of $\beta\geq -1$ are reasonable choices.

Reference implementation #

Below is a reference implementation of NBCE:

Github: https://github.com/bojone/NBCE

As the demo code shows, implementing NBCE is quite simple — you only need to modify how the logits are constructed inside the decoding function, and this doesn't conflict with the choice of decoding algorithm.

Illustration of Naive Bayes-based Context Extension (NBCE)

The provided demo contains 12 different Context segments, totaling over 9,000 characters, which are fed into the model along with 8 questions all at once (the model's training length is 2048, its parameter count is 7B, and it can be downloaded from OpenBuddy). The model is able to correctly answer all 8 questions one by one based on the given Contexts. Notably, the Contexts, questions, and answers combined add up to over 10,000 characters! In addition, some friends have tried simple applications like résumé matching and essay grading, with decent results as well — I'd strongly encourage everyone to try it out themselves.

Related work #

There has already been a fair amount of work on extending LLM context length, but most of it shortens long-context samples via retrieval or summarization, such as Unlimiformer. Since these approaches don't directly process the long Context, they typically can't perform fine-grained reading comprehension, and they usually need to be baked in during training rather than being plug-and-play with an already-trained LLM afterward.

Before NBCE, the approach that could extend context length without fine-tuning was Parallel Context Window (hereafter PCW), from the papers Parallel Context Windows for Large Language Models and Structured Prompting: Scaling In-Context Learning to 1,000 Examples. These two papers are contemporaneous works by different authors, but the proposed methods differ only slightly, so we'll refer to both as PCW here.

PCW is designed for Self-Attention models, and mainly modifies the Position Encoding and the Attention Mask, as shown in the figure below:

Parallel Context Window

First, determine the maximum length $L$ of a Context (6, in the figure). Then, the last position in each Context is encoded as $L-1$, the second-to-last as $L-2$, and so on — we call this scheme "right-aligned" (or "left-indented"). On the other hand, for the Task Tokens part (Prompt + generated content), the position encoding is $L,L+1,L+2,\cdots$. Each Context is encoded independently, so the corresponding Attention Mask is a block-diagonal matrix — and since this is an LM, it's a block-diagonal lower-triangular matrix. As for the Task Tokens part, it needs to attend to all the Contexts, so it attends to all Contexts (as well as to itself). This way, if you take any single Context and pair it with the Task Tokens, the resulting Attention pattern matches that of the original LM.

Some readers may have already noticed that NBCE shares quite similar properties with PCW — for instance, both treat the Contexts as unordered and equally weighted. In fact, if we apply NBCE to a single-layer, single-head attention model, the result is roughly equivalent to PCW. To show this, let's write out the single-layer, single-head attention language model as

\begin{equation}p(x_t|x_{< t}) = softmax\left(\sum_{i=1}^t a_{t,i}v_i W\right)\end{equation}

so we roughly have $\log p(x_t|x_{< t}) \sim \sum\limits_{i=1}^t a_{t,i}v_i W$. Substituting this into Eq. $\eqref{eq:nbce-2}$ and taking $\beta=0$, we get

\begin{equation}\log p(T|S_1, S_2,\cdots,S_n) \sim \frac{1}{n}\sum_{k=1}^n\left(\sum_{i\in S_k} a_{T,i}v_i\right) W = \left(\sum_{i\in S_1\oplus\cdots\oplus S_n} \frac{a_{T,i}}{n}v_i\right) W \end{equation}

Here we're assuming $T$ is a single token, but this loses no generality; $\oplus$ denotes concatenation. In the expression above, $S_k\oplus T$ is treated as a continuous segment for inference (as NBCE assumes), so their position encodings are adjacent, while $a_{T,i}/n$ forms an overall Attention (again summing to 1) between $T$ and all the $S_i$. These properties are indeed consistent with PCW — PCW simply integrates this more elegantly into every layer via the Attention Mask. Thus, PCW is roughly the Average-Pooling version of NBCE, and indeed, in our experiments, it shares similar shortcomings with the Average-Pooling version of NBCE: as the number of Contexts increases, the output starts becoming less accurate — typically staying on-topic, but wrong as an actual answer to the question.

Further thoughts #

A major drawback of NBCE is its lack of order-awareness — it cannot recognize the input order of the Contexts, which may hurt performance in scenarios like story continuation. To mitigate this, one could consider adding a prefix indicating sequence information before each Context, much like "Chapter One," "Chapter Two," etc. in a novel.

Overall, my current tests of NBCE are limited to "reading comprehension" scenarios — that is, "understanding" long text. Whether this method can be used to "generate" long text remains an open question, and I look forward to seeing everyone's results.

There's also an interesting question worth pondering: if Naive Bayes can find a place in the LLM era, could other traditional probabilistic models (such as HMMs) also carve out a niche in this domain?

Summary #

This article proposes NBCE (Naive Bayes-based Context Extension), which extends the context-processing length of LLMs based on the idea of Naive Bayes. It has the advantages of being plug-and-play, model-agnostic, requiring no fine-tuning, having linear efficiency, and being simple to implement — and its results look pretty good. Feel free to give it a try.