RAPT: Retrieval-Augmented Post-hoc Thresholding for Multi-Label Classification

TL;DR

RAPT is a versatile, post-hoc thresholding method that enhances multi-label classification accuracy by adaptively selecting label thresholds based on retrieval of similar cases, without retraining classifiers.

Contribution

It introduces a model-agnostic, retrieval-augmented wrapper that improves label set selection in industrial multi-label document understanding tasks.

Findings

RAPT outperforms static thresholding baselines across multiple benchmarks.

Achieves 0.87 Macro-F1 with metric learners, surpassing transformer variants and LLMs.

Requires significantly less inference time and GPU memory than LLM baselines.

Abstract

Industrial multi-label document understanding pipelines score candidate labels and threshold or rank them to form a label set per document. This early selection step directly affects the accuracy of downstream information extraction from the document, as well as the associated verification effort. In practice, OCR noise, label imbalance, instance-dependent label cardinality, and asymmetric error costs make global score thresholds brittle and hard to maintain as document formats evolve. We present RAPT, a deployment-oriented retrieval-augmented score thresholding wrapper, applied post-hoc to improve label set selection without retraining the underlying classifier. RAPT is a model-agnostic wrapper: any predictor that provides document representations for similarity search and per label confidence scores can be used, including metric learning encoders and fine-tuned transformer classifiers. For each query document, given a classifier's score vector, RAPT retrieves similar document thresholding situations (cases) and adapts the query's label set selection threshold using their outcomes. The adaptation selects the final label set by locally aggregating neighbour solutions (e.g. average label count, cutoff calibration). Evaluation compared multi-label classifiers (metric learners and transformers) combined with RAPT against global and label-wise thresholding baselines, and against few-shot LLMs. Across an industrial dataset and six public benchmarks, RAPT consistently outperformed global and label-wise static thresholding baselines. In the industrial setting, RAPT achieved its best predictive performance with metric learners, reaching 0.87 Macro-F1, while fine-tuned transformer variants on average achieved 0.775 Macro-F1, outperforming fewshot LLM baselines (K = 5) by 2x and requiring at least 115x less inference time and 13.5x less GPU memory.