MATrIX - Modality-Aware Transformer for Information eXtraction

Following prior work[22][23][2], we tokenize the input words using WordPiece[21]. Given a word $w$, the tokenizer $T$ will output $n$ tokens which are then passed through an embedding $W$, giving $n$ representations which are averaged at the word-level to create the encoder input text token $t$. Formally:

This approach allows us to train with a maximum sequence length of 512 word tokens as opposed to 512 sub-word tokens in previous work, allowing us to process entire, large documents with their entire text content.

To enrich the document representation we add visual tokens to the input. Given an image, we resize the long side to 512 pixels, preserving the aspect ratio, we define a reduction ratio of 32 which gives us a maximum of 16x16 visual tokens. In a departure from previous work, we do not enforce a visual token sequence length and we can handle arbitrary resolutions given padding to the nearest reduction ratio multiple. To each vision token, we add a spatial embedding based on its feature map element size with respect to the non-padded image area. The weights of the visual-spatial embedding layer are shared with the text spatial embedding.

To each text and visual token, we add a spatial embedding to provide the encoder with 2D spatial information. As opposed to DocFormer which embedded the bounding box and LayoutLMv2 which used two embeddings, we encode the token index, top-corner $x$ and $y$ relative positions, and the relative width and height of the token. Formally $\text{pos}=(i,x,y,w,h)$. The embedding itself is a 3-layer feed-forward network with leaky ReLU activation.
We use the same embedding for visual and text tokens. The intuition behind this choice is that both types of features exist in the same 2D space.

Vision and text features are usually heavily spatially correlated in VDU. In order to include this inductive bias, LayoutLMv2[23] introduced a learn-able relative attention bias, learned independently for each attention head.

This introduces a simple mechanism for the model to increase the self-attention of close-by tokens and reduce the attention of far-away tokens. We modified slightly this implementation to condition the relative attention bias jointly on the $x$ and the $y$ axis instead of two independent biases.

One of the difficult challenges for a multi-modal transformer is to build token representations that can be compared and combined across modalities through self-attention. Even if vision and text modalities both belong in the 2D space and share a common spatial representation, it doesn’t mean that their spatial correlations are identical. For example, vision features might help the model uniformly around a given text token, while for text, the 1D ordering of a sentence that is over-represented in the $x$ axis relationship, is the most meaningful. We adapt the previously described relative attention bias by conditioning the learned bias based on the pair-wise modality of the tokens being considered.

In this novel task, we infer the line bounding box of the input text tokens. Given the output of the encoder, we apply a linear projection for each text token to predict its line coordinates as $(x,y,w,h)$. The concept of line is defined as returned by the OCR engine which is a group of words separated by spaces. We use a mean-square error loss denoted $L_{LReg}$.

The model can achieve this objective only by learning to group words together semantically and visually to combine their spatial information in order to output the correct line for each word in a given line.

In this novel VDU task, we switch text tokens in the input and use a projection to predict whether a token was switched. We use this task to encourage the model to learn a richer semantic understanding of the text sequence. We optimize the objective using a cross-entropy loss represented by $L_{TS}$.

In this task, we try to predict masked text tokens. We mask or replace tokens in the original sequence $s$ and use the encoder to predict the correct token $t$ at a given position. In a similar fashion to DocFormer[2], we use a modified formulation of the original BERT[7] MM-MLM in which the visual region is not masked, as opposed to prior work by Xu et al.[22][23]. Our implementation differs from DocFormer as it accounts for the mean bag that is applied to the output of the tokenizer. We apply an L1 loss between the input and the output text embeddings, however, this approach is prone to collapsing as the embedding converges to a null vector $\vec{0}$. To prevent this side-effect, we rely on the token switch task which forces the model to converge towards a solution that also allows for the differentiation of the text embeddings and thus makes the $\vec{0}$ solution invalid.

In this task, introduced by Appalaraju et al.[2] we seek to reconstruct the input image from the multi-modal representation generated by the encoder. We use a simplistic decoder based on convolutional and up-sampling layers to infer the input image. We compute the error using a smooth L1 loss denoted as $L_{LTR}$. This task is expected to force the model to leverage textual features to reconstruct characters from the original image.

In this task, we create mismatched image and text modality pairs and introduce them into our dataset with a 20% sampling rate. The encoder is then tasked with identifying mismatched pairs. We add this task because requires a global understanding of the document as opposed to the other pre-training tasks which rely on local features. We use a cross-entropy loss denoted as $L_{TDI}$.

In this task, introduced by Xu et al.[23] we redact randomly selected lines from the input image by setting the pixels in the region to 0. The encoder is then tasked with classifying whether a given token was redacted in the input image. With this task, we seek to instill a strong sense of the spatial relationship between text and image. We use a cross-entropy loss denoted as $L_{LRed}$.


The final loss for pre-training is as follow:

With $\alpha_{1..n}=\left\{6.0,0.5,1.0,1.0,1.0,1.0\right\}$. The coefficient prevent early divergence by over-weighting line regression and under-weighting line redaction.

COnsolidated Receipt Dataset for Post-OCR Parsing[18] is a receipt dataset containing a total of 1000 samples with 800 for training, 100 for validation, and 100 for test. It defines 30 classes grouped in 5 superclasses. F1 score is reported on the 30 classes classification task.

Form Understanding in Noisy Scanned Documents[8] is a form dataset with a total of 199 samples, making it much smaller than other datasets we evaluated. 149 samples are used for training and the reported F1-score is reported on the remaining 50. We evaluate entire documents and do not truncate them to the first 512 sub-word tokens.

We report our results in Table 1. We achieve a competitive F1 score on CORD, falling short of the previous SOTA by 0.94 with MATrIX which has about 3 times fewer parameters than LayoutLMv2-large and DocFormer-large. We do not achieve a high F1 score on FUNSD, which we attribute to its smaller number of samples.

We conduct an ablation study to determine the impact of using pre-trained BERT weights for the attention layer and sub-word token embeddings, and modality-aware relative attention on the final results for the CORD downstream task. This shows that modality-aware relative attention offers a significant improvement over regular multi-modal self-attention.

We conduct an ablation study to determine the impact of each pre-training task on the final results for the CORD downstream task. To minimize resource usage, these pre-trainings only ran for a single epoch on the 5M dataset. In table 4, MM-MLM was always trained with the token switch task to prevent collapse. Appalaraju et al.[2] showed that the learn to reconstruct and text describe image tasks were beneficial for this task, therefore we attribute this regression to insufficient training.

We conduct an ablation study to analyze the impact of each modality on the final results for each downstream task. We do not do additional pre-training and simply input 0 values for the missing modalities during the fine-tuning step on CORD. This shows that the additional modalities are being leveraged.

We conduct an ablation study to measure the impact of increasing and decreasing the input resolution when fine-tuning, given an encoder pre-trained at a maximum resolution of 512x512. We observe indeed that the model is resilient to resolution change.