Naive Bayes classifier

Deep dive

How it works

In: a labelled item (feature ids, class) to train; an item's feature ids to predict · Out: per-class log scores, or a normalised posterior · State: count[class][feature], tokens per class, items per class · Cost: train O(tokens), predict O(tokens x classes)

1. Train: for every feature of a labelled item add 1 (or its count) to count[c][w] and to the class totals. One pass, any order.
2. Merge: add the count tables and totals of two models. Integer addition is exact and order-free.
3. Smooth: P(w|c) = (count[c][w] + a) / (tokens[c] + a·V). a = 1 is Laplace, a < 1 Lidstone; this equals the posterior mean under a Dirichlet(a) prior (bayesian-rate-estimation), and keeps unseen words finite.
4. Freeze: tabulate ln P(w|c) and ln P(c) once.
5. Score: s[c] = ln P(c) + sum over tokens ln P(w|c); skip out-of-vocabulary words.
6. Decide: argmax of s, or posterior exp(s[c] - m) / sum exp(s - m) with m = max s (log-sum-exp).
7. Variants change only the likelihood: Bernoulli also scores absent words (precompute sum ln(1-p), correct for present ones); Gaussian keeps mean and variance per feature and class; complement NB (Rennie et al.) weights by counts from all other classes.

Code

Multinomial NB, add-alpha smoothing, precomputed log table; synthetic 3-class corpus of 1,000 words (70 % of tokens from a shared background, so it is noisy). Asserts: accuracy, bit-exact merging, finite scores and posteriors summing to 1. Compiled and run, --release.

const V: usize = 1000; // vocabulary size
const K: usize = 3;    // classes
#[derive(Clone)]
struct NaiveBayes {
    count: Vec<u32>,  // count[c*V + w]: token w seen in class c
    tokens: [u64; K], // total tokens per class
    docs: [u64; K],   // documents per class
    alpha: f64,       // add-alpha (Lidstone) smoothing; 1.0 = Laplace
}
impl NaiveBayes {
    fn new(alpha: f64) -> Self {
        NaiveBayes { count: vec![0; K * V], tokens: [0; K], docs: [0; K], alpha }
    }
    /// Incremental training: counts are the sufficient statistics.
    fn train(&mut self, doc: &[u32], c: usize) {
        for &w in doc { self.count[c * V + w as usize] += 1; }
        self.tokens[c] += doc.len() as u64;
        self.docs[c] += 1;
    }
    /// Merging two shards = adding their counts (integer adds: exact, order-free).
    fn merge(&mut self, o: &NaiveBayes) {
        for (a, b) in self.count.iter_mut().zip(&o.count) { *a += *b; }
        for c in 0..K { self.tokens[c] += o.tokens[c]; self.docs[c] += o.docs[c]; }
    }
    /// Precompute ln P(w|c) as lp[w*K + c] and ln P(c).
    fn freeze(&self) -> Frozen {
        let n: u64 = self.docs.iter().sum();
        let mut lp = vec![0.0; V * K];
        for c in 0..K {
            let denom = (self.tokens[c] as f64 + self.alpha * V as f64).ln();
            for w in 0..V { lp[w * K + c] = (self.count[c * V + w] as f64 + self.alpha).ln() - denom; }
        }
        let prior = std::array::from_fn(|c| ((self.docs[c] as f64 + 1.0) / (n as f64 + K as f64)).ln());
        Frozen { lp, prior }
    }
}
struct Frozen { lp: Vec<f64>, prior: [f64; K] }
impl Frozen {
    /// Joint log-scores ln P(c) + sum_i ln P(w_i|c); words >= V (out of vocabulary) are skipped.
    fn log_joint(&self, doc: &[u32]) -> [f64; K] {
        let mut s = self.prior;
        for &w in doc.iter().filter(|&&w| (w as usize) < V) {
            for c in 0..K { s[c] += self.lp[w as usize * K + c]; }
        }
        s
    }
    /// Posterior via log-sum-exp, so long documents do not underflow to 0/0.
    fn posterior(&self, doc: &[u32]) -> [f64; K] {
        let s = self.log_joint(doc);
        let m = s.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
        let z: f64 = s.iter().map(|x| (x - m).exp()).sum();
        s.map(|x| (x - m).exp() / z)
    }
}
// ---- deterministic synthetic corpus ----
struct Rng(u64);
impl Rng {
    fn next(&mut self) -> u64 { self.0 ^= self.0 << 13; self.0 ^= self.0 >> 7; self.0 ^= self.0 << 17; self.0 }
    fn unit(&mut self) -> f64 { (self.next() >> 11) as f64 / (1u64 << 53) as f64 }
}
/// Each class: Zipf-like weights 1/(rank+1) over a class-specific permutation of the vocabulary.
fn class_cdfs(rng: &mut Rng) -> Vec<(Vec<f64>, Vec<u32>)> {
    (0..K).map(|_| {
        let mut perm: Vec<u32> = (0..V as u32).collect();
        for i in (1..V).rev() { perm.swap(i, (rng.next() % (i as u64 + 1)) as usize); }
        let h: f64 = (1..=V).map(|r| 1.0 / r as f64).sum();
        let cdf = (1..=V).scan(0.0, |a, r| { *a += 1.0 / r as f64 / h; Some(*a) }).collect();
        (cdf, perm)
    }).collect()
}
/// 70 % of tokens come from a shared background (class 0's table), 30 % from the class's own table.
fn make_doc(rng: &mut Rng, models: &[(Vec<f64>, Vec<u32>)], c: usize, len: usize) -> Vec<u32> {
    (0..len).map(|_| {
        let m = if rng.unit() < 0.7 { &models[0] } else { &models[c] };
        let u = rng.unit();
        m.1[m.0.partition_point(|&p| p < u).min(V - 1)]
    }).collect()
}
fn main() {
    let mut rng = Rng(0x9E3779B97F4A7C15);
    let model = class_cdfs(&mut rng);
    let mut gen = |n: usize, len: usize| -> Vec<(Vec<u32>, usize)> {
        (0..n).map(|i| (make_doc(&mut rng, &model, i % K, len), i % K)).collect()
    };
    let (train, test) = (gen(30_000, 10), gen(3_000, 10));
    let (a, b) = train.split_at(12_345);
    let fit = |docs: &[(Vec<u32>, usize)]| { let mut m = NaiveBayes::new(0.1); for (d, c) in docs { m.train(d, *c); } m };
    let (all, ma, mb) = (fit(&train), fit(a), fit(b));
    // (1) accuracy on held-out documents
    let f = all.freeze();
    let argmax = |s: &[f64; K]| (0..K).max_by(|&i, &j| s[i].partial_cmp(&s[j]).unwrap()).unwrap();
    let acc = test.iter().filter(|(d, c)| argmax(&f.log_joint(d)) == *c).count() as f64 / test.len() as f64;
    println!("accuracy {acc:.3}");
    assert!(acc > 0.8); // chance is 0.333
    // (2) shard A merged with shard B == trained on A+B, bit for bit
    let mut merged = ma.clone(); merged.merge(&mb);
    let fm = merged.freeze();
    assert!(f.lp.iter().zip(&fm.lp).all(|(x, y)| x == y) && f.prior == fm.prior);
    // (3) smoothing: a word never seen in a class (or at all) gives a finite score; posteriors sum to 1
    let mut m = NaiveBayes::new(0.1);
    m.train(&[1, 2, 3], 0); m.train(&[4, 5], 1); m.train(&[6], 2);
    let tiny = m.freeze();
    let p = tiny.posterior(&[999, 1, 999_999]); // 999 never seen, 999_999 out of vocabulary
    assert!(tiny.log_joint(&[999]).iter().all(|x| x.is_finite()));
    assert!((p.iter().sum::<f64>() - 1.0).abs() < 1e-12 && p[0] > p[1] && p[0] > p[2]);
    let long = vec![7u32; 100_000]; // the raw product of 1e5 probabilities underflows to 0
    assert!((f.posterior(&long).iter().sum::<f64>() - 1.0).abs() < 1e-12);
    // throughput; rep-dependent indices stop the compiler reusing earlier reps
    let tok: usize = train.iter().map(|(d, _)| d.len()).sum();
    let t0 = std::time::Instant::now();
    let mut m = NaiveBayes::new(0.1);
    for rep in 0..20 { for (d, c) in &train { m.train(d, (*c + rep) % K); } }
    println!("train {:.0} M tokens/s (checksum {})", 20.0 * tok as f64 / t0.elapsed().as_secs_f64() / 1e6, m.count[5]);
    let t0 = std::time::Instant::now();
    let mut s = 0;
    for rep in 0..200 { for i in 0..test.len() { s += argmax(&f.log_joint(&test[(i + rep) % test.len()].0)); } }
    println!("predict {:.2} M docs/s (checksum {s})", 200.0 * test.len() as f64 / t0.elapsed().as_secs_f64() / 1e6);
}

Output, measured on this Mac (Apple M4 Pro), one core, three runs: accuracy 0.846 on 3,000 held-out documents (chance 0.333); train about 2.6-2.8 G tokens/s; predict about 0.94-1.0 x 10^8 documents/s. Ten-token documents and a 24 KB table in L1: best case.

Complexity and constants

• Train O(N) for N tokens. Predict O(n·K) for n tokens and K classes; sparse inputs touch only present words.
• Memory: counts K·V·4 B (12 KB here); log table V·K·8 B (24 KB). 1 M words x 100 classes is 400 MB of u32 counts: use hashing or prune rare words.
• Cost drivers: tokenising and hashing usually dominate; the arithmetic is K adds per token.

Efficiency

• Cost: about 0.4 ns per training token and about 10 ns per 10-token, 3-class prediction on one M4 Pro core (measured above); bound by dependent adds, and by cache misses for large tables.
• Levers:
  • precomputed log table — removes one ln per token and class (not separately measured) — always
  • word-major table ([word][class]) — one cache line per token instead of K misses — many classes, large vocabulary
  • sharded counting plus merge — merge is one vector add, bit-exact (asserted above) — large corpora
  • hashed features — fixed memory, no vocabulary pass, some collisions — open vocabularies
• Naive → standard → state of the art:
  • multiply raw probabilities · underflows on long documents
  • log table, argmax of sums · about 10 ns per short document (measured above)
  • hashed sparse features, sharded counts merged · O(tokens) single pass, mergeable

Implementation notes

• Underflow: 100,000 factors of 0.001 vanish in f64; sum logs, normalise with log-sum-exp.
• Smoothing: tune a on held-out data; a = 0 gives -inf for any unseen word.
• Priors: from training frequencies. If deployment class rates differ, replace ln P(c) with the deployed rates; no retraining.
• Length: the log score grows with document length, so long documents give posteriors near 0 or 1; compare scores, not probabilities, across documents.
• Normalise text first (text-normalization); Rennie et al. also transform term frequencies and normalise length.
• Thresholds on probabilities need calibration on held-out data.

Pitfalls

• Copies of one feature (word and its stem, both n-gram orders) count as independent evidence and inflate confidence.
• Gaussian NB with a near-constant feature gets variance near 0: add a variance floor.

Variants and alternatives

In production

• Spam filtering (Graham's 2002 essay describes a per-token Bayesian filter), language identification (language-identification), topic routing, text-pipeline baselines.

Implementations

Sources

• Rennie, Shih, Teevan, Karger, Tackling the Poor Assumptions of Naive Bayes Text Classifiers, ICML 2003. paper
• McCallum, Nigam, A Comparison of Event Models for Naive Bayes Text Classification, AAAI/ICML-98 Workshop on Learning for Text Categorization, 1998. paper
• Domingos, Pazzani, On the Optimality of the Simple Bayesian Classifier under Zero-One Loss, Machine Learning 29, 1997. paper
• Ng, Jordan, On Discriminative vs. Generative Classifiers: A comparison of logistic regression and naive Bayes, NeurIPS 2001. paper
• Graham, A Plan for Spam, 2002. essay
• scikit-learn, Naive Bayes. docs