【論文解説】Donut：OCR無しのEnd-to-Endな文書理解モデル

本記事は，京都大学人工知能研究会KaiRAのAdvent Calender 10日目の記事です．

https://qiita.com/advent-calendar/2023/kyoto-kaira

# 本記事の概要

本記事では，2021年に発表された論文「OCR-free Document Understanding Transformer」の解説を行います．特に断らない限り，画像はすべてこの論文からの引用です．

https://arxiv.org/abs/2111.15664

この論文では表題にあるとおり，OCRを使わない文書理解モデルである**Donut**を提案しました．

本記事では，以下の3本立てで解説していきます．

- 従来の文書理解モデルのしくみと問題点
- Donutのしくみ
- 精度検証

# 1. 従来の文書理解モデルのしくみと問題点

Donutの説明に入る前に，従来の文書理解モデルおよびその問題点について見ていきます．

## 「画像文書理解」とはどういうタスクか
**Visual Document Understanding（VDU; 画像文書理解）** は，画像形式の文書から構造化情報を抽出するタスクです．

例えば「レシートの画像から商品名・単価・個数・合計金額を読み取りたい」というタスクは画像文書理解に該当します．

## 従来のOCRを使ったやりかた

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1900599/3215092d-1cf5-32c6-1b41-857b836407c3.png)

さて，画像文書理解をやろうと思ったときにまず思いつきそうなのが上図のようなやり方です．

1. 画像をOCRにツッコむ【矢印(b)】
1. OCRからバウンディングボックス（＝bbox）の座標とbbox内の文字が出力される【矢印(c)】
1. 出力から情報を取り出して構造化してくれるモデルにツッコむ【矢印(d)】
1. 完成

3番目のモデルには，例えばBERTやRoBERTa，LayoutLM[^1]などが考えられます．なおBERTやRoBERTaの場合はテキストの位置情報は使っていない（と思う）ですが，LayoutLMでは位置情報も使っています．

[^1]:Microsoft社．3つのバージョン（[LayoutLM](https://arxiv.org/abs/1912.13318)，[LayoutLMv2](https://arxiv.org/abs/2012.14740)，[LayoutLMv3](https://arxiv.org/abs/2204.08387)）があるようです．

要は

- テキスト抽出
- 情報抽出

という2ステージで行うわけです（下図参照）．

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1900599/57c6d0a5-81a2-48e3-9d34-307502f2fd34.png)

## 従来手法の問題点
従来手法はわかりやすく，OCR+LayoutLMで精度も十分出てはいるものの，次のような問題点があります．

- OCRの計算コストが高い
- 言語やドメインの変更に対する柔軟性が低い
- OCRでミスってしまうと後に響く


1つめ「OCRの計算コストが高い」ですが，高品質のOCR結果を得たいとなると推論コストが高くつくそうです．

2つめ「言語やドメインの変更に対する柔軟性が低い」は，他言語に対応しようとすると再学習の必要があり，**計算コストと大規模データ**が必要となるということです．

3つめの「OCRでミスってしまうと後に響く」は，多ステージ制のモデル全般に言えることなのですが，上流プロセスでミスが起きてしまうと，後続プロセスでどうあがいても良い精度は出せないわけです．とくに，**精度が比較的良くないとされる韓国語や中国語のOCRでは結構問題になるようです**．

# 2. Donutのしくみ
こういう従来の問題点を解決するぜ，というのが**Donut**という手法です．

使用モデルはTransformerオンリーであることから「**Do**cume**n**t **u**nderstanding **t**ransformer」，略して**Donut**と論文では呼んでいます[^2]．

[^2]:「n」がちょっと無理矢理感...


Donutの一番の特徴はやはり，OCRを使わない**End-to-End**な構造であるという点です．
画像を入力したら，出力に半構造化データが現れる，という感じです．

End-to-Endにしたことで従来手法の問題点である「OCRでミスってしまうと後に響く」を解決できています．

さらに，計算時間も従来のほぼ半分で済むので「計算コストが高い」問題も解決し，「言語やドメインの変更に対する柔軟性が低い」という問題も**SynthDoG**（後述）を利用することで解決できるとしています．

## モデル構造

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1900599/175df1c9-d3e0-04bf-7c22-d82858782925.png)

Donutのモデル構造は極めてシンプルで，上図のようになっています．

まず，Transformerのエンコーダーで（$n$パッチに分割した）文書画像を潜在変数ベクトル$\\{\boldsymbol{z}_i\mid\boldsymbol{z}_i\in\mathbb{R}^d,1\leq i \leq n\\}$にエンコードし，デコーダーへ渡します．

※エンコードすりゃあいいのでCNN系モデルで特徴抽出してもいいのですが，この論文ではSqin TransformerというViTの改良版を使っているそうです．

デコーダーは，潜在変数ベクトル$\\{\boldsymbol{z}\\}$を入力として，トークン系列$(\boldsymbol{y}_i) _{i=1}^{m}$を出力します．

デコーダーの方はBARTのデコーダーを使っており，多言語で事前学習されたBARTモデル（既存のやつ）を初期値として学習させたようです．

## 学習①：事前学習
まず，モデルが「どのように画像文書を読むか」を理解させるための学習である**事前学習(pre-training)** を行います．

エンコーダーには画像，デコーダーにはこれまでのテキストを入力すると，後続するテキストを生成させるタスクを解かせることで学習します．（OCRモドキなことをやっている）

学習データは**IIT-CDIP**という英語文書画像1100万枚を使用したそうです．ラベリングはCLOVA OCR APIを使って行ったみたいです．

しかし注意なのは"英語の"文書である点です．しかし，本論文では**Synthetic Document Generator(SynthDoG)** を使うことで多言語対応が可能であると主張しています．

SynthDoGで生成される画像は下図のようなものです．
もしかするとGoogle翻訳アプリでこういうのを見たことがある人もいるかもしれませんね．

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1900599/600d2f3a-87a0-0f54-e6c7-daa864d002a1.png)

具体的には，

- 背景画像：ImageNetからランダムに抽出
- 紙地部分：紙画像からランダムに抽出
- テキスト：Wikipediaから抽出
- 合成方法：シンプルなルールベースアルゴリズム（グリッドをランダムに積み重ねる）＋画像レンダリング技術

です．

## 学習②：ファインチューニング

事前学習ができたら，各タスクに適合させる**ファインチューニング（fine-tuning）** を行います．

本論文では，次の3つのタスクでファインチューニングを行っています（詳細は後述）．

- 文書クラス分類
- 文書情報抽出
- 文書質問応答


## ちょっと細かいところ

### teacher-forcing
推論する際にデコーダーに入力するのは，これまでにモデル自身が生成した系列ですが，学習時にはモデルが生成した系列ではなくGround-truthを入力します．こういう学習手法を**teacher-forcing**と呼ぶそうです．

eacher-forcingによって，モデルが前段階で誤ったものを生成してしまっても，Ground-truthで強制的に修正されるので学習が安定しやすいというメリットがあります．

### 特殊トークン
transformerのデコーダーはいきなりJSONを出力するわけではなく，特殊トークンを含んだトークン列が出力されます．

例えば，

[START_name]食パン[END_name]

というトークンが出力されれば，これは簡単なルールベースプログラミングで

{"name" : "食パン"}

のようにJSON形式へ変換することができますよね．

# 3. 精度検証

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1900599/88aae82a-b472-92df-376e-35c88c9fe432.png)


先述したように，本論文では次の3タスクを行っています．上図は入力（&出力）データの例です．

- (a) 文書クラス分類（Document Classification）
- (b) 文書情報抽出（Document Information Extraction）
- (c) 文書質問応答（Document Visual Question Answering；DocVQA）


**文書クラス分類**では，モデルがいろいろな種類の文書を区別できるかを見ます．ただしこれは単純な分類タスク（Softmax通して確率...）ではなく，JSON形式で出力させて精度を検証しています．

**文書情報抽出**は，レシートから情報抽出するという，冒頭でも述べたアレです．

**文書質問応答**では，文書画像と質問のペアが与えられ，画像と質問文の両方から情報を集めて回答を生成します．一瞬，別のモデルが必要ではないかと思われるかもしれませんが，文書画像はエンコーダーに入力し，質問文はデコーダーの先頭につけることで実現できます．


## 文書クラス分類の性能

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1900599/76b50b26-5550-ef25-51c6-80066b1d366e.png)

従来手法に比べて推論時間がほぼ半分になり，精度も最高です．しかもパラメータ数も比較的少なく抑えられています．

## 文書情報抽出の性能

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1900599/efb99c54-a449-4e9f-ada9-d04c47206708.png)

4タイプのドメインで検証されています．

- CORD：ラテン語のレシート．公開ベンチマークデータセット．
- Ticket：中国語のチケット．公開ベンチマークデータセット．
- Business Card：日本語の名刺．
- Receipt：韓国語のレシート．

いずれのドメインにおいても，Donutが最高精度かつ最速推論を行っています．

## 文書質問応答

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1900599/e65fd77f-79fd-afc0-bcf0-69613a76922a.png)

こちらでも，一部を除いて最速かつ最高精度をたたき出しています．

DonutとLayoutLMの出力を比べているのが下図です．質問応答タスクにおいては，Donutは印刷文字よりも手書きに強いといったところでしょうか．

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1900599/80b86dff-9db8-7d5f-9d70-14b8d567337f.png)

# おわり

以上がDonut論文の簡単な解説でした．

たぶん今ならGPT-4V投げれば，いろんなタスク対応できそうだし精度高くなりそうだなーとか思いました（計算速度でいうとGPT-4Vは劣ると思うけど）．