DocLLM: A layout-aware generative language model for multimodal document understanding

|

Code와 Model checkpoint가 공개되지 않은 논문

1. Introduction

DocLLM은 특정 양식을 갖춘 자료 혹은 각종 청구서를 처리하기 위 모델로, JPMorgan에서 개발하였다. 해당 모델은 아래 그림과 같이 요약될 수 있다.

DocLLM: A layout-aware generative language model for multimodal document understanding


  1. 학습에 사용되는 데이터는 OCR을 통해 추출한 text token과 그에 대한 bounding box가 포함된 데이터이다. 해당 모델은 기존의 Multimodal Model들이 시각 정보를 처리할 때 image encoder를 사용하는 것과는 달리 문서 내 시각 정보(layout 구조)를 이해할 때 OCR을 통해 얻어진 bounding box 정보를 사용한다. (image encoder를 사용하지 않기 때문에 모델 크기가 타 Multimodal Model에 비하여 작고, 이에 따라 추론 시간도 단축된다.)

  2. 해당 모델은 text semantic과 spatial layout 간의 관계를 포착할 때 확장된 attention mecanism을 도입한다. 본 모델은 문서 내 text 와 공간적(spatial) 정보를 각각 따로 처리하기 위해 두 modality 간의 cross-alignment를 포착할 때 classical transformer의 self-attention mechanism을 분해하여 사용한다. 구체적으로, 각 modality에 대한 attention score 뿐만 아니라 두 modality간의 관계를 포착하는 attention score도 계산한다.

  3. 다양한 layout과 문서 내 시각적 정보를 잘 파악하기 위해 사전 학습 과정에서 infill text segment task를 학습한다. infill text segment task는 다양한 layout을 지닌 문서에서 이전 text 정보가 현재 text와 관련 없을 가능성이 있기 때문에 이러한 경우를 파악하기에 적합한 과제이다.

  4. 사전 학습 후에는 네 가지 핵심 문서에 대해 intelligence task에 대한 fine-tuning(instruction-tuning)을 수행한다.

본 연구의 contribution은 아래와 같다:

  • visual document를 이해하기 위해 설계된 light-weight LLM 개발
  • text and layout modality 간의 cross-alignment를 위한 분리된 attention mechanism 제안
  • 불규칙적인 layout을 효과적으로 파악하기 위한 infilling pre-training objective 도입
  • visual document intelligence task 수행을 위해 특수 생성된 instruction-tuning dataset 구축

2. Related Work

2.1 LLM Architectures

Autoregressive Infilling

autoregressive infilling approach에는 두 가지 방법이 있다.

  • FIM(fill-in-the-middle): single span 예측
  • GLM(General language model pretraining with autoregressive blank infilling): multiple span 예측

OpenAI의 FIM approach는 문서를 (prefix, middle, suffix) 이렇게 세 부분으로 나눈다. 나누어진 문서는 (prefix, suffix, middle)로 재구성되어 모델이 middle segment를 예측할 수 있게 한다. 이 과정은 [PRE], [SUF] 그리고 [MID] token을 기준으로 수행된다. 즉, 실제 수행 시 [MID] token을 예측 시작 기준으로 삼는다. 이와 같은 방법은 autoregressive model이 middle part가 누락된 텍스트를 채우는 방법을 학습하게 한다.

GLM은 multiple span을 채우는 과제이다. 빈칸을 채우기 위해 해당 과제에서는 [blank_mask]와 [start_to_fill] 쌍을 사용한다.

Disentangled attention

Disentangled attention은 DeBERTa에서 처음 도입되었다. 해당 mechanism은 token embedding과 relative positional encoding을 합치지 않고 별도로 유지하며, attention weight 를 계산할 때 분리된 matrix를 사용하여 각각 독립적으로 처리한다. 이는 content와 position이 각각 분리된 attention alignment를 학습할 수 있게 하는데, 이것이 NLU benchmark에서 DeBERTa가 RoBERTA-large와 T5를 능가하게 한 지점이라고 알려져 있다.

3. DocLLM Framework

DocLLM: A layout-aware generative language model for multimodal document understanding


3.1 Model Architecture

DocLLM은 auto-regressive transformer language model의 decoder 구조를 기반으로 개발되었다. transformer decoder 모델을 구성하는 각 transformer block은 multi-head self-attention layer 과 fully connected feed forward network를 포함하며, transformer block은 여러개로 구성될 수 있다.

Standard language model은 일반적으로 unimodal로, text만 입력으로 받는다. 그리고 Visual Language Model은 text와 image data를 입력받아 각각의 encoder(model)를 통해 얻어진 modlity 정보를 통합시킨다. 반면 DocLLM은 OCR을 사용하여 text token에 대한 공간적(spatial)위치와 크기를 얻고 이를 활용하여 시각 정보를 통합한다.

3.2 Disentangled Spatial Attention

classical transformers attention mechanism

input toeken sequence $x = (x_1, …, x_i, …, x_T )$가 있을 때, 일반적인 attention에서는 학습된 embedding matrix와 token position에 대한 학습된 parameter set을 사용하여 input token을 우선 hidden vector $H\in \mathbb{R}^{T\cdot d}$ self-attention head은 toekn $i$와 $j$간의 attention score를 계산한다.

$W^{q}\in \mathbb{R}^{d\cdot d}$과 $W^{k}\in \mathbb{R}^{d\cdot d}$는 projection matrix이고 윗 첨자 $t$는 text modality를 나타낸다. attention score $A \in \mathbb{R}^{T\cdot T}$와 또다른 projection matrix $W^{v}$는 hidden vector $H^{‘}$를 계산하는데 사용되며, 이 hidden vector는 다음 layer의 입력으로 사용된다.

DocLLM attention mechanism

DocLLM은 text token과 그에 대한 bounding box 정보가 쌍을 이룬 data를 input으로 받는다. 본 모델은 bounding box를 separate hidden vector로 encoding하고 attention mechanism을 네 가지 score로 분해한다: text-to-text, text-to-spatial, spatial-to-text, spatial-to-spatial. 각각의 점수의 균형을 맞추기 위해 projection matrix와 hyperparameter를 사용한다. spatial information에 대한 hidden vector는 layer 전반에 걸쳐 재사용된다.

DocLLM의 input은 $\mathrm{{(x_i,b_i))}}_{i=1}^{T}$으로 표현되는데 여기에서 $b_i$은 $x_i$에 해당하는 bounding box (left, top, right, bottom)이다. LLM이 spatial information을 포착할 수 있도록 하기 위해 bounding box를 hidden vector $S\in \mathbb{R}^{T\cdot d}$로 encoding한 후 attention matrix를 네 가지 다른 score로 분해한다: text-to-text, text-to-spatial, spatial-to-text, spatial-to-spatial.

여기에서 $W^{s,q}\in \mathbb{R}^{d\cdot d}$와 $W^{s,k}\in \mathbb{R}^{d\cdot d}$는 classical attention과 달리 spatial modality에 해당하는 새로 도입된 projection matrix이다. 그리고 $\lambda$들은 각 점수들(4 가지)의 상대적 중요성을 제어하는 hyperparameter이다. input hidden vector $H^{‘}$는 classical attention과 동일하게 계산된다.

hidden vector $S$는 모든 layer에 걸쳐 재사용되지만, 각 layer는 다른 projection matricx를 사용할 수 있다. bounding box정보를 encodding하는데 필요한 추가 parameter 수는 VLM과 같이 image model을 따로 두는 방식보다 훨씬 적다.

3.3 Pretraining

DocLLM은 unlabeled document에 대해 self-supervised 방식으로 사전학습 한다. 학습 과제는 text infilling objective로, 이를 수행하기 위해 사전 학습이 진행하는 동안 OCR 정보로부터 식별된 text block음 무작위로 masking되며, masking된 block이 shuffle된(GLM 방식과 동일). block은 본 블로그 DocLLM Framework 아래 있는 그림에서 “Name”, “John Doe”, “Doctor”과 같은 것들이다. (이 block들의 span 범위는 더 커질 수 있다.) block infilling은 시작과 끝을 나타내는 special token과 함께 autoregressively하게 수행된다. 해당 task는 모델이 문맥적으로 빈칸을 적절히 완성할 수 있게한다. 이와 같은 방법은 사전학습 단계에서만 사용되고, fine-tuning 단계에서는 사용되지 않는다.

3.4 Instruction Tuning

DocLLM은 16개 dataset과 그에 해당하는 OCR을 사용하여 4가지 DocAI task를 처리할 수 있도록 instruction-tuning된다. 4가지 DocAI task는 아래 표에 기재된 바와 같이 Visual Question Answering (VQA), Natural Language Inference (NLI), Key Information Extraction (KIE), Document Classification (CLS)이다.

DocLLM: A layout-aware generative language model for multimodal document understanding


본 연구에서는 DocAI dataset의 탬플릿으로부터 파생된 다양한 instruction으로 instruction tuning을 수행했다. 또한 supervised fine tuning (SFT) 수행 시 입력되는 instruction의 다양성이 높을수록 해당 모델의 zero-shot generalization 능력이 향상되기 때문에 최대한 다양한 template을 사용하려고 했다(“Finetuned language models are zero-shot learners”, “Scaling instruction-finetuned language models”, “ Training language models to follow instructions with human feedback”). 각 template은 다른 질문을 제시하고, 몇몇 경우 다른 유형의 답변을 반환하도록 했다. dataset의 탬플릿은 “Modularized multimodal large language model for document understanding”과 “Universal ocr-free visually-situated language understanding with multimodal large language model” 연구에서도입된 탬플릿을 가능한 재사용하였다.

Visual Question Answering(VQA)

DocVQA, WikiTableQuestions(WTQ), VisualMRC, DUDE, BizDocs2를 수집하여 VQA instruction-tuning data mix를 구축한다. VQA 학습 시 input은 위 표에 나타난 바와 같이 하나의 instruction template이다.

{document} {question}

e.g. {document} What is the deadline for scientific abstract submission for ACOG - 51st annual clinical meeting?

Natural Language Inference(NLI)

본 연구에서 구축한 instruction-tuning data mix를 DocAI NLI dataset으로 사용하기에는 수량 부족한 관계로, 여기에 TabFact dataset을 추가하여 사용한다.

{document}\"{statement}\", Yes or Ne?

e.g. {document} \"The UN commission on Korea include 2 Australians.\", Yes or No?

Key Information Extraction(KIE)

Kleister Charity(KLC), CORD, FUNSD, DeepForm, PWC, SROIE, VRDU ad-buy (with random train-test splitting), and BizDocs를 수집하여 KIE instruction-tuning data를 구축하였다. instruction template 유형은 위 표에 기재된 바와 같이 총 3가지를 사용하였다: extraction, internal classification, MCQ.

[Extraction] -> {document} What is the value for the \"{key}\"?

e.g. {document} What is the value for the \"charity number\"?

[Internal classification] -> {document} What is \"{value}\" in the document?

[MCQ] -> {document} What is \"{value}\" in the document? Possible choices: {choices}.

Document Classification(CLS)

RVL-CDIP, BizDocs를 수집하여 CLS instruction-tuning data를 구축하였다. 이 작업에는 두 가지 유형의 template이 사용되었다.

[MCQ] -> {document} What type of document is this? Possible choices: {choices}.

e.g. {document} What type of document is this? Possible answers: [budget, form, file folder, questionnaire].

[Internal classification] -> {document} What type of document is this?

4. Experiments

4.1 Model Setup and Training Details

실험에 사용된 모델은 Falcon-1B를 기반으로 한 DocLLM-1B, Llama2-7B를 기반으로 한 DocLLM-7B로 두 가지 버전이 있다. 실험 시 사용된 hardware specification은 24GB A10g GPU 하나고 fully sharded data parallelism으로 학습이 진행되었다. 실험에 사용된 모델의 각종 config와 hyperparameter는 아래 표와 같다:

DocLLM: A layout-aware generative language model for multimodal document understanding


DocLLM-1B의 경우에는 1 epoch의 pre-training, 10 epoch의 instruct-tuning을 수행했고, DocLLM-7B의 경우에는 1 epoch의 pre-training, 3 epoch의 instruct-tuning을 수행했다.

4.2 Downstream Evaluation

본 연구에서는 두 가지 설정 하에 실험을 진행한다.

  • Same Datasets, Different Splits (SDDS)

    일반적인 방법이다. 16개의 데이터셋에 대해 trainset을 학습하고, 각 trainset을 학습한 모델을 상응하는 testset으로 평가하는 방식이다.

    이와 같은 방식 아래에서 zero-shot setting으로 실험한 결과, GPT-4와 Llama2를 포함한 다른 모델들보다 16개 데이터셋 중 12개에서 더 좋은 성능을 보였다. 특히 KIE와 CLS와 같이 레이아웃 집약적인 task에서 특히 좋은 성능을 보였다.

    아래는 SDDS 실험 결과이다.

DocLLM: A layout-aware generative language model for multimodal document understanding


  • Same Tasks, Different Datasets (STDD)

    동일한 task를 수행하는 모델이지만 testdata와 traindata의 출처를 상이하게 설정한 실험이다. 예를 들어 16개의 dataset에서 11개의 dataset을 train에 사용하고 나머지 3개의 dataset을 평가에 활용한다. 이때 (1 )각 task에 대해 최소 하나의 dataset을 SFT 수행에서 제외시키고, (2) data split 시 각 task마다 충분한 trainset을 확보하고, (3) data leakage를 방지한다. 예를 들어 동일한 출처로부터 data가 중복되지 않도록 한다. 마지막으로 (4) 가능한 유명한 benchmark를 사용한다.

    실험 결과, 5개의 데이터셋 중 4개에서 Llama2를 능가했으며, 특히 KIE 작업에서 두 개의 데이터셋에서 최고 점수를 기록했다. 하지만 classification accuracy가 낮은데, 이는 하나의 분류 dataset을 학습했기 때문에 새로운 데이터셋에 대한 generalization이 떨어진 상황으로 파악된다.

    아래는 STDD 실험 결과이다.

DocLLM: A layout-aware generative language model for multimodal document understanding


5. Ablation Studies

5.1 Disentangled Spatial Attention

spatial-to-spatial interaction에 집중할 때 문서 이해력이 가장 높은 것으로 확인되었다. 이 결과는 문서 분석 시 공간적 feature가 중요하다는 것을 보여준다.

DocLLM: A layout-aware generative language model for multimodal document understanding


5.2 Autoregressive Block Infilling

Pretraining task로 block infilling + spatial modality를 학습한 모델이 가장 높은 성능을 보이는 것으로 확인되는데 이 또한 공간적 정보가 중요하다는 것을 나타낸다.

DocLLM: A layout-aware generative language model for multimodal document understanding


5.3 Prefix Decoder and Causal Decoder

해당 실험은 prefix decoder와 causal decoder를 비교하는 실험이다. 실험 결과, 여러 조합에서 두 decoder 간의 차이는 미미했지만, causal decoder가 약간 더 좋은 성능을 보였따.

DocLLM: A layout-aware generative language model for multimodal document understanding


DocLLM: A layout-aware generative language model for multimodal document understanding


Reference

DocLLM: A layout-aware generative language model for multimodal document understanding

https://llava-vl.github.io/