GV doc - Graph-based Visual DO cument Classification (2023)

Abstract: The robustness of a model for real-world deployment is decided by how well it performs on unseen data and distinguishes between in-domain and out-of-domain samples. Visual document classifiers have shown impressive performance on in-distribution test sets. However, they tend to have a hard time correctly classifying and differentiating out-of-distribution examples. Image-based classifiers lack the text component, whereas multi-modality transformer-based models face the token serialization problem in visual documents due to their diverse layouts. They also require a lot of computing power during inference, making them impractical for many real-world applications. We propose, GVdoc, a graph-based document classification model that addresses both of these challenges. Our approach generates a document graph based on its layout, and then trains a graph neural network to learn node and graph embeddings. Through experiments, we show that our model, even with fewer parameters, outperforms state-of-the-art models on out-of-distribution data while retaining comparable performance on the in-distribution test set.

## introduction documents digitization and their intelligent processing in various industries such as finance, insurance, and medicines has resulted in the rapid development of structured document understanding methods, a.k.a. document ai. document classification is one of the essential tasks in document ai for labeling documents. a number of deep convolutional neural network (cnn) and transformerbased models have achieved superior performance on many document-ai tasks @xcite @xcite . however, they tend to employ bigger models with hundreds of millions of parameters, subsequently increasing computational demand that can be a challenge in real-world applications. yet many of them fail to perform well on out-of-distribution (ood) data @xcite . this is because, in many cases, training and testing examples are from a fixed distributionsuch as a particular language, time frame, and industry. however, the layout of the documents evolves over time, and the model should perform well on such out-of-distribution data. further, the model is expected to be able to differentiate between known and unknown categories of documents, thus minimizing false-positive predictions during testing. initial work on document classification employed off-the-shelf image classifiers @xcite and models pre-trained on imagenet @xcite or similar datasets. these methods struggle to label documents having similar layouts but different text contexts. later, focus shifted towards language models @xcite and multimodality models @xcite @xcite . these models also incorporated layout information obtained from optical character recognition (ocr). therefore, the performance of these methods, particularly transformer-like models, degrades due to the imperfection of the ocr engine, such as errors in parsed text or the order of tokens sequence. almost all of these methods tried to improve the performance on the in-distribution test set, neglecting the generalization for real-world applications. to confirm, recently @xcite collected an ood version of rvlcdip dataset @xcite and evaluated several image and multi-modal classifiers. however, none of them performed well on the ood dataset. our method, called gvdoc (for graph-based visual document classification), studies document classification as a graph classification problem, where we take text words as nodes and the relationship between words as edges in a graph. we generate a document-level graph using that layout information from ocr (see figure 1 ) and learn the embedding using graph neural networks (gnns). gvdoc is more robust to changes in the test set; hence it shows improved performance on out-ofdistribution data. we make the following contributions: ## related work visual document classification cnns have achieved excellent performance on natural scene images, so they became the first obvious choice for visual document classification @xcite @xcite . however, documents have overlapping intra-class visual and structural characteristics @xcite , which makes visual features less discriminative for classification. the semantics of text in the document and the layout are essential to understand the visual documents. a second line of work studies document classification as a sequence classification problem @xcite @xcite . they follow language modeling strategies, but aside from text, they also incorporate layout information. such approaches parse text and layout information by applying ocr on document images. then, they train transformer-like models. structurallm @xcite ) adds text and layout embeddings and trains a transformer model (similar to bert @xcite ) on specialized pretraining tasks. some of the recent works employ multi-modal features including visual, text and layout @xcite @xcite . these models train a single transformer on concatenations of text and visual tokens @xcite or train a separate transformer branch for both text and visual modalities @xcite . the methods that utilize text consider serialized tokens from ocr as an input, so their performance varies with the correctness of the ocr engine. for examples, if we replace the proprietary microsoft azure ocr in layoutlmv2 @xcite with tesseract foot_0 , an open source ocr, its performance drops for visual document classification @xcite . transformer-based models consider input sequence based on ocr reading order @xcite , which may not reflect tokens in their actual reading order @xcite . therefore, a few recent studies model the document as a graph by suggesting several possible edge types. zhang et al. (2020) proposed k-nearest neighbors graphs, but these may contain connections with isolated tokens. fully connected graphs employed by @xcite do not leverage the sparsity of the document, hence their approach is similar to transformers. on the other hand, @xcite relied on a proprietary ocr technology to identify "text fields", then utilized a 360-degree line-ofsight (los) graph. we initially used los graphs but that did not show very good performance. form-net @xcite models a document as a graph using a β-skeleton graph (kirkpatrick and @xcite and tries to minimize the serialization error by learning localized super-token embeddings using graph convolutions before a transformer. however, they used etc transformer (ainslie et al., 2020) for schema learning from gcn-encoded structure-aware super-tokens. our approach differs from prior graph-based work in two important ways: graph generation and learning embeddings. our unique document-level sparse graph incorporates both spatial layout and ocr reading order, leveraging the document's sparsity and making our model less sensitive to common mistakes in ocr reading order. moreover, we solely use a gnn to learn embeddings. thus, we do not require a transformer component, making our approach more memory-efficient than models that incorporate a transformer @xcite @xcite . our approach also uses more expressive edge embeddings than that of @xcite . feature fusion initial research simply added together the text and layout embedding @xcite , incorporated position bias in attention mechanism @xcite , designed cross-modality attention layers @xcite @xcite , and explored 1d position and 2d layout aware attention weights using a disentangled matrix @xcite . lilt @xcite adds attention weights from layout and text embeddings and updates both types of embeddings through two separate transformers. however, adding attention weights does not fully leverage the cross-domain features. selfdoc @xcite took the value (v) of one modality as key (k) for the other modality while computing crossattention in transformer layers to learn dependency between language and vision features. finally, it added features of both text and visual modalities. ## gvdoc document graph we now describe our approach for representing document using both textual and layout features. we represent a document d as a graph where each token is a node and edges reflect the spatial relationship between them. nodes we define vertices for all tokens as @xmath0, v 2 , ..., v n } where features of v i are a fusion of the text and layout embeddings defined later in equation (5). in addition, we define a virtual super node that summarizes the graph, similar to the cls token in bert. edges token sequence can be important in understanding text, but this information provided by ocr is noisy. we therefore generate edges in the document graph reflecting two types of relationships between vertices: (a) "ball-of-sight" using β-skeleton graph @xcite and (b) paragraph-based neighborhood. a β-skeleton graph @xcite defines an edge between two bounding boxes if both intersect a circle that does not intersect any other bounding box; the resulting "ball-ofsight" graph is sparser than one using line-of-sight edges @xcite . @xcite found this useful for message passing in gnns. the paragraph-based neighborhood connects tokens within the same paragraph and connects paragraphs based on ocr's reading order predictions. while we could fully connect all tokens in the same paragraph, we aim to reduce computation by increasing sparsity; therefore, we add edges for each token with the k nearest neighbors within the same paragraph. then, for each pair of paragraphs that are adjacent in the ocr's reading order, we define an edge between the last token of the prior paragraph and the first token of the following paragraph. finally, we define a super-node and connect it with the first and last token of each paragraph, considering them as representative tokens of the paragraph. to construct the final graph, we take the union of the edges from the β-skeleton graph and the paragraph-based neighborhood as shown in figure 1. thus, we generate a graph that is sparse but also has enough connections for learning node embeddings through message passing in the gnn (as evident in ## gvdoc model our gvdoc model, shown in figure 2 , consists of input embedding, feature fusion, and task-specific prediction modules. we learn node embeddings in an end-to-end fashion through various unsupervised pre-training tasks. then, we fine-tune the model for downstream tasks. ## experiments we hypothesize that our gvdoc model will be more robust to changes in the test distribution than other models. we therefore designed experiments to measure how our model performed on two tasks: (a) classifying in-domain but out-of-distribution documents, and (b) distinguishing out-of-domain documents from in-domain documents. ## conclusion in this paper, we address the limitation of existing visual document classification models by modeling a document as a graph and learning its embeddings using a graph attention network. by defining two types of edges (β skeleton and paragraph-based), we leverage the benefit of layout information while minimizing the effects of the errors from ocr reading order. thus, effectively embracing coarse and fine-grained layout information, gvdoc generalizes better for different layouts. while most visual document classifiers tend to perform well on in-distribution data, they fail or struggle on outof-distribution data; our model does not drop its performance on ood data. through experiments, we demonstrate the generalization of our model on out-of-distribution data.