基于Transformer架構(gòu)的文檔圖像自監(jiān)督預(yù)訓(xùn)練技術(shù)

本文簡(jiǎn)要介紹ACM MM 2022錄用論文“DiT： Self-supervised Pre-training for Document Image Transformer”［1］的主要工作。該論文是2022年微軟亞研院發(fā)表的LayoutLM V3［2］的前身工作，主要解決了文檔領(lǐng)域中標(biāo)注數(shù)據(jù)稀少和以視覺為中心的文檔智能任務(wù)骨干網(wǎng)絡(luò)的預(yù)訓(xùn)練問題。

一、研究背景

近年來(lái)自監(jiān)督預(yù)訓(xùn)練技術(shù)已在文檔智能領(lǐng)域進(jìn)行了許多的實(shí)踐，大多數(shù)技術(shù)是將圖片、文本、布局結(jié)構(gòu)信息一起輸入統(tǒng)一的Transformer架構(gòu)中。在這些技術(shù)中，經(jīng)典的流程是先經(jīng)過(guò)一個(gè)視覺模型提取額外文檔圖片信息，例如OCR引擎或版面分析模型，這些模型通常依賴于有標(biāo)注數(shù)據(jù)訓(xùn)練的視覺骨干網(wǎng)絡(luò)。已有的工作已經(jīng)證明一些視覺模型在實(shí)際應(yīng)用中的性能經(jīng)常受到域遷移、數(shù)據(jù)分布不一致等問題的影響。而且現(xiàn)有的文檔有標(biāo)注數(shù)據(jù)集稀少、樣式單一，訓(xùn)練出來(lái)的骨干網(wǎng)絡(luò)并非最適用于文檔任務(wù)。因此，有必要研究如何利用自監(jiān)督預(yù)訓(xùn)練技術(shù)訓(xùn)練一個(gè)專用于文檔智能領(lǐng)域的骨干網(wǎng)絡(luò)。本文針對(duì)上述問題，利用離散變分編碼器和NLP領(lǐng)域的常用預(yù)訓(xùn)練方式實(shí)現(xiàn)了文檔圖像的預(yù)訓(xùn)練。

圖1具有不同布局和格式的視覺豐富的業(yè)務(wù)文檔，用于預(yù)培訓(xùn)DiT

二、DiT原理簡(jiǎn)述

2.1總體結(jié)構(gòu)

圖2 DiT的總體架構(gòu)

Fig 2是DiT的整體結(jié)構(gòu)。DiT使用ViT［3］作為預(yù)訓(xùn)練的骨干網(wǎng)絡(luò)，模型的輸入是圖像Patch化后的Embedding特征向量，Patch的數(shù)量和離散變分編碼器的下采樣比例有關(guān)。輸入經(jīng)過(guò)ViT后輸出到線性層進(jìn)行圖像分類，分類層的大小是8192。預(yù)訓(xùn)練任務(wù)和NLP領(lǐng)域的完型填空任務(wù)一致，先對(duì)輸入的Patch隨機(jī)掩膜，在模型輸出處預(yù)測(cè)被遮蓋的Patch對(duì)應(yīng)的Token，Token由Fig 2 中左側(cè)的離散變分編碼器生成，作為每個(gè)Patch的Label，預(yù)訓(xùn)練過(guò)程使用CE Loss監(jiān)督。

2.2 離散變分編碼器dVAE

離散變分編碼器作為Image Tokenizer，將輸入的Patch Token化，來(lái)源于論文DALL-E［4］，在預(yù)訓(xùn)練任務(wù)開始前需要額外訓(xùn)練。本文使用數(shù)據(jù)集IIT-CDIP［5］重新訓(xùn)練了DALL-E中的離散變分編碼器以適用于文檔任務(wù)。在預(yù)訓(xùn)練任務(wù)中只使用到編碼器的部分，解碼器不參與預(yù)訓(xùn)練，編碼器將輸入圖片下采樣到原來(lái)的1/8，例如輸入尺度為112*112，那編碼后的Token Map為14*14，此時(shí)的Map大小，應(yīng)與ViT輸入Patch數(shù)保持一致。

2.3 模型微調(diào)

圖3在不同檢測(cè)框架中應(yīng)用DiT作為骨干網(wǎng)絡(luò)的圖示

模型預(yù)訓(xùn)練完成后，需針對(duì)下游任務(wù)進(jìn)行微小的結(jié)構(gòu)改動(dòng)，針對(duì)分類任務(wù)，輸入經(jīng)過(guò)平均池化和線性層進(jìn)行分類。針對(duì)檢測(cè)任務(wù)，如Fig 3所示，在ViT的特定層進(jìn)行下采樣或上采樣，然后輸入到FPN和后續(xù)的檢測(cè)框架中。

三、主要實(shí)驗(yàn)結(jié)果及可視化效果

表1.RVL-CDIP上的文檔圖像分類精度（%），其中所有模型都使用224×224分辨率的純圖像信息（無(wú)文本信息）。

表2.PubLayNet驗(yàn)證集上的文檔布局分析mAP@IOU［0.50:0.95］。ResNext-101-32×8d縮短為ResNext，級(jí)聯(lián)為C。

表3.ICDAR 2019 cTDaR的表檢測(cè)精度（F1）

表4.文本檢測(cè)精度（IoU@0.5）在FUNSD任務(wù)#1中，掩碼R-CNN與不同的主干（ResNeXt、DeiT、BEiT、MAE和DiT）一起使用。“+syn”表示使用包含1M文檔圖像的合成數(shù)據(jù)集訓(xùn)練DiT，然后使用FUNSD訓(xùn)練數(shù)據(jù)進(jìn)行微調(diào)。

圖4使用不同標(biāo)記器進(jìn)行圖像重建

從左到右：原始文檔圖像，使用自訓(xùn)練dVAE標(biāo)記器進(jìn)行圖像重建，使用DALL-E標(biāo)記器進(jìn)行的圖像重建從表1、表2、表3、表4

來(lái)看，文章所提方法在各種下游任務(wù)中取得了state-of-the-art的結(jié)果，驗(yàn)證了該方法在文檔領(lǐng)域的有效性。Fig 4中展示了重新訓(xùn)練的離散變分編碼器的可視化輸出，結(jié)果顯示本文中的離散變分編碼器效果更好。

四、總結(jié)及討論

本文設(shè)計(jì)了一個(gè)利用大量無(wú)標(biāo)簽文檔圖像預(yù)訓(xùn)練ViT的自監(jiān)督方法，該方法的核心是利用離散變分編碼器對(duì)圖像Patch進(jìn)行Token化，再使用NLP領(lǐng)域的掩碼重建任務(wù)進(jìn)行預(yù)訓(xùn)練。從實(shí)驗(yàn)結(jié)果可以看出，該方法在多個(gè)下游任務(wù)的有效性，探索了自監(jiān)督任務(wù)在文檔領(lǐng)域的可能性。

審核編輯：郭婷