Text Classification for Document AI

Short answer: classification is the cheapest accuracy you will buy in a document pipeline, because knowing what a document is lets every downstream step be narrow. And the technique that wins is not always the newest — a logistic regression over TF-IDF still beats an LLM on cost, latency and stability for high-volume, stable, well-labelled categories.

The instinct in 2026 is to prompt a large model for everything. Sometimes right, often wasteful, and occasionally worse.

Why classify before you extract

Running one general extraction model across every document type is the reliable way to get mediocre results everywhere. An invoice, a passport and a lab report share no structure.

Classify first and each downstream extractor becomes narrow enough to tune properly. You also get a clean place to route unknown types to a human rather than guessing — which is the difference between a system that degrades gracefully and one that silently produces nonsense on a document it has never seen.

It makes the pipeline extensible too. A new document type is a new label plus a new extractor, not a modification to a model everything else depends on. We covered the surrounding architecture in intelligent document processing.

The three families

Classical: TF-IDF or bag-of-words into linear models or gradient boosting. Unfashionable and frequently correct. Trains in seconds, runs in microseconds on CPU, costs nothing per document, and is completely deterministic — the same input gives the same output forever, which matters more than people admit when you are debugging a production incident. Needs labelled data and a reasonably stable vocabulary. For high-volume routing across a fixed set of categories, this is often the right answer and it will still be running unchanged in three years.

Fine-tuned transformers. A BERT-family encoder fine-tuned on your labels. Handles paraphrase and context that bag-of-words misses, needs meaningfully less labelled data than training from scratch, and runs cheaply enough at inference to sit in a high-volume path. The cost is a training pipeline, a model to version, and somewhere to serve it. This is the sweet spot when categories are subtle but volume is high.

Zero-shot and few-shot LLM. Describe the categories in a prompt and classify with no training data at all. Unbeatable for getting started, for long-tail categories with few examples, and for label sets that change often. The costs are real: per-document inference cost, latency measured in hundreds of milliseconds rather than microseconds, and non-determinism — the same document can classify differently across runs or model versions.

Which family fits depends on how many classes you have and how much labelled data exists, and it is one of the first decisions in a document AI build.

Choosing between them

Four questions settle it:

Do you have labels? No labels means starting zero-shot, whatever the end state.
What is the volume? At millions of documents, per-document LLM cost dominates everything and a small model pays for itself quickly.
How stable are the categories? Frequently changing label sets favour prompting, because retraining has a cost every time.
Does it need to be deterministic? Regulated or audited pipelines often need the same input to classify identically every time, and that argues for a classical or fine-tuned model.

The pattern we use most: bootstrap with an LLM, distil into a small model. Use zero-shot classification to label a few thousand documents, have a human correct a sample, then train a small fast classifier on that corrected set. You get LLM-quality labels at classical-model cost and latency, and you keep the LLM for the long tail the small model flags as uncertain.

Confidence is the part that makes it operable

A classifier returning a label and nothing else is hard to run in production, because you cannot tell the reliable predictions from the guesses.

Return calibrated probabilities and route on them. High confidence proceeds automatically; low confidence goes to a human. Then calibrate against real outcomes rather than intuition — measure how often a prediction at 0.8 is actually wrong, and set the routing threshold from that measurement. Teams typically set thresholds by feel and end up either drowning reviewers or passing errors silently.

Note that raw LLM token probabilities are not well-calibrated confidence, and treating them as such is a common mistake. If you need real calibration, a small model with a proper probability output is easier to trust.

Multi-label, hierarchy and the unknown class

Real document sets are messier than a flat single-label problem.

Multi-label. A document can be an invoice and a dispute. Forcing a single label loses information the downstream workflow needed.

Hierarchy. Coarse routing (financial vs clinical vs legal) then fine-grained within it. Easier to train, easier to extend, and it degrades sensibly — getting the coarse level right and the fine level wrong is far less damaging than the reverse.

The unknown class. Most important and most often missing. A classifier trained on eight types will confidently assign one of those eight to a ninth type it has never seen. An explicit “unknown” route, triggered by low confidence across all classes, is what stops a new document type quietly flowing into the wrong workflow for a month.

Measure it separately

Classification errors are cheap to fix and expensive to leave, because everything downstream inherits them. Track per-class precision and recall, not overall accuracy — an overall number looks fine while your rarest and most important class is failing.

Watch for drift. Vocabulary changes, suppliers change templates, new document types appear. A classifier that was excellent at launch degrades quietly, and the symptom is a rising exception rate that looks like an extraction problem.

The takeaway

Classify before you extract, and pick the technique on volume, label availability and determinism rather than novelty. Bootstrap with an LLM and distil into a small model when volume justifies it. Return calibrated confidence, add an explicit unknown class, and measure per-class rather than in aggregate.

EpochC builds intelligent document processing and OCR, KYC and identity verification, and the workflow automation around them. See the KYC OCR automation case study — 98% field accuracy on 70% less compute — or start a project.

Text classification techniques for document AI: what still works