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 proposes a new unified framework for pre-training image/text encoders for capturing both visual and textual features in documents. The main Transformer model is pre-trained with multimodal embeddings as input features (visual & textual features) through three pre-training tasks: Masked Sentence Modeling, Visual Contrastive Learning, and Vision-Language Alignment. The pre-trained model can be fine-tuned on downstream tasks on which the proposed UniDoc model outperforms baselines.