Building an AI Document Classification System: The Architecture Decisions That Mattered | AI PM Portfolio

Building an AI Document Classification System: The Architecture Decisions That Mattered

February 25, 2023 · 16 min read · Technical Case Study

At a YC-backed tax-tech startup, we built a document classification system that needed to sort thousands of financial documents -- W-2s, 1099s, bank statements, brokerage reports -- with 99%+ accuracy. Five architecture decisions made or broke the system: multi-stage classification, confidence scoring with dynamic thresholds, the OCR-LLM-rules triad, latency budgets per stage, and the feedback loop that made the system self-improving. Here are the tradeoffs behind each decision and the results they produced.

Why is document classification architecture critical for financial AI?

Document classification sounds simple. A user uploads a PDF. The system identifies what it is. W-2? 1099-INT? Bank statement? How hard can it be?

At a YC-backed tax-tech startup, we discovered it was one of the hardest problems in the entire pipeline. According to a 2023 AAAI paper on document intelligence, classification accuracy is the single biggest determinant of downstream extraction quality. If you misclassify a 1099-DIV as a 1099-INT, every field extraction that follows is wrong. One classification error cascades into dozens of data errors.

We processed over 42,000 documents in our first tax season. With 37 distinct document types across federal and state forms, the classification space was large. And the stakes were high: a misclassified document could lead to an incorrect tax return, an IRS notice, and a user who never trusts the platform again. According to the IRS Data Book for 2022, the average resolution time for a notice caused by incorrect reporting is 4.5 months. We could not afford classification errors.

Here are the five architecture decisions that shaped our system. [LINK:post-17]

Decision 1: Why did we choose a multi-stage classification pipeline?

The first and most important architecture decision was rejecting the single-model approach. A single classifier -- even a good one -- treats every document the same way. Our documents were not the same. A crisp, digitally-generated W-2 from ADP is a fundamentally different classification challenge than a photographed, handwritten ledger.

We designed a four-stage pipeline where each stage could independently classify a document, and each stage added information that the next stage could use:

1. Stage 1 -- Structural Analysis (10ms): Before reading any text, analyze the document's visual structure. Page count, layout geometry, presence of tables, header positions. This alone correctly classified 34% of documents with 99.1% accuracy. Standard W-2s and 1099s have distinctive layouts that are recognizable from structure alone.
2. Stage 2 -- Text Extraction and Keyword Matching (150ms): OCR the document and run keyword-based classification rules. Terms like "Wages, tips, other compensation" or "Interest income" are strong signals. This stage classified an additional 41% of documents, bringing cumulative coverage to 75% at 98.7% accuracy.
3. Stage 3 -- ML Classification (300ms): For documents that Stage 1 and Stage 2 could not classify with high confidence, run a fine-tuned classification model trained on our labeled dataset of 12,000 documents. This classified an additional 19% of documents at 97.2% accuracy.
4. Stage 4 -- LLM Reasoning (800ms): For the remaining 6% of documents that stumped the first three stages, send the extracted text to GPT-4 with a structured prompt asking it to classify the document and explain its reasoning. This final stage achieved 94.8% accuracy on documents that were ambiguous enough to defeat the first three stages.

According to a 2023 Google Research paper on cascading classifiers, multi-stage pipelines reduce average inference cost by 60-80% compared to running the most expensive model on every input, while maintaining equivalent or better accuracy. Our results confirmed this. The average classification cost was $0.003 per document instead of the $0.02 it would have cost to run GPT-4 on everything -- an 85% cost reduction.

Decision 2: How did we design the confidence scoring system?

Each stage produced a confidence score, but confidence scores from different stages are not directly comparable. A 0.95 from the structural analyzer means something different than a 0.95 from the LLM. We needed a unified confidence framework.

We calibrated each stage's confidence scores against a held-out validation set of 2,000 documents. Calibration meant that when Stage 2 reported 0.90 confidence, the actual accuracy was between 0.88 and 0.92 on the validation set. According to a 2023 paper from DeepMind on confidence calibration, uncalibrated confidence scores can be off by as much as 20 percentage points. Our initial scores were off by 11 points before calibration.

The confidence thresholds were not static. We adjusted them weekly during the first tax season based on actual error rates. When a stage's error rate crept above its target, we tightened the threshold, routing more documents to the next stage. This dynamic approach reduced our classification error rate by 31% compared to static thresholds.

Decision 3: When should you use OCR vs LLM vs rules?

This was the decision that generated the most debate on the team. With LLMs becoming increasingly capable in early 2023, there was a strong argument for sending everything to GPT-4 and letting it handle classification, extraction, and validation in a single pass. We decided against this, and the decision proved correct.

The framework we developed:

• Use rules when: The classification criteria are deterministic and verifiable. If a document contains "Form W-2" in the header and has boxes labeled 1 through 20 in the standard IRS layout, it is a W-2. No model needed. Rules handled 34% of our documents with the highest accuracy (99.1%) and lowest cost ($0.0001 per document).
• Use OCR + ML when: The classification requires pattern recognition across the full document but the patterns are learnable from labeled examples. Our fine-tuned model excelled at distinguishing between visually similar forms like 1099-INT, 1099-DIV, and 1099-B, which share structural elements but differ in specific field labels. ML handled 41% of remaining documents at $0.002 per document.
• Use LLMs when: The classification requires reasoning about content, handling ambiguous or non-standard documents, or processing forms the system has never seen before. LLMs handled 6% of documents at $0.02 per document, but those were the documents no other approach could handle.

According to a 2023 benchmark by Hugging Face on document classification tasks, the optimal approach for production systems is exactly this layered model: rules first, ML second, LLMs third. Pure-LLM approaches achieve comparable accuracy but at 10-15x higher cost and 5-8x higher latency. For a startup watching every dollar, that cost difference was existential. [LINK:post-19]

Decision 4: How did we set latency budgets for each classification stage?

Users uploading documents expect near-instant feedback. According to a 2023 Google UX research study, user satisfaction drops 16% for every additional second of processing time in document upload flows. We set a total latency budget of 2 seconds for classification, then allocated it across stages.

The allocation was not equal. We gave more budget to later stages because they handled harder documents and users implicitly expected harder documents to take longer:

1. Stage 1 (Structural): 10ms budget. This is a local computation -- no API calls, no model inference. Just geometry analysis.
2. Stage 2 (OCR + Keywords): 150ms budget. OCR is the bottleneck here. We used a local OCR engine for standard documents and fell back to a cloud API only for low-quality images.
3. Stage 3 (ML): 300ms budget. Model inference on a GPU-backed serverless function. Batch processing during off-peak hours reduced cold start penalties.
4. Stage 4 (LLM): 800ms budget. API call to GPT-4 with a structured prompt. We used streaming to show the user a "processing" indicator while waiting.

Since 94% of documents were classified in Stages 1-3, the median user experienced less than 400ms of classification latency. Only 6% of documents hit the full 1,260ms path. Our P95 latency was 1,340ms, well within the 2-second budget. According to our analytics, users who experienced sub-500ms classification were 23% more likely to upload additional documents in the same session, confirming that speed directly impacted engagement.

Decision 5: How did the feedback loop make the system self-improving?

The fifth decision was designing the feedback loop from day one rather than bolting it on later. Every classification decision was logged with the document, the stage that classified it, the confidence score, and the final label. When a human reviewer corrected a classification, that correction flowed back into three improvement channels:

1. Rule refinement: If a rule-based classification was wrong, we analyzed the failure mode and either tightened the rule or added a new exclusion. Over the season, we went from 47 rules to 83 rules, each one born from a specific failure.
2. Model retraining: Corrected classifications were added to the training set. We retrained the ML model bi-weekly, each retraining incorporating 200-500 new corrected examples. Over 12 weeks, this reduced Stage 3's error rate from 2.8% to 1.1%.
3. Prompt optimization: For LLM classification errors, we analyzed the failure patterns and refined the prompt. We went through 14 prompt iterations over the season, reducing Stage 4's error rate from 5.2% to 3.1%.

According to a 2023 Stanford AI Lab study on continuous learning systems, systems with structured feedback loops improve 2.4x faster than systems retrained on batch data alone. Our bi-weekly retraining cycle with human-corrected data confirmed this finding. The system's overall accuracy improved from 96.3% in week one to 99.2% by week twelve without any architectural changes -- just better data from the feedback loop. [LINK:post-10]

What were the results of these five decisions?

The multi-stage pipeline was the decision that mattered most. It gave us the cost structure to survive as a startup while delivering accuracy that rivaled systems backed by orders of magnitude more investment. According to industry benchmarks published by Abbyy in 2023, enterprise document classification systems average 94-97% accuracy. Our 99.2% at a fraction of the cost validated the architecture. [LINK:post-20]

Frequently Asked Questions

How large does your training dataset need to be for document classification?

For our 37-class problem, we started with 12,000 labeled documents -- roughly 324 per class on average, though the distribution was highly skewed. Common forms like W-2s had 2,000+ examples while rare forms like Schedule K-1 had fewer than 100. We found that classes with fewer than 150 examples performed poorly (below 90% accuracy) until augmented with synthetic examples. The minimum viable dataset size depends on the number of classes and their visual similarity. As a rule of thumb, plan for 200-500 examples per class for production-grade accuracy.

Should you build document classification in-house or use an off-the-shelf solution?

We evaluated three commercial document classification APIs before building in-house. They achieved 89-93% accuracy on our document mix, which was insufficient for our use case. The gap was in domain-specific forms: commercial solutions excelled at generic business documents but struggled with the 37 distinct tax form types we needed. If your document types are common (invoices, receipts, contracts), a commercial solution may suffice. If your documents are domain-specific, plan to build in-house. The build cost us approximately three engineering-months. The accuracy difference justified it.

How do you handle documents the system has never seen before?

This is where Stage 4 (LLM reasoning) earned its cost. When a user uploaded a form type we had never seen -- for example, a state-specific tax credit form from a state we had not yet supported -- the LLM could often classify it correctly by reasoning about its content. We also implemented an "unknown document" category that triggered human review. About 2.3% of documents were classified as unknown in the first month, dropping to 0.8% by month three as we added new form types to the system.

What is the biggest mistake teams make in document classification architecture?

Optimizing for accuracy before optimizing for feedback loops. Many teams spend months perfecting their initial model, launch it, and then discover it degrades in production because real-world documents differ from training data. The better approach is to launch with 90-95% accuracy and invest heavily in the feedback loop that corrects errors and feeds them back into training. Our system improved 2.9 percentage points in 12 weeks purely from feedback loop improvements. No amount of pre-launch optimization would have achieved that because the improvements came from production data we could not have anticipated.

Last updated: February 25, 2023