TextMonkey: An OCR-Free Large Multimodal Model for Understanding Document

OCR 없는 문서 이해를 위한 혁신적 AI 모델. 고해상도 이미지 처리, 토큰 최적화, 텍스트 스팟팅 기능 탑재. 12개 벤치마크에서 우수한 성능 입증. 문서 이해 AI의 새 지평을 열다

Inc Lomin

Apr 25, 2024

TextMonkey: An OCR-Free Large Multimodal
Model for Understanding Document

Contents

Paper

TextMonkey: An OCR-Free Large Multimodal Model for Understanding Document

We present TextMonkey, a large multimodal model (LMM) tailored for text-centric tasks, including document question answering (DocVQA) and scene text analysis. Our approach introduces enhancement...

https://arxiv.org/abs/2403.04473

TextMonkey demo

vlrlab-monkey.xyz

http://vlrlab-monkey.xyz:7684/

2024.03.07 자 논문

Introduction

Motivation

기존의 LMM (Large Multimodal Models) 분야의 vision encoder는 낮은 해상도 (224~336)의 이미지만으로 학습된 CLIP에 의존하여 많은 텍스트를 가진 document task 에 잘 맞지 않음 (a)

이를 해결하기 위해 이미지를 여러 patch로 나눠 연산하는 crop-based method가 제시되었는데, 이는 작은 단어들을 분할할 가능성이 있기 때문에 patch들을 합한 후에 의미가 달라질 경우가 발생함 (b)

Frequency-based method는 frequency를 활용하기 때문에 높은 해상도의 이미지를 처리할 수 있지만, pre-trained된 vision encoder를 사용하기 어렵다는 단점이 있음 (c)

Contribution

고해상도 문서 이미지 처리: TextMonkey는 Shifted Window Attention과 Zero-Initialization을 도입하여 높은 입력 해상도에서의 cross-window 연결성을 달성하고 초기 훈련을 안정화시킵니다. 이를 통해 고해상도 문서 이미지를 효율적으로 처리할 수 있습니다.

토큰 길이 최적화: Token Resampler를 사용하여 이미지 내의 중복 토큰을 필터링하고 토큰 길이를 줄입니다. 이는 모델의 성능을 향상시키고 계산 비용을 줄이는 데 기여합니다.

텍스트 스팟팅 및 그라운딩: TextMonkey는 텍스트 스팟팅과 그라운딩을 포함한 다양한 기능을 지원합니다. 위치 정보를 응답에 포함시켜 모델의 해석 가능성을 높이고 환각을 최소화합니다.

Prompt 방식으로 여러 subtask 수행 가능

추가적으로 표 추출이 가능하고 대부분의 벤치마크에서 높은 성능 달성

추가적으로 몇 가지 데모를 돌려본 결과 zero-shot 성능이 실제로 아주 높은 것을 확인

Related work

OCR-Model-Driven Methods (layoutLM, DocFormerv2)

다른 모델의 Text extraction 에 의존 (성능, 시간적 문제), OCR 성능에 bound된다는 단점

OCR-Free Methods

특정 task에 fine-tuning이 필요하다는 단점 존재하고, LLM을 도입한 모델들 (MLLM) 이 나왔지만 motivation과 같은 단점

Methodology

Overview

입력 이미지를 448*448 크기의 겹치지 않는 patch로 분리

이 patch는 14*14 크기의 32*32 patch들로 다시 나눠지고, 각각을 token으로 취급

CLIP pre-trained model의 transformer block을 이용해서 각각의 토큰 처리

Image Resampler, Token Resampler를 이용해서 token 수 압축

Image token 과 Instruction token을 Multi-modal LLM에 입력

1. Shifted Window Attention

text, vision object가 분할되는 약점이 존재 이를 위해 CLIP 모델의 vision ability를 위해 shifted window attention (swin transformer) 도입

이미지를 겹치지 않는 window로 자른 뒤, 각각의 window를 CLIP transformer block을 거치게 함

Sliding window를 좌상단 방향으로 이동해서 self attention 적용 (Shifted Attention Layers) 이 때 window 사이의 관계가 형성된다고 함

이를 반복하여 (그림의 Vit Encoder Layers - Shifted Attention Layers 를 N번 반복)

추가적으로 MLP 단계에서 두 linear matrix A,B 의 initialization을 각각 다르게 (A: Gaussian init, B: Zero init) 함으로서 학습 초기에 너무 큰 변화를 겪지 않게 하였음

2. Token Resampler

이 논문에서는 높은 해상도를 활용하기 때문에 많은 토큰 수가 필요

이전의 논문에서 자연어에 존재하는 중복 정보들을 활용하여 유효 토큰 수를 압축하는 것이 가능함을 보임

이와 비슷하게 이미지 해상도가 높아진다면 시각적 토큰 정보에도 중복이 존재할 거라는 가설을 세웠음

저자들은 image token간의 중복성을 확인하기 위해, Image resampler 이후 임의로 선택된 20개의 feature를 cosine similarity를 이용해서 비교하였음

또한 여러 해상도에서 토큰의 중복성을 정량적으로 비교했을 때, 해상도가 증가함에 따라 반복되는 토큰의 발생률이 높아진다는 것을 확인

해상도가 각각 448, 896, and 1334일 때, 중복되는 토큰의 비율이 68/256 (26.6%), 571/1024 (55.8%), and 1373/2304 (59.5%) 로 해상도가 커질 수록 증가함을 보였음 (Threshhold 0.8 기준)

image resampler는 논문에 정보가 없지만 Detail에서 256 learnable query라고 기재되어 있는 것으로 보아 Vit layer 결과 벡터를 (256*D) 크기의 벡터로 변환하기 위해 (token화하기 위해) 사용한 것으로 추정 ( )

중요한 토큰을 식별하고, 중복된 토큰을 제거하는 방법 (Token filter)은 아래 알고리즘을 따름

각 토큰에 대한 유사성 계산: 모든 토큰에 대해, 다른 모든 토큰들과의 유사성을 계산하고 그 중 최대값을 찾음. 이는 토큰이 다른 토큰들과 얼마나 유사한지를 나타냄

중요도 점수 계산: 각 토큰의 중요도를 1에서 해당 토큰의 최대 유사성을 뺀 값으로 계산합니다. 토큰이 다른 것들과 매우 유사하면, 최대 유사성 값이 높을 것이고, 중요도 점수는 낮아질 것입니다. 반대로, 토큰이 유니크하면 최대 유사성 값이 낮고 중요도 점수는 높아질 것입니다.

중요한 토큰 선택: 계산된 중요도 점수를 바탕으로, 상위 'r'개의 토큰을 선택합니다. 여기서 'r'은 유지하고 싶은 토큰의 수입니다.

원래 순서로 토큰 정렬: 선택된 토큰들을 그들이 원래 있던 순서대로 정렬합니다.

3. Position-Related Task

LLM 모델에서 자주 발생하는 hallucination issue를 해결하기 위해 (여기서는 input image와 관련 없는 응답을 생성하는 문제) 모델이 시각 정보를 분석하고, 답변에 visual information을 통합하도록 하였음

모델의 text position 인식을 향상시키기 위해 아래와 같은 task를 수행하도록 함

Text와 localtion data의 관계를 위해 텍스트 정보는 위치 정보의 앞에 나오도록 하였음

다양한 비율의 이미지를 처리하기 위해 x,y 좌표를 (0, 1000) scale로 정규화하였음

4. Dataset Construction

이 연구에서는 오픈소스 데이터만을 사용했음

총 409.1k 쌍의 dialogue data, 2.1M QA pair 를 사용

추가적으로, 표와 그래프 등을 포함하는 structured data에 특화된 모델을 만들기 위해 위의 데이터의 10%와 문서, 표, 그래프를 포함하는 structured data의 일부를 활용하여 fine-tuning 진행 (TextMonkey+)

5. Loss

I: input image, Q: question sequence, : output sequence, s: input sequence, L: length of output

Experiments

Implementation Details

Model Configuration

Vit transformer block: Vit-BigG

LLM: Qwen-VL (중국어 pre-trained large multimodal model)

Image resampler: 256 learnable queries

Token resampler ratio: 512 for images with a 896px, 1024 for images with a 1344px

학습은 주로 896px에 맞춰져 있음

Training

Optimizer: AdamW optimizer

Learning rate: 1e-5, to 5e-6 (cosine learning rate scheduler)

0.9, 0.95

Warmup: 150 steps

Batch size: 128

Weight decay: 0.1

모델 전체 학습은 A800 12개를 활용하여 1epoch에 하루가 걸렸다고 함

Evaluation

다른 연구와의 비교를 위해서 모델 response가 정답을 포함하고 있다면 맞았다고 처리하도록 했음

Results

Quantitative Results

TextMonkey+: document, table 등을 활용하여 fine-tuning한 모델

Qualitative Results

Structuralization

Conclusion

고해상도 문서 이미지 처리: TextMonkey는 Shifted Window Attention과 Zero-Initialization을 도입하여 높은 입력 해상도에서의 cross-window 연결성을 달성하고 초기 훈련을 안정화시킵니다. 이를 통해 고해상도 문서 이미지를 효율적으로 처리할 수 있습니다.