[Paper] LayoutLMv2: Multi-modal Pre-training for Visually-rich Document Understanding

Introduction

최근 Visually-rich Document Understanding (VrDU) task은 두 가지 방향으로 진행되고 있다.

 

1. textual과 visual/layout/style 정보들의 shallow fusion

pre-trained NLP와 CV 모델을 사용하여 좋은 성능을 달성하였지만 하나의 document type knowledge가 다른 type에서 쉽게 전달되기 어렵다. 그래서 보통 이러한 경우 document type에 따라 re-train을 해서 모델을 변경한다.

이에 따라 일반적인 document layout에서 변하지않는 것들(local invariance) (key-value 쌍으로 되어있는 left-right layout, grid layout에서의 table 등)이 충분히 이용되지 않는다.

 

2. 상당히 많은 다른 도메인에서 unlabeled document로 texual, visual, layout 정보들의 deep fusion

pre-trained 모델이 다른 document type들에서 변하지 않는(local invariance) layout과 style이 유지되어 cross-modal knowledge를 흡수한다. 또한 다른 document format으로 바꾸려고 할 때 약간의 labeld sample만 있어도 fine-tune을 통해 좋은 성능을 얻을 수 있다.

그래서 본 논문에서 제안된 모델은 두 번째 방향을 따라 VrDU task에서 pre-training starategy를 개선하는 방법에 대해 살펴보았다.

 

기본적인 LayoutLM에서 fine-tuning stage에서 visual embedding을 결합했지만 LayoutLMv2에서는 pre-training stage에서 visual information을 합쳤다. 이것은 visual과 textual 정보 간의 cross-modality interaction을 학습할 수 있는 Transformer 구조의 장점을 사용했다.

 

추가로 1-D relative position representation에서 착안하여 LayoutLMv2에서는 토큰 쌍에 대한 2-D relative position representation을 제안한다. 기존 LayoutLM의 absolute 2-D position embedding과 다른 점은 relaive position embedding이 contextual spatial modeling에 대한 더 넓은 시야를 가지게 한다.

 

LayoutLMv2에서는 masked visual-language modeling 외에도 두 가지 pre-training strategy를 사용한다.

• text-image alignment strategy: text line과 해당하는 image region을 정렬
• text-image matcing strategy: document image와 texual context가 서로 연관되어있는지 아닌지를 학습

 

Model Architecture

Text Embedding

text embedding은 세 개의 embedding을 더한다.

• token 자신을 표현하는 Token embedding
• token index를 표현하는 1D positional embedding
• 다른 text segment들을 구분하는데 사용되는 segment embedding

$t_i = TokEmb(w_i) + PosEmb1D(i) + SegEmb(s_i)$

 

Visual Embedding

visual encoder backbone으로 ResNeXt-FPN을 사용했다.

1. document page image I
2. 224x224 resize
3. visual backbone에 넣고 output feature map을 얻음
4. output feature map은 average-pool을 통해 고정된 사이즈 W(width)와 H(height)가 됨
5. WxH length로 flatten하여 visual embedding sequence($VisTokEmb(I)$)가 됨
6. projection layer를 통해 각 visual token embedding은 text embedding과 합쳐짐
7. CNN기반의 visual backbone은 potisional 정보를 못잡아내기 때문에 visual token embedding에 1D positional embedding을 추가 (이 1D positional embedding은 text embedding layer와 공유됨)
8. segment embedding의 경우 모든 visual token들에 visual segment $[C]$를 붙임

$v_i = Proj(VisTokEmb(I)_i) + PosEmb1D(i) + SegEmb([C])$

 

Layout Embedding

i번째 text/visual token의 nomailze된 bounding box($box_i = (x_{min}, x_{max}, y_{min}, y_{max}, width, height)$)가 주어지면 layout embedding layer는 token-level의 2D posiional embedding으로 만들고 6개 bouding box feature들을 결합한다.

$I_i = Concat(PosEmb2D_x(x_{min}, x_{max}, width), PosEmb2D_y(y_{min}, y_{max}, height))$

 

Multi-modal Encoder with Spatial-Aware SelfAttention Mechanism

encoder는 visual embedding과 text embedding을 concat하고 layout embedding을 더한 sequence를 첫 번째 layer에 넣는다.

기존 self-attention mechanism은 absolute posiion hint로 input token들 간의 관계 implicitly하게 잡는다. document layout에있는 local invariance를 효율적으로 모델링하기 위해서는 relative position 정보를 명시적으로 집어넣어주어야한다. 그래서 본 논문에서는 self-attention layer로 spatial-aware self-attention mechanism을 소개한다.

 

편의를 위해 single self-attention laye와 single head를 기준으로 설명하겠다.

기존 self-attention mechanism은 두 vector를 projection하고 attention score를 계산하는 것으로 query($x_i$)와 key($x_j$)간의 연관성을 잡는다.

아주 큰 범위의 position을 고려하여 semantic relative position과 spatial relative position을 bias로 사용하여 parameter가 많아지는 것을 방지했다.

learnable 1D, 2D realtive position bias를 각각 $b^{(1D)}, b^{(2D_x)}, b^{(2D_y)}$라고 했을 때 이 bias들은 attention head들에서 다르지만 모든 encoder layer에서 공유된다.

i번째 bounding box의 top left corner의 좌표를 $(x_i, y_i)$라고 했을 때 spatial-aware attention score를 계산한다면 다음과 같다.

 

Pre-training Tasks

Text-Image Alignment

Text-Image Alignment(TIA)은 token line을 랜덤으로 선택해서 이것들의 image region은 document image에서 cover된다. 본 논문에서는 이것을 covering이라 부른다. pre-training에서는 classification layer가 covered인지 아닌지를 예측하고 binary cross-entropy loss로 계산한다.

Text-Image Matching

Text-Image Matching(TIM)은 document image와 textual content간의 일치(?)를 학습하는 것을 돕기위해 적용되었다. 같은 document page에서 나온 image와 text인지를 예측하기 위해 classifier에 $[CLS]$를 넣는다.

nagarive sample을 만들기위해서는 image를 다른 document로 대체하거나 drop시킨다.