Understanding Long Document with Different Position-Aware Attentions
Hai Pham†∗, Guoxin Wang‡ , Yijuan Lu‡ , Dinei Florencio‡ , Cha Zhang‡†
Language Technologies Institute, Carnegie Mellon University‡
Microsoft Azure AI
htpham@cs.cmu.edu {guow,yijlu,dinei,chazhang}@microsoft.com

Abstract
Despite several successes in document understanding, the practical task for long document understanding is largely under-explored
due to several challenges in computation andhow to efficiently absorb long multimodal
input. Most current transformer-based approaches only deal with short documents andemploy solely textual information for attention
due to its prohibitive computation and memory
limit. To address those issues in long document understanding, we explore different approaches in handling 1D and new 2D positionaware attention with essentially shortened context. Experimental results show that our proposed models have the advantages for this task
based on various evaluation metrics. Furthermore, our model makes changes only to theattention and thus can be easily used for any
transformer-based architecture.

1

Introduction

The task of document understanding has recently
gleaned many successes (Xu et al., 2020, 2021b;
Appalaraju et al., 2021). This task requires multimodal input that makes it heavier than the text-onlyones, resulting in most models only being capableof dealing with short documents, i.e. having up to
512 tokens. However, there exist long documents
almost everywhere, e.g. contracts, scientific papers,
newsletters, or Wikipedia articles, which are typically longer than 1,000 words. To automaticallysummarize and understand such long documents
urges long document understanding to become an
important task in both NLP and AI.
Long document understanding faces several big
challenges. 1) Recent document understanding
approaches heavily rely on transformer (Vaswani
et al., 2017). However, transformer suffers fromthe quadratic attention that usually limits the input
to 512 words. Therefore, the correlation across∗Work done during an internship at Microsoft.

Figure 1: Distribution of document length in RVLCDIP (Harley et al., 2015), a subset of IIT-CDIP usedpredominantly in the document understanding pretraining task. Most of them are longer than 1,000 words.

long paragraphs/pages is yet to be learned. 2)
Understanding long documents requires power tomodel all long information available, not only justin text but also in other modalities such as spatial
information. For example, LayoutLM (Xu et al.,2020) showed that short document understandingis largely improved by additionally embedding spatial into text information. How to efficiently makeuse of spatial information for long document understanding, however, is still an open and challenging
problem regarding computation cost and adaptability.
For long documents as shown in Figure 1, it
is reasonable to assume that useful information isspanned across their lengths. Especially current
OCR technology, which is essential for data preprocessing, only supports extracting spatial information on every page basis, without the knowledgeof other pages. This behavior poses yet another big
challenge in dealing with long documents, which
requires a proper method to connect information
across pages for all input modalities given.
In this paper, we discover new approaches indealing with long document understanding, whichaddress the aforementioned challenges. We carefully preprocess OCR data to establish the proper
linkages across pages. Then we explore approaches
for directly reducing the heavy attention cost whileachieving high performance, flexibly using thetypical 1D (textual) and/or novelly, 2D (spatial)
reduced contextual information, without adding
more components into the already-heavy transformer (Appalaraju et al., 2021; Nguyen et al.,2021). Despite being simple, we show throughexperiments that both 1D and 2D information canenhance the practicality of transformer-based models while achieving the needed power of handling
long documents.
Our contributions 1) We newly motivate thesimplistic, flexible use of spatial input in attention,
making it plug-able to transformer. 2) We are able
to tackle the document understanding task with input data up to 4096 words. 3) Experimental results
prove the advantages of our approaches on various long-document datasets in comparison to shortmodels for both 1D and 2D contextual information.

2

Related Work

Transformer Attention For Long Documents
There are several methods that address thequadratic cost transformer attention.
Longformer (Beltagy et al., 2020) uses sliding window to
reduce the context, only retrains some sparse global
connections. Similarly, ETC (Ainslie et al., 2020)
embeds relative positions and adds contrastive predictive encoding. Bigbird (Zaheer et al., 2020)
optimizes Longformer’s sliding window by adding
random connections. Our model similarly uses
sliding window but differs in that it exploits layout
input along with the typical text input flexibly anddirectly into attention.
Multimodal Document Pretraining Document understanding largely inherits from multimodal pretraining (Li et al., 2020; Chen et al.,2020; Luo et al., 2020) with the successes from
LayoutLM (Xu et al., 2020, 2021a). Recently,
Docformer (Appalaraju et al., 2021) and StrucuturalLM (Li et al., 2021) introduced a two-pronged
approach: having new pretrain tasks and suitablechanges to the processing or embedding. Probably
Skim-Attention (Nguyen et al., 2021) has the mostrelated motivation for long documents, although
we have a more memory-efficient, and faster wayof handling layout input directly into attention andnot from after the embedding like theirs, and con-sequently support longer input (4096 vs. 2048).

3

Our Model

3.1

Pretrain Model Architecture

We employ Masked Lanuage Model (MLM) architecture as in other document intelligence work, e.g.
Xu et al. (2020, 2021a); Appalaraju et al. (2021),
but make proper changes to enable the capabilityof long documents. Different from a typical MLM,
we have multimodal–instead of text-only input–
which makes the model heavier and thus cannot
deal with long documents without proper changes,
as described in our model shown in Figure 2.
First, we use the sliding-window inspired
from Beltagy et al. (2020), given its lightweightand elegance in limiting the context window, making it significantly more memory friendly. Second,
we introduce new spatial-based attention masks,
in which each context window to a bounding box
is determined by calculating its spatial neighbors,
instead of the given neighboring words. Likewise,
our model does not only use spatial input in theembedding but also in attention directly with preserved spatial correlation. Section 3.3 will elaborate on the establishment and usage of these new
distance masks in comparison with others.
3.2

Post-OCR Processing

This processing is crucially important for long documents that usually have multiple pages because
current OCR engines only generate single-page results, without any connections among pages. Morecurrent models are short models that support up to
512 tokens, and thus discard the rest of valuable
information. Consequently, to make our model
capable of long documents, we process the postOCR data to establish the connections for all input
components among the pages. For example, thebounding boxes on page n need adjusting the coordinates to include the previous n − 1 pages.
3.3

Different Attention Masks

We begin to describe the original transformer attention and our different approaches for long documents, using 1D and 2D input data.
Original Attention Masks (Vaswani et al.,2017) (also known as full attention masks) For
each layer, it is calculated by Equations (1) and (2),
QKTscore(Q, K) = softmax √dk
attn_score(Q, K, V) = score(Q, K) · V,


(1)
(2)

MASKEDLANGUAGE
MODEL

Pre-Training
Target

word2

<CLS>

word1

[MASK]

word3

…

V

Q

TRANSFORMERLAYERS

LINEARADD
NORMALIZE

LINEAR<SEP>

word4094

K

ATTENTIONEMBEDDINGS

bboxPAD

Ebbox1

Ebbox2

Ebbox3…

Ebbox4094

bboxPAD2D (bbox)
Embedding

<CLS>

E1

[MASK]

E3

…

E4094

<SEP>

1D (Text)
Embedding

0

1

2

3

4095

4096

Position
Embedding

bboxPAD

bbox1

bbox2

bbox 3…

bbox4094

bboxPAD

Bounding
Boxes

<CLS>

word1

[MASK]

word3…

word4094

<SEP>

Text

OCR

Figure 2: Our MLM pretrain model architecture. Unlike LayoutLM, we use 1D and 2D input for not only inembeddings but also in transformer attention.

where Q, K, V stand for the learnable Query, Key,and Value matrices respectively. Given the lengthsof these three matrices are all N (input length), thecomplexity of each step is O(N 2 ).
Sliding-Window Masks (Figure 3a) We use thesliding-window approach as inspired from Beltagy
et al. (2020), which limits the context for eachtoken from N down to a smaller M , e.g. N =4096, M = 512, and so the complexity is reduced
to O (N M ).
Kw = get_window(K)
QKT
score(Q, K) = softmax √ w
dkVw = get_window(V)attn_score(Q, K, V) = score(Q, K) · Vw

(3)
(4)
(5)
(6)

Equations 3 and 5, with extraction of random tokens.
Kw = get_2D_window(K)Vw = get_2D_window(V)Spatial Distance Masks (Figure 3c) Different from previous attention types, the M context
tokens for each token is not decided by textual(1D) but instead by spatial (2D) input with somesteps. First, we calculate the centers of all bounding boxes. Second, we fit the kNN algorithm tothe sequence of those points based on L2 distance,
resulting in a 2D distance matrix (having the same
shape N M as sliding window). In summary, we
replace Equations 3 and 5 with Equations 7 and 8.
3.4

Using that intuition, the calculations are now
changed to Equations (3–6), with the added
get_window steps in Equations (3) and (5) 1 .
Sliding-Window plus Random Token Masks(Figure 3b) On top of sliding windows, we add a
few random tokens to establish more connections
to the attention, similarly to Zaheer et al. (2020).
This operation essentially makes changes only to
1
To enable fast calculations in Equations (4) and (6) with
now-changed matrix shapes, one has to extract and chunk thecontexts for all tokens in a way that can exploit fast matrix
multiplication (e.g. by using einsum)

(7)
(8)

Pretrain Model Variants

We build out MLM pretrain architecture with various attention mechanisms for long documents as
described in Section 3.3 and compare their performances in several tasks. Since this change is only
made directly to the attention, our method can beused off-the-shelf for transformer-based architecture with multimodal input.
SW Model This model directly uses SlidingWindow (SW) masks for attention, which significantly reduces the computation and was shown tobe effective for long documents in text-based tasks.

(a) SW(b) SW+R AND(c) D ISTANCE(d) D ISTANCE +SW

Figure 3: Visualization of our models’ different types of attention mask for real samples from RVL-CDIP
dataset (Harley et al., 2015) with limit length of 2048 and context size 512 (for both textual and spatial cases).
Fig 3a is sliding window (SW), Fig 3b is sliding window in blocks with 1-per-block random blocks (SW+R AND),Fig 3c is spatial-based distance mask, and Fig 3d is the combination of sliding window and distance modes. Legend: Attention mask may only have values of 0 and 1, which are represented as the light-yellow background anddark-blue foreground colors, respectively.

SW+R AND Model This model uses SlidingWindow plus Random Token Masks.
D ISTANCE Model This model uses Spatial Distance Masks, with all neighboring contexts being
preemptively computed using kNN, and is implemented in the data loading instead of transformer
encoding phase, not to slow down the main process.
D ISTANCE +SW Model. In this model, we
combine the spatial and textual attention masks
together in a single attention pass, with the motivation of combining the benefits of those two. In
detail, it is done by Equations (3–6), with Equation(3) now being replaced by Equation (7).

4Experiments

4.1

Tasks and Datasets

Pretraining We use IIT-CDIP Test Collection
1.0 2 dataset for our MLM pretraining task. Thisis a large-scale dataset with over 6M multi-page
documents and around 11M pages in total (eachpage is stored as an image and is preprocessed by
an OCR engine).
Document Classification This document classification task uses RVL-CDIP (Harley et al., 2015)
dataset, which is a subset of the pretraining dataset
IIT-CDIP. It comprises 16 classes and each classequally has 25K grayscale images. The document
length distribution is shown in Figure 1.
Sequence Labeling There are two datasets forthis task, namely Kleister-NDA and FunSD.
1) FunSD (Guillaume Jaume, 2019)3 This isa lightweight dataset that has 199 noisy scanned
forms, which contain around 31K words and 9.7K
2
3

https://ir.nist.gov/cdip/
https://guillaumejaume.github.io/FUNSD

entities with 7 given token classes. Though it is
not a long-document dataset (all documents have <512 words), it is useful for ablation studies.
2) Kleister-NDA (Graliński et al., 2020;
Stanisławek et al., 2021)4 This dataset has 540 documents in total (254 train, 83 val, and 203 test) with
2,160 entities annotated and an average of 2,540
words per document. Due to the difficulty in reproducibility with unclear results post-processing, thistask is cast similarly to FunSD with 4 classes.
4.2

Baselines

We compare our 4 model variants (Figure 3) withthe following baselines:
Text: This group consists of models that only
accept text input including BERT (Devlin et al.,2019), RoBERTa (Liu et al., 2019), and other longmodels including Bigbird (Zaheer et al., 2020) andLongformer (Beltagy et al., 2020) 5 .
Text+Layout: This group contains models that
accept both text and layout information, including
LayoutLM (Xu et al., 2020) variants.
4.3

Results and Discussions

Document Classification As shown in Table 1,long models (SeqLen 4096) clearly outperform
short ones in both baseline groups, with or without
layout information added to the input. Furthermore,
all our 4 model variants outperform all the baselines.
This result concurs with our observation thatlong documents have valuable information spanned
45

https://github.com/applicaai/kleister-nda

Our SW and SW+R AND models share the similarity withthose last two ones, with difference of handling multimodal
input for document intelligence.

ModelSeqLen Acc (%) ↑BERT-base
512
89.81
RoBERTa-base
512
90.06
BERT-large
512
89.92
RoBERTa-large
512
90.11
Text
Bigbird-base
4096
93.48
Longformer-base
4096
93.85
Bigbird-large
4096
93.34
Longformer-large
4096
93.73
LayoutLM-base
512
91.88
LayoutLM-large
512
91.90
Ours SW
4096
94.50
Text+Layout
Ours SW+R AND
4096
95.25
Ours D ISTANCE
4096
94.79
Ours D ISTANCE +SW 4096
94.69
Type

4.4

The purpose of this study is to explore how longmodels perform on short documents, which also appear in practice, to see whether they can generalize
their performance to shorter data.
Model
SeqLen F1 ↑BERT-base
512 60.3
RoBERTa-base
512 66.5
BERT-large
512 65.6
RoBERTa-large
512 70.7
Text
Bigbird-base4096 45.8
Longformer-base4096 71.4
Bigbird-large
4096 46.8
Longformer-large
4096 73.5
LayoutLM-base
512 78.7
LayoutLM-large
512 79.0
Ours SW
4096 69.9
Ours SW+R AND
4096 77.1
Text+Layout
Ours D ISTANCE
4096 64.0
Ours D ISTANCE +SW 4096 61.8
Type

Table 1: Classification accuracy for RVL-CDIP. For
this long-document dataset, the models capable of using 4096 words uniformly beat other models and layout
information helps with the task compared with using
Text input. All our long models show their advantages
on this long dataset.
Model
SeqLen F1 ↑BERT-base
512 47.06
BERT-large
512 52.66
Text
Longformer-base4096 61.78
Bigbird-base4096 46.98
LayoutLM-base
512 55.69
LayoutLM-large
512 61.95
Ours SW
4096 64.06
Text+Layout
Ours SW+R AND
4096 58.92
Ours D ISTANCE
4096 57.01
Ours D ISTANCE +SW 4096 44.70
Type

Table 2: Results on Kleister-NDA. Although thisdataset is challenging, long models still show advantages over short ones.

across the length. And importantly, our models
show advantages of handling long multimodal input, and hence are more practical with real datathat are usually longer than 512 tokens.
Sequence Labeling with Kleister-NDA6 Comparing the “base” versions (separated from their“large” counterparts), Table 2 shows that most of
our models, which are also the “base” ones, clearly
have better scores. Particularly, our SW model isthe best performer.
Furthermore, our D ISTANCE +SW is not performing equally well. Our hypothesis is that the
OCR engine cannot understand the decoying annotation in this dataset, and thus generates spatial
results that do not correlate well with the text. Consequently, the combination of textual and spatial
information does not result in the benefits of thosetwo.
6
The results are from the validation split due to no annotation for the test split provided in the dataset.

Ablation: Long Models on Short Dataset

Table 3: Results on FunSD dataset. As usual, layout
information is helpful in boosting performance. However, long models do not perform well compared withshort models on this small, short-document dataset.

Table 3 shows that on FunSD, we see again thatlayout information generally helps in the case ofmultimodal input. However, long models do not
perform well compared to short ones, although
the gap between the best of ours and the baselines
are not very far away (77.1 vs. 79.0). The main
reason is that long models essentially have much
more parameters than short ones. And not only is
FunSD short, it is also very small. As a result, thelimited phase of fine-tuning on only 199 samples
can hardly tune parameters well for good results.
Especially, since all documents are short, most long
input to the model is zero padding and thus notenough for contributing for better scores.
Another reason is that long models have theirembedding representations trained for the lengthof 4096 tokens and hence are hard to adapt to 512token input with just a few fine-tuning steps. As
a result, analyzing the data well to design suitable
pretraining and fine-tuning models is very important.
The next 2 studies will explore the implicationsof the newly-added spatial attention masks in ourmodels.
4.5

Ablation: Different-Length Documents

This study aims to explore how the models work
if we do not cut any information from documents(the models take input up to their maximum length

ACCURACY (%)

96

Figure 4: RVL-CDIP performance on different document types based on their original lengths (i.e. without purging) with LayoutLM (with the best “large”
version) and our spatial models (D ISTANCE and D IS TANCE +SW). Our models are consistently better.

limit). Out of 40K test samples in RVL-CDIP, there
are 9268 samples with length ≥ 512, 2312 withlength ≥ 1024, and only 106 with length ≥ 2048.
Figure 4 shows the consistent observation that
our models are much better than LayoutLM, andyet perform slightly worse as the original document
length increases. There could be several possible
reasons for this behavior: the models are not well
pretrained and/or fine-tuned, many long documents
have lots of confusing parts, or there are many
noises in OCR results.
4.6

Ablation: Different Max Input Lengths

Given the pretrained models that can accept input
up to 4096 tokens, we finetune them with the input
of different maximum lengths, i.e. excess will bepurged. As a result, we use all 40K test samples in
RVL-CDIP for this study.
As shown in Figure 5, our models are better andbetter as more tokens are absorbed, thus once again
confirming our intuition that valuable informationis spanned across the length. As a result, if themodel capacity permits, we should not limit thecapacity to 512 tokens as in most current models
in the literature.
4.7

Further Discussion on Spatial Masks

As seen in the above experimental results, direct
usage of 2D layout context information in the transformer attention has some advantages. However,its performance does not match the typical usageof 1D textual information. This might be discouraging at first since introducing spatial masks brings
heavier computation compared to textual ones. We
hypothesize the drawbacks are due to some objective limitations. First, the kNN suffers someinaccuracy compared with normal (and slow) cal-95

Ours DISTANCE+SW
Ours DISTANCE

94

93.81

93.14

93

92.89

94.04

94.79
94.69

93.29

9291.85
91 5121024

2048

MAX LENGTH

4096

Figure 5: RVL-CDIP performance on different
maximum lengths using our D ISTANCE and D IS TANCE +SW models. For each of lengths 512, 1024,
2048, and 4096, the test set contains the same 40K samples. A longer maximum length gives better results.

culations. Second, the performance of the whole
pipeline heavily depends on OCR quality, e.g. in
Kleister-NDA with decoy design, OCR results are
not well aligned with the text. Consequently, we
conjecture that with future development in OCR
technologies, the use of spatial masks would be
more and more helpful in practice.

5

Conclusion and Discussion

We propose a versatile solution for long document
understanding, in which the shortened context canbe used in the form of textual and/or layout input for the attention mechanism in a flexibly pluggable manner. We keep our approach simple by
not putting extra overhead on complicated embedding or encoding methods. Despite its simplicity, our solution has shown promising experimental results on document understanding tasks withlong, multimodal input. In the future, we will further reduce the memory consumption of models
with given multimodal input and speed up the pretraining. Similar to LayoutLM, pretraining usuallytakes 80 hrs/epoch with 8 V100 GPUs. Thus there
are certainly lots of room for improvement to makethese models more efficient and practical.

6

Acknowledgements

The authors would like to thank Microsoft for providing GPU training resources for the projects, thecolleagues in Microsoft, and last but not least, theanonymous reviewers for helpful discussions andfeedback.

References
Joshua Ainslie, Santiago Ontanon, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham, Anirudh
Ravula, Sumit Sanghai, Qifan Wang, and Li Yang.
2020. Etc: Encoding long and structured inputs
in transformers. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language
Processing (EMNLP), pages 268–284.
Srikar Appalaraju, Bhavan Jasani, Bhargava Urala
Kota, Yusheng Xie, and R Manmatha. 2021. Docformer: End-to-end transformer for document understanding. ICCV.
Iz Beltagy, Matthew E Peters, and Arman Cohan.
2020. Longformer: The long-document transformer.
arXiv preprint arXiv:2004.05150.
Yen-Chun Chen, Linjie Li, Licheng Yu, Ahmed
El Kholy, Faisal Ahmed, Zhe Gan, Yu Cheng, andJingjing Liu. 2020. Uniter: Universal image-text
representation learning. In European Conference on
Computer Vision, pages 104–120. Springer.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, andKristina Toutanova. 2019. BERT: Pre-training ofdeep bidirectional transformers for language understanding. pages 4171–4186.
Jonathan Frankle and Michael Carbin. 2019. The lottery ticket hypothesis: Finding sparse, trainable neural networks. ICLR.
Filip Graliński,
Tomasz Stanisławek,
Anna
Wróblewska, Dawid Lipiński, Agnieszka Kaliska,Paulina Rosalska, Bartosz Topolski, and Przemysław Biecek. 2020. Kleister: A novel task
for information extraction involving long documents with complex layout.
arXiv preprint
arXiv:2003.02356.
Jean-Philippe Thiran Guillaume Jaume, Hazim Kemal Ekenel. 2019. Funsd: A dataset for form understanding in noisy scanned documents. In Acceptedto ICDAR-OST.
Adam W Harley, Alex Ufkes, and Konstantinos G Derpanis. 2015. Evaluation of deep convolutional nets
for document image classification and retrieval. International Conference on Document Analysis andRecognition.
Chenliang Li, Bin Bi, Ming Yan, Wei Wang, Songfang Huang, Fei Huang, and Luo Si. 2021. Structurallm: Structural pre-training for form understanding. ACL.
Xiujun Li, Xi Yin, Chunyuan Li, Pengchuan Zhang, Xiaowei Hu, Lei Zhang, Lijuan Wang, Houdong Hu,
Li Dong, Furu Wei, et al. 2020. Oscar: Objectsemantics aligned pre-training for vision-languagetasks. In European Conference on Computer Vision,
pages 121–137. Springer.

Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis,
Luke Zettlemoyer, and Veselin Stoyanov. 2019.
Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692.
Huaishao Luo, Lei Ji, Botian Shi, Haoyang Huang, Nan
Duan, Tianrui Li, Xilin Chen, and Ming Zhou. 2020.
Univilm: A unified video and language pre-training
model for multimodal understanding and generation.
arXiv preprint arXiv:2002.06353.
Laura Nguyen, Thomas Scialom, Jacopo Staiano,and Benjamin Piwowarski. 2021. Skim-attention:
Learning to focus via document layout. EMNLP
Findings.
Tomasz Stanisławek,
Filip Graliński,
Anna
Wróblewska, Dawid Lipiński, Agnieszka Kaliska,Paulina Rosalska, Bartosz Topolski, and Przemysław Biecek. 2021. Kleister: Key information
extraction datasets involving long documents withcomplex layouts. arXiv preprint arXiv:2105.05796.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé
Jégou. 2021. Training data-efficient image transformers & distillation through attention. In International Conference on Machine Learning, pages
10347–10357. PMLR.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob
Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz
Kaiser, and Illia Polosukhin. 2017. Attention is all
you need. Proceedings of the 31st International
Conference on Neural Information Processing Systems (NeurIPS).
Yang Xu, Yiheng Xu, Tengchao Lv, Lei Cui, Furu
Wei, Guoxin Wang, Yijuan Lu, Dinei Florencio, Cha
Zhang, Wanxiang Che, et al. 2021a. Layoutlmv2:
Multi-modal pre-training for visually-rich document
understanding. ACL.
Yiheng Xu, Minghao Li, Lei Cui, Shaohan Huang,
Furu Wei, and Ming Zhou. 2020. Layoutlm: Pretraining of text and layout for document image
understanding. In Proceedings of the 26th ACM
SIGKDD International Conference on Knowledge
Discovery & Data Mining, pages 1192–1200.
Yiheng Xu, Tengchao Lv, Lei Cui, Guoxin Wang, Yijuan Lu, Dinei Florencio, Cha Zhang, and Furu Wei.
2021b. Layoutxlm: Multimodal pre-training for
multilingual visually-rich document understanding.
arXiv preprint arXiv:2104.08836.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava
Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang,
Li Yang, et al. 2020. Big bird: Transformers forlonger sequences. Advances in Neural Information
Processing Systems (NeurIPS), 33.

Parameter Name
Value
do_lower_case
true
fp16
true
fp16_backend
amp
gradient_accumulation_steps
4
max_seq_length
4096
max_2d_position_embeddings1024
max_steps
1000000
model_name_or_path
allenai/longformer-base-4096
dataloader_num_workers
64
tasks
mask_lm
optimizer
transformers_AdamW
learning_rate
5e-5
warmup_ratio
0.1
weight_decay
0.01
whole_word_masking
false
add_prefix_spacetrueattention_window
512

Table 4: Main pretrain hyperpameters on the MLM
pretraining task for the ITT-CDIP large-scale dataset.
There are 3 variants share this set of parameters that are
Ours SW, Ours D ISTANCE and Ours D ISTANCE +SW
models. All of them use the pretrained weights from
Longformer-base (Beltagy et al., 2020) model.

A

More Information on the Pretrain
Task

Pretrain Data Preprocessing As described, forpretrain model we retain the same OCR engine forgenerating and aligning layout and text informationfrom LayoutLM (Xu et al., 2020). The task is also
the same, which is Masked Language Modeling(MLM). To deal with long documents, we have toimplement the additional sliding-window, randomblock and distance-based masks.
Pretrain Model Implementation Our solution
only makes changes to the attention module, inwhich uses can choose to use any types of attention
masks from the 4 variants illustrated in Figure 3.
For the SW and SW+R AND models which arealso our new pretrain models, we implement thelayout-related part on top of the original BigBird 7
and Longformer 8 implementations from Huggingface’s transformers, respectively. Otherwise thedistance-based masks, which are employed in D IS TANCE and D ISTANCE +SW models, are newly
implemented as a pluggable module.
Training MLM We pre-train the task on the
IIT-CDIP datasets, using a single-node multi-GPU
mode. Each job was run on a server with 8 V100
7

https://huggingface.co/transformers/model_doc/
bigbird.html
8
https://huggingface.co/transformers/model_doc/
longformer.html

Nvidia GPUs, each of which has 32GB memory
and fast processors. For text-only models, pleaserefer to LayoutLM’s github 9 .
For SW model, we use the public pretrained
weights from Lomgformer (Beltagy et al., 2020).
Other of our models employ the same set of parameters, except for the pretrained weights, inwhich SW+R AND model uses the weights from
Bigbird (Zaheer et al., 2020) and the last twomodels having distance masks (D ISTANCE andD ISTANCE +SW models) use the same pretrained
weights as SW model, as demonstrated in Table 4.
It is also worth noting that the pretrained weightsfrom Longformer and Bigbird models are usefuleven for the models using distance masks because
those two model families support documents withlength 4096, so the position embeddings are helpful. For speed and memory tradeoff, we limit thecontext for distance masks to only 128 (vs. 512 in
textual contexts), without sacrificing much performances, as reported in Section 4.3.
Training Notes Although not reported in themain content, we note some lessons learned fromthe pretraining task. As we observe, the Ours SW
model consistently achieves the best results, while
consuming the least GPU memory. For the basemodel, it only consumes about 7 GB GPU memory and Ours D ISTANCE +SW that uses slidingwindow attention on its half processing also consumes about 9 GB memory. Both models, as aresult, can be deployed well on a broad range ofGPUs in the market.
Unlike those conveniences, Ours SW+R AND
and Ours D ISTANCE do not share the same advantages. In fact, they consumes about more than
30GB GPU memory each, limiting their practicality. We hypothesize the main reason for suchdrawbacks is that they have random, inconsistentpatterns, and hence there is no efficient way to takeadvantage of fast memory-efficient and fast matrix
operations.
Finally, although showing promising practicalbehaviors, all baselines and our models, and almost any transformer-based ones are certainly notlightweight models. And although there are advancements in compressing those heavy models(e.g. (Touvron et al., 2021; Frankle and Carbin,2019), there seems to be a considerable way to go
for making these model run on mobile devices in
the near future.
9

https://github.com/microsoft/unilm

B

More information on Finetuning Tasks

As described in the main content, after pretraining,
the saved models are the backbone for the respective fine-tuning model types. For that reason, theparameters are mostly shared with their pretrain
counter-part models, e.g. Table 4 for Ours SW
models. Generally, we keep the same optimizer
and batch size of 32 (combined across all used
parallel GPUs).
For RVL-CDIP in the document classificationtask, we use the SequenceClassification
model type. On top of the pretrain skeleton, we
add a small classifier with 2 fully-connected layers
and a drop-out layer in between. The final output isthe single class for the whole sequence/document.
For FunSD and Kleister-NDA datasets, we instead use the TokenClassification model
type, which is designed to classify all-document
entities. The similar classifier is added to the pretrained skeleton, now with a different usage inwhich each token/entity is to be classified into 1 of
the number of given classes.
What’s more, to preprocess these two datasets,
we have to ingest all available document tokens.
Likewise, with documents longer than the maximum lengths, we need to cut those documents, andrecursively treat the overflowing parts in the sameway. In terms of implementation, unlike FunSD
that is lightweight, we always want to avoid loading
the whole dataset into the memory but rather takeadvantage of the data buffering in feeding to themodels. As a result, we pre-process all data first,
save them to disks and only load the respective
parts when needed.
Additional Information for Kleister-NDA It
is worth recalling that the evaluation of it is tricky
if using the provided official GEval evaluation
script (Graliński et al., 2020)10 . In detail, given
the predited tokens, one has to retrieve the associated texts in a group. For example, the beginningof an entity group usually starts with a class beginning with "B-", followed by a series of "I-" tokens.
However, there is no guarantee that the prediction
will always return a group having this meaningful
pattern, let alone many other complicated casesthat can happen. Such complications make thepost-processing of the prediction– before feeding
to GEval–very difficult and importantly, not easily
reproducible. In fact, amongst recent papers thatreport performance on this dataset (e.g. in Xu
10https://github.com/applicaai/kleister-nda

Figure 6: More distance masks from RVL-CDIP samples with the limit length of 2048 and 512 neighbors
each.

et al. (2021a); Appalaraju et al. (2021)), there isreference code with which for us to compare.
Consequently, we treat this dataset the same as
FunSD, given their similarity in annotation. In
addition, because this dataset is larger and much
more difficult (due to decoying texts) compared to
FunSD, we analyze the train dataset and employ theweighted loss based on the distribution the givenlabels. As a result, our method is more transparent
and reproducible.

C

Additional Samples on Distance Masks

Complementary to Figures 3c and 3d, we present
some more distance masks based on real samples
taken from RVL-CDIP with the same setting in
Figure 6.