# UniHDSA: Unified Relation Prediction for Hierarchical Document Structure Analysis > A research note on UniHDSA (Wang et al., 2025), a unified relation prediction framework that consolidates page-level and document-level document structure sub-tasks into a single label space, reporting the best published results on the Comp-HRDoc benchmark. - Authors: Hunter Heidenreich - Published: 2026-04-02 - Paper: Preprint (arXiv) - Paper: Published (Pattern Recognition 2025) - Code: UniHDSA (code + Comp-HRDoc benchmark) - Source: Bevaya Labs, --- ## TL;DR Wang et al. extend their earlier Detect-Order-Construct work by proposing UniHDSA, a unified relation prediction framework for hierarchical document structure analysis. Rather than running separate modules for detection, reading order, TOC extraction, and reconstruction, UniHDSA merges all sub-tasks at both the page-level and the document-level into a single label space handled by a shared relation prediction head. A multimodal Transformer system (vision backbone + BERT) built on this framework reports the best published results across all sub-tasks on the Comp-HRDoc benchmark and competitive results on DocLayNet. ## What kind of paper is this? **Dominant:** $\Psi\_{\text{Method}}$: The paper's central contribution is a new architectural principle (unified label space) and the multimodal Transformer system built on it. Most of the paper describes the page-level module, the document-level module, the type-wise query design, the loss functions, and extensive ablations verifying each design choice. This is a journal extension of DLAFormer (Wang et al., ICDAR 2024), which introduced unified relation prediction for page-level tasks only. UniHDSA extends that to document-level tasks and adds a language model for multimodal fusion. ## What is the motivation? Hierarchical Document Structure Analysis (HDSA) encompasses multiple interdependent sub-tasks: page object detection, text region detection, logical role classification, reading order prediction, table-of-contents extraction, hierarchical list extraction, and cross-page grouping. Prior methods fell into one of two camps: 1. **Single-task models**: Effective in isolation but unable to exploit cross-task interactions. 2. **Multi-branch unified systems** (e.g., Detect-Order-Construct): Each branch handles a distinct sub-task, sharing feature maps but not operating on a common representation. Sequential execution introduces cascading errors: an error in the detection stage propagates to the reading order stage, which then propagates into TOC extraction. The authors also note that multi-branch designs are hard to extend. Adding a new sub-task requires engineering a new branch and possibly a new decoding algorithm. UniHDSA addresses these problems by framing every sub-task at both the page-level and the document-level as a relation prediction problem and collapsing all relation labels into a single label space. A single prediction head then handles all sub-tasks concurrently, reducing the potential for cascading errors and making the system easier to extend. ## What is the novelty? ### Unified label space The core idea is to define a label matrix $M \in \mathbb{Z}^{H \times W}$, where rows correspond to source queries and columns to target queries. Each cell can take one of a small number of values encoding the relation type (or no relation). At the page-level, three relation types are defined: - **Intra-region**: reading order between text-lines within a text region, or a self-referential marker for single-line regions. - **Inter-region**: reading order between adjacent text regions, or semantic association between a text region and a graphical object (e.g., a caption pointing to a table). - **Logical role**: assignment of each text-line or graphical object to a predefined logical role class (title, section heading, paragraph, caption, etc.). At the document-level, two relation types are defined: - **Intra-region**: cross-page grouping (e.g., sub-tables on consecutive pages belonging to the same logical table, or paragraphs split across page boundaries). - **Inter-region**: hierarchical relationships between section headings or list items, supporting TOC and list extraction. This encoding means the same relation prediction head, with a relation prediction module and a relation classification module, handles all tasks at both levels. ### Relation prediction head Given a set of queries $\\{q\_1, \ldots, q\_H\\}$ (text-line and graphical object queries) and logical role queries $\\{q\_{H+1}, \ldots, q\_W\\}$, the relation prediction module computes: $$ f\_{ij} = FC^r\_q(q\_i) \circ FC^r\_k(q\_j), \quad i \le H,\ j \le W $$ $$ s\_{ij} = \begin{cases} \dfrac{\exp(f\_{ij})}{\sum\_{j'=1}^{H} \exp(f\_{ij'})}, & j \le H \\\\[6pt] \dfrac{\exp(f\_{ij})}{\sum\_{j'=H+1}^{W} \exp(f\_{ij'})}, & H < j \le W \end{cases} $$ The separate softmax normalizations enforce single-successor constraints for intra-region and inter-region relations independently of the logical role relations. The relation classification module then uses a bilinear classifier to predict the relation type for pairs with a predicted relation: $$ p\_{ij} = \text{BiLinear}(FC^c\_q(q\_i),\, FC^c\_k(q\_j)), \quad c\_{ij} = \text{argmax}(p\_{ij}) $$ ### Type-wise query initialization The page-level module is built on a Deformable DETR-like encoder-decoder. The authors introduce a type-wise query selection strategy: instead of using a binary foreground classifier to select the top-$K$ encoder features as object proposals, they use a multi-class classifier. The predicted category is then used to initialize the content query from a set of learnable per-category embeddings. This gives the decoder a semantically grounded starting point rather than a generic learnable token. Text-line queries are initialized directly from BERT embeddings, while predefined logical role queries use learnable embeddings with auxiliary supervision from the union boxes of all queries in that role class. ### Two-stage efficiency for long documents The framework divides the problem into page-level analysis (run per page) and document-level analysis (run once over the entire document, taking only the relevant entity types as input). The document-level Transformer encoder uses RoPE to encode the predicted reading order sequence. Because only section headings, list items, tables, and boundary paragraphs are passed to the document-level module, the quadratic cost of full cross-page attention is avoided. Training uses a random sampling strategy (6-8 consecutive pages per sample), which also introduces diversity. The authors report that models trained on 6-8 page samples produced promising TOC extraction results on a 44-page document in their qualitative evaluation. ### Loss functions The overall objective combines detection and relation prediction losses: $$ L\_{\text{overall}} = \lambda\_{\text{graphical}} \cdot (L\_{\text{graphical}}^{\text{enc}} + L\_{\text{graphical}}^{\text{dec}}) + \lambda\_{\text{relation}}^{\text{page}} \cdot L\_{\text{relation}}^{\text{page}} + \lambda\_{\text{relation}}^{\text{doc}} \cdot L\_{\text{relation}}^{\text{doc}} $$ All three $\lambda$ weights are set to 1.0. The graphical detection losses follow Deformable DETR: focal loss for classification and a linear combination of $L\_1$ and GIoU for box regression. Relation losses are softmax cross-entropy over both the relation prediction scores and the relation classification logits: $$ L\_{\text{rel}} = \frac{1}{N} \sum\_{l} \sum\_{i} L\_{\text{CE}}(s\_i, s\_i^{\ast}) $$ $$ L\_{\text{rel.cls}} = \frac{1}{N} \sum\_{l} \sum\_{i} L\_{\text{CE}}(c\_i, c\_i^{\ast}) $$ where $l$ indexes decoder layers (auxiliary loss at each layer) and $N$ is the number of queries. ## What experiments were performed? ### Datasets - **Comp-HRDoc**: 1,000 training and 500 test documents drawn from the HRDoc-Hard dataset. Four sub-tasks evaluated: page object detection (segmentation mAP), reading order prediction (REDS), table-of-contents extraction (Semantic-TEDS), and hierarchical structure reconstruction (Semantic-TEDS). OCR files provide pre-filtered text-line bounding boxes. - **DocLayNet**: 69,375 training pages and 6,489 test pages, 11 object classes, diverse document types. Evaluated with COCO-style box mAP at IoU $\\{0.50, 0.55, \ldots, 0.95\\}$ and mean F1 score at the same thresholds. ### Baselines On Comp-HRDoc: Mask2Former (detection only), Lorenzo et al. (reading order), MTD (TOC), DSPS Encoder (reconstruction with ground-truth reading order and bounding boxes provided), and the prior best-reported end-to-end system, Detect-Order-Construct (DOC). On DocLayNet: Faster R-CNN, Mask R-CNN, YOLOv5, SwinDocSegmenter, DINO, VSR, DOC, DLAFormer (a vision-only predecessor from the same group), and M2Doc. ### Ablations - Attention mask ablation: disabling interaction between text-line and graphical object queries in the page-level decoder to test the value of unified modeling. - Larger sampling window (14-16 pages vs. 6-8 pages) for document-level tasks. - Semantic feature ablation: comparing vision-only variants with the late-fusion multimodal variant across different backbone strengths (DOC-R18, DOC-R50, DLAFormer-R18, UniHDSA-R18, UniHDSA-R50). - Backbone comparison: ResNet-18, ResNet-50, ResNet-101, Swin-Tiny, InternImage-Tiny, InternImage-Small. ## What are the outcomes/conclusions? ### Comp-HRDoc end-to-end results UniHDSA outperforms DOC (the prior best end-to-end system) across all sub-tasks: | Sub-task | Metric | DOC-R18 | UniHDSA-R18 | UniHDSA-R50 | |---|---|---|---|---| | Page object detection | Seg. mAP | 88.1 | 90.9 | 91.2 | | Reading order (text) | REDS | 93.2 | 96.4 | 96.7 | | Reading order (graphical) | REDS | 86.4 | 90.6 | 91.0 | | TOC extraction | Micro-STEDS | 86.1 | 87.9 | 88.3 | | TOC extraction | Macro-STEDS | 87.9 | 89.5 | 88.8 | | Hierarchical reconstruction | Micro-STEDS | 83.7 | 88.0 | 88.9 | | Hierarchical reconstruction | Macro-STEDS | 83.7 | 87.8 | 88.6 | Reading order shows the largest gains: +3.2 points for text regions and +4.2 points for graphical regions versus DOC, despite DOC using a two-stage reading order prediction requiring additional parameters. ### DocLayNet page object detection UniHDSA-R50 achieves 87.0% mAP and 87.9% mean F1 score. This is competitive with M2Doc (87.3% mAP, 87.5% F1 after correcting for unreasonable post-processing). The authors note that M2Doc's early-fusion strategy allows one bounding box to receive multiple category predictions, which inflates AP but is problematic for downstream tasks such as markdown reconstruction. After correcting M2Doc's post-processing to enforce a single category per box, UniHDSA remains competitive. UniHDSA particularly outperforms at the strictest IoU threshold (0.95), suggesting the hybrid approach produces tighter bounding box predictions. ### Ablation findings - **Unified modeling of text and graphical elements**: Masking cross-attention between text-line and graphical object queries drops formula detection by 26.6 mAP points and table detection by 2.0 points, confirming that text-lines within graphical objects provide critical signal for top-down detection. - **Sampling window size**: Training on 14-16 pages versus 6-8 pages yields no significant improvement on the test set (which contains full documents of roughly 20 pages). The smaller window provides sufficient context while improving training efficiency. - **Semantic features**: Late fusion boosts semantically ambiguous categories (author, mail, affiliation, caption) substantially. However, for visually distinctive but semantically confusable categories (header vs. title), semantic features can hurt slightly. Gains from semantic features diminish as visual backbone quality increases. - **Backbone scaling**: ResNet-50 delivers the best overall balance on Comp-HRDoc. InternImage-Small, despite using a smaller sampling window due to GPU memory constraints, achieves 91.7% page object detection mAP, suggesting further gains are possible with sufficient compute. ### Relation prediction performance Page-level intra-region and logical role relations exceed 97% macro F1, while inter-region relations exceed 93%. Document-level inter-region relations also exceed 93%, validating TOC and hierarchical list extraction. Document-level intra-region relations (cross-page paragraph grouping) achieve around 87% F1, which the authors attribute to limited contextual connections across pages and content ambiguity in this task. ### Limitations - No suitable public datasets exist to evaluate cross-page table grouping and hierarchical list extraction in isolation. Without held-out test splits for these specific tasks, it is difficult to characterize model performance on them independently. - The vision and language backbones are pretrained separately, without joint visual-language pretraining. Features may not be optimally aligned for cross-modal reasoning. - The quadratic relation matrix may create scalability challenges for pages with very large numbers of text-lines or graphical regions. - The authors acknowledge a broader gap: there is no comprehensive benchmark and agreed-upon evaluation metrics for end-to-end document digitization, making fair comparison between specialist models and general multimodal large language models difficult. ## Reproducibility ### Models - **Vision backbone**: ResNet-18 or ResNet-50 (pretrained on ImageNet). Swin-Tiny, ResNet-101, InternImage-Tiny, and InternImage-Small also tested. UniHDSA-R18 and UniHDSA-R50 total approximately 150M and 162M parameters, respectively. - **Language model**: BERT-Base (for Comp-HRDoc) or LayoutXLM (for multilingual settings, mentioned but not primary in the experiments). - **Page-level Transformer**: 3-layer Deformable Transformer encoder and decoder; 8 attention heads; 256 hidden dim; 1024 FFN dim. Top-$K$ graphical object proposals: 50 (Comp-HRDoc), 100 (DocLayNet). - **Document-level Transformer**: 3-layer encoder; 8 heads; 256 hidden dim; 2048 FFN dim; RoPE positional encoding. - No pretrained weights for the full trained system are reported as publicly released. Configurations and evaluation code are available at [https://github.com/microsoft/CompHRDoc/tree/main/UniHDSA](https://github.com/microsoft/CompHRDoc/tree/main/UniHDSA). ### Algorithms - **Framework**: PyTorch v1.11. - **Optimizer**: AdamW; betas (0.9, 0.999); epsilon 1e-8. - **Comp-HRDoc**: mini-batch 1 (one sampled sub-document), 40 epochs. CNN backbone: lr 4e-5, weight decay 1e-2. BERT: lr 8e-5, weight decay 1e-2. All other parameters: lr 4e-4, weight decay 1e-2. - **DocLayNet**: mini-batch 4, 24 epochs. Same lr settings. Logit-adjusted softmax cross-entropy loss used to handle the long-tail distribution of text-lines inside graphical objects. - **Multi-scale training**: shorter side randomly chosen from $\\{320, 416, 512, 608, 704, 800\\}$; longer side capped at 1024. Test: shorter side 512. - **Document sampling**: 6-8 consecutive pages sampled per training example. - Text-line grouping in DocLayNet uses a heuristic merging step (grouping separated words/characters by horizontal and vertical distance relative to average character height) before feeding into the model. ### Data - **Comp-HRDoc**: 1,000 train / 500 test documents (HRDoc-Hard extension). Available at `https://github.com/microsoft/CompHRDoc` under MIT license. Provides pre-filtered OCR with text-lines inside graphical objects removed. - **DocLayNet**: 69,375 train / 6,489 test pages, IBM/CC-BY-4.0 license. OCR files have many fragmented text-lines requiring the heuristic grouping step described above. - Annotation for document-level tasks (TOC hierarchy, list hierarchy) is derived from the structural metadata of the source documents (academic PDFs and technical documents). ### Evaluation - **Page object detection**: COCO-style segmentation mAP (Comp-HRDoc) and box mAP + mean F1 at IoU 0.50:0.05:0.95 (DocLayNet). - **Reading order**: REDS (Comp-HRDoc), evaluated separately for text regions and graphical regions. - **TOC extraction and hierarchical reconstruction**: Semantic-TEDS (Micro and Macro), which measures tree-edit distance weighted by semantic content. - No error bars or significance tests are reported for the main comparison tables. The authors do note average performance standard deviations for UniHDSA-R18 (0.16) and UniHDSA-R50 (0.23) on Comp-HRDoc without specifying the number of runs. - M2Doc comparisons required replicating M2Doc with corrected post-processing (one box, one category). The corrected numbers differ from the M2Doc paper's reported figures. ### Hardware - 16 Nvidia Tesla V100 GPUs (32 GB each) for all experiments. - Training time and inference latency are not reported. ## BibTeX ```bibtex @article{wang2025unihdsa, title={UniHDSA: A Unified Relation Prediction Approach for Hierarchical Document Structure Analysis}, author={Jiawei Wang and Kai Hu and Qiang Huo}, journal={Pattern Recognition}, volume={165}, pages={111617}, year={2025}, doi={10.1016/j.patcog.2025.111617} } ```