UniDoc: Unified Pretraining Framework for Document Understanding

1 Introduction . Document intelligence is a broad research area that includes techniques for information extraction and understanding . Unlike plain-text documents in natural language processing ( NLP ) [ 1 , 2 ] , a physical document can be composed of multiple elements : tables , figures , charts , etc . In addition , a document usually includes rich visual information , and can be one of various types of documents ( scientific paper , form , resume , etc . ) , with various combinations of multiple elements and layouts . Complex content and layout , noisy data , font and style variations make automatic document understanding very challenging . For example , to understand text-rich documents such as letters , a system needs to focus almost exclusively on text content , paying attention to a long sequential context , while processing semi-structured documents such as forms requires the system to analyze spatially distributed short words , paying particular attention to the spatial arrangement of the words . Following the success of BERT [ 3 ] on NLP tasks , there has been growing interest in developing pretraining methods for document understanding [ 4 , 5 , 6 ] . Pretrained models have achieved state-of-the-art ( SoTA ) performance across diverse document understanding tasks [ 7 , 8 ] . Huge training datasets help pretraining models to learn a good representation for downstream tasks . However , we observe three major problems with the current pretraining setup : ( 1 ) documents are composed of semantic regions . Most of the recent document pretraining works follow BERT and split documents into words . However , unlike the sequence-to-sequence learning in NLP , documents have a hierarchical structure ( words form sentences , sentences form a semantic region , and semantic regions form a document ) . Also , the importance of words and sentences are highly context-dependent , i.e. , the same word or sentence may have different importance in a different context . Moreover , current transformer-based document pretraining models suffer from input length constraints . Also , input 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . length becomes a problem for text-rich documents or multi-page documents . ( 2 ) documents are more than words . The semantic structure of the document is not only determined by the text within it but also the visual features such as table , font size and style , and figure , etc . Moreover , the visual appearance of the text within a block are often overlooked . Most of recent BERT-based pretraining works only take the words as input without considering multimodal content and alignment of multimodal information within semantic regions . ( 3 ) documents have spatial layout . Visual and layout information is critical for document understanding . Recent works encode spatial information via 2D position encoding and model spatial relationships with self-attention , which computes attention weights for long inputs [ 4 , 5 ] . However , for semi-structured documents , such as forms and receipts , words are more related to their local surroundings . This corresponds strongly with human intuition – when we look at magazines or newspapers , the receptive fields are modulated by our reading order and attention . Based on the above observations , we ask the following question : Can unified document pretraining benefit all of these different kinds of documents ? We propose a unified pretraining framework for document understanding , shown in Fig . 1 . Our model integrates image information in the pretraining stage by taking advantage of the transformer architecture to learn cross-modal interactions between visual and textual information . To handle textual information , we encode sentences using a hierarchical transformer encoder . The first level of the hierarchical encoder models the formation of the sentences from words . The second level models the formation of the document from sentences . With the help of the hierarchical structure , UDoc learns how words form sentences and how sentences form documents . Meanwhile , it reduces model computation complexity exponentially and increases the number of input words . This also mimics human reading behaviors since the sentence/paragraph is a reasonable unit for people to read and understand—people rarely check the interactions between arbitrary words across different regions in order to understand an article . Convolution has been very successful in the extraction of local features that encode visual and spatial information [ 9 ] , so we use convolution layers as a more efficient complement to self-attention for addressing local intra-region dependencies in a document image . Meanwhile , self-attention uses all input tokens to generate attention weights for capturing global dependencies . Thus , we combine convolution with self-attention to form a mixed attention mechanism that combines the advantages of the two operations . We depart from previous vision-language pretraining [ 10 , 11 ] by extracting both the textual and visual features for each semantic region . We propose a novel gated cross-attentional transformer that enables information exchange between modalities . A visually-rich region ( figure , chart , etc ) may have stronger visual information than textual information . Instead of treating outputs from both modalities identically , we design a gating mechanism that can dynamically control the influence of textual and visual features . This approach enables cross-modal connections and allows for variable highlight the relevant information in visual and textual modality and enables cross-modal connections . During pretraining , the CNN-based visual backbone and multi-layer gated cross-attention encoder are jointly trained in both pretraining and fine-tuning phases . Our contributions are summarized as follows : ( 1 ) We introduce UDoc , a powerful pretraining framework for document understanding . UDoc is capable of learning contextual textual and visual information and cross-modal correlations within a single framework , which leads to better performance . ( 2 ) We present Masked Sentence Modeling for language modeling , Visual Contrastive Learning for vision modeling , and Vision-Language Alignment for pretraining . ( 3 ) We present extensive experiments and analyses to validate the effectiveness of the proposed UDoc . Extensive experiments and analysis provide useful insights on the effectiveness of the pretraining tasks and show outstanding performance on various downstream tasks . 2 Related Work . Self-supervised learning has shown great success in producing generic representations that learn from large-scale unlabeled corpora [ 3 ] . Like the development of pretraining in computer vision [ 12 ] and NLP [ 3 ] , there has been a surging interest in self-supervised learning for Vision-Language ( VL ) tasks [ 10 , 13 , 14 , 11 ] . Transformers [ 3 ] are the key technology that enables learning contextualized representations from large-scale unlabeled training data . The unique characteristics of document images ( spatial layout and multiple elements ) distinguish document image pretraining from pretraining works in NLP and VL domains . In the NLP domain , the inputs are pure texts without spatial layouts ( bounding boxes ) . In the VL domain , the inputs are the visual objects and captions . While for document images , the input elements are spatially distributed , and the visual and textual information co-occur within the semantic regions . Several recent works have explored pretraining on document images [ 4 , 5 , 15 ] . LayoutLM [ 4 ] extends BERT to learn contextualized word representations for document images through multi-task learning . It takes a sequence of OCR words as input during pretraining and incorporates the 2D position embedding as input for each token . However , LayoutLM only considers textual information during pretraining without modeling the alignment between visual and textual information–visual information is only incorporated into the model during the fine-tuning stage . The most recent version , LayoutLMv2 [ 5 ] , improves on this by incorporating the image encoder into pretraining and jointly training the image encoder along with the BERT model . LayoutLMv2 splits the document image into several parts and concatenates the visual embeddings and text embeddings into a single sequence . Apart from masked language learning ( MLM ) , LayoutLMv2 also considers image-text alignment and image-text matching during pretraining . The most related work to ours is SelfDoc [ 6 ] , which proposes a multimodal document pretraining framework . It first extracts the document object proposals from pre-trained Faster R-CNN [ 16 ] and then applies OCR for each proposal to get the words . It takes the pre-extracted RoI features and sentence embeddings as input , and models the perform learning over the textual and visual information using the cross-modality encoder . There is a noticeable difference between our proposed method , UDoc , and other concurrent works in document image pretraining . UDoc is a multimodal end-to-end pretraining framework for document images . Unlike the fixed document object detector in [ 6 ] , the parameters of the image encoder with RoI align , which derive the visual features for semantic regions , are also updated in UDoc . In contrast to [ 5 ] , our visual features come from the semantic regions instead of splitting the image into fixed regions . Like the object-level semantic elements in natural images , for document images , we represent the typical document layout elements such as paragraph , title , figure , and table as semantic regions . Moreover , to learn the contextualized visual representations , UDoc masks visual information in the latent space and learns contextualized representations by solving a contrastive learning task defined over a quantization of the latent visual embeddings . 3 Method . 3.1 Model Architecture . Fig . 1 illustrates our approach , UDoc , which consists of four components : feature extraction , feature embedding , multi-layer gated cross-attention encoder , and pretraining tasks . Given a document image and the locations of document elements ( sentence or RoI ) , UDoc takes image regions and words that correspond to each document elements as inputs , and extracts their respective embeddings through a visual feature extractor and a sentence encoder . These embeddings are then fed into a transformer-based encoder to learn the cross-modal contextualized embeddings that integrate both visual features and textual features . In the feature extraction step , we first employ an off-the-shelf OCR tool [ 17 ] to extract text from a document image I , where the words are grouped into sentences S = { s1 , . . . , sN } whose corresponding bounding boxes are P = { p1 , . . . , pN } . For each sentence bounding box pi , we use a ConvNet-based backbone fImEnc and RoI Align [ 18 ] fRoIAlign to extract the pooled RoI features vi . To obtain a feature embedding , we extract the sentence embedding si for each sentence si via a pretrained sentence encoder fSentEnc . Each region ’ s RoI feature vi is discretized into a finite set of visual representations vQi ∈ VQ via product quantization [ 19 ] . The multi-layer Gated Cross-Attention encoder takes the position information , masked visual features Ṽ and masked textual features S̃ as inputs , and then it generates the contextualized multimodal representations ( HlV and H l S , l ∈ [ 1 , L ] ) and outputs the predicted features ( V̂ and Ŝ ) , where L is the number of stacked transformer blocks . More formally , the pretraining procedure can be decomposed into the following steps : I OCR−−→ ( P S ) fImEnc+fRoIAlign−−−−−−−−→ fSentEnc ( V , VQ S ) −−−→ fMask ( Ṽ S̃ ) −→ ( HlV HlS ) −→ ( V̂ Ŝ ) −→ LPretraining ( 1 ) where fMask denotes the masking function that randomly masks RoI features and sentence embeddings with the respective probabilities pvMask and p s Mask . LPretraining is composed of three pretraining tasks : Masked Sentence Modeling ( MSM ) , Visual Contrastive Learning ( VCL ) , and Vision-Language Alignment ( VLA ) . Next , we provide details mentioned in Eq . 1 . Feature Extraction and Embedding . Formally , a document image I ∈ RW×H consists of N regions , where each region ’ s bounding box is characterized by a 6-d vector , as pi = { xLTW , yLT H , xRB W , yRB H , w W , h H } , where w and h are of the width and height the region , W and H are the width and height of I , while ( xLT , yLT ) and ( xRB , yRB ) denote the coordinates of the top-left and bottom-right corners respectively . The 6-d vector is mapped into a high-dimensional representation via a linear mapping function . The visual embedding is the sum of the mapped RoI feature and position embedding . Likewise , textual embedding is the sum of sentence embedding and position embedding . We also have different types of segments to distinguish different modalities . The input sequence to the transformer-based encoder starts with a special start element ( [ CLS ] and full visual features ) , then it is followed by multimodal elements , and it ends with a special ending element ( [ SEP ] +full visual features ) . For the special elements ( [ CLS ] and [ SEP ] ) , the corresponding full visual features are features extracted from the whole input image , by applying fImEnc to an RoI covering the whole input image . Quantization Module . Unlike the fixed image encoder in [ 6 ] , we jointly learn the image encoder in an end-to-end fashion alongside the multimodal model . A visual representation can be learned by predicting the visual features of the masked regions , but it is challenging to predict such features exactly , since they are unconstrained and of continuous representation . To constrain the representation space of the visual features and facilitate the end-to-end learning of image encoder ( see Task # 2 in Sec . 3.2 ) , we follow [ 20 , 21 ] and use vector quantization to discretize the visual features V = { v1 , . . . , vN } into a finite set of representations VQ = { vQ1 , . . . , v Q N } . Specifically , we define latent embedding spaces e ∈ RC×E , where C is the number of codebooks , and E is the number of entries for each codebook . For each vi , we first map it to logits v ` i ∈ RC×E , and calculate the probability for the j-th codebook entry in i-th group as pc , e = exp ( ( v ` c , e + ge ) /τ ) / ∑E k=1 exp ( ( v ` c , k + gk ) /τ ) , where τ is a non-negative temperature , g1 : E are i.i.d samples drawn from Gumbel ( 0,1 ) distribution . During the forward pass , we choose one entry vector from each codebook by ẽi ∼ argmaxepc , e and generate the quantized representation vQi by a concatenation of { ẽ1 , . . . , ẽG } which is then followed by a linear transformation . During the backward pass , the gradients are computed through a Gumbel-Softmax estimator [ 22 ] . Gated Cross-Attention . To model the interactions among multimodal inputs , we introduce a multimodal transformer with gated cross-attention to model the cross-modality relationships . Let Hl+1m be output features at the l-th layer for one modality m , and let n be another modality ( m , n ∈ { V , S } ) . We obtain the features at ( l + 1 ) -th layer as : Hl+1m = fLN ( fLN ( Hlm + f l Cross-Att ( H l m|Hln ) ) + f lFF ( fLN ( H l m + f l Cross-Att ( H l m|Hln ) ) ) ) ( 2 ) where fLN denotes layer normalization [ 23 ] . The feed-forward sub-layer fFF in Eq . 2 is further composed of two fully-connected sub-layers , both wrapped in residual adds and fLN . The core part of Eq . 2 is the cross-attention fCross-Att ( · ) . Given the intermediate representations Hlm and Hln , the cross-attention output for modality m is computed as : fCross-Att ( H l m|Hln ) = [ Cross-Att1 ( Hlm|Hln ) ; . . . ; Cross-Atth ( Hlm|Hln ) ] U ( 3 ) Cross-Atti ( Hlm|Hln ) = softmax ( f iq ( H l m ) f i k ( H l n ) T / √ d ) f iv ( H l n ) ( 4 ) where f iq ( H l m ) , f i k ( H l n ) , and f i v ( H l n ) are the query , key , and value calculated by linear mapping layers for the i-th head . d is the model dimension , h is the number of heads , and U is the weight matrix that combines the outputs of the heads . Considering the substantial diversity of document images and the different information needs of differing document types , we use a gating mechanism [ 24 ] to dynamically weight the outputs of the visual and textual branches . Specifically , we feed the concatenated the visual and textual features to a non-linear network fGate ( [ Hl+1m ; H l+1 n ] ) , which generates the modality-specific attention weights α l m and αln , and returns the weights separately to their respective modality-specific branches to perform element-wise products . We multiply the features for modality m with its modality-specific attention weight , and compute the updated feature as : Hl+1m = H l+1 m ( 1 + α l m ) , same that for modality n .

Summary: This paper presents UniDoc as a multimodal pre-trained model for document understanding. UniDoc consists of 3 layers: a feature extraction layer, which extracts sentence embeddings and region embeddings from the original document; a fusion layer, which fuses textual and visual information based on a gated cross-attention mechanism; a task layer, which uses 3 tasks as pre-training objectives to optimize model parameters. Evaluations are conducted on 4 downstream tasks, including Form Understanding, Receipt Understanding, Document Classification and Document Object Detection. Comparing to several SOTA baselines, UniDoc performs better on Form Understanding and Receipt Understanding, but worse on Document Classification. The paper also includes some ablation studies to investigate the effects of V/L/V+L, pre-training tasks and visual backbones, and some interesting future directions.