Scaling document classification to 100k+ labels

Scaling document classification to 100k+ labels | Databricks Blog Skip to main content

Summary

Mapping text to large taxonomies of 100,000+ labels, whether that's biomedical entity linking, vendor normalization, or company deduplication, is a common production problem where regex, trained classifiers, and direct LLM calls all struggle on cost, maintenance, and context limits.

Our solution pairs vector search with the Databricks AI Classify function: retrieve a shortlist of candidate labels per document, then let the AI Classify function pick from that shortlist instead of the full taxonomy.

Across three benchmarks spanning these use cases, SQL-native vector search plus AI Classify beat the best cost-efficient frontier model by five points of accuracy at roughly a hundredth of the token cost.

Across Databricks, thousands of customers build production workloads that map freeform text to normalized taxonomies of 100k+ labels. A few common use cases include: Biomedical entity linking. Clinical notes and research papers mention diseases, drugs, and procedures that must be matched to a concept in the Unified Medical Language System , a vocabulary with thousands of biomedical concept identifiers. Vendor normalization. To analyze spending patterns, financial companies normalize raw transaction strings like "SBUX #4471 SEATTLE WA" against a set of 100,000+ merchants. Company normalization. Enterprises deduplicate customer records across different internal systems by matching company descriptions against a set of 100,000+ companies.

Each of these use cases defines a large set of labels, also called a taxonomy , and every input must be mapped to one or more of its labels. Customers typically evaluate their production solution based on these criteria: Quality: Produce classifications that are accurate enough to drive downstream decisions. Cost: Keep the per-document classification cost low, especially at scale. Throughput: Process large workflows, often on the order of hundreds of thousands of documents, within a reasonable timeframe.

Meeting all three requirements makes production large taxonomy classification difficult, and we prioritize researching and building solutions to them at Databricks. Why existing approaches fall short Historically, organizations approached large taxonomy classification by relying on custom regex rules, keyword matching, and supervised machine learning classifiers. However, these approaches fall short on a few dimensions: Brittle pattern matching. Regex and keyword rules depend on exact matches, so they break on unpredictable formatting from real life data. Companies like YipitData needed to continuously update their regex patterns to address edge cases and taxonomy changes which are difficult to maintain at the scale of thousands of labels. Sparse and skewed ground truth. Real label distributions are long-tailed: a handful of labels cover most documents, while thousands appear rarely. Most customers lack ground truth documents for every label in their taxonomy, which makes supervised classifiers hard to train. A classifier cannot learn to classify labels that are not included in the training data, and class imbalance causes classifiers to over-predict commonly represented labels and under-predict rare labels. Taxonomy drift. Labels and their descriptions are constantly added, retired, and rewritten to account for new business use cases and to improve classification quality. On each new taxonomy version, regex patterns need to be updated and classifiers retrained and deployed. Recently, we've seen more customers use large language models (LLMs) for classification. With LLMs, customers no longer need to train models or maintain brittle regex patterns. However, at the scale of hundreds of thousands of labels, LLMs struggle to fit and reason over the taxonomy within their context window. With thousands of candidates in a single prompt, models also begin to hallucinate, returning labels that do not exist in the taxonomy at all. Passing the entire taxonomy for hundreds of thousands of documents to a frontier model is also costly at scale, even when taking advantage of prompt caching. How we approach large label classification We evaluate three different methods to find the best balance of cost and quality for large taxonomies: Method 1: Vector search The first method we test is using vector search to retrieve the best label given the input. We embed every label, with its description when present, and the input document using the Qwen3-Embedding-8B model, a top-ranked open-weight embedding model as of July 2026. We then score each label’s relevance to the document with a hybrid score that weighs both semantic and lexical similarity. The semantic score is determined by calculating the cosine similarity between the document and label embeddings; a higher cosine similarity means the label matches the document closely in meaning even if the text does not exactly match. The lexical score is computed with the BM25 algorithm , which weights each shared term by its inverse document frequency across the label set, so a generic verb like "use" that occurs in thousands of labels receives a low weight while a rare token like "ETL" receives a high one. Each search method returns its own ranked list of labels. We merge the two lists with Reciprocal Rank Fusion, which scores each label by its position in each list. We take the top k, sweeping k over 1, 5, 10, 20, 50, 100, and 200 to find the value with the best average accuracy and use the top-ranked label as the prediction. At the scale of hundreds of thousands of labels, embedding the taxonomy once and building an in-memory index is more cost effective than creating a hosted vector search index. Embedding 100k labels uses ~1.6 GB of storage assuming Qwen3-8B 4096-dim float32 embeddings and takes approximately 1-3 minutes to embed using the Databricks AI Query function . Then at the start of a workload, we create a vector search instance in memory, ingest the persisted embeddings, and use the index for all documents throughout the duration of the workload. You can find example code for our vector search implementation in this tutorial notebook . Method 2: Vector search + AI Classify The second method we test is a two-step workflow combining the vector search approach and the Databricks AI Classify function. AI Classify is a Databricks AI function that takes a document and a map of labels to descriptions and...