Retrieval Augmented Classification for Confidential Documents

2026-04-09

AI Summary

This study addresses the challenges of classifying confidential documents in dynamic document streams, particularly class imbalance and the infeasibility of embedding sensitive information into model weights. To overcome these issues, the authors propose a novel retrieval-augmented classification (RAC) approach that leverages an external vector store for similarity-based matching to make classification decisions, thereby preventing sensitive content from being固化 in model parameters and enabling rapid updates without retraining. The work introduces the first RAC pipeline and evaluation framework tailored for confidential document classification. Experiments on the WikiLeaks diplomatic cables dataset demonstrate that RAC achieves approximately 96% accuracy (with F1 scores up to 94%) under both original and balanced settings, significantly outperforming fine-tuning baselines and exhibiting greater robustness in context-constrained and class-imbalanced scenarios. Practical design guidelines for managed deployments are also provided.

Technology Category

• Machine Learning: Multi-class/Multi-label Learning & Extreme Classification
• Natural Language Processing: Safety and Robustness
• Data Mining & Knowledge Management: Conversational Systems for Recommendation & Retrieval

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

Unauthorized disclosure of confidential documents demands robust, low-leakage classification. In real work environments, there is a lot of inflow and outflow of documents. To continuously update knowledge, we propose a methodology for classifying confidential documents using Retrieval Augmented Classification (RAC). To confirm this effectiveness, we compare RAC and supervised fine tuning (FT) on the WikiLeaks US Diplomacy corpus under realistic sequence-length constraints. On balanced data, RAC matches FT. On unbalanced data, RAC is more stable while delivering comparable performance--about 96% Accuracy on both the original (unbalanced) and augmented (balanced) sets, and up to 94% F1 with proper prompting--whereas FT attains 90% F1 trained on the augmented, balanced set but drops to 88% F1 trained on the original, unbalanced set. When robust augmentation is infeasible, RAC provides a practical, security-preserving path to strong classification by keeping sensitive content out of model weights and under your control, and it remains robust as real-world conditions change in class balance, data, context length, or governance requirements. Because RAC grounds decisions in an external vector store with similarity matching, it is less sensitive to label skew, reduces parameter-level leakage, and can incorporate new data immediately via reindexing--a difficult step for FT, which typically requires retraining. The contributions of this paper are threefold: first, a RAC-based classification pipeline and evaluation recipe; second, a controlled study that isolates class imbalance and context-length effects for FT versus RAC in confidential-document grading; and third, actionable guidance on RAC design patterns for governed deployments.

Problem

• confidential document classification
• class imbalance
• information leakage
• retrieval-augmented classification
• security-preserving machine learning

Innovation

• Retrieval Augmented Classification
• Confidential Document Classification
• Class Imbalance Robustness
• Low-Leakage AI
• Vector Store Grounding