Document AIで注目を浴びているLayoutLMv3の論文を読んだので、備忘録も兼ねて簡単に整理します。

Huang, Yupan, et al. "Layoutlmv3: Pre-training for document ai with unified text and image masking." Proceedings of the 30th ACM International Conference on Multimedia. 2022.

LayoutLMv3とは

簡単にまとめると、

テキスト中心、画像中心のDocument AIタスク両方でSOTAを達成する、CNNやFaster R-CNNのようなバックボーンに依存しないDocument AI初のマルチモーダルモデル
MLM（Masked Language Model）, MIM（Masked Image Modeling）によりテキスト・画像モダリティの事前学習の粒度を統一
さらに、WPA（Word Patch Alignment）による事前学習で両モダリティ間の連携を明示的に学習

LayoutLMv3のコードとモデルは公開されています。

Document Understanding

Document AIは、図1で示すようにスキャンされたフォームや学術論文などの様々な文書のレイアウト解析や情報の抽出を目指す分野です。

Document Understandingタスクは急速な進歩を遂げており、様々な事前学習技術が議論されています。LayoutLMv3では、テキストモデルと画像モデルの事前学習目的の不一致に着目し、より効果的なマルチモーダル表現学習の実現を目指しています。
また、LayoutLMv2よりも簡潔なモデル構造となっています。

LayoutLM v3

Layoutv3のモデル構造を上図に示します。

文書画像に対してOCRによる前処理を適用し、RoBERTaによる単語の埋め込みを行います。位置埋め込みには、1D Position Embeddingと2D Position Embeddingがあり、前者はテキスト列内のトークンのインデックス、後者はテキスト列のbounding box座標です。

画像の埋め込みは、画像パッチの線形射影で表現します。具体的には、文書画像をリサイズ、パッチ列への分割、線形射影します。その後、1D Position Embeddingと2D Position Embeddingを追加します。

事前学習目的の不一致

既存のほとんどのマルチモーダルDocument AIモデルにおけるテキストモデル部分では、BERTで提案されたMLMでマスクされたトークンを予測する事前学習しています。一方で、画像モデルの部分では事前学習の目的がテキストモデルと異なっています。例えば、DocFormerでは画像ピクセルの再構成を学習しており、文書レイアウトのような高レベルの構造よりも、ノイズの多い細部を学習してしまう傾向があります。画像（高密度の画像ピクセルや連続領域特徴）とテキスト（=離散トークン）の目的の粒度が異なるため、整合性のあるマルチモーダル表現学習が困難になっています。

LayoutLMv3では、テキストと画像のそれぞれに対して、MLM, MIMという同じ粒度の事前学習を行っています。

MLM（Masked Language Modeling）
Span Maskingというマスク手法を用いており、テキストトークンの30%をマスクします。Span Maskingでは、ランダムなトークンではなく、連続したランダムなスパンをマスクします。
MIM（Masked Image Modeling）
MIMは、BEiTで提案された事前学習手法です。画像をグリッドで分割して作成したパッチのうち約40%をランダムにマスクし、Tokenizerによって生成したトークンに戻すことを学習しています。DeepSquare mediaの解説記事がわかりやすいです。

さらに、それぞれのモダリティを明示的に連携して学習するために、WPA（Word-Patch Alignment） を行います。
WPAは、単語に対応する画像パッチがマスクされているかを予測する二値分類タスクです。損失関数はCross Entropy Lossで、2層からなるMLPを使用します。

モデルアーキテクチャ

LayoutLMv3のモデルアーキテクチャはLayoutLMおよびLayoutLMv2に従っており、BaseモデルとLargeモデルがあります。それぞれの違いは以下の表に示す通りです。

モデル名	Transformerエンコーダの層数	Self Attentionのヘッド数	隠れサイズ (D)	Feed Forward Networkの中間サイズ
LayoutLMv3BASE	12	12	768	3,072
LayoutLMv3LARGE	24	16	1,024	4,096

実験とデータセット

CORD、FUNSD、RVL-CDIP、DocVQAデータセットにおける既存の公開モデルとの比較実験を行います。
それぞれのデータセットの概要は以下の通りです。

CORD
レシート理解のデータセット。数千のインドネシアのレシートで構成されており、OCR 用の画像とボックス/テキスト注釈、構文解析のためのマルチレベルの意味ラベルが含まれている。
FUNSD
フォーム理解のデータセット。注釈が付けられた199個の実際のスキャンされたフォームで構成。
RVL-CDIP
文書画像分類のためのデータセット。手紙、フォーム、電子メール、履歴書、メモなどの 16 のクラスに属するスキャンされた文書画像で構成されている。
DocVQA
文書画像に関するVisual Quession Answing(VQA)のためのデータセット。

実験結果

結果を以下に示します。また、表内のModalityのT, L, IはそれぞれText, Layout, Imageを表しており、Layoutであればレイアウト情報を活用することを示しています。
また、R, G, Pは順に、Faster R-CNNで画像領域の特徴を取得、バックボーンをCNNとして画像グリッドの特徴を取得、CNNを使ったバックボーンを単純な線形埋め込みに置き換えることを指しています。

実験の結果、複数のデータセットにおいてLayoutLMv3が他の手法を上回っていることが確認できます。

さらに、画像埋め込みの有無と事前学習に関するアブレーションは以下の通りです。

ここでは、PubLayNetデータセットを用いて実験を行っています。このデータセットには、テキスト、タイトル、リスト、図、表の5つの文書レイアウトのカテゴリにわたって、バウンディングボックスと多角形セグメンテーションで注釈が付けられた研究論文画像が含まれています。

MLM, MIM, WPAすべて行う場合に性能が最大となっていることがわかります。

視覚タスクでのファインチューニング

文書レイアウト解析タスクでのLayoutLMv3の汎用性を確かめます。文書レイアウト解析タスクは、非構造化デジタル文書のレイアウトを検出するタスクです（参考：図1）。

実験結果は以下の通りです。

LayoutLMv3はすべてのカテゴリにおいて他を上回るmAPスコアを達成しています。特にタイトルカテゴリで他のモデルと差をつけています。これは事前学習で言語モダリティを組み込んだことにより、「タイトル」という他のカテゴリよりも小さく、テキスト内容をみることで識別できるカテゴリに対して優位性があったのではないかと考察されています。

まとめ

テキスト部分と画像部分の事前学習の目的の不一致の解決と各モダリティ間の効果的な融和をMLM, MIM, WPAで実現し、様々なデータセットで高い性能を出している手法でした。
Document AI関連の分野はあまり詳しくないので他の手法についても今後調べて知識を得ていこうと思います。

参考文献