본문 바로가기

논문/Multimodal

ERNIE-Layout: Layout Knowledge Enhanced Pre-trainingfor Visually-rich Document Understanding

Abstract

최근 많은 visuall-rich document understanding의 pre-training 기법이 많은 발전을 이루었지만 여전히 layout-centered knowledge의 사용에는 부족함이 있다. 이 논문에서는 ERNIE-Layout이라는 text와 layout 및 image 정보를 잘 합쳐서 layout 지식을 향상시키는 모델을 제안한다. 구체적으로 먼저 input sequence를 재정렬한 다음 문서의 적절한 읽기 순서를 학습하기 위해  pre-training task인 읽기 순서 예측을 제시한다. 모델의 layout 인식을 향상시키기 위해 spatial-aware disentengled attion을 multi-modal transformer에 합치고 replaced region prediction task를 pre-training 했다. 실험결과는 다양한 downstremtask에서 SOTA이다.

 

1. Introduction

 Visually-rich document Understanding(VrDU)는 다양한 유형의 스캔 또는 디지털 기반 비즈니스 문서(예: 양식, 송장)를 처리하는 것을 목표로 하는 중요한 연구 분야로, 다양한 응용 프로그램으로 인해 산업 및 학계에서 큰 관심을 끌었다. Natual language understanding(NLU) tasks가 text만 사용하는 것과는 달리 VrDU는 가장 원시적인 데이터 특성을 사용한다. 

 The preliminary works for VrDU

  • Uni-modal or shallow multi-modal fusion approaches
  • LayoutLM, LayoutLMv2: pre-training language models

 Document understanding model의 특징은 layout knowledge를 얼마나 잘 사용하는가에 달려있다는 점에서 다른 모델과 차이가 있다.

 일반적으로 단순히 2D 좌표를 model layout 인식을 위해 부여하는 함정에 빠진다. VrDU의 특징을 고려하면 layout-centered 지식은 체계적으로 수집되어야하고 두 가지 방면에서 이용되어야 한다.

  1. Layout 정보는 문서의 읽는 순서를 잘 반영해야한다. 기존에는 단순 Optical Character Recognition (OCR) 결과를 좌측상단에서 우측 하단으로 serialization 하여 사용했다.
  2. 한편, Layout은 세 번째 modality인 반면, 현재 모델은 layout을 특수 위치 feature로 사용한다. Layout 정보를 입력 layer에 embedding하는 것과 (Xu et al., 2020) attention layer에 bias로 사용(Xu et al., 2021). 

즉 text/image 와의 cross-modal intereaction이 부족하다.

 

 ERNIE-Layout 모델 요약

  1. Off-the-shelf layout-based document parser를 serialization 단계에서 사용 사람의 읽기 습관과 더 일치하도록 함
  2. 각 textual/visual token을 position/layout embedding 후 stacked multimodal transformer layer에 보낸다.
  3. Cross-modal interaction 향상을 위해 spatial-aware disentangled attention mechanism을 이용 (DeBERTa - 토큰 간 attention 가중치는 hidden states와 relative position을 기반으로 disentangled matrices를 사용하여 계산, He et al.,2021)

 학습 방식

  1. Reading order prediction: 다음 token을 예측
  2. Masked visual-language modeling
  3. Text-image alignment task 
  4. Replaced region prediction

모델 평가에 사용된 VrDU downstream tasks

  1. The key information extraction task
    : FUNSD (Jain and Wigington,2019), CORD (Park et al., 2019), SROIE (Huanget al., 2019), Kleister-NDA (Grali ́nski et al., 2021) datasets
  2. The  document  question  answering  task with the DocVQA (Mathew et al., 2021) dataset
  3. The document image classification task with the RVL-CDIP (Harley et al., 2015) dataset

Contribution

  • 입력 token 재배열과 reading order prediction task을 pre-training 단계에서 수행
  • disentangled attention mechanism을 사용하며 replaced prediction pre-training task을 설계
  • 다양한 VrDU Task에서 SOTA

 

2. Related Work

Layout-aware Pre-trained Model.

  • LayoutLM (Xu et al.,2020)은 최초로 2D 좌표로 layout embeddings 를 각 token에 했고masked language modeling pre-training task (De-vlin et al., 2019) 를 적용시켰다.
  • LayoutLMv2 (Xuet al., 2021) 문서 image patches를 textual tokens과 연결시키고, 두 pre-training tasks, text-image matching과 text-image alignment를 제안했다.
  • StructralLM (Li et al., 2021a) 단어 수준이 아닌 segment 수준의 layout feature를 활용하여 어떤 단어가 동일한 cell에서 오는지 모델이 인식하도록 한다
  • DocFormer (Appalaraju et al., 2021) 학습된 spatial embedding을 모든 modality가 공유하여 모델이 text를 visual token과 쉽게 연관시키고 그 반대도 가능하게 한다.
  • TILT (Powalski et al.,2021) sequence labelling limitation을 해결하기 위해 input sequence에 명시적으로 포함되지 않은 결과를 생성하는 encoder-decoder model을 제안한다.
  • 그러나 위 방법들은 raster-scanning 방법을 사용하여 사람이 읽는 방식과는 차이가 있다. 이 문제를 해결하기 위해 LayoutReader (Wang et al., 2021c) 는 적절한 순서로 문서를 읽기 위한 seq2seq framework를 개발했다. 그러나 순서에 집중한 나머지 문서 이해 능력은 떨어졌다. 

Knowledge-enhanced  Representation.

BERT (Devlin et al., 2019) 이후로 유익한 representation 학습을 위한 pre-trained language model이 여럿 개발 되었다.

  • 일부 연구에서는 WikiData 및 WordNet의 사실과 같은 추가 지식이 pre-trained language model에 더 도움이 될 수 있지만 text의 word embedding과 knowledge graph의 entities는 동일한 벡터 공간에 있지 않으므로 이를 해결하기 위한 모듈이 필요하다. 
  • ERNIE (Sunet al., 2019)는 인간 knowledge를 language model에 통합하기 위해 pre-training에서 entity-level mask를 제안
  • SpanBERT(Joshiet al., 2020)는 text span을 더 잘 표현하고 예측하기 위해 schema and training objectives를 수정
  • BERT-wwm (Cui et al., 2021)는 Chinese language models을 위한 whole word masking strategy 적용
  • ERNIE-ViL (Yu et al., 2021) scene graphs로 부터 얻은 구조화된 knowledge를 이용하여 vision-language의 joint representation을 학습

위 방법들에 영감을 받아 사람이 읽는 순서를 이용하고자 했다

 

3. Methodology

Figure 1: The architecture and pre-training objectives of ERNIE-Layout. The serialization module is introduced to correct the order of raster-scan, and the visual encoder extracts corresponding image features. With the spatial- aware disentangled attention mechanism, ERNIE-Layout is pre-trained with four tasks

Figure 1은 ERNIE-Layout의 개요를 보여준다. Layout knowledge를 이용하여 token sequence를 재배열하고 visual encoder로부터 visual feature를 뽑아낸다. Textual 과 layout embedding은 linear projection을 통해 textual feature로 합쳐지고 visual feature 역시 linear projection을 한다. 이후 feature들은 concatenated 되어 spatial-aware disentangled attention mechanism이 포함된 transformer에 입력된다. ERNIE-Layout은 reading order prediction, replaced region prediction task, masked visual-language modeling, text-image alignment task 4가지의 pre-training을 수행한다.

 

3.1 Serialization Module

우선, OCR을 사용하여 문서의 단어와 해당 좌표를 얻는다. 우선 raster-scan방식으로 input sequence를 생성한다. 

Figure 2: The effect of layout knowledge enhanced serialization compared with vanilla raster-scanning order. By using Document-Parser, the perplexity of the document with a complex layout is significantly reduced.

하지만 이 방식은 사용하기는 쉽지만 Figure 2와 같은 문서에서 raster-scan 방식을 사용하면 각 cell 내의 정보들이 분리되므로 적절하지 못하다.

 사람이 읽는 방식과 유사한 입력을 얻기 위해 Layout-Parser 기반의 Document-Parser를 사용했다. Figure1에서 우선 OCR을 이용하여 인식된 boxes와 words를 기반으로 문서의 elements(문단, 목록, 표, 그림)를 찾아내고 특정 알고리즘을 사용하여 서로 다른 element의 특성을 기반으로 단어 간의 논리적 관계를 얻고 적절한 읽기 순서를 얻는다.

 Layout knowledge 향상 serialization의 이점을 분석하기 위해 perlexity(PPL)을 평가 지표로 사용했다. Figure2에서 보면 Document-Parser를 이용한 방식이 raster-scanning order보다 낮은 PPL을 보여주는 것을 확인할 수 있다. 

 

3.2 Input Representation

Text Embedding.

 BERT와 같이 [CLS]와 [SEP] special token이 text sequence의 처음과 끝에 추가된다. Token sequence $T$에 대해

$$T = E_{tk}(T) + E_{1p}(T) + E_{TP}(T)$$

여기서 $E_{tk}$, $E_{1p}(T)$, $E_{TP}(T)$는 각각 token embedding, 1D position embedding 그리고 token type embedding layer를 의미한다. 

Visual Embedding.

 Visual feature를 뽑기 위해 Faster-RCNN(Ren et al.,2015) 을 backbond으로 한 visual encoder로 사용했다. 문서의 image는 $224 \times 224$로 resize 되었으며, adaptive pooling layer를 적용하여 output 을 고정된 width(W)와 height(H)의 feature로 변환했다. (해당 연구에서는 7을 값으로 사용) 이후 feature를 visual sequence $V$로 flatten한 후 text embedding과 같은 크기로 linear layer $F_{vs}( \cdot )$을 이용하여 project한다. 유사하게 1D position 과 token type $[\texttt{V}]$ 가 추가적으로 고려되어 visual embedding을 생성한다.

$$V = F_{vs}(V) + E_{1p}(V) + E_{TP}([\texttt{V}])$$

Layout Embedding.

 Textual token에 대해 OCR을 사용하면 각 textual token에 대해 2d좌표 $(x_0, y_0, x_1, y_1, w, h)$ 가 생성된다. 모든 좌표는 $[0, 1000]$ 으로 normalized 한다. Visual token에 대해 유사한 계산 과정이 수행된다. Embedding layer는 수평과 수직 두 방향으로 생성했다.

$$ L=E_{2x}(x_0, x_1, w)+E_{2y}(y_0, y_1, h) $$

여기서 $2x$와 $2y$는 각각 x축과 y축 embedding layer를 의미한다. 

 

 결과적으로 input representaton $H$를 얻기 위해 Layout embedding을 각각 Text, Visual embdding에 더해주고 concatanate한다. 이렇게 얻어진 $H$의 길이는 $N+HW$로 $N$는 textual part의 max length이다.

$$H=[T+L;V+L]$$

3.3 Multi-modal Transformer

 생성된 input representation에서 modality 간의 연관성을 잘 고려하기 위해서 disentangled attention of DeBERTa (He et al., 2021)를 참고했다. DeBERTa는 Attention 가중치는 content의 상대적 위치에 대한 disentangled matrices를 사용하여 계산된다. 해당 논문에서는 spatal-aware disentangled attention을 적용한 multi-modal transformer를 통해 layout feature를 잘 활용하고자 한다. 

 우선 1D 위치를 예로들어 $i$와 $j$ token에 대한 상대적 거리 $\delta_{1p}(\cdot )$을 정의하면, 2d layout에 대한 x-axis와 y-axis의 정의도 같다. 

$$ \delta_{1p}(i,j)=\left\{\begin{matrix}
0 & \texttt{for} \quad i-j\leqslant -k\\
2k-1 & \texttt{for} \quad i-j\geq k \\
i-j+k & \texttt{others,} \\
\end{matrix}\right.$$

DeBERTA : Decoding-Enhanced BERT with Disentangled Attention (slideshare.net)

코드 상에서 k의 default 값은 max_rel_pos=128, max_rel_2d_pos=256이다.

 다음 input 차원과 같은 relative position vectors를 얻기 위해 relative position embedding table $E^{'}_{1p}$, $E^{'}_{2x}$, $E^{'}_{2y}$를 1d, 2d x축, 2d y축에 대해 세 개를 사용한다. 일련의 Projection matricies은 contents vectors 뿐만 아니라 relative position vectors를 $Q^ \star$, $K^ \star$, $V^ \star$ 로 매핑한다. 여기서 $\star \in \{ct, 1p, 2x, 2y\}$. Attention 계산 과정에서, 1D/2D feature와 contents의 관계를 알기 위해 4개의 파트로 raw score를 decouple한다. 

Figure 3: The internal working principle of spatial-aware disentangled attention.

$$\begin{matrix}
A_{ij}^{ct, ct}=Q_i^{ct}{K_j^{ct}}^T\\
A_{ij}^{ct, 1p}=Q_i^{ct}{K_{\delta _{1p}(i,j)}^{1p}}^T + K_j^{ct}{Q_{\delta _{1p}(k,i)}^{1p}}^T\\
A_{ij}^{ct, 2x}=Q_i^{ct}{K_{\delta _{2x}(i,j)}^{2x}}^T + K_j^{ct}{Q_{\delta _{2x}(k,i)}^{2x}}^T \\
A_{ij}^{ct, 2y}=Q_i^{ct}{K_{\delta _{2y}(i,j)}^{2y}}^T + K_j^{ct}{Q_{\delta _{2y}(k,i)}^{2y}}^T \\
\end{matrix}$$

 마지막으로 이 모든 attention score를 더해서 attention matrix $\hat A$을 얻는다. 마지막으로 scaling과 normalization후 spatial-aware disentangled attention의 출력은 다음과 같다.

$$\hat A_{ij}=A_{ij}^{ct,ct}+A_{ij}^{ct,1p}+A_{ij}^{ct,2x}+A_{ij}^{ct,2y}$$

$$H_{out} = \texttt{softmax}(\frac{\hat A}{\sqrt{3d}})V^{ct}$$

3.4 Pre-training Tasks

Reading Order Prediction (ROP).

Serialization 결과는 단어와 2D 좌표를 포함하여 여러 text segment로 구성한다. 그러나 transforemer에서 받은 input sequence의 text segment 간에는 명확한 경계가 없다. 모델이 layout 지식을 잘 활용하여 읽는 순서를 이해하기 위해서 잘못된 순서로 입력을 받았을 때도 잘 작동하기를 원한다. Reading Order Prediction (ROP)을 통해 $\hat A$가 읽는 순서에 대한 지식을 함유하고 있기를 원한다. $j$번째 token이 $i$번째 token의 다음 token일 확률을 $\hat A_{ij}$로 표현한다. ground truth 행렬 $G$는 0과 1의 행렬로 1은 읽는 순서에 관련이 있다는 뜻이다. 마지막 token에 대해 다음 token은 자기 자신이다. 최종적으로는 아래의 Cross-Entropy식을 계산한다.

$$ L_ROP = - \sum_{0\leq i< N}\sum_{0\leq j< N}G_{ij} \texttt{log} (\hat A_{ij}) $$

 

Replaced Region Prediction (RRP).

10%의 패치는 임의로 다른 image로 교체되어 multi-modal transformer의 입력으로 encode되어 들어간다. [CLS] 의 vector output은 어떤 패치가 변했는지 예측하는데 사용되며 손실은 다음과 같다. 

$$L_{ROP} = - \sum_{0\leq i< HW}[G_i\texttt{log} (P_{i})+(1-G_i)\texttt{log} (1-P_{i})$$

여기서 $G_i$는 교체된 패치의 golden label이고 $P_i$는 예측의 nomalized probability이다.

 

Masked Visual-Language Modeling (MVLM).

Text token을 maksing한 후 복원

 

Text-Image Alignment (TIA).

몇 text 줄을 임의로 선택하고 관련된 region의 이미지가 맞는지 아닌지 분류기를 통해 예측한다. 

 

최종적으로 pre-training objective는 다음과 같다.

$$L =L_{ROP}+L_{RRP}+L{MVLM}+L{TIA}$$

 

4 Experiments

4.1 Datasets

실험의 공정성을 위해서 layout knowledge enhaced serialization 을 pre-training data의 읽기 순서 재배열에만 사용한다. 즉 ERNIE-Layout 은 다른 모델과 동일한 입력을 받는다.

 

Pre-training.

IIT-CDIP Test Collection (Lewis et al., 2006) 에서 임의로 1천만장을 추출하여 학습에 이용

Fine-Tuning.

  1. The key information extraction task
    : FUNSD (Jain and Wigington,2019), CORD (Park et al., 2019), SROIE (Huanget al., 2019), Kleister-NDA (Grali ́nski et al., 2021) datasets
  2. The  document  question  answering  task
    : DocVQA (Mathew et al., 2021) dataset
  3. The document image classification task
    : RVL-CDIP (Harley et al., 2015) dataset

Table 1: Statistics of datasets for downstream tasks

4.2 Settings

Pre-training.

ERNIE-Layout은 1024 hidden unit과 16개의 attention head로 이루어진 24개의 transformer layer로 구성된다. Maximum sequence length는 각각 textual은 512 visual은 49 tokens 이다. Transformer는 RoBERTa-large (Liu et al., 2019)와 visual encoder는 Faster-RCNN (Ren et al., 2015)으로 초기화 되고 나머지 parameter는 임의로 초기화 된다. 

Adam optimizer 사용 learning rate 1e-4 / weight decay 0.01 / learning rate는 10% step까지 warm up 하고 0으로 decay

학습은 24개의 Tesla A100 GPU로 batch size 576으로 20 epochs 수행

 

Fine-tuning.

Key information extraction task (FUNSD, CORD, SROIE, Kleister-NDA)는 sequence labeling framework로 도입하여 token-level classificatino layer를 추가하여 BIO labels 예측하도록 했다. 

Document qustion answering task (DocVQA)는 extrative question-answering paradigm을 따라서 token-level classifier 를 통해 output representation으로 정답의 start 와 end position을 예측하도록 한다. 

Document image classification (RVL-CDIP) [CLS]의 representation을 이용하여 fully-connected network가 document label을 예측하도록 한다. 

모든 downstream task에 대해 Adam optimizer 사용 learning rate 2e-5 / weight decay 0.01 이외는 pre-training과 상동

Table 2: Hyper-parameters for downstream tasks

4.3 Results

Table 3: Results (Entity-level F1 score) of ERNIE-Layout and previous methods on the Key Information Extraction task (FUNSD, CORD, SROIE, Kleister-NDA). The highest and second-highest scores are bolded and underlined.

 

Table 9: Results of ERNIE-Layout (base-level model) and previous methods on various downstream VrDU tasks. &dagger; marks the results without any description of fine-tuning set (train or train+dev), The bold and underlined scores indicate the best and second results, respectively.

Key Information Extraction.

FUNSD, CORD, Kleister-NDA에서 SOTA를 달성했으며, FUNSD의 경우 7.98 포인트 향상된 값을 보여준다.

Document Question Answering.

Table 4: Results (Average Normalized Levenshtein Similarity, ANLS) of ERNIE-Layout and previous methods on the Document Question Answering task (DocVQA). "-" means the fine-tuning set is not clearly described in the original paper. &Delta;ANLS means ANLS difference between the multi-modal model and its corresponding text-only model, where ERNIE-Layout is initialized from RoBERTa and LayoutLMv2 is initialized from UniLMv2.

Table 4에서와 같이 DocVQA에 대해 Average Normalized Levenshtein Similarity (ANLS) 점수를 사용. 주목할 점은 LayoutLMv2의 경우 UniLMv2 기반으로 개발되었고 LayoutLM을 뛰어넘는 놀랄만한 성능을 보여주었으나 pre-training 코드를 공개하지 않았다. 따라서 해당 parameter로 RoBERTa를 initialize하여 ERNIE-Layout을 학습했더니 top-1 점수를 기록하였다.

Document Image Classification.

Table 5: Results (Accuracy) of ERNIE-Layout and previous methods on the Document Image Classification task (RVL-CDIP).

Image 분류는 문서의 macro 한 인식이 중요하다. 비록 ERNIE-Layout 이 세분화된 cross-modal 매칭을 집중적으로 했지만 이런 작업에서도 여전히 좋은 성능을 보여준다. 

4.4 Analysis

Pre-training task의 효율을 보고자 했다. FUNSD와 CORD를 사용했으며 모든 hyper-parameter는 동일한 값을 사용했다. 보고된 값은 5개의 다른 임의 시드의 평균값이다.

Effectiveness of Pre-training Tasks.

 

Table 6: Performance analysis with different pre-training tasks and attention mechanisms, in which SADA refers to the spatial-aware disentangled attention in ERNIE-Layout, SASA refers to the spatial-aware self-attention pro- posed by LayoutLMv2. &dagger; indicates the added module is proposed in this paper.

Table 6에서 RRP가 0.95% FUNSD 성능을 향상 시킨 것으로보아 해당 학습이 세부적인 cross-modal 상호작용의 이점을 잘 보여준다.  ROP를 사용했을 때 FUNSD 성능은 1.3% 향상된다. 따라서 읽는 순서를 학습시키는 것이 문서 이해를 위한 더 나은 representation을 학습하게 된다.

Effectiveness of Attention Mechanism.

Table 6에서 #4와 #6, #4와 #5를 비교하면 LayoutLMv2에 적용된 spatial-aware self attention과 disentangleing attention 이 단순 매커니즘의 변경만으로 많은 성능 향상을 가져오는걸 알 수 잇다. 

Table 7: Performance analysis with different serialization modules, in which Raster-Scan means serialization with vanilla OCR results, while Layout-Parser and Document-Parser arrange the recognized words with the help of layout knowledge.

Effectiveness of Serialization Modules.

비록 fine-tuning task에서는 적용하지 않았지만 table 7에서 확인할 수 있는 것처럼 읽는 순서를 재배열 하는 것만으로 성능향상을 가져올 수 있다. 

5 Conclusion

본 논문에서는 serialization과 attention이라는 두 가지 측면에서 layout knowledge를 document pre-training model에 통합하기 위해 ERNIE-Layout을 제안한다.
  ERNIE-Layout은 문서의 인식된 단어를 재정렬하려고 시도하여 원래 raster-scan 순서에 비해 downstream tasks에서 상당한 개선을 달성한다.
  게다가 ERNIE-Layout이 text/image와 layout feature 간의 더 나은 상호 작용을 구축할 수 있도록 새로운 attention 메커니즘을 설계한다.