Source-linked AI summary
Memory-based Language Models: An Efficient, Explainable, and Eco-friendly Approach to Large Language Modeling
Antal van den Bosch, Ainhoa Risco Patón, Teun Buijse, Peter Berck, Maarten van Gompel
TL;DR
Memory-based language modeling offers an alternative to neural language models, but naive k-NN inference is computationally costly. The paper implements OLIFANT using fast k-NN approximations and compares it with Transformer-based models; its estimates indicate substantially lower emissions and inference costs.
Problem
Naive k-NN classification has upper-bound complexity O(nd), motivating fast approximations for memory-based language-model inference.
Source: §1 · p. 1Method
OLIFANT implements memory-based language modeling with fast approximations of k-NN classification and runs on CPUs.
Source: §1 · p. 1, §2.2 · p. 3Results
Training with OLIFANT is estimated to cause about 1,000 times fewer CO2 emissions than neural LM training, with inference costs 10-100 times lower.
Source: §7 · p. 12Takeaways & Limitations
Memory-based models may be useful where relatively small and fast foundation models are needed as a component, such as in speculative decoding.
Source: §7 · p. 11OLIFANT is confined to a position-specific context whose width is microscopically small compared with usual neural-LM context width.
Source: Limitations and future work · p. 13
Abstract
from arXiv · showhide
We present memory-based language modeling as an efficient, eco-friendly alternative to deep neural network-based language modeling. It offers log-linearly scalable next-token prediction performance and strong memorization capabilities. Implementing fast approximations of k-nearest neighbor classification, memory-based language modeling leaves a relatively small ecological footprint both in training and in inference mode, as it relies fully on CPUs and attains low token latencies. Its internal workings are simple and fully transparent. We compare our implementation of memory-based language modeling, OLIFANT, with GPT-2 and GPT-Neo on next-token prediction accuracy, estimated emissions and speeds, and offer some deeper analyses of the model.
1 Introduction
The introduction presents memory-based language modeling as a potentially lower-cost alternative to Transformer-based models, while identifying inference speed and memory use as practical constraints. It outlines comparisons of accuracy, emissions, speed, and memorization.
1 Introduction
- Naive k-NN inference has upper-bound complexity O(nd), but prefix-trie approximations offer faster classification and compression.The reported log-linear accuracy trend with more data may change at larger scales or plateau because of memory limitations.Source: §1 · p. 1
- OLIFANT is explored as a potentially eco-friendly alternative because memory-based modeling has relatively low learning and inference costs.The implementation uses TiMBL as its classification engine.Source: §1 · p. 1
- The largest pre-trained OLIFANT models currently available require about 256 GB of RAM to run.Scaled training also eventually requires substantial RAM.Source: §1 · p. 2
- The paper compares OLIFANT variants with GPT-2 and GPT-Neo on next-token prediction, estimated emissions, and speed, and examines memorization and token distributions.The later analyses also compare validation perplexities and assess limitations.Source: §1 · p. 2
2 OLIFANT: Background and architecture
The section situates OLIFANT in memory-based classification and describes its trie-based architecture, three classification variants, and data-driven scaling framework.
2.1 Previous and related work
- Memory-based language modeling draws on k-NN and decision-tree classification, with earlier language-modeling work and applications including text completion and spelling correction.A related neural-LM approach interpolates a pretrained model with k-NN, which was reported to help surface low-frequency events.Source: §2.1 · p. 2
- Prior work also studies neural-LM memorization largely through adversarial elicitation of recitations, which are described as sparse and difficult to trigger.The cited studies associate such recitations with large models, greedy inference, and frequent training-data occurrences.Source: §2.1 · p. 2
2.2 Architecture
- OLIFANT stores context-to-next-token instances in a prefix trie, ordering context positions by gain ratio and retaining next-token distributions at trie nodes.For next-token prediction, the immediately preceding context token is checked first, followed by earlier tokens.Source: §2.2 · p. 3, Figure 1 · p. 3
- IB1-IG retrieves nearest neighbors, TRIBL2 combines trie traversal with k-NN, and IGTree returns a last-match distribution while omitting majority-consistent nodes.IGTree's omitted nodes and subtrees typically yield about 95% compression relative to losslessly storing all training instances.Source: §2.2 · pp. 3–4
2.3 Scaling dimensions in training
- Transformer training scales across data, model size, and compute, whereas memory-based training couples the model to the data and uses learning curves to study added data.The prefix trie is the model, and the amount of data linearly determines the compute needed to create it.Source: §2.3 · p. 4, Figure 3 · p. 4
3 Learning curve analyses
The analyses compare OLIFANT’s next-token accuracy, latency, emissions, and model size with GPT systems. TRIBL2’s extrapolated accuracy approaches or may surpass older small GPT systems, while OLIFANT’s faster variants and lower estimated emissions accompany substantially smaller IGTree structures.
3 Learning curve analyses
- The learning-curve analyses test whether each n-fold increase in training data yields a roughly constant accuracy gain, using EduFineWeb training data and a 512,660-token validation set.Training data comes from the corpus’s first shards, up to five billion tokens.Source: §3 · pp. 4–5
3.1 Context width
- Context width 1 begins with the highest accuracy but plateaus near 16%; widths 3–5 continue rising similarly, so later experiments use width 4.The authors note that width 4 uses substantially less context than typical Transformer-based predictors.Source: Figure 4 · p. 5, §3.1 · p. 5
3.2 Comparison with GPT variants
- The comparison plots OLIFANT accuracy across increasing training-set sizes against GPT systems with fixed training sets, whose curves therefore remain vertical.The GPT systems were trained iteratively over multiple epochs on their respective datasets.Source: §3.2 · p. 5, §3.3 · p. 6
3.3 Token prediction accuracy
- At 1 trillion tokens, TRIBL2 projects 47.7% accuracy under the post-1-billion-token fit versus 39.2% under the whole-range fit; four-token context may approach or surpass older small GPT systems.The paper attributes GPT systems’ measured accuracy surplus to their attention mechanism and wider context.Source: §3.3 · p. 6, Figure 5 · p. 7
3.4 Token prediction latency
- At 100 million training tokens, IGTree generates 739 tokens per second and TRIBL2 118, both faster than the GPT systems at roughly similar accuracy.The latency comparison runs GPT systems on 16 CPU cores.Source: §3.4 · p. 6, Figure 6 · p. 8
3.6 Model sizes and CO2 emissions during training
- Assuming about 140 grams of CO2 per billion training tokens, TRIBL2 and IGTree would emit about 49 kilograms over BLOOM 176B’s 350-billion-token training volume, three orders of magnitude less.This is an estimate based on the OLIFANT systems’ assumed electricity use, not a direct training run at that scale.Source: §3.6 · p. 8
3.5 Estimated CO2 emissions
- TRIBL2 and IGTree remain below 200 grams of CO2 equivalents on the full validation set, versus about 2.7 kilograms for GPT-Neo 2.7B.The OLIFANT emission estimates show non-monotonic variation and greater variability at the largest training-set sizes.Source: §3.5 · p. 7, Figure 7 · p. 9
3.6 Model sizes and CO2 emissions during training
- IGTree trees are 90%–95% smaller than TRIBL2’s at equal data, while TRIBL2’s 1-billion-token tree requires 53.23 gigabytes of memory.TRIBL2’s tree size scales slightly sub-linearly between 1 billion and 2 billion training tokens.Source: Figure 8 · p. 9, §3.6 · p. 9
4 Memorization
Memory-based models can reproduce training tokens accurately, but memorization is imperfect because four-word contexts can support competing next words. Errors become more common when the tested material is part of a larger training set.
4 Memorization
- 3.7% of memorization predictions were incorrect for TRIBL2 trained only on the test set itself, with ambiguous four-word contexts producing competing next words.The table compares TRIBL2 and IGTree on the same 486,655-token test material, either alone or as part of larger training sets.Source: §4 · pp. 9–10, Table 3 · p. 10
- Close to 20% of reproductions were incorrect when the 100-million-instance training set included the tested first 10,000 lines, which represented 0.5% of that set.Exactly matching four-word contexts elsewhere in the larger training set could make alternative next words more likely than the original token.Source: §4 · p. 10
- Memory-based language models are reasonable but imperfect memorizers, and their reproduction accuracy declines as training data grows.The section attributes errors to competing outcomes for limited four-word contexts; equally likely matches may also be resolved by random tie-breaking.Source: §4 · pp. 9–10
5 Explainability
OLIFANT exposes nearest-neighbor evidence behind next-token predictions, making some classifications traceable to stored training instances. The amount of visible evidence depends on the classifier variant.
5 Explainability
- In the illustrated TRIBL2 example, the model predicts “by” instead of “through” from 20 equally distant neighbors, and “through” is absent from their outcomes.The neighbors match two of the four context tokens; “by” is the most frequent outcome among them.Source: §5 · p. 10, Figure 9 · p. 10
- For TRIBL2, k = 1 can use all equidistant nearest neighbors together to form the next-word distribution.The example’s 20 neighbors are at the same distance, despite differing in which context tokens match.Source: §5 · p. 10, Figure 9 · p. 10
- TRIBL2 and IB1-IG predictions can be traced to individual instances stored in the prefix trie, and hidden identifiers could link those instances to source locations.IGTree exposes less: its output shows the next-word distribution and path depth.Source: §5 · pp. 10–11
6 Perplexity and a distributional analysis of OLIFANT predictions
OLIFANT’s memory-based variants produce sparse next-token distributions, and their reported perplexities are lower than those of IGTree and GPT systems on tokens assigned non-zero probability. Their predicted tokens also differ in frequency profile from GPT-2-small’s.
6 Perplexity and a distributional analysis of OLIFANT predictions
- TRIBL2’s next-token distribution has a median of 6 candidates and a mean of 87.5, while IGTree’s has a median of 13 and a mean of 356.The authors attribute IGTree’s larger distributions to its earlier stopping higher in the tree, where nodes represent larger sets of memorized instances.Source: §6 · p. 11
- IB1-IG and TRIBL2 have considerably lower reported perplexity than IGTree and the GPT systems, calculated only over tokens receiving non-zero probability.The authors note that IB1-IG and TRIBL2 often assign zero probability to the actual next token, whose −inf logprob was excluded from the calculation.Source: §6 · p. 11
- At 100 million training tokens, IGTree assigns 49.6% of actual next tokens non-zero probability and TRIBL2 assigns 41.9%; at 5.5 billion tokens, the reported share rises to 63%.GPT2-small assigns non-zero weight to 100% of actual next tokens.Source: §6 · p. 11
- GPT2-small predicts fewer low-frequency and more high-frequency tokens than the memory-based models in the frequency-binned comparison.The comparison uses GPT2-small trained on the 8.9-billion-token WebText corpus and memory-based models trained on 100 million tokens.Source: §6 · p. 11
7 Discussion and conclusions
The authors present OLIFANT as a low-emission, transparent memory-based alternative with scalable next-token prediction and strong memorization. They also identify its narrow context and outline directions for extending the approach and its integration.
7 Discussion and conclusions
- The best reported performances only approach GPT-2 from six years earlier, though the authors suggest memory-based models may suit small, fast components such as speculative decoding.The authors note that performance may improve with more training data.Source: §7 · p. 11
Limitations and future work
- The four-token context may constrain later learning-curve growth; the authors propose standard decoder benchmarks such as GLUE and wider-context representations as future work.Wider contexts enlarged tries; IGTree gained little, while TRIBL2 slowed for small accuracy gains.Source: §7 · pp. 12–13, Limitations and future work · p. 13
7 Discussion and conclusions
- TRIBL2 supports incremental learning and checkpointed finetuning; TiMBL multiprocessing can potentially provide speedups up to n CPU cores in training and inference.The potential speedup depends on how TiMBL is integrated into other software.Source: §7 · pp. 12–13
- Performance rises approximately log-linearly or better with more data while model size scales linearly; IGTree models are 90%-95% smaller than TRIBL2.IGTree processes in the order of hundreds of tokens per second.Source: §7 · p. 12
- OLIFANT runs on CPUs; estimates put training emissions about 1,000 times below neural-LM training and inference costs 10-100 times lower.The authors describe these comparisons as estimates.Source: §7 · p. 12
- TRIBL2 offers nearest-neighbor explanations for predictions using individual examples, and neighbors can be source-tagged for provenance.Source: §7 · p. 12
- At context size 4, OLIFANT models can recite 80% or more of training tokens faithfully.Source: §7 · p. 13
Limitations and future work
- Potential extensions include exemplar weighting, larger k with distance weighting, richer MVDM similarity, and better compatibility with model platforms.The authors describe raising k as preliminarily beneficial for TRIBL2; broader platform compatibility could ease integration into larger setups.Source: Limitations and future work · p. 13