[논문 리뷰] DiT: Self-supervised Pre-training for Document Image Transformer

Junlong Li, Yiheng Xu외 4명|arXiv (Cornell University)|2022. 03. 04.

Handwritten Text Recognition Techniques인용 수 9

한 줄 요약

DiT는 인간의 레이블이 없는 대규모 문서 이미지를 활용하여 일반적인 시각적 표현을 학습하는 자기지도 학습 기반 사전 훈련된 문서 이미지 트랜스포머를 제안한다. 문서 전용 이산형 VAE(디스crete VAE)를 사용한 마스킹 이미지 모델링을 통해 다양한 문서 AI 작업에서 최신 기술 수준의 성능을 달성한다. 이는 문서 이미지 분류(92.69% top-1 정확도), 레이아웃 분석(94.9% F1), 표 탐지(96.55% wF1), 텍스트 탐지(94.29% F1) 등에서 성과를 보였다.

ABSTRACT

Image Transformer has recently achieved significant progress for natural image understanding, either using supervised (ViT, DeiT, etc.) or self-supervised (BEiT, MAE, etc.) pre-training techniques. In this paper, we propose extbf{DiT}, a self-supervised pre-trained extbf{D}ocument extbf{I}mage extbf{T}ransformer model using large-scale unlabeled text images for Document AI tasks, which is essential since no supervised counterparts ever exist due to the lack of human-labeled document images. We leverage DiT as the backbone network in a variety of vision-based Document AI tasks, including document image classification, document layout analysis, table detection as well as text detection for OCR. Experiment results have illustrated that the self-supervised pre-trained DiT model achieves new state-of-the-art results on these downstream tasks, e.g. document image classification (91.11 $ ightarrow$ 92.69), document layout analysis (91.0 $ ightarrow$ 94.9), table detection (94.23 $ ightarrow$ 96.55) and text detection for OCR (93.07 $ ightarrow$ 94.29). The code and pre-trained models are publicly available at \url{https://aka.ms/msdit}.

연구 동기 및 목표

• 문서 AI에서 비전 백본을 사전 훈련하기 위한 대규모 인간 레이블이 부여된 문서 이미지 데이터셋의 부족을 해결한다.
• 템플릿 및 포맷 불일치로 인한 도메인 이탈과 실세계 문서 AI 응용 프로그램에서의 성능 저하 문제를 해결한다.
• 다양한 문서 레이아웃과 포맷에 일반화되는 자기지도 학습 기반 사전 훈련 프레임워크를 개발한다.
• 감독 학습 레이블에 의존하지 않고도 비전 기반 문서 AI 작업의 최종 성능을 향상시킨다.
• 컴퓨터 비전 및 자연어 처리 작업을 동시에 지원하는 통합 트랜스포머 아키텍처를 통해 문서 AI를 위한 기초 모델을 구축한다.

제안 방법

• 대규모 인간 레이블이 없는 문서 이미지에서 마스킹 이미지 모델링(MIM)을 사용해 비전 트랜스포머 백본을 사전 훈련한다.
• 일반적인 모델(예: DALL-E)보다 더 관련성이 높은 시각적 토큰을 생성하기 위해 회색조 문서 이미지에서 도메인 전용 이산형 VAE(dVAE)를 훈련한다.
• 입력 문서 이미지의 일부 패치를 마스킹하고, dVAE에서 생성된 원래의 시각적 토큰을 트랜스포머가 재구성하도록 훈련한다.
• 학습된 위치 임베딩과 다중 헤드 자기주의 어텐션을 사용한 표준 이미지 트랜스포머 아키텍처를 활용해 문서 레이아웃 내 장거리 의존성을 모델링한다.
• 표준 분류 및 탐지 헤드를 사용해 최종 문서 AI 작업에 대해 사전 훈련된 DiT 모델을 미세 조정한다.
• 실제 데이터 외에도 100만 장의 합성 문서 이미지로 추가로 미세 조정하여 제로샷 일반화 성능을 더욱 향상시킨다.

실험 결과

연구 질문

• RQ1 대규모 인간 레이블이 없는 문서 이미지에서 자기지도 학습 기반 사전 훈련이 인간 레이블 없이도 문서 AI 작업의 일반화 성능을 향상시킬 수 있는가?
• RQ2 문서 이해를 위한 의미 있는 시각적 표현을 학습하는 데 있어, 일반적인 이미지 토크나이저(예: DALL-E)에 비해 문서 전용 dVAE는 어떻게 성능을 냈는가?
• RQ3 DiT는 감독 학습 및 자기지도 학습 기반 기준 모델 대비 문서 이미지 분류, 레이아웃 분석, 표 탐지, 텍스트 탐지 등의 최종 작업에서 얼마나 성능을 향상시켰는가?
• RQ4 제한된 또는 도메인 특화된 데이터셋으로 사전 훈련된 모델에 비해 DiT는 다양한 문서 포맷과 레이아웃 간에 더 잘 일반화되는가?
• RQ5 DiT는 레이아웃 기반 모델(예: LayoutLM)을 포함한 다중 모odal 문서 AI 시스템의 강력한 기초 모델로 활용될 수 있는가?

주요 결과

• 문서 이미지 분류 작업에서 RVL-CDIP에 대해 DiT는 92.69% top-1 정확도로 이전 기준 모델의 91.11%보다 향상된 새로운 최신 기술 수준 성능을 달성했다.
• 문서 레이아웃 분석 작업에서 PubLayNet에 대해 DiT는 94.9% F1을 기록했으며, 이는 이전 최신 기술 수준의 91.0%보다 뚜렷한 향상이다.
• ICDAR 2019 cTDaR에서 표 탐지 작업에 대해 DiT-L은 96.55% wF1 스코어를 기록했으며, 이는 이전 최고 성능인 94.23%를 2.3% 이상 초월했다.
• FUNSD에서 텍스트 탐지 작업에 대해 DiT-L은 IoU@0.5 기준 94.29% F1을 달성했으며, 상용 OCR 엔진 기준 모델 및 모든 다른 ViT 및 CNN 백본보다 뛰어난 성능을 보였다.
• 100만 장의 합성 문서 이미지로 구성된 데이터셋으로 DiT를 추가로 미세 조정함으로써 성능이 더욱 향상되었으며, FUNSD에서 94.29% F1, 표 탐지에서 96.55% wF1을 기록했다.
• 모델는 특히 색상 제약이 있는 dVAE로 인해 자원이 제한된 환경(예: cTDaR의 고전적 서류 서브셋)에서도 강력한 소수의 샘플 일반화 성능을 보였으며, 대부분의 기준 모델보다 뛰어났다.