Contrastive Training Improves Zero-Shot Classification of Semi-structured Documents
Year: 2023

Abstract:
We investigate semi-structured document classification in a zero-shot setting. Classification of semi-structured documents is more challenging than that of standard unstructured documents, as positional, layout, and style information play a vital role in interpreting such documents. The standard classification setting where categories are fixed during both training and testing falls short in dynamic environments where new document categories could potentially emerge. We focus exclusively on the zero-shot setting where inference is done on new unseen classes. To address this task, we propose a matching-based approach that relies on a pairwise contrastive objective for both pretraining and fine-tuning. Our results show a significant boost in Macro F1 from the proposed pretraining step in both supervised and unsupervised zero-shot settings.

heading: 1 Introduction
text: Textual information assumes many forms ranging from unstructured (e.g., text messages) to semistructured (e.g., forms, invoices, letters), all the way to fully structured (e.g., databases or spreadsheets). Our focus in this work is classification of semi-structured documents. A semi-structured document consists of information that is organized using a regular visual layout, and includes tables, forms, multi-columns, (nested) bulleted lists, and that is either understandable only in the context of its visual layout or that requires substantial more work to understand without the visual layout. Automatic processing of semi-structured documents comes with a unique set of challenges including a non-linear text flow (Wang et al., 2021), layout inconsistencies, and low-accuracy optical character recognition. Prior work has shown that integrating the two-dimensional layout information
∗Work done while at AWS AI Labs †Corresponding author
of such documents is critical in models for analyzing such documents (Xu et al., 2020, 2021; Huang et al., 2022; Appalaraju et al., 2021). Due to these challenges, methods for unstructured document classification, such as static word vectors (Socher et al., 2013) and standard pretrained language models (Devlin et al., 2019; Reimers and Gurevych, 2019; Liu et al., 2019) perform poorly with semi-structured inputs as they model text in a one-dimensional space and ignore information about document layout and style (Xu et al., 2020).
Past work on semi-structured document classification (Harley et al., 2015; Iwana et al., 2016; Tensmeyer and Martinez, 2017; Xu et al., 2020, 2021) has focused exclusively on the full-shot setting, where the target classes are fixed and identical across training and inference, neglecting the zero-shot setting (Xian et al., 2018), which requires generalization to unseen classes during inference.
Our work addresses zero-shot classification of semi-structured documents in English using the matching framework, which has been used for many tasks on unstructured text (Dauphin et al., 2014; Nam et al., 2016; Pappas and Henderson, 2019; Vyas and Ballesteros, 2021; Ma et al., 2022). Under this framework, a matching (similarity) metric between documents and their assigned classes is maximized in a joint embedding space. We extend this matching framework with two enhancements. First, we use a pairwise contrastive objective (Rethmeier and Augenstein, 2020; Radford et al., 2021; Gunel et al., 2021) that increases the similarity between documents and their ground-truth labels, and decreases it for incorrect pairs of documents and labels. We augment the textual representations of documents with layout features representing the positions of tokens on the page to capture the twodimensional nature of the documents. Second, we propose an unsupervised contrastive pretraining procedure to warm up the representations of documents and classes. In summary, (i) we study
7499
zero-shot classification of semi-structured documents, which, to the best of our knowledge, has not been explored before. (ii) we use a pairwise contrastive objective to both pretrain and fine-tune a matching model for the task. This technique uses a layout-aware document encoder and a regular text encoder to maximize the similarity between documents and their ground-truth labels. (iii) Using this contrastive objective, we propose an unsupervised pretraining step with pseudo-labels (Rethmeier and Augenstein, 2020) to initialize document and label encoders. The proposed pretraining step improves F1 scores by 9 and 19 points in supervised and unsupervised zero-shot settings respectively, compared to a setup without this pretraining.

heading: 2 Approach
text: This section describes our proposed architecture (§ 2.1), pretrained model (§ 2.2), as well as the contrastive objective used for pretraining (§ 2.3) and fine-tuning (§ 2.4).

heading: 2.1 Model
text: Our goal is to learn a matching function between documents and labels such that similarity between a document and its gold label is maximized compared to other labels, which can be seen as an instance of metric learning (Xing et al., 2002; Kulis et al., 2012; Sohn, 2016). This requires encoding documents and class names1 into a joint documentlabel space (Ba et al., 2015; Zhou et al., 2019; Chen et al., 2020; Hou et al., 2020). In this work, documents and class names are of different nature— documents are semi-structured (§ 1), while class names are one or two-word fragments of text.
We use two encoders to account for this difference: a document encoder Φdoc suitable for semistructured documents, and a label (class) encoder Φlabel suitable for the natural language representations of the class labels. Φlabel is simply a vanilla pretrained BERTBASE model (Devlin et al., 2019). Φdoc, as in prior work (Xu et al., 2020; Lockard et al., 2020), is a pretrained language model that encodes the text and the layout of the document using coordinates of each token. The next section explains this model, LayoutBERT, in detail. We choose this model for its simplicity, but our proposed approach can be combined with more sophisticated
1We use class names as the natural language representation of a class, but more descriptive representations can be used if available (e.g. dictionary definitions) (Logeswaran et al., 2019)
document encoders that incorporate layout and visual information in different ways (Huang et al., 2022; Xu et al., 2021; Appalaraju et al., 2021).

heading: 2.2 LayoutBERT
text: LayoutBERT is a 6-layer Transformer based on BERTBASE (Devlin et al., 2019) and is pretrained using masked language modeling on a large collection of semi-structured documents (§ 3). Unlike prior work, LayoutBERT has a simpler architecture that decreases model footprint while maintaining accuracy. Specifically, there are three main architectural differences between LayoutBERT and LayoutLM, which is the most comparable architecture in the literature (Xu et al., 2020): (a) LayoutLM uses 12 transformer layers while LayoutBERT uses only 6 layers (b) LayoutLM uses four positions per token, namely upper-left and bottom-right coordinates, while LayoutBERT use only two positions viz. the centroid of the token bounding box. (c) Unlike LayoutLM, LayoutBERT does not use an image encoder to obtain CNN-based visual features.2

heading: 2.3 Unsupervised Contrastive Pretraining
text: Φlabel and Φdoc are models that have been pretrained independently. To encourage these models to produce similar representations for documents and their labels, we continue pretraining Φlabel and Φdoc via an unsupervised procedure based on a pairwise contrastive objective. The unsupervised objective can learn from large amounts of unlabeled semi-structured documents. This also allows us to directly use the pretrained encoders in an unsupervised zero-shot setting (§ 3.3.1).
2The results in Xu et al. (2020) show that image features are not always useful. To keep things simple, we do not include the CNN component in our model.
Since we do not assume access to ground truth labels for this step, our pretraining procedure relies solely on self-supervision via pseudo-labels (Rethmeier and Augenstein, 2020). These pseudo-labels are generated by sampling a continuous block of tokens from the document with a length drawn from a shifted geometric distribution. A pseudo-label extracted from a document is treated as the positive label for that document and is encoded using Φlabel.
We now describe our contrastive objective which is based on the multi-class n-pair loss (Sohn, 2016; Radford et al., 2021). Let B be a training batch that consists of training documents D and their pseudo-labels L, such that D = (d1, d2, ..., d|B|) and L = (l1, l2, ..., l|B|). Let Φdoc and Φlabel be the document and label encoders, respectively. We start by encoding each document and pseudolabel in the batch and then computing a matching matrix MB ∈ R|B|×|B| of pairwise dot products between every document-label pair, such that MBij = Φlabel(li)
T · Φdoc(dj). Our objective is to increase the value of diagonal elements Mij , where i = j, as compared to all other elements. More precisely, the loss function for a batch is a symmetric loss, LB , that can be expressed with the equation:
LB = 1 2 [LBrow + LBcol]. (1)
Hhere LBrow and LBcol are the per-batch row-wise and column-wise losses, respectively, with
LBrow = |B|∑
i=1
 − log(exp(MBii )) + log( |B|∑
j=1
exp(MBij ))   .
(2)
The first term in Eq. 2 maximizes the diagonal elements, while the second term minimizes the offdiagonal elements. The column-wise loss is the same with i and j swapped. We directly optimize the raw dot products rather than cosine similarity as we observed dot-products to perform much better, which also agrees with Karpukhin et al. (2020).

heading: 2.4 Contrastive Fine-tuning
text: For the supervised zero-shot setting (§ 3.3.2), we fine-tune the model using the same objective as the pretraining step (Equation 1), except that the labels L = (l1, l2, ..., l|B|) for a batch B are ground-truth labels and not pseudo-labels.

heading: 3 Experiments and Results

heading: 3.1 Data
text: We evaluate our approach on the RVL-CDIP dataset (Harley et al., 2015), which consists of 400K documents balanced across 16 classes such as letter, advertisement, scientific report, form, etc. Since zero-shot performance can vary depending on which classes are used for train and test, we follow previous work (Ye et al., 2020) and create four zero-shot splits of the data with non-overlapping test classes. Thus, each split has 8 training classes (200K documents), 4 validation classes (100K documents), and 4 test classes (100K documents).3
Our document encoder is pretrained on documents from CommonCrawl (see Appendix B for more details).4 While this pretraining corpus is different from the one used for LayoutLM, our objective is not to compare directly with this model but to explore zero-shot classification. Our contrastive pretraining corpus consists of 800K documents sampled from this pretraining corpus. We first sample l ∼ Geometric( 120), and then sample a block of l tokens from each document to obtain a pseudo-label for that document. We run contrastive pretraining for 50K steps with batch size of 256.

heading: 3.2 Experimental Setup
text: LayoutBERT is a 6-layer model initialized using BERTBASE weights and further pretrained using the MLM loss with layout information for 50K steps with a batch size of 2048 and a peak learning rate of 10−4. Unlike LayoutLM, where the extra position embeddings are initialized from scratch, we initialize them from BERT positional embeddings, which we found to speed up convergence. We used dynamic subtoken masking (Liu et al., 2019) with pmask = 0.15 and preplace = 0.80.
The representation of the [CLS] token is used as the encoding of input documents and an affine layer with a dimension of 768 is applied to the output of both encoders. We fine-tune the matching model on the data from the train classes for 30 epochs with a batch size of 40 and a learning rate of 3× 10−5. The model with the best macro F1 on the validation set is used for evaluation on the held out test set.

heading: 3.3 Results
text: We experiment with two settings — unsupervised zero-shot, and supervised zero-shot. In the former,
3The exact classes used for each split are in Appendix A. 4https://commoncrawl.org/
no fine-tuning is involved and all models are directly used for inference. In the latter, all models are fine-tuned on data from classes different than those present in the test set. Thus, the latter is strictly more challenging.

heading: 3.3.1 Unsupervised Zero-shot
text: We start with the more challenging unsupervised setup and compare three models (Table 1). The first model uses a vanilla pretrained BERTBASE as both the document and label encoders. The second model replaces the BERTBASE document encoder with LayoutBERT model. For these two models, we remove the affine layer after both encoders (§ 3.2) since in the absence of pretraining/finetuning, they will not be trained. The third model uses the same components as the second model but is pretrained using the unsupervised contrastive loss (§ 2.3).
The results yield three key observations. First, the vanilla BERT model performs the worst with an F1 score of 13. This is unsurprising as BERT does not capture any layout information. Second, the value of layout information can be verified by replacing the BERTBASE document encoder with LayoutBERT. This improves the average F1 by ~8 points. Finally, contrastive pretraining (CPT) is critical to produce better initialization for the encoders and it improves the average performance of the previous model by ~19 F1 points.

heading: 3.3.2 Supervised zero-shot
text: Next, we turn to the supervised zero-shot setup, where models are finetuned on data from classes different than those in the test set. We only experiment with the LayoutBERT (doc), BERT (label) setup since it performed the best in unsupervised
settings. Table 2 shows the Macro F1 with our in-batch contrastive training objective as well as a standard cross-entropy loss (Dauphin et al., 2014; Ye et al., 2020). We also show the fine-tuning performance with contrastive pretraining (§ 2.3).
We observe that the in-batch contrastive objective yields comparable F1 to the cross-entropy loss on average (with and without pretraining). However, the in-batch loss also has higher variance across different runs compared to the cros-entropy loss,5 possibly due to the stochastic nature of inbatch contrastive training. Crucially, though, we observe a strong F1 boost in almost all cases with contrastive pretraining, and in some cases as much as ~21 F1 points. This reemphasizes the importance of pretraining in producing similar representations for related documents and labels.
Finally, comparing Tables 1 and 2 shows that the zero shot performance is better in the unsupervised case than the supervised case. This is likely due to the fact that in the latter, the model is fine-tuned towards a specific type of documents (i.e. those present in the training/validation) classes, which hinders generalization to unseen inference classes. More sophisticated approaches (Finn et al., 2017; Nichol et al., 2018) can potentially improved the supervised setup, but we leave this to future work.

heading: 4 Conclusion
text: This work explores zero-shot classification of semistructured documents. We proposed two contrastive techniques for pretraining and fine-tuning
5Tables 4 and 5 in Appendix C show means and standard deviations with three random seeds. Experiments with more random seeds did not yield any meaningful differences.
of a matching model. Our fine-tuning objective showed comparable results to the standard crossentropy loss used widely in the literature and our contrastive pretraining significantly boosted zeroshot F1 in supervised and unsupervised scenarios.

heading: 5 Limitations
text: The current work is an initial attempt at studying the problem of zero-shot classification of semistructured documents. There are two key aspects that this work does not cover and we encourage future work to explore.
First, as pointed out in §2.1, we choose LayoutBERT as our document encoder, Φdoc. This work does not experiment with the variety of encoding strategies in the literature that combines textual, visual, and layout information (Appalaraju et al., 2021; Xu et al., 2021; Huang et al., 2022). It is likely that richer document representations derived from these diverse encoders will further push the limits of zero-shot classification when combined with our proposed unsupervised contrastive pretraining procedure.
Second, results in this paper are on a single dataset, i.e. the RVL-CDIP dataset. While we mitigate this to a large extent by creating four nonoverlapping test splits (see §3.1 and Appendix A), results on more datasets might yield more useful insights. In practice, the lack of datasets for this task (of semi-structured document classification) is what makes this exploration difficult and might require creation of new resources

heading: Acknowledgements
text: We thank our colleagues at AWS AI Labs and the ACL reviewers who have helped improve the paper through their feedback.

heading: A Data Splits
text: As stated in section 3, we split the RVL-CDIP dataset into four splits with non-overlapping test classes. Table 3 shows the classes used in each split.

heading: B Pre-training data from Common Crawl
text: We build our pre-training corpus by first extracting all documents from CommonCrawl with a ‘.pdf‘ extension. We then remove duplicate documents based on the MD5 hash using fdupes.6. The resulting documents are then passed through PDFPLUMBER7 to extract both the text as well as the co-ordinates of the tokens in the documents, and any documents that cannot be processed byPDFPLUMBER are discarded. We analyzed a sample of the crawled documents and found a large amount of structured information in the documents, so we use all documents at this stage without additional filtering. This leaves us with 2.3 million documents with approximately 850 million tokens.

heading: C Supervised Zero-shot Results
text: Tables 4 and 5 shows the full results of the supervised zero-shot finetuning with macro F1 means and standard deviations across three different runs. While in-batch contrastive fine-tuning outperforms the standard loss in many cases, we can see that, in general, the contrastive loss exhibits higher F1 variance. For example, in Table 4, the standard deviation when evaluating on the test set of the split II is 10.28, which is very high.
6https://github.com/adrianlopezroche/fdupes 7https://github.com/jsvine/pdfplumber
ACL 2023 Responsible NLP Checklist

heading: A For every submission:
text: 3 A1. Did you describe the limitations of your work?
5
A2. Did you discuss any potential risks of your work? Not applicable. Left blank.
3 A3. Do the abstract and introduction summarize the paper’s main claims? 1
7 A4. Have you used AI writing assistants when working on this paper? Left blank.
B 3 Did you use or create scientific artifacts? 3.1

heading: 3 B1. Did you cite the creators of artifacts you used?
text: 3.1
7 B2. Did you discuss the license or terms for use and / or distribution of any artifacts? Left blank.

heading: 3 B3. Did you discuss if your use of existing artifact(s) was consistent with their intended use, provided
text: that it was specified? For the artifacts you create, do you specify intended use and whether that is compatible with the original access conditions (in particular, derivatives of data accessed for research purposes should not be used outside of research contexts)? 3.1
7 B4. Did you discuss the steps taken to check whether the data that was collected / used contains any information that names or uniquely identifies individual people or offensive content, and the steps taken to protect / anonymize it? Data used as in prior work
7 B5. Did you provide documentation of the artifacts, e.g., coverage of domains, languages, and linguistic phenomena, demographic groups represented, etc.? Data used as in prior work
3 B6. Did you report relevant statistics like the number of examples, details of train / test / dev splits, etc. for the data that you used / created? Even for commonly-used benchmark datasets, include the number of examples in train / validation / test splits, as these provide necessary context for a reader to understand experimental results. For example, small differences in accuracy on large test sets may be significant, while on small test sets they may not be. Left blank.
C 3 Did you run computational experiments? Section 4

heading: 3 C1. Did you report the number of parameters in the models used, the total computational budget
text: (e.g., GPU hours), and computing infrastructure used? 3.2
The Responsible NLP Checklist used at ACL 2023 is adopted from NAACL 2022, with the addition of a question on AI writing assistance.

heading: 3 C2. Did you discuss the experimental setup, including hyperparameter search and best-found
text: hyperparameter values? 3.2

heading: 3 C3. Did you report descriptive statistics about your results (e.g., error bars around results, summary
text: statistics from sets of experiments), and is it transparent whether you are reporting the max, mean, etc. or just a single run? 3.3
C4. If you used existing packages (e.g., for preprocessing, for normalization, or for evaluation), did you report the implementation, model, and parameter settings used (e.g., NLTK, Spacy, ROUGE, etc.)? Not applicable. did not use existing packages
D 7 Did you use human annotators (e.g., crowdworkers) or research with human participants? Left blank.
D1. Did you report the full text of instructions given to participants, including e.g., screenshots, disclaimers of any risks to participants or annotators, etc.? No response.
D2. Did you report information about how you recruited (e.g., crowdsourcing platform, students) and paid participants, and discuss if such payment is adequate given the participants’ demographic (e.g., country of residence)? No response.
D3. Did you discuss whether and how consent was obtained from people whose data you’re using/curating? For example, if you collected data via crowdsourcing, did your instructions to crowdworkers explain how the data would be used? No response.
D4. Was the data collection protocol approved (or determined exempt) by an ethics review board? No response.
D5. Did you report the basic demographic and geographic characteristics of the annotator population that is the source of the data? No response.

heading: Future Work
text: There are two key aspects that this work does not cover and we encourage future work to explore. First, as pointed out in §2. 1, we choose LayoutBERT as our document encoder, Φdoc. This work does not experiment with the variety of encoding strategies in the literature that combines textual, visual, and layout information (Appalaraju et al. , 2021; Xu et al. , 2021; Huang et al. , 2022). It is likely that richer document representations derived from these diverse encoders will further push the limits of zero-shot classification when combined with our proposed unsupervised contrastive pretraining procedure. Second, results in this paper are on a single dataset, i. e. the RVL-CDIP dataset. While we mitigate this to a large extent by creating four nonoverlapping test splits (see §3. 1 and Appendix A), results on more datasets might yield more useful insights. In practice, the lack of datasets for this task (of semi-structured document classification) is what makes this exploration difficult and might require creation of new resources