Fast-StrucTexT: An Efficient Hourglass Transformer with Modality-guided Dynamic Token Merge for Document Understanding

As the first heuristic work, NLP-based approaches Devlin et al. (2019); Liu et al. (2019) adopt the language model to extract the semantic structure. Various works Hong et al. (2022); Lee et al. ; Li et al. (2021a); Xu et al. (2020) then jointly leverage layout information by spatial coordinates encoding, leading to better performance and extra computations simultaneously. After that, some researchers realize the effectiveness of deep fusion among textual, visual, and layout information from document images. A quantity of works Gu et al. (2021); Li et al. (2021b, c); Xu et al. (2021); Gu et al. (2022); Appalaraju et al. (2021) rely on text spotting to extract semantic region features with a visual extractor. However, the transformer-based architectures are inefficient for long sequence modeling because of computationally expensive self-attention operations. To this end, we propose an efficient multi-modal model, Fast-StrucTexT, for visual document understanding.

A well-known issue with self-attention is its quadratic time and memory complexity, which can impede the scalability of the transformer model in long sequence settings. Recently, there has been an overwhelming influx of model variants proposed to address this problem.

Given a visual document and its text content extracted from OCR toolkits, the feature extractor first extracts both textual and visual features from the region proposals of text segments. These features are then fed into a transformer-based encoder to learn multi-modal contextualized representations via alternating self-attention and cross-attention mechanisms. By leveraging redundancies across input tokens, the proposed encoder consisting of several Merging-Blocks and Extension-Blocks is designed as a hierarchical structure for shortening the sequence length progressively. In particular, the Merging-Block dynamically merges nearby tokens and the Extension-Block recovers the shortened sequence to the original scale according to the merging information to support token-level downstream tasks. Besides, we further adopt the alignment constraint strategy to improve the model ability by introducing a CAE Chen et al. (2022) head in the pre-training. The generated contextual representations can be used for fine-tuning downstream tasks of visual document understanding.

We employ an off-the-shelf OCR toolkit to a document image $I\in\mathbb{R}^{W\times H}$ to obtain text segments with a list of sentences $S$ and corresponding 2D coordinates $B$. We then extract both segment-level visual features and word-level textual features through a ConvNet and a word embedding layer. For visual features $f_{V}$, the pooled RoI features of each text segment extracted by RoIAlign are projected to a vector. For textual features $f_{L}$, we utilize the WordPiece to tokenize text as sub-words and convert them to the ids.

We add a special start tag [CLS] with coordinates $B_{0}=[0,0,W,H]$ at the beginning of the input to describe the whole image. Besides, several [PAD] tags with zero bounding boxes append to the end of $f_{V}$, $f_{L}$ to a fixed length. We ensure the length of $f_{L}$ is an integral multiple of $f_{V}$.

We propose an hourglass transformer as the encoder module and progressively reduce the redundancy tokens. Instead of shortening the sequence length by a fixed pattern of merging tokens, the model performs a dynamic pooling operation under the guidance of semantic units (text sentence/segment), which leads to a shorter hidden states without harming model capability. The encoder consists of several Merging- and Extension-blocks. The Merging-block (M-Block) merges the tokens of hidden states and the Extension-block (E-Block) conducts up-sampling to make up for the original length.

Merging. The M-Block suggests merging nearby $k$ tokens with weighted 1D average pooling, where $k$ is the shortening factor. In view of the multi-model hidden states, a referred weighting is predicted from another modality by a linear layer. The process is denoted in Figure 4.

where $i$ is the stage index of M-Block, $m$ and $n$ denotes textual and visual modality $(m,n\in\{V,L\})$. Notably, $\text{Linear}^{*\rightarrow*}$ share their parameters for all M-Blocks.

Extension. The E-Block is required to transform shortened sequence of hidden states back to the entire token-level state. In detail, we simply apply repeat up-sampling method Dai et al. (2020b) to duplicate the vector of each merged token for $k$ times. This method is computationally efficient. For maintaining the distinct information of tokens, we fed the hidden states from the corresponding M-Block into the E-Block through a skip-connection.

Symmetry Cross-Attention. Cross-Attention has shown the effectiveness of multi-modal learning on the vision-and-language tasks. To enhance the interactions among textual and visual modalities, we introduce Symmetry Cross-Attention (SCA) module to model the cross-modality relationships, which consists of two dual cross-attentions to handle the text and visual features in this work. We also leverage SCA to provide cross-modality guidance for token merging in the M-Blocks. The SCA is defined as follows:

where $F_{q},F_{k},F_{v}$ are the linear project for query, key and value. $\sigma$ is the Softmax function, $d$ is the hidden size and $W_{o}$ is the output weight. The multi-head settings of attention are omitted for brevity.

Yet it’s worth noting that our SCA incorporates semantic embedding as an additional input that gives the identical index for each segment and its corresponding words. SCA can provide multi-granularity interaction information for the subsequent token merging in addition to ensure the multi-modal interaction. Furthermore, a global context information be taken into account by Self-Attention (SA). Therefore, SCA and SA are adopted in turn to build transformer layers.

We adopt four self-supervised tasks simultaneously during the pre-training stage, which are described as follows. Masked Visual-Language Modeling (MVLM) is the same as LayoutLMv2 Xu et al. (2021) and StrucTexT Li et al. (2021c). Moreover, a CAE Chen et al. (2022) head is introduced to eliminate masked tokens in feature encoding and keep the consistency of document representation between pre-training and fine-tuning.

Graph-base Token Relation (GTR) constructs a ground truth matrix $G$ with 0$\sim$9 to express the spatial relationship between each pairwise text segments. We give $G_{ij}$ a layout knowledge, i.e., $0$ means the long distance (exceeding half the document size) between text segment $i$ and $j$, and 1$\sim$9 indicate eight buckets of positional relations (up, bottom, left, right, top-left, top-right, bottom-left, bottom-right). We apply a bilinear layer in the segment-level visual features to obtain the pairwise features and fed them into a linear classifier driven by a cross-entropy loss.

Sentence Order Prediction (SOP) uses two normal-order adjacent sentences as positive examples, and others as negative examples. SOP aims to learn fine-grained distinctions about discourse-level coherence properties. Hence, the encoder is able to focus on learning semantic knowledge, and avoids the influence of the decoder.

Text-Image Alignment (TIA) is a fine-grained cross-modal alignment task. In the TIA task, some image patches are randomly masked with zero pixel values and then the model is pre-trained to identity the masked image patches according to the corresponding textual information. It enable the combine information between visual and text. Masked text token is not participating when estimating TIA loss.

We fine-tune our model on visual information extraction tasks: entity labeling and entity linking.

Entity Labeling. The entity labeling task aims to assign each identified entity a semantic label from a set of predefined categories. We perform an arithmetic average on the text tokens which belong to the same entity field and get the segment-level features of the text part. To yield richer semantic representations of entities containing the multi-modal information, we fuse the textual and visual features by the Hadamard product operation to get the final context features. Finally, a fully-connected layer with Softmax is built above the fused features to predict the category for each entity field. The Expression is shown as follow,

where $W^{c}\in\mathbb{R}^{d\times d}$ is the weight matrix of the MLP layer and $\sigma$ is the Softmax function. For entity $e^{i}$, $Z_{F}^{i}$ indicates the final fused contextual features, $k$ is the token length and $P_{i}$ is the probability vector.

Entity Linking. The entity linking task desires to extract the relation between any two semantic entities. We use the bi-affine attention Zhang et al. (2021) for linking prediction, which calculates a score for each relation decision. In particular, for entity $e^{i}$ and $e^{j}$, two MLPs are used to project their corresponding features $Z_{F}^{i}$ and $Z_{F}^{j}$, respectively. After that, a bilinear layer $W^{b}$ is utilized for the relation score $Score^{i,j}$.

where $W^{k},W^{v},W^{b}$ is the parameter weights that $\in\mathbb{R}^{d\times d}$, $b^{k}$ and $b^{v}$ is the bias. $\sigma$ denotes the sigmoid activation function.

We followed the typical pre-training and fine-tuning strategies to train our model. For all pre-training and downstream tasks, we resized the images along their longer side and padded them to a size of $512\times 512$. The input sequence is set to a maximum length of 640, with text tokens padded to a length of 512 and image tokens padded to a length of 128.

Model Configurations. The hourglass encoder in our model comprises three M-Blocks and three E-Blocks. Each block consists of a SA and a SCA layer, with 12 heads, in addition to Merging or Extension operations. The encoder architecture is a 12-layer transformer with a hidden size of 768 and an intermediate size of the feed-forward networks, $D_{f}$, of 3072. The shortening factor $k$ is set to 2 for all stages. The input sequence lengths of the three blocks in M-Blocks are ${256,128,64}$, successively, and vice versa in E-Blocks. We tokenize segment-level text by BERT Devlin et al. (2019), and trans to text sequence using One-Hot embedding with the vocabulary size $L_{v}$=30522. To pre-process the image sequence, we use ResNet18 He et al. (2016) pre-trained on ImageNet Deng et al. (2009) to extract ROI features with the size of $C\times H\times W=128\times 4\times 64$, followed by a linear project of $D_{t}$=1024.

Pre-training. We pre-train our model on overall IIT-CDIP dataset. The model parameters are randomly initialized. While pre-training, we apply AdamX optimizer with a batch size of 64 for 1 epoch on 8 NVIDIA Tesla A100 80GB GPUs. The learning rate is set as $1\times 10^{-5}$ during the warm-up for 1000 iterations and the keep as $1\times 10^{-4}$. We set weight decay as $10^{-4}$ and $(\beta_{1}=0.9,\beta_{2}=0.999)$.

Fine-tuning on Downstream Tasks. We fine-tune two downstream tasks: entity labeling and entity linking. For entity labeling task on FUNSD, CORD, and SROIE, we set the training epoch as 100 with a batch size of 8 and learning rate of $5\times 10^{-5}$, $1\times 10^{-4}$, $1\times 10^{-4}$, respectively.

We compare Fast-StrucTexT with BASE scale multi-modal transformer pre-trained models on public benchmarks. We evaluate our model on three benchmark datasets for entity labeling and entity linking tasks with metrics such as Frames Per Second (FPS), Floating Point Operations Per Second (FLOPs), Parameters, and F1 score.

Entity Labeling. The comparison results are exhibited in Table 1. Our Fast-StrucTexT achieves state-of-the-art FPS and outperforms the performance of the current state-of-the-art model by 156%. Fast-StrucTexT achieves a 1.9$\times$ throughput gain and comparable F1 score on the CORD and SROIE entity labeling tasks.

To demonstrate the effectiveness of Fast-StrucTexT in Chinese, we pre-train the model in a self-built dataset which consists of 8 million document images in Chinese, and fine-tune the pre-trained model on EPHOIE. Table 2 illustrates the overall performance of the EPHOIE dataset, where our model obtains the best result with 98.18% F1-score.

Entity linking. As shown in Table 3, we compare Fast-StrucTexT with several state-of-the-art methods on FUNSD for entity linking. Compared with BROS Hong et al. (2022), our method achieves a comparable performance with 2$\times$ speed in Table 1.

We conduct ablation studies on each component of the model on the FUNSD and SROIE datasets, including the backbone, pre-training tasks, pooling strategy, and shorten factor. At last, we evaluate the cost of our proposed hourglass transform with various sequence lengths.

Backbone. To prove the pre-trained Fast-StrucTexT can obtain state-of-the-art efficiency and performance. We replaced a variety of popular lightweight backbones Guo et al. (2021); Verma (2021); Wang et al. (2022) for evaluation. As shown in Table 4, Fast-StrucTexT can achieve the highest performance and FPS. The point here is that transformer takes responsibility for token feature representation, which could benefit from large model size. It is the reason why those lightweight architectures lead to worse performance. Particularly, we consider that those methods only reduce time complexity, but it is not considered that multi-modal interaction for visual document understanding. Therefore, we integrate multi-modal interaction into the design of our model.

In Table 4, the model shows better efficiency and performance than other lightweight encoders and achieves FPS = 103.50, F1 = 90.35%. Comparing “Ours w/o M&E” and “Ours”, we can obverse that merging and extension operations are not only efficient but also can improve performance. Specifically, it obtains $1.5\times$ the throughput of ”w/o M&E” settings, and 0.41% improvement in the labeling task.

Pooling. Table 5 gives the results of different merging strategies for entity labeling task on FUNSD. GlobalPool is a form to directly merge all token-level features into segment-level before encoding. AvgPool is our merging strategy without cross-modal guidance. DeformableAttention Xia et al. (2022) attempts to learn several reference points to sample the sequence. The experimental results show the effectiveness of our token merging method.

Shorten factor. We extend the ablation study on the hyper-parameter $k$ in Table 6. We have evaluated multiple benchmarks and have established that the setting $k$=2 is the best trade-off for document understanding. In addition, we can adjust the $k$ in the fine-tuning stage. Nevertheless, ensuring consistency of $k$ during pre-training and fine-tuning can take full advantage of knowledge from pre-trained data. Our framework supports multi-scale pre-training with a list of $k$ factors to handle different shortening tasks.

Sequence length. Our token merging can adapt to the arbitrary length of the input. The ratio of merging is adjustable and determined by the value of $k$ and the number of M-blocks and E-blocks. Referring to Table 7, we investigate the ability of our method with various sequence lengths. The experimental results show that the speed and computation gains become more pronounced with the sequence length increasing.

As shown in Fig 5, we conducted an ablation study on the hourglass architecture and SCA module. We calculated the number of more than 0.7 in each layer of the attention map and then accumulated it layer by layer to observe the utilization of each token. Compared with the yellow line and the blue line, the number of highly responsive areas in the later layer is significantly increased, which can show that after aggregating the multi-granularity semantics, each token has sufficient ability to express more information. The blue line is slightly higher than the red line, which indicates that our SCA has better modal interaction ability and lower computational capacity than SA.