Sayna Ebrahimi, Sercan Ö. Arik, Tomas Pfister
Google Cloud AI Research
发表于 Transactions on Machine Learning Research (06/2023)
研究背景与目的
随着企业数字化进程的加速,视觉文档理解(Visual Document Understanding, VDU)已成为人工智能领域的关键应用技术,广泛应用于发票处理、身份验证、医疗记录解读及合规管理等场景。目前的机器学习解决方案大多依赖“自监督预训练+监督微调”的范式,这虽然在特定任务上表现优越,却面临一个核心挑战:领域偏移(Domain Shift)。当模型在某一源域(Source Domain)数据上训练后,若直接应用于未见过的、具有不同分布的目标域(Target Domain),性能往往大幅下降。由于现实世界中的文档类型、模板、排版、扫描质量等千变万化,这种分布差异极为普遍,严重制约了VDU模型在实际应用中的稳定性与可靠性。
针对该问题,过往的无监督领域自适应(Unsupervised Domain Adaptation, UDA)方法通常需要同时访问带标签的源域数据和无标签的目标域数据,但这在实际部署中可能因隐私保护、法规限制或存储成本而不可行。测试时自适应(Test-time Adaptation, TTA),也称无源领域自适应(Source-free Domain Adaptation),仅凭目标域的无标签数据对源域训练的模型进行在线调整,更具现实意义。然而,现有TTA研究集中于图像分类和语义分割,忽略了VDU任务独有的多模态特性(文本、图像、布局)及长序列输出带来的噪声累积问题。因此,本文首次提出了专为文档理解设计的测试时自适应方法——DocTTA,并构建了相应的基准测试集。
研究方法与工作流
该研究设计并实现了一个名为DocTTA的测试时自适应框架,其核心流程是:给定一个已在源域上训练好的VDU模型,在仅拥有无标签的目标域文档数据且不访问源域数据的条件下,通过优化几个特定的损失函数来在线更新模型参数,完成模型自适应。
DocTTA的工作流主要包含以下三个并行的训练目标,它们构成了方法的独到创新之处:
掩码视觉语言建模(Masked Visual Language Modeling, MVLM)
受BERT等模型在自然语言处理领域成功的启发,DocTTA在测试时引入了掩码视觉语言建模的自监督学习任务。具体操作是:对于目标域输入的文本序列,随机掩盖(mask)掉15%的文本标记(token),其中80%替换为特殊标记[MASK],其余替换为随机词汇。模型的任务是,在视觉信息(文档图像)和布局信息(每个词对应的边界框(Bounding Box), 以6维向量表示)保持完整的条件下,通过一个分类器预测被掩盖掉的原始词汇。该目标的直觉在于,迫使模型学习目标域数据的文本与视觉、布局模态间的内在联系,从而弥合不同模态间的鸿沟,强化对目标文档的内容理解。
基于不确定性感知的伪标签自训练(Self-training with Uncertainty-Aware Pseudo-Labeling)
该方法在自适应过程中,使用模型当前版本对目标数据进行在线预测,生成伪标签(Pseudo Label),并将其作为监督信号进行训练。与传统的离线聚类生成伪标签或每epoch更新一次不同,DocTTA在每个批次(batch)数据上都生成最新的伪标签,实现了模型的自纠正。更为关键的是,为解决伪标签可能存在的噪声(尤其是训练初期预测错误率高)问题,研究提出了一种基于不确定性(Uncertainty)的选择机制:仅选取预测结果的香农熵(Shannon’s Entropy)低于特定阈值的伪标签样本参与交叉熵损失(Cross-entropy Loss)的计算。这有效避免了因累积错误而导致模型性能恶化,并改善了模型的校准度。
多样性最大化目标(Diversity Objective)
为防止模型在伪标签自训练过程中,预测结果盲目偏向某个出现频率最高的“主导类别”,该方法引入了一个多样性损失函数。通过最小化模型在整个目标数据上平均预测概率分布的熵,鼓励模型对各类别进行更为均衡和多样化的预测,从而提升整体泛化能力。
最终的训练总损失为上述三部分损失的加权和。此外,该框架亦可扩展为一个名为DocUDA的UDA方法,即在训练时额外加入源域数据的监督损失。DocUDA同样优于对比的UDA方法。
为了严谨评估DocTTA,研究团队基于现有的公开数据集,构建了三个面向不同VDU任务的TTA基准测试集:
所有实验均采用LayoutLMv2-base架构(约2亿参数)作为骨干模型。对比基准包括:源域直接测试(Source-only)、仅在目标域训练测试(Train-on-target)作为上下限;UDA方法DANN和CDAN;TTA方法Batch Normalization自适应(BN)、TENT和SHOT。
主要实验结果
实验结果表明,DocTTA在所有任务和基准上均显著优于其他TTA基线方法,甚至在一些场景下超越了需要源域数据的UDA方法。
在实体识别(FUNSD-TTA)任务上,DocTTA将源模型的F1分数从80.80%提升至84.23%,相对提升了3.43%,而所有TTA基线方法(BN、TENT、SHOT)的提升幅度均不足0.5%。在键值对提取(SROIE-TTA)任务上,DocTTA将源模型的F1分数从92.45%提升到了94.34%,优于所有TTA基线。
在更具挑战性的文档视觉问答(DocVQA-TTA)任务中,DocTTA展现了最显著的性能提升。例如,在从“电子邮件与信函”域向“图表与图形”域适应时,DocTTA将ANLS分数从37.79%提升至40.36%,提升2.57%;而在差距最大的从“图表与图形”域向“电子邮件与信函”域适应时,ANLS分数更是从5.23%大幅跃升至22.91%,提升高达17.68%。DocTTA在12个适应场景中的11个上,性能超过了需要源域数据的CDAN方法。
值得注意的是,在涉及“图表与图形”域的某些适应场景(如“电子邮件与信函”到“图表与图形”),DocTTA的性能甚至超越了其对应的UDA版本(DocUDA)。论文分析这得益于两方面因素:一是目标域数据集相对较小,导致UDA联合训练时被较大的源域数据主导;二是两个域之间差距巨大,使得联合表征产生偏差。这反而凸显了DocTTA在极端领域偏移下的鲁棒性优势。
消融实验在DocVQA-TTA基准上进行,验证了DocTTA各组件的有效性。结果显示:单独使用基于置信度的伪标签选择机制会导致性能极差,甚至不如不使用伪标签;联合使用置信度和不确定性能带来接近完整方法的性能;去掉MVLM损失或多样性损失,以及完全去掉伪标签损失,都会使模型性能显著下降,证明了每个损失函数对最终性能的关键贡献。
结论与价值
本文首次将测试时自适应引入视觉文档理解领域,并提出了创新的DocTTA方法。该方法通过巧妙结合掩码视觉语言建模和不确定性感知的在线伪标签学习,成功解决了VDU模型在面对分布漂移时性能严重下降的关键瓶颈问题。其在多个VDU核心任务上的突破性表现,为在医疗、金融、法律等对准确性和鲁棒性要求极高的场景中,实现可靠、无监督的模型部署开辟了新途径,具有重大的学术研究价值和深远的应用前景。
研究亮点