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

1 Introduction

Visual document understanding technology aims to analyze visually rich documents (VRDs), such as document images or digital-born documents, enables to extract key-value pairs, tables, and other key structured data. At present, multimodal pre-training transformer models [Lee et al., 2022; Gu et al., 2022; Huang et al., 2022; Hong et al., 2022] have shown impressive performance in visual document understanding. The inside self-attention mechanism is crucial in modeling the long-range dependencies to capture contextual information. However, its quadratic computational complexity is the limitation of the transformer involved in visual document understanding with long sequences directly.

∗Corresponding author

To improve the computational efficiency, there are three typical solutions. The first solution [Kim et al., 2022; Hongxu et al., 2022] simply reduces input sequence length by efficient sampling or new tokenization process. The second solution attempts to redesign transformer architectures, such as reformulating the self-attention mechanism [Verma, 2021; Wang et al. et al., 2020] or producing a lightweight model [Jiao et al., 2020; Wu et al., 2022; Zhang et al., 2022]. The third solution merges tokens by MLP [Ryoo et al., 2021], grouping [Xu et al., 2022a], or clustering [Marin et al., 2023] to prune unnecessary information related to several works [Pan et al., 2022; Beltagy et al., 2020] have proved the redundancies in attention maps.

Although those methods significantly reduce the computational complexity of the transformer model, they do not take into account the multiple granularity expressions in the visual documents, including words, lines, paragraphs, and so on. Furthermore, some token reduction methods learn multigranularity information to a certain extent, but without considering the correlation between modality and granularity.

In Figure 2, we visualize the attention maps and token interactions of a standard transformer in the inference phase. In particular, Figure 2(a) displays the aggregated attention map of all Transformer layers. The left half shows a selfattention map inside words and the right half shows the crossattention map between words and segments. We re-sample two regions with highlighted red boxes and zoomed in Figure 2(b). Moreover, for closer observation, Figure 2(c) gives the correlation visualization (curves and lines) of words and segments based on attention scores related to Figure 2(b). Almost the words and segments belonging to one semantic entity have significant correlations to their counterparts, indicating a high redundancy between irrelevant tokens in attention scores, which illustrates two key viewpoints:

1. Strong correlations are existed across granularities.
2. There is rich redundancy in attention computation.

In this paper, we propose an efficient multi-modal transformer called Fast-StrucTexT based on StrucTexT [Li et al., 2021c], which is not only devoted to improving the model efficiency but also enhancing the expressiveness. We design an hourglass transformer that consists of merging and extension blocks, which receives the full-length sequence as input and compresses the redundant tokens progressively. A weighted correlation between two granularity features is produced by each merging block to dynamically guide the merging process. Since massive redundant information is eliminated, our model yields great efficiency gain and higher-level semantics. Further, we develop a Symmetry Cross-Attention (SCA) module, a dual cross-attention mechanism with multi-modal interaction, to conduct multi-modal feature fusion and provide modal semantic guidance for token merging. We pretrain Fast-StrucTexT with four task-agnostic self-supervised tasks for learning a good representation, and then fine-tune the pre-trained model in three benchmark datasets. Experiment results demonstrate that our model achieves state-ofthe-art performance and FPS.

The main contributions are summarized below:

• We propose Fast-StrucTexT with an efficient hourglass transformer by performing modal-guided dynamic token merging to reduce the number of tokens.
• We develop a dual multi-modal interaction mechanism named Symmetry Cross-Attention, which can enhance the multi-modal feature fusion from visual documents.
• Extensive experiments on four benchmarks show our model achieve state-of-the-art speed and performance.

2 RelatedWork

2.1 Multi-Modal Pre-trainingModel

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., 2022; 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; Li et al., 2021c; 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.

2.2 Efficient Transformers

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.

Fixed Patterns. The earliest modification to self-attention simply specifies the attention matrix by limiting the field of view, such as local windows and block patterns of fixed

𝑓𝑓𝑉𝑉

Figure 3: The overall Fast-StrucText architecture. Encoder is an hourglass architecture with the input of language and visual features. The merging operation can effectively reduce redundancy computations, and multi-modal interaction information is obtained through SCA & SA. The skip-connection exits from before merging to after extension.

strides. Sparse Transformer [Child et al., 2019] converts the dense attention matrix to a sparse version by only computing attention on a sparse number of qi, kj pairs. GroupingViT [Xu et al., 2022a] divides tokens into multiple groups, and then group-wised aggregates these tokens. Chunking input sequences into blocks that reduces the complexity from N 2 to B2 (block size) with B ≪ N, significantly decreasing the cost.

Down-sampling. Down-sampling methods that narrow the resolution of a sequence can effectively reduce the computation costs by a commensurate factor. Zihang Dai et al. [Dai et al., 2020a] have highlighted the much-overlooked redundancy in maintaining a full-length token-level representation. To solve this problem, they compress the sequence of hidden states to a shorter length, thereby reducing the computation cost. TokenLearner [Ryoo et al., 2021] uses MLP to project the tokens to low-rank space. The recent Nystr¨omformer [Xiong et al., 2021] is a down-sampling method in which the "landmarks" are simply strided-based pooling in a similar spirit to Set Transformer [Lee Juho, 2019], Funnel Transformer [Dai et al., 2020a], or Perceiver [Jaegle Andrew, 2021]. Inspired by those works, we design the hourglass transformer composed of modalityguided dynamic token merging and extension.

3 Method

In this section, we provide the framework of the proposed Fast-StrucTexT. First, we introduce the model architecture and describe the approach for generating multi-modal input features. Next, we present the details of the hourglass transformer with a hierarchical architecture. Finally, we explain the pre-training objectives and fine-tuning tasks. The overall Fast-StrucTexT architecture is depicted in Figure 3.

3.1 Model Architecture

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 ExtensionBlocks is designed as a hierarchical structure for shortening the sequence length progressively. In particular, the MergingBlock dynamically merges nearby tokens and the ExtensionBlock 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 finetuning downstream tasks of visual document understanding.

3.2 Feature Extraction

We employ an off-the-shelf OCR toolkit to a document image I ∈ RW×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 fV , the pooled RoI features of each text segment extracted by RoI-Align are projected to a vector. For textual features fL, we utilize the WordPiece to tokenize text as subwords and convert them to the ids.

We add a special start tag [CLS] with coordinates B0 = [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 fV , fL to a fixed length. We ensure the length of fL is an integral multiple of fV .

3.3 Hourglass Encoder

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 ∈ {V,L}). Notably, Linear∗→∗ 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-andlanguage tasks. To enhance the interactions among textual and visual modalities, we introduce Symmetry CrossAttention (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 Fq, Fk, Fv are the linear project for query, key and value. σ is the Softmax function, d is the hidden size and Wo 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.

3.4 Pre-training Objectives

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∼9 to express the spatial relationship between each pairwise text segments. We give Gij a layout knowledge, i.e., 0 means the long distance (exceeding half the document size) between text segment i and j, and 1∼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.

3.5 Fine-tuning

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 ∈ Rd×d is the weight matrix of the MLP layer and σ is the Softmax function. For entity ei , Z iF indicates the final fused contextual features, k is the token length and Pi 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 ei and ej , two MLPs are used to project their corresponding features Z iF and Z jF , respectively. After that, a bilinear layer W b is utilized for the relation score Scorei,j .

whereW k ,W v ,W b is the parameter weights that ∈ Rd×d , bk and bv is the bias. σ denotes the sigmoid activation function.

4 Experiment

In this section, we introduce several datasets used for visual document understanding. We then provide implementation details, including our pre-training and fine-tuning strategies for downstream tasks. We conclude with evaluations of FastStrucTexT on four benchmarks, as well as ablation studies.

4.1 Datasets

IIT-CDIP [Harley et al., 2015] is a large resource for various document-related tasks and includes approximately 11 million scanned document pages. We pre-trained our model on this dataset.

FUNSD [Jaume et al., 2019] is a form understanding dataset designed to address the challenges presented by noisy scanned documents. It comprises fully annotated training samples (149) and testing samples (50), and we focused on the semantic entity labeling and linking tasks.

CORD [Park et al., 2019] is typically utilized for receipt key information extraction, consisting of 800 training receipts, 100 validation receipts, and 100 test receipts. We used the official OCR annotations and an entity-level F1 score to evaluate our model's performance.

SROIE [Huang et al., 2021] is a scanned receipt dataset consisting of 626 training images and 347 testing images. Each receipt contains four predefined values: company, date, address, and total. We used the official evaluation tools to evaluate our model's performance.

EPHOIE [Wang et al., 2021] is a collection of 1,494 actual Chinese examination papers with a rich text and layout distribution. The dataset is divided into 1,183 training images and 311 testing images, and each character in the document is annotated with a label from ten predefined categories.

4.2 Implementation

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 × 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. We apply the PaddleOCR to extract the text segment.

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, Df , 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 Lv=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 × H × W = 128 × 4 × 64, followed by a linear project of Dt=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 × 10-5 during the warm-up for 1000 iterations and the keep as 1×10-4 . We set weight decay as 10-4 and (β1 = 0.9, β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×10-5 , 1×10-4 , 1×10-4 , respectively.

4.3 Comparison with the State-of-the-Arts

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-theart model by 156%. Fast-StrucTexT achieves a 1.9× 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 selfbuilt 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 FastStrucTexT 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× speed in Table 1.

4.4 Ablation Studies

We conduct ablation studies on each component of the model on the FUNSD and SROIE datasets, including the backbone, 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 believe that those methods do not take into multi-modal interaction for visual document understanding. 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× the throughput of "w/o M&E" settings, and 0.41% improvement in the labeling task.

FPS‡ FUNSD

120.33 86.62

107.46 88.01

DeformableAttention 83.14 88.57

103.50 90.35

k SROIE-F1 FUNSD-F1 CORD-F1 FPS FLOPs

97.15 81.85 46.48G

97.15 103.50 19.85G

96.85 117.69 12.46G

93.85 133.48 10.14G

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 segmentlevel 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 study the hyper-parameter k as shown 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

8192 1.22 3.82 (+213%) 1933.49G 527.95G (-72%)

Table 7: Ablation study of sequence length, where V.T. represents

Vanilla Transformer.

more pronounced with the sequence length increasing.

Modal interaction ability. 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 than SA.

5 Conclusion

In this paper, we present Fast-StrucTexT, an efficient transformer for document understanding task. Fast-StrucTexT significantly reduces the computing cost through hourglass transformer architecture, and utilizes multi-granularity information through modality-guided dynamic token merging operation. Besides, we propose the Symmetry Cross-Attention module to enhance the multi-modal interaction and reduce the computational complexity. Our model shows state-of-the-art performance and efficiency on four benchmark datasets.

Acknowledgments

This work was supported by the Natural Science Foundation of China (NSFC) under Grants No. 62172041 and No. 62176021.