Text classification

Deep dive

How it works

Parts:

• Featurizer — words, char n-grams or subword ids (BPE), hashed to a fixed index space (feature-hashing).
• Model rung — rules; TF-IDF + logistic regression / linear SVM / fastText; frozen embeddings + linear head; fine-tuned encoder; prompted LLM.
• Output layer — softmax for one label, a sigmoid per label for several.
• Decision layer — per-label thresholds, abstain threshold, human queue.
• Monitor — label-rate and confidence drift.

Operations:

1. Define labels a human can apply consistently; inter-annotator agreement caps the achievable score.
2. Rung 1: one keyword/regex rule per label, with precedence. It is the baseline and often seeds the first labelled set.
3. Rung 2: word unigrams + bigrams (or char 2-5-grams for typos, URLs, log lines), TF-IDF weights, logistic regression or linear SVM; or fastText, which averages n-gram embeddings into a linear classifier.
4. Rung 3: embed each text once with a frozen sentence encoder; train a linear head on the vectors.
5. Rung 4: fine-tune a BERT-family encoder end to end, linear head on the first token. SetFit, for few labels: contrastive fine-tuning of a sentence encoder on text pairs, then a linear head.
6. Rung 5: prompt an LLM with label definitions and a few examples; constrain output to the label set (grammar, enum schema, or compare label log-probabilities).
7. Multi-label: binary cross-entropy per label, each threshold chosen on validation data to maximise that label's F1; imbalance via class weights or threshold shifts. Calibrate before treating scores as confidence.
8. Serve: abstain when max probability (or every sigmoid) is below a threshold taken from the accuracy-vs-coverage curve; humans label the abstentions; retrain.

text -> normalise -> features -> [ rules | tfidf+LR | embed+head | encoder | LLM ]
                       scores -> softmax (1 label) / sigmoid per label (many)
                       p_max >= tau ? --yes--> label   --no--> human queue --> new labels --> retrain

Hierarchical labels: a flat classifier over leaves plus a path check, or one classifier per node top-down (an early error stays). Label noise: review the highest-loss training examples; keep a clean validation set. Drift: new vocabulary or abuse tactics shift inputs; watch label and abstain rates (distribution-drift-detection).

Code

Hashed word unigram+bigram features, softmax regression by SGD, synthetic 3-class corpus (5 % cross-class words, so a ceiling below 100 % is intended). Compiled and run, --release.

const D: usize = 1 << 14; // hash buckets per class
const K: usize = 3;

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 below(&mut self, n: usize) -> usize { (self.next() % n as u64) as usize }
    fn unit(&mut self) -> f64 { (self.next() >> 11) as f64 / (1u64 << 53) as f64 }
}

const CLASS_WORDS: [[&str; 8]; K] = [
    ["invoice", "refund", "charge", "card", "billing", "payment", "receipt", "price"],
    ["parcel", "delivery", "tracking", "courier", "shipped", "address", "customs", "package"],
    ["password", "login", "account", "email", "reset", "profile", "verify", "username"],
];
const NOISE: [&str; 12] = ["the", "my", "is", "not", "please", "help", "have", "since", "again", "order", "you", "still"];

fn make_doc(rng: &mut Rng, class: usize) -> String {
    let mut words = Vec::new();
    for _ in 0..10 {
        let r = rng.unit();
        // 25 % own-class word, 5 % other-class word (confusion), rest shared noise
        let w = if r < 0.25 { CLASS_WORDS[class][rng.below(8)] }
            else if r < 0.30 { CLASS_WORDS[(class + 1 + rng.below(K - 1)) % K][rng.below(8)] }
            else { NOISE[rng.below(NOISE.len())] };
        words.push(w);
    }
    words.join(" ")
}

fn fnv(bytes: &[u8], seed: u64) -> u64 {
    let mut h = 0xcbf29ce484222325u64 ^ seed;
    for &b in bytes { h ^= b as u64; h = h.wrapping_mul(0x100000001b3); }
    h
}

/// Hashed word unigrams + bigrams -> bucket indices (all values 1.0, binary features).
fn featurize(text: &str, out: &mut Vec<usize>) {
    out.clear();
    let mut prev: Option<&str> = None;
    for w in text.split_whitespace() {
        out.push(fnv(w.as_bytes(), 1) as usize % D);
        if let Some(p) = prev {
            let mut b = Vec::with_capacity(p.len() + w.len() + 1);
            b.extend_from_slice(p.as_bytes()); b.push(b' '); b.extend_from_slice(w.as_bytes());
            out.push(fnv(&b, 2) as usize % D);
        }
        prev = Some(w);
    }
}

struct Model { w: Vec<f32>, b: [f32; K] } // w[f*K + k]

impl Model {
    fn proba(&self, feats: &[usize]) -> [f32; K] {
        let mut z = self.b;
        for &f in feats { for k in 0..K { z[k] += self.w[f * K + k]; } }
        let m = z.iter().cloned().fold(f32::MIN, f32::max);
        let mut s = 0.0;
        for v in z.iter_mut() { *v = (*v - m).exp(); s += *v; }
        for v in z.iter_mut() { *v /= s; }
        z
    }
    fn sgd_step(&mut self, feats: &[usize], y: usize, lr: f32, l2: f32) {
        let p = self.proba(feats);
        for k in 0..K {
            let g = p[k] - (k == y) as u8 as f32; // d loss / d logit_k for softmax + NLL
            self.b[k] -= lr * g;
            for &f in feats { let w = &mut self.w[f * K + k]; *w -= lr * (g + l2 * *w); }
        }
    }
}

fn argmax(p: &[f32; K]) -> usize { (0..K).max_by(|&a, &b| p[a].partial_cmp(&p[b]).unwrap()).unwrap() }

fn macro_f1(pred: &[usize], gold: &[usize]) -> f64 {
    let n = |f: &dyn Fn(usize, usize) -> bool| pred.iter().zip(gold).filter(|(p, g)| f(**p, **g)).count() as f64;
    (0..K).map(|c| {
        let (tp, fp, fneg) = (n(&|p, g| p == c && g == c), n(&|p, g| p == c && g != c), n(&|p, g| p != c && g == c));
        if tp == 0.0 { 0.0 } else { 2.0 * tp / (2.0 * tp + fp + fneg) } // F1 = 2TP / (2TP + FP + FN)
    }).sum::<f64>() / K as f64
}

fn main() {
    let mut rng = Rng(0x9E3779B97F4A7C15);
    let gen = |rng: &mut Rng, n: usize| -> Vec<(String, usize)> { (0..n).map(|i| (make_doc(rng, i % K), i % K)).collect() };
    let (train, test) = (gen(&mut rng, 3000), gen(&mut rng, 3000));
    let (mut m, mut buf) = (Model { w: vec![0.0; D * K], b: [0.0; K] }, Vec::new());
    let mut order: Vec<usize> = (0..train.len()).collect();
    for epoch in 0..5 {
        for i in (1..order.len()).rev() { order.swap(i, rng.below(i + 1)); } // Fisher-Yates shuffle
        let lr = 0.5 / (1.0 + epoch as f32);
        for &i in &order { featurize(&train[i].0, &mut buf); m.sgd_step(&buf, train[i].1, lr, 1e-4); }
    }
    let (mut pred, mut gold, mut conf) = (vec![], vec![], vec![]);
    for (t, y) in &test {
        featurize(t, &mut buf);
        let p = m.proba(&buf);
        pred.push(argmax(&p)); gold.push(*y); conf.push(p[argmax(&p)]);
    }
    let acc = pred.iter().zip(&gold).filter(|(p, g)| p == g).count() as f64 / gold.len() as f64;
    println!("accuracy {acc:.3}  macro-F1 {:.3}", macro_f1(&pred, &gold));
    // abstain below threshold: keep only confident predictions
    let tau = 0.9;
    let kept: Vec<usize> = (0..pred.len()).filter(|&i| conf[i] >= tau).collect();
    let kept_acc = kept.iter().filter(|&&i| pred[i] == gold[i]).count() as f64 / kept.len() as f64;
    println!("tau {tau}: kept {:.1} %, accuracy on kept {kept_acc:.3}", 100.0 * kept.len() as f64 / pred.len() as f64);
    assert!(acc > 0.8 && kept_acc > acc && kept.len() < pred.len());
    // throughput: featurize + predict, single thread
    let t0 = std::time::Instant::now();
    let mut s = 0;
    for _ in 0..20 { for (t, _) in &test { featurize(t, &mut buf); s += argmax(&m.proba(&buf)); } }
    println!("{:.0} docs/s (checksum {s})", 20.0 * test.len() as f64 / t0.elapsed().as_secs_f64());
}

Output (Apple M4 Pro, one thread, 10-word documents, about 19 features each; three runs gave 2.59, 2.59, 2.61 M docs/s): accuracy 0.877 macro-F1 0.877, tau 0.9: kept 91.2 %, accuracy on kept 0.914. The assert checks that abstaining raises accuracy on the kept set.

Complexity and constants

• Train O(epochs x documents x nonzeros x classes); predict O(nonzeros x classes). Weights here: 2^14 x 3 x 4 B = 192 KiB, resident in L2; 2^22 buckets x 100 classes is 1.7 GB and cache-misses on every feature.
• Linear rung: cost is tokenizing and hashing, not the dot product; 200 tokens give order of 400 nonzeros with bigrams.
• Encoder: about 2 x parameters x tokens FLOPs per document (order of; BERT-base is 110 M parameters per Devlin et al.) plus quadratic attention; padding wastes the rest of a batch.
• LLM: input tokens (prompt + examples + document) x price; the few-shot prefix dominates unless cached.

Efficiency

• Cost: rung 2: about 0.4 µs per 10-word document on one M4 Pro core (measured, 2.6 M docs/s with featurizing), CPU-bound. Rung 4: order of 20 GFLOPs for 128 tokens on BERT-base (2 x ~85 M non-embedding parameters x 128, my arithmetic), GPU-class. Rung 5: a network call priced per token (not measured).
• Levers:
  • hash n-grams incrementally, no per-bigram string allocation — saves one allocation per bigram (not measured separately) — any n-gram model.
  • cascade: linear model answers confident texts, encoder or LLM the rest — cost scales with the escalated share (9 % at tau 0.9 in the toy run) — when most texts are easy.
  • embed once, cache by text hash — one forward pass per distinct text — rung 3, duplicate-heavy streams such as logs.
  • bucket by length and batch — less padding — rung 4.
  • SetFit — "an order of magnitude faster to train" than comparable few-shot methods (Tunstall et al. abstract) — a few labels per class.
• Naive → standard → state of the art:
  • LLM call per document · per-token price, network latency.
  • n-gram + linear model or fastText · millions of short docs/s per core (measured above).
  • cascade: linear first, encoder/LLM on abstentions, humans on the residue · rung-2 cost for the confident majority.

Implementation notes

• Split by time or source, not at random: near-duplicates leak (near-duplicate-detection-minhash).
• Tune thresholds on validation data; report per-class F1 on a separate test set (evaluation-metrics).
• Keep rules as a pre-filter for must-catch cases; log which rung decided; store hash seed and tokenizer version with the model.
• Add "none of the above" when inputs can fall outside the label set.

Pitfalls

• 98 % accuracy by always predicting the majority class: use macro-F1 and per-class recall.
• Fine-tuned encoders are over-confident; tau 0.9 keeps a different share per model. Calibrate, then choose tau.
• Softmax forces exactly one label; sigmoid can emit none.
• Leakage through signatures, ticket IDs, template headers.
• Prompted LLMs change between model versions and with label order; pin both and test on the same held-out set.

Variants and alternatives

In production

• Help desks route tickets, mail systems filter spam, platforms moderate content with encoders behind cheap pre-filters, back offices tag scanned invoices after OCR, SRE tools categorise log lines; a cascade with a human queue is the common shape.

Implementations

Sources

• Joulin, Grave, Bojanowski, Mikolov, Bag of Tricks for Efficient Text Classification, EACL 2017. paper · code
• Devlin, Chang, Lee, Toutanova, BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding, NAACL 2019. paper
• Tunstall et al., Efficient Few-Shot Learning Without Prompts, 2022 (SetFit). paper · code
• scikit-learn, Text feature extraction. docs
• Vowpal Wabbit. code