Some Supplementary Notes and Analysis on the NBCE Method
Last week, in NBCE: Extending the Context-Processing Length of LLMs Using Naive Bayes, we introduced NBCE (Naive Bayes-based Context Extension), a scheme for extending the context length of LLMs based on naive Bayes. Since it is plug-and-play, model-agnostic, and requires no fine-tuning, it has gained some recognition from readers, and overall the feedback on its test performance so far has been fairly positive.
Of course, some readers also raised a number of questions while using it. This post combines readers' questions with the author's subsequent thoughts to provide some supplementary notes and analysis on the NBCE method.
Recap of the Method
Suppose $T$ is the token sequence to be generated, and $S_1,S_2,\cdots,S_n$ are the given Contexts. We need to generate $T$ conditioned on $S_1,S_2,\cdots,S_n$, which requires estimating $p(T|S_1, S_2,\cdots,S_n)$. Following the naive Bayes idea, we obtain
\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}more
where $\beta = n - 1$, $\overline{\log p(T|S)} = \frac{1}{n}\sum\limits_{k=1}^n \log p(T|S_k)$; see the previous post for details. NBCE made two modifications: 1) treating $\beta$ as a hyperparameter to be tuned; 2) replacing $\overline{\log p(T|S)}$ with a general Pooling method $\mathcal{P}$. The result becomes
\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}
Finally, the Pooling scheme that NBCE selected is "take the one with minimum entropy":
\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}\label{eq:min-h}\end{equation}
Truncated Prediction
Equation $\eqref{eq:nbce-2}$ is the standard naive Bayes result, but when the author implemented it as written, it was found that as $n$ increases, the performance of equation $\eqref{eq:nbce-2}$ gradually deteriorates, eventually turning into complete gibberish. So after repeated adjustments, "take the one with minimum entropy" was finally chosen as the Pooling scheme for NBCE. But thinking about it carefully afterward, this behavior of equation $\eqref{eq:nbce-2}$ is not normal, because the only assumption made by naive Bayes is that the Contexts are mutually independent, and the Contexts I tested with were several randomly selected news articles, which satisfy this assumption to a reasonable extent. So no matter how bad equation $\eqref{eq:nbce-2}$ performs, it shouldn't produce complete gibberish.
While I was racking my brain over this, @孔某人 reminded me on WeChat: the training labels for language models are always one-hot, so aside from the head (the part with the highest probability), the rest of the prediction is essentially untrustworthy. This hint really hit the nail on the head and immediately cleared up the mystery: since equation $\eqref{eq:nbce-2}$ contains the term $-\beta\log p(T)$, it amplifies the predictions in the tail, and if the tail predictions are unreliable, this amplification effect can even completely overturn the accurate results in the head. Why doesn't this affect "take the one with minimum entropy"? Because for the minimum-entropy result, the head probability tends to be larger and the tail probability smaller, so even if the term $-\beta\log p(T)$ amplifies the tail, it still cannot outweigh the head. But for equation $\eqref{eq:nbce-2}$, which is an average over all predictions, the head gets weakened, so that after multiplying by $-\beta\log p(T)$, the tail ends up outweighing the head.
With this clue, the fix becomes obvious: just apply Top-P or Top-K truncation to each prediction. In the code on GitHub, the author chose Top-P truncation.
Handling Infinities
However, the story doesn't end there. After truncation, the tails of $\log p(T|S_k)$ and $\log p(T)$ both become $-\infty$, at which point equation $\eqref{eq:nbce-2}$ or equation $\eqref{eq:nbce-3}$ may run into the meaningless operation $(-\infty)-(-\infty)$. In general, there are the following cases:
$$\begin{array}{c|cc|c} \hline & \log p(T|S_k) & \log p(T) & \log p(T|S_k) - \log p(T) \\ \hline \text{case 1} & > -\infty & > -\infty & > -\infty\\ \text{case 2} & > -\infty & = -\infty & = +\infty \\ \text{case 3} & = -\infty & > -\infty & = -\infty \\ \text{case 4} & = -\infty & = -\infty & \text{NaN}\\ \hline \end{array}$$
Among these, "Case 1" and "Case 3" compute normally; "Case 2" also computes normally, but its result of positive infinity is unreasonable; "Case 4" is an ill-defined, meaningless operation. That is, we need to find a way to fix "Case 2" and "Case 4," and these two cases correspond exactly to $\log p(T)=-\infty$, so we modify equation $\eqref{eq:nbce-3}$ as follows:
\begin{equation}\log p(T|S_1, S_2,\cdots,S_n) =\left\{ \begin{aligned} &\color{red}{\mathcal{P}[\log p(T|S)]}, \quad \text{if}\color{green}{\log p(T) = -\infty} \\[5pt] &\color{red}{(\beta + 1)\mathcal{P}[\log p(T|S)]} - \color{green}{\beta\log p(T)}, \quad \text{other}\\ \end{aligned}\right\} + \color{skyblue}{\text{const}}\label{eq:nbce-4}\end{equation}
After this treatment, the standard naive Bayes equation $\eqref{eq:nbce-2}$ can also produce normal output (although its final performance is still not as good as taking the minimum entropy, at least it no longer produces gibberish), and the modified code is also more robust with respect to the Pooling method and $\beta$.
Transition Probability
When used to answer certain opinion-based questions or questions leaning toward free-form creative writing, NBCE can suffer from repeatedly jumping back and forth between Contexts. Specifically, this happens because the model isn't confidently attending to any particular Context, so the values of $H_1,H_2,\cdots,H_n$ are not very different from one another, causing the $\mathop{\text{argmin}}$ result in equation $\eqref{eq:min-h}$ to be unstable — each generation step ends up selecting a different Context, which leads to semantic discontinuity in the generated result, or even results that are entirely unrelated to the Context, exacerbating the LLM's "hallucination" phenomenon.
To alleviate this issue, we can borrow the idea of transition probabilities and appropriately upweight the Context chosen in the previous step, so that the model only switches Context "when necessary." Concretely, we introduce a parameter $\eta > 0$ and modify equation $\eqref{eq:min-h}$ as follows:
\begin{equation}\color{red}{k} = \mathop{\text{argmin}} \big\{H_1,\cdots,H_{k'-1},H_{k'}\color{red}{-\eta},H_{k'+1},\cdots,H_n\big\}\end{equation}
where $k'$ is the index of the Context selected in the previous generation step. This way, a Context switch only occurs when $H_k < H_{k'} - \eta$, which lowers the probability of switching.
All of the modifications mentioned above have already been synced to GitHub:
GitHub: https://github.com/bojone/NBCE
Applicable Scenarios
Given the independence assumption made by naive Bayes, many readers may wonder: when there is significant semantic overlap between Contexts, will NBCE's performance drop noticeably? Or put differently, what are the appropriate scenarios for using NBCE?
In fact, it is the standard naive Bayes formulation — i.e., equation $\eqref{eq:nbce-2}$ — that is constrained by the independence assumption. After generalization, equations $\eqref{eq:nbce-3}$ and $\eqref{eq:min-h}$ are essentially no longer bound by this independence assumption. In fact, the "minimum entropy" version of NBCE is essentially using the LLM's entropy as a similarity measure to retrieve Contexts, updating the retrieval result at every generation step. So the applicable scenario for NBCE is: assume the answer to be predicted can be divided into several segments, each of which depends on only one Context.
Based on this conclusion, when we only have a single long text as Context (e.g., a novel), we can use overlapping sliding windows to automatically split the long Context into multiple short Contexts, rather than necessarily having to manually split it into relatively independent segments — because, as the conclusion above tells us, NBCE's applicability has nothing to do with the overlap between Contexts. As for why we use overlapping sliding windows at all, that's simply to make it as likely as possible that a complete result can be produced by relying on a single Context.
The following two scenarios are unlikely to work well with NBCE:
1. Ordered Context: This refers to cases where the generated result strongly depends on the input order of the Contexts (or some more complex nested structure). NBCE typically doesn't work well here, because it preserves the order-invariance property of naive Bayes. A typical example of this scenario is writing a summary of a novel (where the novel is split into multiple Contexts). A partial workaround is to manually tag each Context segment with an order-identifying marker, such as "Chapter XX";
2. Coupled Context: This refers to cases where the output must be constructed by combining two or more Contexts, in which NBCE also performs poorly, since NBCE only selects one Context at a time. @孔某人 gave a typical example: "Given $x > 1$, and $x < 0$, find the solution set of $x$." Suppose the two conditions are split into two separate Contexts — then one must combine both Contexts to output the correct answer "empty set"; looking at either Context alone cannot determine that the set is empty.
If NBCE is to be developed further, the improvements will largely need to focus on these two scenarios.
Summary
This post introduced some follow-up updates and analysis of the context-length extension scheme NBCE, and further discussed the scenarios in which NBCE is applicable.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.