Named entity recognition

Deep dive

How it works

Parts:

• Tokenizer — words or subwords with character offsets.
• Features or encoder — hashed word/shape/context features, or a transformer.
• Scorer — per-token scores over BIO tags.
• Decoder — Viterbi over tag transitions; I-x only after B-x/I-x.
• Gazetteer and rules — Aho-Corasick matches as features or recall fallback.
• Span builder — (start, end, type); a linker then maps each to a KB id.

Operations:

1. Tokenize, keep offsets; with subwords label only each word's first subword and mask the rest (Devlin et al.).
2. Score every token for every tag.
3. Decode the best valid path, O(n x tags^2); per-token argmax plus repair is close for strong encoders.
4. Collect spans; apply rules, drop overlaps, normalize values (text normalization).
5. Train with a sequence loss (perceptron, CRF) or per-token cross-entropy; score entity P/R/F1.

text -> tokens+offsets -> features/encoder -> tag scores (n x tags)
        gazetteer, rules ------------------> extra features / fallback
scores -> Viterbi (BIO-valid) -> B-ORG I-ORG O O B-LOC -> spans -> linker

Schemes: BIO, or BIOES/BILOU with end and single-token tags; Ratinov and Roth: BILOU 90.57 vs BIO 89.15 F1, CoNLL-2003 test. CRFs (Lafferty et al.) learn tag transitions; Lample et al. put one on a BiLSTM; transformer taggers often drop it.

Code

Averaged structured perceptron, BIO-constrained Viterbi, hashed features, a location gazetteer and exact-match F1 on synthetic news, clinical and contract sentences; test surnames, firms and cities are unseen.

const D: usize = 1 << 16; // hashed feature buckets
const NT: usize = 9; // O, then B/I for PER ORG LOC AMT: B=1+2k, I=2+2k
const TR: usize = NT * D; // transitions p*NT+t follow the emissions
type Span = (usize, usize, usize); // [start, end), type
fn ok(p: usize, t: usize) -> bool { t == 0 || t % 2 == 1 || p == t - 1 || p == t } // I-k only after B-k/I-k; start = O
fn spans(tags: &[usize]) -> Vec<Span> {
    let mut v: Vec<Span> = vec![];
    for (i, &t) in tags.iter().enumerate().filter(|(_, &t)| t != 0) {
        match v.last_mut() {
            Some(s) if t % 2 == 0 && s.1 == i && s.2 == (t - 2) / 2 => s.1 = i + 1,
            _ => v.push((i, i + 1, (t - 1) / 2)), // B-k, or an orphan I-k repaired to B-k
        }
    }
    v
}
fn f1(gold: &[Vec<Span>], pred: &[Vec<Span>]) -> (f64, f64, f64) { // exact match: start, end and type (CoNLL)
    let tp: usize = gold.iter().zip(pred).map(|(g, p)| p.iter().filter(|s| g.contains(s)).count()).sum();
    let (np, ng) = (pred.iter().map(Vec::len).sum::<usize>(), gold.iter().map(Vec::len).sum::<usize>());
    let (pr, re) = (tp as f64 / np.max(1) as f64, tp as f64 / ng.max(1) as f64);
    (pr, re, if tp == 0 { 0.0 } else { 2.0 * pr * re / (pr + re) })
}
fn h(s: &str) -> usize { (s.bytes().fold(0xcbf29ce484222325u64, |h, b| (h ^ b as u64).wrapping_mul(0x100000001b3)) >> 20) as usize % D }
const GAZ: [&str; 8] = ["paris", "lima", "oslo", "london", "new", "york", "san", "diego"]; // locations only
fn feats(w: &[String], gaz: bool) -> Vec<Vec<usize>> {
    (0..w.len()).map(|i| {
        let (c, lo) = (&w[i], w[i].to_lowercase());
        let shape: String = c.chars().map(|x| if x.is_ascii_digit() { 'd' } else if x.is_uppercase() { 'X' } else { 'x' }).take(2).collect();
        let nb = |j: Option<usize>| j.and_then(|j| w.get(j)).map_or("^".to_string(), |s| s.to_lowercase());
        let mut f = vec![h("b"), h(&format!("w={lo}")), h(&format!("s={shape}")), h(&format!("x={}", &lo[lo.len().saturating_sub(3)..])),
            h(&format!("p={}", nb(i.checked_sub(1)))), h(&format!("n={}", nb(Some(i + 1)))), h(&format!("p+s={}{shape}", nb(i.checked_sub(1))))];
        if gaz && GAZ.contains(&lo.as_str()) { f.push(h("gaz=loc")); }
        f
    }).collect()
}
struct Tagger { w: Vec<f32>, u: Vec<f32>, c: f32 } // averaged structured perceptron
impl Tagger {
    fn decode(&self, f: &[Vec<usize>]) -> Vec<usize> { // Viterbi, O(n * NT^2)
        let n = f.len();
        let (mut best, mut bp) = (vec![[f32::MIN; NT]; n], vec![[0usize; NT]; n]);
        for i in 0..n { for t in 0..NT {
            let e: f32 = f[i].iter().map(|&k| self.w[t * D + k]).sum();
            for p in (if i == 0 { 0..1 } else { 0..NT }).filter(|&p| ok(p, t)) {
                let prev = if i == 0 { 0.0 } else { best[i - 1][p] };
                if prev == f32::MIN { continue; }
                let s = prev + self.w[TR + p * NT + t] + e;
                if s > best[i][t] { best[i][t] = s; bp[i][t] = p; }
            }
        }}
        let mut t = (0..NT).max_by(|&a, &b| best[n - 1][a].partial_cmp(&best[n - 1][b]).unwrap()).unwrap();
        let mut tags = vec![0; n];
        for i in (0..n).rev() { tags[i] = t; t = bp[i][t]; }
        tags
    }
    fn add(&mut self, i: usize, v: f32) { self.w[i] += v; self.u[i] += self.c * v; } // u: weighted sum for averaging
    fn train(&mut self, f: &[Vec<usize>], gold: &[usize]) {
        let pred = self.decode(f);
        if pred != gold {
            for i in 0..gold.len() {
                let (gp, pp) = (if i > 0 { gold[i - 1] } else { 0 }, if i > 0 { pred[i - 1] } else { 0 });
                if gold[i] != pred[i] { for &k in &f[i] { self.add(gold[i] * D + k, 1.0); self.add(pred[i] * D + k, -1.0); } }
                if (gp, gold[i]) != (pp, pred[i]) { self.add(TR + gp * NT + gold[i], 1.0); self.add(TR + pp * NT + pred[i], -1.0); }
            }
        }
        self.c += 1.0;
    }
    fn average(&mut self) { for (w, u) in self.w.iter_mut().zip(&self.u) { *w -= u / self.c; } }
}
struct Rng(u64); // xorshift64
impl Rng { fn pick<'a>(&mut self, v: &[&'a str]) -> &'a str { self.0 ^= self.0 << 13; self.0 ^= self.0 >> 7; self.0 ^= self.0 << 17; v[(self.0 % v.len() as u64) as usize] } }
const TEMPLATES: [&str; 5] = ["0 joined 1 in 2 last week", "1 acquired 1 for 3", "patient 0 received 3 of ibuprofen at 1", "the contract names 0 as director of 1", "0 and 0 flew from 2 to 2"];
fn sentence(r: &mut Rng, test: bool) -> (Vec<String>, Vec<usize>) {
    let (mut w, mut t) = (vec![], vec![]);
    let (last, pre, loc): (&[&str], &[&str], &[&str]) = if test { (&["Kowalski", "Nguyen", "Okafor", "Larsen"], &["Orbit", "Vantara", "Kestrel"], &["Kyiv", "Hanoi", "Oslo", "New Delhi"]) }
        else { (&["Smith", "Jones", "Brown", "Miller", "Moore"], &["Zenta", "Nordex", "Helion", "Acme"], &["Paris", "Lima", "London", "New York", "San Diego"]) };
    for tok in r.pick(&TEMPLATES).split(' ') {
        let ent: Vec<String> = match tok { // 0 PER, 1 ORG, 2 LOC, 3 AMT, else a plain word
            "0" => vec![r.pick(&["Alice", "Bob", "Carol", "David", "Emma"]).into(), r.pick(last).into()],
            "1" => vec![r.pick(pre).into(), r.pick(&["Bank", "Pharma", "Ltd", "Hospital"]).into()],
            "2" => r.pick(loc).split(' ').map(String::from).collect(),
            "3" => format!("{} {}", r.pick(&["5", "20", "250", "3"]), r.pick(&["mg", "million euros", "dollars"])).split(' ').map(String::from).collect(),
            _ => { w.push(tok.to_string()); t.push(0); continue; }
        };
        for (j, e) in ent.into_iter().enumerate() { w.push(e); t.push(1 + 2 * tok.parse::<usize>().unwrap() + (j > 0) as usize); }
    }
    (w, t)
}
fn main() {
    let mut r = Rng(0x9E3779B97F4A7C15);
    let prep = |r: &mut Rng, n, test, gaz| (0..n).map(|_| { let (w, t) = sentence(r, test); (feats(&w, gaz), t, w) }).collect::<Vec<_>>();
    for gaz in [true, false] { // same data, with and without the gazetteer feature
        let (train, test) = (prep(&mut r, 1000, false, gaz), prep(&mut r, 400, true, gaz));
        let mut m = Tagger { w: vec![0.0; TR + NT * NT], u: vec![0.0; TR + NT * NT], c: 1.0 };
        for _ in 0..5 { for (f, t, _) in &train { m.train(f, t); } }
        m.average();
        let (g, p): (Vec<_>, Vec<_>) = test.iter().map(|(f, t, _)| (spans(t), spans(&m.decode(f)))).unzip();
        let (pr, re, f) = f1(&g, &p); println!("gazetteer={gaz}: held-out entity P {pr:.3} R {re:.3} F1 {f:.3}");
        if gaz {
            assert!(f > 0.9);
            let s: Vec<String> = "Alice Kowalski joined Vantara Bank in New Delhi last week".split(' ').map(String::from).collect();
            assert_eq!(spans(&m.decode(&feats(&s, true))), vec![(0, 2, 0), (3, 5, 1), (6, 8, 2)]);
            let (t0, mut n) = (std::time::Instant::now(), 0);
            for _ in 0..50 { for (_, _, w) in &test { n += spans(&m.decode(&feats(w, true))).len(); } }
            println!("{:.0} sentences/s, one thread (checksum {n})", 50.0 * test.len() as f64 / t0.elapsed().as_secs_f64());
        }
    }
    let (pr, re, f) = f1(&[vec![(0, 2, 0), (5, 6, 2)]], &[vec![(0, 2, 0), (5, 7, 2)]]);
    assert!((pr - 0.5).abs() < 1e-9 && (re - 0.5).abs() < 1e-9 && (f - 0.5).abs() < 1e-9); // near-miss = 1 FP + 1 FN
}

Output on an Apple M4 Pro, one thread, release build: held-out entity F1 0.972 with the gazetteer, 1.000 without; 107 k to 183 k sentences/s over ten runs, typically 165 k. Templates are easy; real corpora score lower.

Complexity and constants

• Emissions O(tokens x features x tags); Viterbi O(tokens x tags^2). Weights 9 x 65,536 x 4 B = 2.3 MiB, cache-resident.
• Transformer: about 2 x parameters x subwords FLOPs; padding is the waste.

Efficiency

• Cost: toy about 6 µs per sentence (measured, one M4 Pro core; format! and hashing dominate). BERT-base, 30 subwords: order of 5 GFLOPs (my arithmetic: 2 x ~85 M parameters x 30). LLM: a priced network call.
• Levers:
  • Greedy instead of Viterbi/beam — comparable F1 (Ratinov and Roth) — strong local features.
  • Aho-Corasick gazetteers — one pass however long the list — drug, gene, sanctions lists.
  • Rules first for fixed-grammar IDs — fewer model calls — logs, KYC.
  • Length-sorted batches, distilled encoder — less padding (no figure) — bulk indexing.
• Naive → standard → state of the art:
  • LLM prompt per document, parse JSON · per-token price, a round trip.
  • Fine-tuned encoder with a token head · one forward pass, CoNLL-2003 test F1 92.8 (Devlin et al.).
  • Rules and lists, encoder, LLM for low confidence or new types · LLM cost on a small share.

Implementation notes

• Evaluate and redact on character offsets.
• Long documents: overlapping windows, then merge, or entities split at window edges.
• Tune PII for recall, not F1.

Pitfalls

• Gazetteer over-trust. Measured on the toy: F1 0.972 with a flag firing on every training city, 1.000 without; unseen cities lacking the flag were read as amounts.
• Train/test name overlap inflates scores.
• LLM-reported offsets are often wrong: search the returned string in the source, reject misses.
• Flat BIO loses nested entities (a bank inside "Bank of China Ltd").

Variants and alternatives

In production

• spaCy, Stanza, Flair, Hugging Face token classification, Presidio (PII): news indexing, clinical de-identification, KYC, biomedical mining, log redaction.
• Domain shift is the main operating cost: news-trained models drop on clinical, legal, biological or product text.

Implementations

Sources

• Tjong Kim Sang, De Meulder, Introduction to the CoNLL-2003 Shared Task: Language-Independent Named Entity Recognition, CoNLL 2003. paper
• Lafferty et al., Conditional Random Fields: Probabilistic Models for Segmenting and Labeling Sequence Data, ICML 2001. paper
• Ratinov, Roth, Design Challenges and Misconceptions in Named Entity Recognition, CoNLL 2009. paper
• Lample et al., Neural Architectures for Named Entity Recognition, NAACL 2016. paper · code
• Devlin et al., BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding, 2018 (CoNLL-2003 NER, Table 7). paper · code
• Zaratiana et al., GLiNER: Generalist Model for Named Entity Recognition using Bidirectional Transformer, 2023. paper · code