Document AIの基盤モデルとして提案された Universal Document Processing (UDOP)についての論文です。
完全な訳ではなく、部分的に抜粋して整理しています。
従来のVision&Languageの研究と大きく異なる点として、文書のレイアウトを効果的にモデル化することが必要です。
論文の内容を3行にまとめると、
- 画像、テキスト、レイアウトを統一的に扱うことのできるDocument AIの基盤モデルUDOPを提案
- 新しい自己教師あり事前学習を提案し、教師あり事前学習も行うことで複数のタスクでSOTAを達成
- Document AIで初めての高品質な文書内容の編集と生成を実現
コードとモデル：https://github.com/microsoft/i-Code/tree/main/i-Code-Doc
UDOP
以下は提案するUDOPの概要です。
UDOPは、Vision-Text-Layout Transformerと複数の事前学習タスクを生成タスクとして解くことにより、画像、テキスト、レイアウトを統合します。
以降で、Vision-Text-Layout Transformerを構成するUnified EncoderとDecoderについて述べた後、事前学習タスクについて説明します。
Vision-Text-Layout Transformer
Unified Encoder
Unified Encoderでは、画像、テキスト、レイアウトのそれぞれのモダリティを一つのTransformerエンコーダに入力しています。
まず、文書画像$v \in \mathbb{R}^{H \times W \times C}$から、$M$個の単語トークン$\lbrace s_i \rbrace^{M}_{i = 1}$と、
レイアウト情報としてBounding Box $\lbrace(x^1_i, y^1_i, x^2_i, y^2_i)\rbrace^{M}_{i = 1}$が抽出されます。ここで、$(x^1_i, y^1_i), (x^2_i, y^2_i)$はそれぞれBounding Boxの左上、右上の座標です。
文書画像$v$は、$\frac{H}{P} \times \frac{W}{P}$個のパッチに分割します。そして、$P \times P \times C$のサイズである各パッチをPatch Embeddingで$D$次元に埋め込み、$\lbrace v_i \in \mathbb{R}\rbrace^{N}_{i = 1}$にまとめます。ここで、$N$は$\frac{H}{P} \times \frac{W}{P}$です。
テキストトークンについても同様に、$D$次元の埋め込み$\lbrace s_i \rbrace^{M}_{i = 1}$に変換されます。
そして、$s_i$のBounding Boxの中心が画像パッチ$v_j$内にある場合、それらを結合した表現を作成します。
テキストトークンを含まない画像パッチ$v_j$については、$v_j$自体を結合表現として扱います。
この操作によって作成した表現をエンコーダに入力することにより、各モダリティ間の相互作用を大幅に向上させることができます。
Vision-Text-Layout Decoder
Vision-Text-Layout Decoderは、Text-Layout DecoderとVision Decoderから構成されています。
Text-Layout Decoderは、テキストとレイアウトトークンを生成するための単方向のTransformerです。
一方、Vision DecoderはMAEのデコーダで、テキストのレイアウト情報を使用して直接画像ピクセルを生成します。後述しますが、これにより文書画像内のテキストを編集することができるようになります。
自己教師あり事前学習タスク
ラベルなし文書に対する自己教師あり事前学習タスクを提案しています。ラベルなし文書には、文書画像とトークンレベルのbounding boxを持つOCRテキストの入力が含まれています。ここで、「Ship Date to Retail: Week of March 14, 1994」という入力テキストを例として、提案する事前学習タスクを説明します。
（1）Joint Text-Layout Reconstruction
欠落しているテキストを再構成し、文書画像でそれらトークンを見つけるタスクです。
具体的には、テキストトークンの一部をマスクして、そのトークンとレイアウトトークンの両方を予測します。例えば、「Ship Date」と「of」をマスクすると入力系列とターゲットの系列は以下のようになります。
ターゲット系列に、<100><350><118><372>といったトークンが含まれていますが、これはレイアウトトークンといいます。レイアウトトークンはレイアウトとテキストの表現を統一するために導入しています。
レイアウトトークンは、Bounding Box $(x^1_i, y^1_i, x^2_i, y^2_i)$が$[0, 1]$で正規化されているとしたとき、各座標を語彙サイズで乗算し四捨五入した結果です。例えば、語彙サイズが500で$(0.1, 0.2, 0.5, 0.6)$のとき、レイアウトトークンは<50><100><250><300>になります。
また、MLMと同様にマスク率は15%としています。
(論文中で、<100><350><118><372>は "Date to "のレイアウトトークンを表すと書いていますが、正しくは"Ship Date"？)
（2）Layout Modeling
文書画像と文脈となるテキストが与えられたときに、該当のテキストトークンを含むグループの位置を特定するタスクです。
マスク率を小さくするとタスクが簡単になってしまうため、マスク率は75%としています。
（3）Visual Text Recognition
レイアウトトークンから対応するテキストトークンを予測するタスクです。
マスク率は50%としています。
（4）Masked Image Reconstruction with Text and Layout
図3に示すように、テキストとレイアウトによる画像の再構成を目的とするタスクです。
ここでMAEをベースにしていますが、以下の二つの工夫をしています。
（MAEについては、【論文5分まとめ】Masked Autoencoders Are Scalable Vision Learners
がわかりやすいです。）
文字埋め込みの活用
トークン単位だけではなく、文字単位の埋め込みも使用します。Vision DecoderにCross Attentionを追加し、トークン単位の埋め込みと文字単位の埋め込み両方とのCross Attentionをとることで画像生成の質を大幅に向上させることができました。
Image Decoding
MAEは、マスクされていない画像パッチのみをエンコーダに入力する仕組みであり、提案手法は画像パッチとテキストトークンを融合させるため、エンコーダの出力を直接Vision Decoderに与えることはできません。
そのため、学習可能な埋め込みをVision Decoderに与え、エンコーダの出力とのCross Attentionを取ることでこの問題に対応しています。
上述した自己教師あり事前学習には、IIT-CDIP Test Collection 1.0データセットを使用しています。これには、OCRによって抽出されたテキストとトークンレベルの境界ボックスを含む1100万件のスキャンされた文書が含まれています。
教師あり事前学習タスク
以下に示すような教師ありタスクを事前学習で解きます。この時、いずれのタスクについても生成タスクとして解きます。
| タスク | タスク概要 | 使用データセット |
|---|---|---|
| 分類 | 文書のタイプを予測するタスク | RVL-CDIP（16のカテゴリ） |
| レイアウト解析 | 文書内のエンティティの位置を予測するタスク | PubLayNet |
| 情報抽出 | テキストクエリのエンティティタイプと位置を予測するタスク | DocBank、KLC、PWC、DeepForm |
| 質問応答 | 文書イメージに関連付けられた質問に答えるタスク | WebSRC、VisualMRC、DocVQA、InfographicsVQA、WTQ |
| 文書NLI | 文書内の2つの文の含意関係を予測するタスク | TabFact |
実験
モデル構成
Unified EncoderとText-Layout Decoderは、T5-largeのアーキテクチャに従い、Vision DecoderはMAE-largeのデコーダを使用します。全体として、UDOPは794Mの訓練可能なパラメータを持っています。tokenizerには、T5のtokenizerを使用し、, のようなSpecial Tokenを追加しています。
Curriculum Learning
低い解像度では文書中のテキストを認識し辛くなる一方、高解像度では学習に長い時間がかかってしまいます。そこで、訓練を比較的小さい解像度で始め、1epochごとに次第に1024の解像度まで大きくするCurriculum Learningを行います。
実験結果
様々なタスクに対するUDOPとその他モデルとの性能比較を表2、3に示します。
以下の表2では、UDOPはDUE-Benchmarkの7つのタスクすべてでSOTAを達成していることがわかります。
さらに、CORDデータセットにおいて他のモデルよりも優れた性能を達成しています。
他のモデルは各データセットに対してタスク固有のネットワークを使用しているにもかかわらず、UDOPはそれらよりも優れた結果を示しています。
分析
以下の図4に、オリジナル文書画像（左）の内容の編集及び追加の実験結果（右）を示します。
図5はマスクされた画像（左）と再構成した結果（中央）です。オリジナルの文書画像（右）と比較すると、75%という高いマスク率でも高品質に再構成できることがわかります。
以下は、事前学習目的のアブレーションです。ベースラインはMLMを行ったBERTです。
いずれの事前学習目的も性能向上に有効であることがわかります。
以下の表5は、各モダリティを一つのエンコーダで処理するモデル（UDOP）とテキストエンコーダと画像エンコーダを別々に用意するモデル（UDOP-Dual）の比較結果です。
一つのエンコーダを使用する方がほとんどのデータセットで良い結果を出していることがわかります。
まとめ
テキスト、画像、レイアウトを統合したDocument AIの基盤となるUDOPについての論文でした。
UDOPは、高品質な文書生成と編集を実現しており、自己教師あり事前学習と教師あり事前学習を併用することで多くのタスクにおいてSOTAを達成しています。