Source-linked AI summary

Infini-gram: Scaling Unbounded n-gram Language Models to a Trillion Tokens

Jiacheng Liu, Sewon Min, Luke Zettlemoyer, Yejin Choi, Hannaneh Hajishirzi

arXiv:2401.17377v4cs.CLcs.AIcs.IR

TL;DR

The paper asks whether classical n-gram LMs remain relevant at neural-LLM data scales. It answers by scaling them to 5 trillion tokens, extending context to unbounded n with backoff, and introducing suffix-array-based infini-gram; the resulting models support text analysis and reduce neural-LM perplexity.

  • Problem

    The paper examines whether classical n-gram LMs remain relevant when neural LLMs are trained on trillion-token corpora.

  • Method

    The paper scales n-gram modeling to 5 trillion tokens, introduces an unbounded-n backoff LM, and serves it with a suffix-array engine.

  • Results

    ∞-gram reaches 47% next-token accuracy on human-written text and reduces neural-LM perplexity by up to 73% when interpolated.

  • Takeaways & Limitations

    ∞-gram supports analyses of human- and machine-generated text and can complement neural LMs, including models already trained on related reference data.

  • Takeaways & Limitations

    The combined model can make irrelevant predictions and digress, so it is not ready to replace neural LMs.

Abstract

from arXiv · show

Are $n$-gram language models still relevant in this era of neural large language models (LLMs)? Our answer is yes, and we showcase their values in both text analysis and improving neural LLMs. This was done by modernizing $n$-gram LMs in two aspects. First, we train them at the same data scale as neural LLMs -- 5 trillion tokens. This is the largest $n$-gram LM ever built. Second, existing $n$-gram LMs use small $n$ which hinders their performance; we instead allow $n$ to be arbitrarily large, by introducing a new $\infty$-gram LM with backoff. Instead of pre-computing $n$-gram count tables (which would be very expensive), we develop an engine named infini-gram -- powered by suffix arrays -- that can compute $\infty$-gram (as well as $n$-gram with arbitrary $n$) probabilities with millisecond-level latency. The $\infty$-gram framework and infini-gram engine enable us to conduct many novel and interesting analyses of human-written and machine-generated text: we find that the $\infty$-gram LM has fairly high accuracy for next-token prediction (47%), and can complement neural LLMs to greatly reduce their perplexity. When analyzing machine-generated text, we also observe irregularities in the machine--$\infty$-gram agreement level with respect to the suffix length, which indicates deficiencies in neural LLM pretraining and the positional embeddings of Transformers.

1 Introduction

The paper argues that n-gram LMs remain useful when modernized with trillion-token training data and unbounded context, supported by the infini-gram engine.

  • The model scales n-gram training to 5 trillion tokens and extends n to infinity, using backoff when longer contexts have zero counts.This is presented as the modernization of both training-data scale and context length.
  • Infini-gram uses suffix arrays instead of explicit count tables, supporting n-gram queries with less than 20 milliseconds latency and other queries under 200 milliseconds.The index remains on disk during inference and uses 7 bytes per token.
  • 47% next-token accuracy shows that ∞-gram can predict human-written text, compared with 29% for conventional small-n n-grams.The paper also reports that accuracy is higher when the effective n is larger.
  • Up to 73% perplexity reduction shows that interpolating ∞-gram estimates with neural LMs can improve models up to 70B parameters.The improvement is measured relative to neural LMs alone.
  • Agreement analyses report that nucleus sampling most closely resembles human-written text, while greedy decoding shows suffix-length fluctuations linked to neural pretraining and Transformer positional embeddings.The paper frames these observations as indications of deficiencies in those components.
  • The authors release web, API, Python, and source-code tools for querying and building indexes over several large text corpora.Supported corpora include Dolma, RedPajama, Pile, and C4.

2 ∞-gram LM: Extending n-gram LMs with Unbounded n

The ∞-gram LM extends classical n-gram modeling by backing off from arbitrarily long contexts according to suffix availability, while interpolation addresses its zero probabilities.

  • Background: n-gram LM: Classical n-gram count tables grow almost exponentially with n, limiting prior systems mostly to n = 5 and discarding richer predictive context.A 5-gram table for a 1.4-trillion-token corpus would require 28 TB of disk space.
  • ∞-gram LM: ∞-gram backs off from n = ∞ only until the prompt’s longest training-set suffix is found, making effective n instance-dependent.The effective n equals one plus the longest matching suffix length.
  • ∞-gram LM: Unlike conventional backoff, the ∞-gram distribution is valid by construction because the effective n depends only on the context, not the candidate next token.The cited passage explains that no additional discounting is required.
  • ∞-gram LM: Sparse estimates assign probability 1 to one possible next token and zero to all others, and are more predictive of actual tokens than non-sparse estimates.Sparsity is defined over the vocabulary for a given context.
  • Interpolating with ∞-gram: Because ∞-gram estimates can contain zero probabilities, the paper interpolates them with neural LMs rather than computing standalone ∞-gram perplexity.The combined model mixes ∞-gram and neural probabilities.

3 Infini-gram: A Performant Engine for n-gram/∞-gram Queries

Infini-gram replaces infeasible unbounded n-gram tables with suffix-array indexing, enabling large-scale counting and language-model queries with low resource use.

  • Infini-gram engine: The system is designed to process n-gram and ∞-gram queries efficiently because explicit unbounded-n count tables are infeasible at trillion-token scale.The paper presents suffix arrays as the engine’s enabling data structure.
  • Suffix array: The engine uses suffix arrays because they support substring counting in O(L + log N) time for a needle of length L in a haystack of length N.This provides the core query mechanism for n-gram and ∞-gram models.
  • Suffix array: The dataset is encoded as a byte array with two bytes per token, document separators, and one suffix-array pointer for each token.Each pointer stores a byte offset into the token array.
  • Building the suffix array: Suffix arrays were built for corpora from 200B to 3T tokens, with RedPajama requiring about 48 hours on one 128-CPU, 1 TiB RAM node.The implementation is adapted and optimized for large-scale construction.
  • Inference with the suffix array: Suffix arrays let infini-gram count n-gram occurrences by locating the consecutive suffix-array interval containing each query string.The count is obtained from the first and last occurrence positions.
  • Inference with the suffix array: Less than 20 milliseconds is the average latency for counting a given n-gram on RedPajama, while the index remains on disk without GPU use.The engine also applies parallelized shard processing, hinted search, prefetching, fast effective-n lookup, and amortized query processing.

4 Analyzing Human-written and Machine-generated Text using ∞-gram

The ∞-gram is highly predictive of human-written and machine-generated text, with accuracy increasing for longer effective contexts, and reveals complementary behavior relative to neural LMs.

  • 4.1 Human-written text: 47% overall agreement with human-written text rises above 75% when the effective n is at least 16.For sparse ∞-gram estimates, agreement reaches 75% overall and exceeds 80% when effective n is at least 14.
  • 4.1 Human-written text: A 5-gram has much lower agreement because over 90% of evaluation tokens require at least five-token context, while median effective n is 7.The mean effective n is 9.1, indicating that small fixed n fails to capture sufficient context.
  • 4.1 Human-written text: ∞-gram often completes multi-token words, phrases, and entity names, but performs poorly when recalling factual knowledge such as an entity name’s first token.The authors attribute this limitation likely to insufficient contextualization.
  • 4.1 Human-written text: ∞-gram and neural LMs predict different human-written tokens: even where neural probabilities are poor, ∞-gram agreement exceeds 20% and reaches 50% for sparse estimates.This supports combining the two model types rather than treating their predictions as redundant.
  • 4.1 Human-written text: Nucleus sampling produces machine-generated text with the most human-like agreement distribution, whereas greedy decoding shows strong effective-n fluctuation.The fluctuation is especially rapid for smaller models and is suspected to relate to Transformer positional embeddings and neural-LM pretraining.

5 Improving Neural LMs with the ∞-gram

The section evaluates interpolating ∞-gram estimates with neural LMs and finds consistent perplexity improvements across model families, while cautioning that open-ended generation can suffer from odd mistakes.

  • 5.1 Experimental setup: The evaluation uses Pile and RedPajama reference data, multiple neural-LM families, tokenizer-specific indexes, and validation, test, and time-shifted evaluations.Perplexity comparisons are only directly comparable among models sharing a tokenizer.
  • 5.2 Results: ∞-gram interpolation greatly and consistently improves neural-LM perplexity, including a 12% improvement for Llama-2 70B using Pile-train alone.The improvement generally shrinks as model size grows within a family.
  • 5.2 Results: The improvement varies across model families: ∞-gram improves GPT-2 1.6B by 34% but GPT-Neo 1.3B by 16%.The authors attribute this difference partly to whether the neural model was trained on Pile.
  • 5.2 Results: Combining Llama-2 13B with ∞-gram trained on Pile-train plus RedPajama outperforms Llama-2 70B, while Llama-2 70B perplexity falls below 4.0.The union of Pile-train and RedPajama yields larger improvements than Pile-train alone for Llama-2 models.
  • 5.2 Results: For SILO, ∞-gram helps more with more restrictive training data and improves perplexity more than SILO’s kNN-LM and RIC-LM retrieval augmentations.The contribution can be traced to reference documents, supporting source-data crediting.
  • 5.2 Results: Preliminary experiments find that ∞-gram interpolation may be unhelpful or harmful for open-ended generation because it can predict irrelevant tokens and cause digressions.The combined model is therefore not ready to replace neural LMs for text generation.

6 Related Work

Related work situates Infini-gram among classical, unbounded, and nonparametric language models, emphasizing its combination of trillion-token scale and unbounded n.

  • n-gram language models: Recent studies report mixed results when interpolating n-grams with neural LMs, ranging from improved RNN perplexity to limited gains with Transformers.Other work finds n-grams competitive with small neural LMs.
  • Unbounded n-grams, suffix arrays, suffix trees: Earlier suffix-array and suffix-tree approaches enabled unbounded-n queries but had limitations involving probability validity, storage overhead, or training-data scale.Infini-gram addresses this line of work with a scalable engine and language-model formulation.
  • Nonparametric language models: The ∞-gram LM is presented as a nonparametric language model whose complexity follows the reference data and can scale with modest resources.The authors position it as the largest nonparametric LM in reference-data size and n-gram order.

7 Conclusion

The paper concludes that Infini-gram modernizes n-gram modeling by combining trillion-token data, unbounded n, and suffix-array-based inference to analyze text and improve neural LMs.

  • 7 Conclusion: The paper modernizes n-gram language modeling by scaling to trillion-token data and extending the model to unbounded n.The conclusion frames this as the central modernization of the classical approach.
  • 7 Conclusion: Infini-gram uses suffix arrays to avoid an infeasible explicit count table while supporting scalable training and inference.The implied count table would contain at least 2 quadrillion unique n-grams, motivating the index design.
  • 7 Conclusion: The engine supports additive and subtractive indexes, enabling combinations or differences of disjoint or nested datasets without rebuilding from scratch.These operations add inference work but can be parallelized to mitigate latency overhead.
  • 7 Conclusion: The index stores document offsets and metadata to support retrieval of documents containing queried n-grams.Occurrence pointers in the suffix-array segment can be followed back into the tokenized dataset.
  • 7 Conclusion: Suffix-array queries find n-gram occurrence ranges with binary search in O(n · log N) time and can parallelize shard processing to recover O(log N) complexity.Reusing prior search results and prefetching further reduce query latency.
  • 7 Conclusion: Dense ∞-gram evaluation can reach amortized O(log N) time per token, while binary lifting reduces the operations needed to find the effective suffix length.The ∞-gram computation otherwise scales with the maximum effective n.

A.5 Supported query types and latency benchmarking

Infini-gram supports counting, probability, distribution, and document-retrieval queries, with sub-second latency even on trillion-token reference data.

  • A.5 Supported query types and latency benchmarking: The engine supports six query categories: n-gram counting, n-gram probabilities, n-gram distributions, ∞-gram probabilities, ∞-gram distributions, and document retrieval.Document retrieval also supports conjunctive and disjunctive combinations of n-gram terms.
  • A.5 Supported query types and latency benchmarking: All benchmarked query types achieve sub-second latency on trillion-token training data.The benchmark uses randomly sampled validation tokens and consecutive-token document workloads.
  • A.5 Supported query types and latency benchmarking: ∞-gram token-probability queries on RedPajama take 135 milliseconds, while full next-token distributions take 39 milliseconds for n-gram LMs and 180 milliseconds for ∞-gram.The reported figures are average query latencies from the benchmark.
  • A.5 Supported query types and latency benchmarking: Counting supports arbitrarily large n with roughly constant 20-millisecond latency, validated experimentally up to n = 1000.Full distribution queries are slower because decoding requires computing the entire next-token distribution.

B Decontamination of Reference Data

The evaluation reference data is decontaminated by removing documents with substantial n-gram overlap with evaluation sets, using whole-document filtering and fixed n-gram thresholds.

  • Documents with excessive n-gram overlap with Pile evaluation sets were filtered from Pile’s training set and RedPajama.The Big Friendly Filter was applied before these corpora were used as ∞-gram reference data.
  • Filtering removes whole documents when at least 80% of their n = 13 n-grams occur in the evaluation set.Pile documents were lowercased before filtering to capture more potential contamination.
  • The decontamination statistics for RedPajama and Pile’s training set are reported in Table 4.

C Additional Analysis

Additional analyses examine agreement patterns, evaluation preprocessing, model families, and time-shifted evaluation data for ∞-gram-based modeling.

  • ∞-gram agreement with human-written text is analyzed by effective n and longest-suffix frequency, with bar height showing token count and color showing agreement.
  • Additional results extend the agreement analyses to Llama-2 13B/7B and GPT-Neo models, while evaluation sequences use 1024-token windows with a 512-token stride.
  • The evaluated neural-model families include GPT-2, GPT-Neo/J, and Llama-2, spanning model sizes from 125M to 70B.
  • Time-shifted evaluation uses newly created Wikipedia articles after the Pile and RedPajama reference-data cutoffs.This evaluation tests performance on documents created after the reference data.
  • On documents from four of five months, interpolating ∞-gram improves Llama-2 13B perplexity, with further gains from a Random Forest interpolation policy.The policy uses suffix lengths and reference-data suffix frequencies as features.

D.3 Ablations

Ablations show that ∞-gram gains grow with reference-data scale and are largely retained when using only in-domain reference data.

  • ∞-gram improvements widen as reference-data size grows, following a roughly log-linear relationship except in the NIH ExPorter domain at small sizes.The scaling study repeatedly downsamples the full reference data by factors up to 256x.
  • Using only in-domain reference data is roughly as powerful as using the full reference data.This indicates that almost all measured improvement comes from in-domain, decontaminated data.

E Extended Discussion

The extended discussion presents infini-gram as a general corpus-analysis and retrieval tool, with applications spanning curation, attribution, contamination detection, factuality, and decoding.

  • Infini-gram can inspect massive pretraining corpora through n-gram lookup, including counting corpus contents and identifying absent text.
  • SEARCHDOC can retrieve documents containing n-gram terms or conjunctive/disjunctive expressions, supporting large-scale retrieval and iterative data curation.Indexes support additive and subtractive updates for repeated removal rounds.
  • Infini-gram may support factuality, contamination, memorization, plagiarism, copyright, and entity-popularity analyses through corpus matching and counting.The discussion also reports evidence that ∞-gram can outperform Llama-2-70B on LAMA factual probing benchmarks.
  • The engine can potentially support attribution by tracing predictive n-grams to related training documents and speculative decoding as a fast decoder.
  • The paper releases a public interface, API, Python package, and source code, while related work situates ∞-gram among high-order n-grams, text indexes, and nonparametric language models.
  • The supported query examples cover counting n-grams, n-gram and ∞-gram probabilities or distributions, and document search.
Loading 2401.17377v4…