Naive Bayes

What it is and when you reach for it

Naive Bayes is a generative classifier: instead of learning the decision boundary directly, it models how each class generates its data, then inverts that model with Bayes' rule to ask which class most plausibly produced the example in front of it. The "naive" part is the factorization of the class-conditional distribution into a product of per-feature distributions, which reduces fitting to estimating a handful of one-dimensional quantities per feature per class. Training is one pass of counting or moment accumulation, there is no iterative optimization at all, and prediction is a dot product against a table of log-probabilities. You reach for it when you need a text-classification baseline in the next five minutes, when the dataset is small enough that a discriminative model would overfit, when you need a classifier that trains in milliseconds on a stream, or when you want a sanity floor that any fancier model must beat before it earns its complexity. Its neighbors are logistic regression, which is the discriminative model of exactly the same linear form and usually wins once data is plentiful, and linear SVMs, which held the text-classification crown between the naive Bayes era and the neural one.

The math

Bayes' rule and the independence assumption

For a feature vector x = (x₁, …, x_d) and a class c, Bayes' rule says P(c | x) = P(x | c) P(c) / P(x). The denominator does not depend on c, so classification only needs the joint score P(x | c) P(c). The problem is P(x | c): a full joint distribution over d features is exponentially large, and no realistic dataset pins it down. Naive Bayes assumes the features are conditionally independent given the class:

P(x | c) = Π_j P(x_j | c)

so the whole model is d one-dimensional distributions per class plus a class prior. The assumption is plainly false for text: the word "san" all but guarantees "francisco" nearby, and no class label explains that away. What the falsehood costs is precision in the probabilities, not necessarily in the decisions. Correlated features get counted as if each were fresh evidence, so the model multiplies in the same information several times and its posteriors saturate toward 0 or 1 far more confidently than the data warrants. But the argmax over classes can survive that distortion: Domingos and Pazzani showed in 1997 that naive Bayes is often the optimal classifier even under strong dependence, because ranking classes correctly is a much weaker requirement than estimating probabilities correctly. That is the single most useful fact about this model: expect bad calibration and decent rankings, and design your system accordingly.

Multinomial and Gaussian variants

The variant names describe the per-feature distribution you plug into the product. The multinomial variant treats a document as a bag of word draws: class c has a probability θ_c,j for each vocabulary word j, with Σ_j θ_c,j = 1, and a document with count vector x has likelihood proportional to Π_j θ_c,j^x_j. The maximum likelihood estimate is exactly what intuition says it should be: θ_c,j = N_c,j / N_c, the number of times word j appears in class-c training documents divided by the total number of word tokens in class c. The Gaussian variant is for continuous features: each feature j under class c is modeled as a normal distribution with mean μ_c,j and variance σ²_c,j, both estimated by the per-class sample mean and variance. Geometrically, Gaussian naive Bayes fits one axis-aligned Gaussian blob per class; it is a diagonal-covariance special case of quadratic discriminant analysis, and if you force all classes to share the same variances the decision boundary becomes linear.

Laplace smoothing

The maximum likelihood estimate has a fatal edge case. If the word "viagra" never appears in any training ham, then P(viagra | ham) = 0, and a single occurrence in a test message multiplies the entire ham score by zero, no matter how hammy the other five hundred words are. One unseen word gets veto power over the whole document. Laplace smoothing fixes this by pretending every word was seen α extra times in every class:

θ_c,j = (N_c,j + α) / (N_c + α·V)

where V is the vocabulary size and α is usually 1 (add-one) or a smaller value tuned on held-out data. This is not just a hack: it is the posterior mean under a symmetric Dirichlet prior with parameter α + 1, so the smoothed estimate is the Bayesian answer to "what do I believe about word probabilities after seeing these counts". The practical effect is that unseen words contribute a small, finite penalty instead of an infinite one.

Log space, or nothing works

A 300-word email under a 50,000-word vocabulary produces a product of 300 factors, most of them around 10^-3 to 10^-5. That product is on the order of 10^-1000, which underflows float64 (smallest normal ≈ 10^-308) somewhere around word 100 and float32 far earlier. Every real implementation therefore works with the log of the joint score:

log P(c) + Σ_j x_j log θ_c,j

which for a whole batch of documents is a single matrix multiply of the count matrix against the transposed log-probability table, plus the log prior. The multinomial coefficient (the number of ways to arrange the words) is constant across classes, so it cancels from the argmax and from the normalized posterior and nobody computes it. When you do want actual posterior probabilities, normalize with log-sum-exp rather than exponentiating the raw joint scores; the subtract-max trick that makes that safe is the same one derived on the softmax page.

The log-space form also exposes the model's true shape: the score for each class is linear in the feature counts, weights log θ_c,j plus a bias log P(c). Multinomial naive Bayes is therefore a linear classifier over exactly the same hypothesis space as multinomial logistic regression; the two differ only in how the weights are chosen. Naive Bayes sets them generatively from per-class counts in one pass, while logistic regression fits them discriminatively to optimize the conditional likelihood, which lets it compensate for correlated features at the cost of an iterative solver and more data to reach its potential. That framing explains most of the empirical folklore about when each one wins.

A worked spam example

Take a toy corpus with a three-word vocabulary {money, free, meeting}, 4 spam messages and 6 ham messages, so the priors are P(spam) = 0.4 and P(ham) = 0.6. Suppose the spam messages contain 15 word tokens total: "money" 6 times, "free" 8 times, "meeting" once. The ham messages also contain 15 tokens: "money" once, "free" twice, "meeting" 12 times. With Laplace smoothing at α = 1 and V = 3, each denominator becomes 15 + 3 = 18:

                     spam                    ham
P(money   | c)   (6+1)/18 = 0.389       (1+1)/18 = 0.111
P(free    | c)   (8+1)/18 = 0.500       (2+1)/18 = 0.167
P(meeting | c)   (1+1)/18 = 0.111      (12+1)/18 = 0.722

Now classify the two-word message "free money". In log space (natural logs):

score(spam) = ln 0.4 + ln 0.500 + ln 0.389
            = -0.916 - 0.693 - 0.944  = -2.554

score(ham)  = ln 0.6 + ln 0.167 + ln 0.111
            = -0.511 - 1.792 - 2.197  = -4.500

Spam wins by 1.946 nats. Normalizing, P(spam | "free money") = 1 / (1 + e^-1.946) ≈ 0.875. Note what smoothing did along the way: had "meeting" never appeared in spam at all, the unsmoothed model would assign any message containing "meeting" a spam probability of exactly zero. With α = 1 it instead contributes ln(1/18) ≈ -2.89 to the spam score, a strong but finite vote for ham.

Implementation, twice

Both variants below follow the same shape discipline. For the multinomial model, fit scatter-adds the (n, V) count matrix into a (C, V) per-class count table in one call, and predict is one matmul. For the Gaussian model, per-class means and variances come from a one-hot matmul (a soft histogram), and the log-likelihood is broadcast over an (n, C, d) grid and summed over features. Nothing loops over samples or classes.

First the multinomial variant, the one you want for text counts.

import torch

class MultinomialNB:
    """Multinomial naive Bayes over nonnegative count features.

    fit is one index_add_ (per-class count table) plus a bincount;
    predict_log_proba is one matmul against the log-prob table.
    """

    def __init__(self, alpha: float = 1.0):
        self.alpha = alpha

    def fit(self, X, y):
        # X: (n, V) float counts, y: (n,) int64 class ids in [0, C)
        n_classes = int(y.max()) + 1
        counts = torch.zeros(n_classes, X.shape[1], dtype=X.dtype)
        counts.index_add_(0, y, X)            # N_cj: per-class word counts
        smoothed = counts + self.alpha        # Laplace: Dirichlet posterior mean
        self.feature_log_prob_ = (
            smoothed.log() - smoothed.sum(dim=1, keepdim=True).log()
        )                                     # (C, V) log theta
        class_count = torch.bincount(y, minlength=n_classes).to(X.dtype)
        self.class_log_prior_ = class_count.log() - class_count.sum().log()
        return self

    def joint_log_likelihood(self, X):
        # log P(c) + sum_j x_j log theta_cj; the multinomial coefficient
        # is constant across classes so it cancels and is never computed
        return X @ self.feature_log_prob_.T + self.class_log_prior_

    def predict_log_proba(self, X):
        joint = self.joint_log_likelihood(X)
        return joint - joint.logsumexp(dim=1, keepdim=True)

    def predict(self, X):
        return self.joint_log_likelihood(X).argmax(dim=1)

import jax
import jax.numpy as jnp
from functools import partial

def fit_multinomial_nb(X, y, n_classes, alpha=1.0):
    """X: (n, V) counts, y: (n,) int class ids. Returns the model as
    a pytree (feature_log_prob (C, V), class_log_prior (C,))."""
    counts = jax.ops.segment_sum(X, y, num_segments=n_classes)  # N_cj
    smoothed = counts + alpha            # Laplace / Dirichlet smoothing
    feature_log_prob = (
        jnp.log(smoothed) - jnp.log(smoothed.sum(axis=1, keepdims=True))
    )
    class_count = jnp.bincount(y, length=n_classes).astype(X.dtype)
    class_log_prior = jnp.log(class_count) - jnp.log(class_count.sum())
    return feature_log_prob, class_log_prior

@jax.jit
def joint_log_likelihood(params, X):
    feature_log_prob, class_log_prior = params
    # one matmul; the count multinomial coefficient cancels across classes
    return X @ feature_log_prob.T + class_log_prior

@jax.jit
def predict_log_proba(params, X):
    joint = joint_log_likelihood(params, X)
    return joint - jax.scipy.special.logsumexp(joint, axis=1, keepdims=True)

@jax.jit
def predict(params, X):
    return jnp.argmax(joint_log_likelihood(params, X), axis=1)

Then the Gaussian variant for continuous features. The variance floor matters: a feature that is constant within a class has zero sample variance, and the log-density would divide by it. Following scikit-learn's convention, the floor is a small fraction of the largest feature variance in the data, which keeps the guard scale invariant.

import math
import torch

class GaussianNB:
    """Gaussian naive Bayes: one axis-aligned Gaussian per class."""

    def __init__(self, var_smoothing: float = 1e-9):
        self.var_smoothing = var_smoothing

    def fit(self, X, y):
        # X: (n, d) float, y: (n,) int64 class ids
        n_classes = int(y.max()) + 1
        onehot = torch.nn.functional.one_hot(y, n_classes).to(X.dtype)
        counts = onehot.sum(dim=0)                       # (C,)
        self.theta_ = (onehot.T @ X) / counts[:, None]   # class means (C, d)
        sq_mean = (onehot.T @ X.pow(2)) / counts[:, None]
        # E[x^2] - E[x]^2, floored at a fraction of the largest data variance
        # so within-class-constant features cannot produce a zero divisor
        self.var_ = (sq_mean - self.theta_.pow(2)
                     + self.var_smoothing * X.var(dim=0).max())
        self.class_log_prior_ = counts.log() - counts.sum().log()
        return self

    def joint_log_likelihood(self, X):
        # broadcast (n, 1, d) against (C, d) -> (n, C, d), sum features
        d2 = (X[:, None, :] - self.theta_).pow(2) / self.var_
        log_lik = -0.5 * (d2 + self.var_.log()
                          + math.log(2.0 * math.pi)).sum(dim=-1)
        return log_lik + self.class_log_prior_

    def predict_log_proba(self, X):
        joint = self.joint_log_likelihood(X)
        return joint - joint.logsumexp(dim=1, keepdim=True)

    def predict(self, X):
        return self.joint_log_likelihood(X).argmax(dim=1)

import jax
import jax.numpy as jnp

def fit_gaussian_nb(X, y, n_classes, var_smoothing=1e-9):
    """Per-class means and variances via one-hot matmuls (soft histograms)."""
    onehot = jax.nn.one_hot(y, n_classes, dtype=X.dtype)   # (n, C)
    counts = onehot.sum(axis=0)                            # (C,)
    theta = (onehot.T @ X) / counts[:, None]               # (C, d) means
    sq_mean = (onehot.T @ X**2) / counts[:, None]
    # variance floor tied to the data scale, matching sklearn's convention
    var = sq_mean - theta**2 + var_smoothing * X.var(axis=0).max()
    class_log_prior = jnp.log(counts) - jnp.log(counts.sum())
    return theta, var, class_log_prior

@jax.jit
def gaussian_joint_log_likelihood(params, X):
    theta, var, class_log_prior = params
    d2 = (X[:, None, :] - theta) ** 2 / var                # (n, C, d)
    log_lik = -0.5 * jnp.sum(d2 + jnp.log(var) + jnp.log(2.0 * jnp.pi),
                             axis=-1)
    return log_lik + class_log_prior

@jax.jit
def gaussian_predict_log_proba(params, X):
    joint = gaussian_joint_log_likelihood(params, X)
    return joint - jax.scipy.special.logsumexp(joint, axis=1, keepdims=True)

@jax.jit
def gaussian_predict(params, X):
    return jnp.argmax(gaussian_joint_log_likelihood(params, X), axis=1)

Using it on a real shape of problem

The realistic text shape is thousands of documents against a vocabulary in the tens of thousands, with a count matrix that is overwhelmingly zeros. The snippet below builds a synthetic two-class corpus of that shape: 10,000 documents over a 30,000-word vocabulary, where each class draws word counts from its own skewed distribution.

import torch

torch.manual_seed(0)
n, V, C = 10_000, 30_000, 2
y = torch.randint(0, C, (n,))
# each class has its own word-probability profile; documents are
# ~80 tokens drawn from the class profile (a true multinomial corpus)
profiles = torch.distributions.Dirichlet(
    torch.full((C, V), 0.01)).sample()
X = torch.stack([
    torch.distributions.Multinomial(80, profiles[c]).sample()
    for c in y.tolist()  # sampling loop only; fit/predict stay vectorized
])

model = MultinomialNB(alpha=1.0).fit(X[:8000], y[:8000])
pred = model.predict(X[8000:])
print((pred == y[8000:]).float().mean())   # ~0.99 on this easy separation

import jax
import jax.numpy as jnp

key = jax.random.PRNGKey(0)
n, V, C = 10_000, 30_000, 2
k1, k2, k3 = jax.random.split(key, 3)
y = jax.random.randint(k1, (n,), 0, C)
profiles = jax.random.dirichlet(k2, jnp.full((V,), 0.01), shape=(C,))
# one multinomial document per row, vectorized over the batch
X = jax.vmap(
    lambda k, c: jax.random.multinomial(k, 80, profiles[c])
)(jax.random.split(k3, n), y).astype(jnp.float32)

params = fit_multinomial_nb(X[:8000], y[:8000], n_classes=C)
pred = predict(params, X[8000:])
print(jnp.mean(pred == y[8000:]))          # ~0.99 on this easy separation

On a modern laptop CPU the fit is a few tens of milliseconds and prediction is one (2,000 × 30,000) by (30,000 × 2) matmul; exact timings are machine-dependent, but the point stands at any scale: there is no training loop to wait on. Expect accuracy in the high nineties on a synthetic corpus this cleanly separated, and expect the predicted posteriors to hug 0 and 1: with 80 tokens each contributing an independent-looking log-likelihood vote, the summed evidence is huge even when many of those votes are correlated. Real corpora behave the same way, only with lower accuracy and equally overconfident probabilities.

Applications

The canonical application is spam filtering, and the history is worth knowing because it shaped a decade of email. Paul Graham's 2002 essay "A Plan for Spam" popularized Bayesian filtering of mail, the SpamAssassin project shipped a Bayes subsystem that millions of servers ran, and Mozilla Thunderbird's junk filter and the Bogofilter and SpamBayes projects were all naive Bayes at the core. The reasons it fit the job are the reasons it still gets used: it trains incrementally one message at a time (counting is trivially online), it personalizes to each user's mail, it handles a vocabulary that grows without bound, and it is cheap enough to run on a 2003 mail server. Adversarial pressure eventually forced production filters toward larger feature sets and ensembles, but the Bayes layer survives inside many of them.

Beyond spam, naive Bayes earns its keep as the fast baseline for any text pipeline: language identification, sentiment tagging, topic routing, support-ticket triage. It is the model you fit in the first hour to learn whether the signal is in the words at all, and its accuracy number becomes the floor every later model must justify itself against. Gaussian naive Bayes plays the same baseline role for low-dimensional sensor and tabular features. One more production pattern follows directly from the theory: because naive Bayes ranks well but calibrates badly, it is better used as a scorer feeding a tuned threshold or a downstream ranker than as a source of literal probabilities; if a consumer needs real probabilities, recalibrate the scores with isotonic or Platt scaling on held-out data.

Against the real libraries

The production reference is scikit-learn, whose sklearn.naive_bayes module ships the same two models implemented above plus three siblings. MultinomialNB and GaussianNB match this page's math exactly. BernoulliNB models word presence rather than counts, and it penalizes absent words explicitly, which sometimes helps on short texts. CategoricalNB handles discrete non-count features. The interesting one is ComplementNB, from Rennie et al.'s 2003 paper "Tackling the Poor Assumptions of Naive Bayes Text Classifiers": instead of estimating each class's word distribution from that class's own documents, it estimates it from the complement (all other classes) and scores against that, which fixes the systematic bias multinomial NB shows toward classes with more training data and is usually the strongest of the family on imbalanced text.

What the library adds over this reference implementation is mostly engineering breadth rather than different math: native support for scipy sparse matrices, so the (n, V) count matrix never materializes its zeros; partial_fit for out-of-core and streaming training, which for a counting model is exact rather than approximate; sample weights; configurable or learned priors; and years of numerical edge-case hardening. There is no GPU story and none is needed; the model is a matmul. The from-scratch version is genuinely enough whenever your counts fit in dense memory or you are embedding the scorer inside a PyTorch or JAX pipeline where a scikit-learn dependency is unwelcome.

Verification is easy and you should actually do it: on a fixed seed, both implementations must agree with the library to float64 tolerance, because the estimators are closed-form.

import numpy as np, torch
from sklearn.naive_bayes import MultinomialNB as SkMNB

rng = np.random.default_rng(0)
X = rng.poisson(1.0, size=(500, 200)).astype(np.float64)
y = rng.integers(0, 3, size=500)

sk = SkMNB(alpha=1.0).fit(X, y)
mine = MultinomialNB(alpha=1.0).fit(
    torch.tensor(X), torch.tensor(y))

assert np.allclose(sk.feature_log_prob_,
                   mine.feature_log_prob_.numpy(), atol=1e-10)
assert np.allclose(sk.predict_log_proba(X),
                   mine.predict_log_proba(torch.tensor(X)).numpy(),
                   atol=1e-8)

The same check for GaussianNB needs a loose tolerance of around 1e-6 rather than 1e-10, because scikit-learn computes variances with an incremental update formula whose rounding differs slightly from the two-moment formula used here.

Traps and misconceptions

Trusting the posteriors. The predicted probabilities are systematically overconfident because correlated features are multiplied in as if independent. A message that is 60% likely to be spam will be reported as 99.99% likely. The ranking of examples by score remains useful; the probabilities themselves need recalibration before anyone consumes them.

Skipping smoothing because the data "covers everything". It never does. One unseen feature value in production gives one class a log-probability of negative infinity and silently vetoes it. Smoothing is not an optional refinement, it is the difference between a working model and a time bomb.

Using the wrong variant for the feature type. Gaussian NB on word counts fits normal distributions to quantities that are mostly zero and badly skewed, and it performs accordingly. Multinomial NB on TF-IDF values violates the count semantics of the model; it often still works in practice, which is its own trap, because ComplementNB or plain counts usually work better and nothing warns you.

Believing "naive" means "weak". On small datasets the bias of the independence assumption is exactly what saves you: naive Bayes reaches its (lower) asymptotic accuracy much faster than logistic regression reaches its (higher) one, a trade-off analyzed in Ng and Jordan's 2001 generative versus discriminative comparison. With hundreds of labeled examples, naive Bayes frequently wins; with hundreds of thousands, it frequently loses. Neither regime makes it a toy.

Computing in probability space. Multiplying hundreds of small probabilities underflows to zero in any float format long before real document lengths. Everything is sums of logs, and posterior normalization goes through log-sum-exp, never through raw exponentials of joint scores.