Loading…
Scaling document classification to 100k+ labels
Jane Zhang, Arnav Singhvi
- Source
- Databricks
- Published
- Added to Yomu
Summary
Large-taxonomy classification must map freeform text to 35,000–100,000 labels while meeting quality, cost, and throughput requirements. The benchmark compares vector search, a two-step vector-search-plus-AI-Classify workflow, and direct frontier-model calls across transaction, company, and biomedical datasets, using hybrid semantic and BM25 retrieval with Reciprocal Rank Fusion for the search stages. The combined workflow performed best, reaching 0.81 average accuracy versus 0.76 for Gemini 3.5 Flash, while costing roughly one-hundredth as much per document; its strongest average setting shortlisted the top 20 labels. Vector search alone was nearly 100 times cheaper than the workflow but scored more than 20 points lower, while full-taxonomy frontier calls could exceed context limits. Taxonomy updates can be handled by embedding and ingesting changed labels rather than retraining or redeploying classifiers.
Context
Production classification workloads must map freeform text to taxonomies containing tens or hundreds of thousands of labels. The evaluation targets three competing requirements: accurate classifications for downstream decisions, low per-document cost, and sufficient throughput for large document volumes. Existing regex, keyword, supervised-classifier, and direct LLM approaches face brittleness, sparse or skewed training data, taxonomy drift, context limits, hallucinations, or high cost.
Approach / What changed
The benchmark compares vector search, vector search followed by the Databricks AI Classify function, and direct frontier-model calls across Transactions, Companies, and MedMentions datasets. Vector search embeds labels and documents with Qwen3-Embedding-8B, combines cosine similarity with BM25 using Reciprocal Rank Fusion, and tunes the shortlist size. The combined workflow passes shortlisted labels to AI Classify, while direct calls provide the full taxonomy when context permits.
Takeaways
- The vector-search-plus-AI-Classify workflow achieved 0.81 average accuracy across the three datasets, compared with 0.76 for the next-best direct model, Gemini 3.5 Flash.
- Shortlisting the top 20 labels from vector search produced the best average quality for the AI Classify workflow; vector search alone was nearly 100 times cheaper but scored more than 20 points lower.
- Embedding the taxonomy once and maintaining an in-memory index avoids recurring taxonomy embedding costs, and changed labels can be re-embedded and ingested without classifier retraining or redeployment.