Going Full-TILT Boogie on Document Understanding with Text-Image-Layout Transformer
Encoder-Decoder Model은 앞에서 언급한 sequence labeling의 여러 문제들을 해결할 수 있습니다.
Apr 19, 2022
Introduction
Motivation
저자들이 이 연구를 진행한 motivation은 다음과 같이 세 가지를 들 수 있습니다.
Importance of Spatio-Visual Relations
- 다음의 figure와 같이 text 정보만으로 문서의 내용을 완벽하게 파악하기란 어렵습니다.
- 따라서, 문서의 내용을 완벽하게 이해하기 위해서는 text 정보뿐만 아니라, 공간 정보(layout)와 이미지 정보가 모두 필요하고 이를 통합하는 시스템이 필요합니다.
Limitations of Sequence Labeling
- NLP에서 Key Information Extraction 등에 많이 사용되는 sequence labeling은 token 수준의 텍스트 라벨링과 span에 대한 라벨링이 있어야만 학습을 시킬 수 있습니다.
- 또한, 추출된 entity 중에 어느것이 추출되고 얼마나 많이 추출될지는 때론 복잡한 rule-based 방식으로 후처리를 해주어야 하는 경우가 많습니다.
- 마지막으로, document classification, non-extractive question answering, abstractive summarization과 같은 task에서는 문서 내에 정답이 나타나지 않기 때문에 Sequence Labeling을 사용할 수 없다는 단점이 있습니다.
- 이러한 이유들로 Sequence Labeling 방식에는 한계가 있다고 볼 수 있습니다.
Resorting to Encoder-Decoder Models
- Encoder-Decoder Model은 앞에서 언급한 sequence labeling의 여러 문제들을 해결할 수 있습니다. (span을 예측하는 sequence labeling가 달리, span에 대한 라벨링이 필요 없으며 encoder-decoder model은 전처리나 후처리가 복잡하지 않고 abstractive task들을 수행할 수 있음)
- 또한, Encoder-Decoder Model은 하나의 framework(또는 model)로 다양한 task를 처리할 수 있다는 장점이 있어 general-purpose model을 구축할 수 있습니다.
위의 motivation들로 인해 저자들은 Multi-modal을 처리할 수 있고 Spatial 정보를 잘 활용하며 Encoder-decoder 구조를 사용하는 Transformer 기반의 모델인 TILT(Text-Image-Layout Transformer)를 제안하였습니다.
Contributions
- Key Information Extraction, Question Answering, Document Classification 등의 task에서 복잡한 전처리 및 후처리가 없이 End-to-End로 학습시킬 수 있는 방법론 제안 및 SOTA 성능 달성
- 풍부한 공간 및 시각 정보를 Transformer에서 사용할 수 있는 방법론을 제안하여 text, layout, image 정보를 전부 활용할 수 있게 함
Proposed Method
Model Architecture
TILT 모델의 구조를 설명하기 전에 기본 Transformer(Vanilla Transformer)의 self-attention에 대해 되짚어볼 필요가 있습니다.
Transformer의 Self-attention은 다음과 같은 수식을 갖습니다. 아래에서 는 개의 inpu tokens에 대한 embedding matrix이고 , , 는 에 대한 query, key, value입니다. 는 attention bias로 기본 Transformer의 경우엔 해당 값이 영벡터가 됩니다. 두번째 수식에서 는 token embedding, 는 position embedding을 의미하고 이들의 합이 최종적인 임베딩 벡터인 가 됩니다.s
TIL+T에서는 Encoder-Decoder 모두를 갖는 T5 모델(GPT-3 이전에 SOTA 성능을 달성한 모델)을 base architecture로 채택하였는데, T5에서는 Relative Attention Bias를 사용하는 것에 영감을 얻어 TILT에서는 Spatial Bias 방식을 도입하였습니다.
T5의 Relative Attention Bias 설명
T5에서는 Position Embedding 대신에 Self-attention에 token의 상대적인 위치를 나타내는 bias 값을 추가합니다. Positioin Embedding은 상대적 위치를 나타내지 않고 절대적 위치를 나타내므로 token pair에 대한 attention을 수행하는 self-attention의 특성과 더 잘 맞을 것으로 생각해 해당 기법을 도입했다고 합니다.
Relative Attention Bias에서는 offset boundary를 설정하고 boundary 내의 token들에 relative position encoding 값을 줍니다.
예를 들어, 아래의 그림에서 4번째 토큰에 해당하는 “I”가 기준이고 offset=2라면 offset boundary를 벗어나는 첫번째 토큰 “I”와 두번째 토큰 “think”에는 index 0을 부여(boundary를 벗어나는 토큰들에 모두 같은 인덱스를 할당)하고 offset boundary 내에 있는 token들에는 순서대로 index를 부여하는 방식입니다.
논문에서 제안하는 Spatial Bias는 아래과 같은 수식으로 표현할 수 있습니다.
아래의 수식에서 는 T5에서 제안한 relative attention bias를 의미하고 와 는 각각의 토큰들끼리의 relative horizontal, vertical distances에 대한 bias입니다.
이러한 방식으로 TILT에서는 self attention에 Spatial Bias를 추가한 구조로 self attention layer를 모델링하였습니다.
다음으로, TILT에서는 word embedding과 image embedding을 결합하여 Multi-modal embedding이 가능하도록 합니다.
논문에서 저자들은 Image embedding을 위해서 512 x 384 크기의 이미지를 input으로 하고 64 x 48 x 128 크기의 feature map을 output으로 하는 U-Net을 사용합니다. U-Net의 경우, skip connection으로 인해 해당 token 근처의 정보(font, style 등)뿐만 아니라, 거리가 먼 곳에 있는 정보(caption 등)까지 고려하기 때문에 informative image feature를 구성할 수 있기에 이 network를 채택했다고 합니다.
최종적으로 논문의 저자들은 U-Net을 통해 생성된 image feature들에서 bounding box에 해당하는 부분만을 ROI pooling으로 추출하고 이를 Transformer 모델의 word embedding demension과 맞추기 위해 linear layer로 projection시킵니다.
따라서, 최종적인 각 토큰에 대한 embedding은 다음과 같은 수식으로 표현할 수 있습니다. 이 때, 는 word embedding, 는 image embedding vector입니다.
다음의 figure는 Vanilla Transformer, T5, TILT 모델들의 architecture를 비교한 것입니다. TILT 모델은 Word Embedding, Image Embedding을 결합하고 1D 및 2D Relative Bias를 self-attention layer에 추가한 것임을 확인할 수 있습니다.
Regularization Techniques
Encoder-Decoder 구조를 갖는 생성 모델들의 경우, overfitting이 발생하기 쉽습니다.
논문의 저자들은 조금 더 robust한 모델을 만들기 위해서 다양한 Augmentation 기법들을 적용해 성능을 개선했습니다.
Case Augmentation
현재 흔히 사용되고 있는 subword tokenization은 모두 대문자로 구성된 경우와 같은 비정상적인 형태에서 성능이 안 좋다는 연구 결과가 있었습니다. 특히, 논문에서 사용한 FUNSD, CORD, DocVQA 같은 데이터셋에서는 문서의 제목처럼 대문자로만 구성된 경우가 많기 때문에 이러한 현상은 심화될 수 밖에 없다고 저자들은 설명하고 있습니다.
따라서, 저자들은 특정한 비율로 token을 모두 대문자로 변환하거나 소문자로 변환하는 augmentation을 적용하였다고 합니다.
Spatial Bias Agumentation
Computer Vision 분야에서 이미지를 transform하여 augmentation하는 것과 비슷한 효과를 내기 위해 저자들은 Relative Bias에 무작위로 각 토큰끼리의 수평 및 수직 거리를 곱하는 augmentation을 추가합니다. 이러한 방식이 가로 또는 세로로 긴 document들에 대해 좋은 효과를 가져올 것이라고 저자들은 생각하였습니다.
Affine Vision Augmentation
논문의 저자들은 기존의 image augmentation 기법과 유사하게 affine transform을 사용하여 augmentation을 수행했습니다.
Experiments
Training Procedure
T5는 unsupervised pre-training, supervised pre-training의 2 stage 학습 방법을 통해 pre-training을 하고 만들어진 모델을 통해 downstream task에 대해 fine-tuning을 하는 방식으로 학습을 합니다. TILT의 저자들도 다음과 같이 T5의 학습 방법과 유사하게 학습을 수행합니다.
- unsupervised pre-training : vanilla T5 model의 weight로 모델을 초기화한 다음에 RVL-CDIP(275k docs), UCSF Industry Documents Library(480k docs), PDF files from Common Crawl(350k docs) 데이터로 MLM(Masked language model) task로 pre-training을 합니다.
- supervised pre-training : 레이아웃 features이 풍부한 범용 모델을 만들기 위해 다양한 task들에 correlation이 높은 데이터셋들(CORD, RVL CDIP, SROIE, docVQA)만 선별하여 학습을 수행합니다. 이 때, 현재 단계에서 overfitting되는 것을 막기 위해 10,000 steps 또는 5 epochs만 학습시킵니다.
- supervised fine-tuning : 특정 데이터셋 하나의 downstream task들에 대한 최적 학습을 수행합니다.
Experiment Results
저자들은 CORD, SROIE, DocVQA, RVL-CDIP에 대한 fine-tuning 성능을 여러 모델들과 비교하여 결과를 설명합니다.
DocVQA의 경우, table이나 form 형태의 데이터를 많이 포함하고 있는데, SOTA 모델에 비해 form에서는 89.5 → 94.6으로, table에서는 87.7 → 89.8로 성능이 개선되었다고 합니다. 이를 통해 TILT 모델이 form이나 table에서도 구조를 잘 파악한다고 이해할 수 있습니다. 또한, 최종 성능 또한 SOTA 성능을 달성하였습니다.
RVL-CDIP의 경우, SOTA 성능에는 0.12 떨어지는 95.52 성능을 달성하였습니다.
CORD의 경우, entity가 아닌 경우 decoder가 “None”을 출력하도록 하고 실험하였으며 성능은 SOTA 모델에 비해 0.3 정도 높은 96.33을 달성하였습니다.
SROIE의 경우, SOTA 성능과는 0.07 정도보다 낮은 성능인 98.10을 달성하였습니다. 하지만 SROIE 데이터셋은 OCR 오류와 전체 entity 라벨링에 문제가 많기 때문에 이는 개선될 여지가 있다고 합니다.
Ablation Study
논문에서는 제안한 방법론들이 효과가 있었는지 확인하기 위해 TILT-Base 모델에서 특정 모듈들을 제거해보는 Ablation Study를 수행합니다.
모델의 Architecture 측면에서,
Spatial Bias와 Visual Embeddings을 각각 제거해본 결과 2프로에 가까운 성능이 드랍되었습니다. 이는 Spatial Bias와 Visual Embeddings이 encoder-decoder 모델의 성능에 도움을 준다고 해석할 수 있습니다.
Reguralization 측면에서,
Case Augmentation, Spatial Augmentation, Affine Augmentation을 각각 제거해봤을 때, 모두 다 성능이 드랍되었지만 그렇게 많은 영향을 끼치지는 않았습니다. 하지만 Case Augmentation의 경우, 성능이 다른 augmentation에 비해 많이 드랍된 것으로 볼 때 Case Augmentation이 영어 데이터셋에서 많이 중요하다고 이해할 수 있습니다.
Training Procedure 측면에서,
Supervised Pretraining을 제거한 결과, 1.7 정도의 성능이 드랍되었습니다. 이는 pre-training 단계를 2 stage로 나누는 것이 효과가 있을 수 있음을 보여줍니다.
Conclusion
논문에서는 Key Information Extraction, Question Answering, Document Classification 등의 task에서 복잡한 전처리와 후처리 없이 End-to-End로 학습할 수 있는 방법론으로 Encoder-Decoder 기반의 모델을 활용하였습니다.
또한, text, layout, visual information들을 모두 활용하기 위해 Spatial Bias, Visual Embeddings을 제안하였고 이를 적용한 TILT 모델을 제안하였습니다.
Visual Document Understanding이라는 task를 수행하기 위한 많은 생각과 노력이 있었던 연구였던거 같아서 개인적으로는 상당히 의미가 있는 논문이라고 생각했습니다.
로민에서도 pre-training task에 대한 고민, multi-modal을 잘 활용할 수 있는 model architecture에 대한 고민, 데이터셋 수집 방법 및 처리에 대한 고민 등을 통해 로민만의 Document Understanding 모델을 만들어 볼 수 있었으면 좋겠습니다.
Share article
