Technical Summary: DocHRL – A Hierarchical Reinforcement Learning Framework for Cost-Optimised Document Classification

1. Problem Statement

Current real-world document classification pipelines typically apply a fixed sequence of models to every incoming document, regardless of the document's complexity or type. This static approach leads to significant inefficiencies:

• Resource Misallocation: Simple documents are over-processed with expensive models, while difficult or ambiguous documents may not receive sufficient scrutiny.
• Operational Cost: The total operational cost is fixed and independent of document type. Over large volumes, small per-document inefficiencies accumulate into substantial costs.
• Suboptimal Trade-offs: Existing methods often treat computational cost and misclassification cost as separate concerns or ignore operational costs entirely, failing to jointly optimize for accuracy and economic efficiency.

The paper addresses the need for an adaptive system that dynamically selects the most cost-effective classification policy on a per-document basis, balancing inference costs, misclassification penalties, and human review costs.

2. Methodology

2.1. Formal Cost Model

The authors define a formal cost model for document classification pipelines comprising three components:

• $C_i$ (Inference Cost): The computational cost of a single model forward pass.
• $C_f$ (Cost of Failure): The expected downstream business loss from a misclassification.
• $C_h$ (Human Labeller Cost): The fully-loaded cost of manual expert review.

The paper analyzes four pipeline architectures:

1. Standalone Classifier: Cost is $C_i + E_r \cdot C_f$, where $E_r$ is the error rate.
2. Ensemble Pipeline: Cost aggregates inference costs of multiple models plus the ensemble's error rate.
3. Human-in-the-Loop (HIL) Pipeline: Uses entropy-based escalation. If a classifier's output entropy exceeds a threshold, the document is routed to a human. The cost includes inference, human review, and residual human error.
4. DocHRL Pipeline: Unlike the static thresholds of HIL, DocHRL learns a policy to select the cheapest sufficient set of actions for each specific document instance.

2.2. The DocHRL Framework

DocHRL formulates document classification as a Markov Decision Process (MDP) solved via Hierarchical Reinforcement Learning (HRL).

• State Representation ($s_t$): The agent state includes:

  • $v$: A frozen vision embedding of the document.
  • $P_t$: A probability grid accumulating soft-label predictions from invoked actions.
  • $O_t$: An OCR grid containing text embeddings from OCR sources.
  • $u_t$: A binary mask tracking used actions to prevent redundancy.

• Policy Hierarchy:

  • Top-Level Policy ($\pi_{top}$): Selects a broad option from a set including classify, llm, ocr, human, or specific terminal class labels.
  • Sub-Policies ($\pi_{\theta}$): Dedicated policies for each non-terminal option that select specific tools (e.g., a specific classifier model, a specific OCR engine, or a specific LLM).

• Reward Function: The reward is the negative total expected cost of a classification episode.

  • Costs are incurred at each step based on the action (inference cost of the policy or execution cost of the tool).
  • A failure penalty ($C_f$) is added if a terminal class prediction is incorrect.
  • The total episode cost is the sum of all step costs, and the reward is its negation, propagated via discounted returns.

• Training Algorithm:

  • Optimizer: Proximal Policy Optimization (PPO).
  • PPO Variant: Replaces standard hard clipping with a smooth tanh-gated trust region ($\beta \tanh(r_t - 1) + 1 + \alpha(r_t - 1)$) to ensure non-zero gradients for extreme ratios.
  • AdaPolyTan: A novel adaptive return normalization technique that blends a polynomial and a tanh transform. Parameters adapt based on the mean return magnitude to stabilize value learning across a wide dynamic range of costs.
  • Value Head: Bounded via a $0.5 \cdot \tanh$ activation to match the normalized return range.

3. Key Contributions

1. Unified Cost Objective: A formal cost model unifying inference, failure, and human labelling costs into a single objective for four distinct pipeline architectures.
2. Hierarchical RL Framework: An adaptive, cost-aware classification framework trained end-to-end using a novel PPO variant (tanh-gated trust region, bounded value head, and AdaPolyTan normalization).
3. Empirical Validation: Comprehensive evaluation on the RVL-CDIP benchmark demonstrating simultaneous improvements in classification performance and operational efficiency compared to static baselines.

4. Experimental Results

The framework was evaluated on the RVL-CDIP dataset (400,000 grayscale images, 16 classes).

• Cost Efficiency:

  • DocHRL achieved an average per-document cost of 2.74 normalized units.
  • This represents a 69% reduction compared to the best standalone classifier (EfficientNet-V2-M, cost 8.74) and a 65% reduction compared to the best ensemble (7.74).
  • It also outperformed the optimal entropy-based Human-in-the-Loop configuration (cost 3.085) by 11.1%.

• Classification Performance:

  • DocHRL achieved a Macro F1 of 0.973 and Macro Accuracy of 97.27%.
  • This outperforms the best standalone classifier (EfficientNet-V2-M, F1 0.913) by 6.0 percentage points and the best ensemble (CNN Ensemble, F1 0.923) by 5.0 percentage points.
  • It achieved perfect accuracy (1.000) on email, questionnaire, and resume classes.
  • Significant gains were observed on difficult classes like form (0.928 vs 0.854 baseline) and scientific report (0.935 vs 0.820 baseline).

• Adaptive Behavior:

  • The agent learned to route easy, high-confidence documents (e.g., email) directly to cheap classifiers.
  • For ambiguous documents (e.g., form, scientific report), the agent dynamically invoked additional evidence (OCR, LLMs, or human review), incurring higher costs only where necessary to maintain high accuracy.

5. Significance and Claims

The paper claims that cost-aware reinforcement learning is a practical and principled approach to building economically efficient document understanding systems. The primary significance lies in demonstrating that accuracy and cost efficiency need not be traded against one another when the routing policy is learned rather than fixed.

By treating document classification as a sequential decision problem with a reward signal explicitly capturing the negative total expected cost, DocHRL learns to allocate computational resources in proportion to the difficulty of each document instance. The authors position DocHRL as a proof-of-concept for scenarios where heterogeneous ensembles (including LLMs and human review) are already operationally justified, showing that hierarchical RL can optimize the total expected cost while improving classification performance.

6. Limitations and Future Directions

The authors modestly frame the work as a proof-of-concept and identify several limitations:

• Fixed Failure Cost: The current model assumes a fixed, domain-independent failure cost ($C_f$). Future work could learn class-specific costs from downstream feedback to handle regulated domains (e.g., healthcare, legal).
• Simplified Human Model: The human labeler is modeled as a stochastic function with a fixed error rate, ignoring realistic behaviors like fatigue or learning effects.
• Scalability: The hierarchical PPO training loop requires multiple model invocations per document. Scaling to larger pools or high-volume pipelines may require off-policy RL or experience replay.
• Benchmark Constraints: Evaluation was conducted on the balanced RVL-CDIP benchmark; performance under real-world class imbalance and distribution shift remains to be validated.