Pix2Struct：一种革命性的视觉语言理解预训练模型

pix2struct

Pix2Struct简介

Pix2Struct是Google Research团队于2022年提出的一种新型视觉语言预训练模型。它的核心思想是通过学习解析网页截图来获得视觉和语言的联合表示，从而在各种视觉语言理解任务中表现出色。

相比传统的视觉语言模型，Pix2Struct具有以下几个显著特点:

预训练目标新颖：学习将网页截图解析为简化的HTML结构，这一任务自然地涵盖了OCR、语言建模、图像描述等常见预训练信号。
输入表示灵活：采用可变分辨率的输入表示，更好地适应不同下游任务的需求。
多模态融合更自然：将语言提示如问题直接渲染到输入图像上，实现更紧密的视觉语言交互。
通用性强：单一预训练模型可以在多个视觉语言任务上微调并取得最佳性能。

Pix2Struct模型架构

Pix2Struct的预训练

Pix2Struct的预训练数据来源于公开的C4数据集中的网页URL。研究人员对这些网页进行截图，并将其中的部分内容遮盖。模型的任务是将这些带有遮盖的截图解析为简化的HTML结构。

这一预训练目标巧妙地结合了以下几个方面:

OCR：模型需要识别图像中的文本内容。
语言建模：需要理解和生成HTML结构中的文本。
布局分析：需要理解网页的视觉布局结构。
图像理解：需要解析图像中的非文本视觉元素。

通过这种方式，Pix2Struct可以学习到丰富的视觉语言联合表示，为下游任务提供了坚实的基础。

Pix2Struct的应用

Pix2Struct在9个不同领域的视觉语言任务中进行了评估，涵盖了文档、插图、用户界面和自然图像等多个方面。这些任务包括:

TextCaps：为包含文本的图像生成描述
ChartQA：回答关于图表的问题
Widget Captioning：为UI部件生成描述
Screen2Words：为移动应用截图生成描述
RefExp：根据自然语言表达式定位UI元素
DocVQA：回答关于文档图像的问题
InfographicVQA：回答关于信息图的问题
OCR-VQA：回答需要OCR的图像问题
AI2D：回答关于科学图表的问题

在这9个任务中，Pix2Struct在6个任务上取得了最先进的性能，展现了其强大的通用能力。

Pix2Struct在多个任务上的性能

Pix2Struct的实现与使用

Google Research已经开源了Pix2Struct的代码和预训练模型。研究人员和开发者可以通过以下步骤来使用Pix2Struct:

克隆GitHub仓库并安装依赖:

git clone https://github.com/google-research/pix2struct.git
cd pix2struct
conda create -n pix2struct python=3.9
conda activate pix2struct
pip install -e ."[dev]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

准备数据集:

Google提供了多个数据集的预处理脚本，可以将原始数据转换为Pix2Struct所需的格式。

微调模型:

可以使用T5X库提供的训练脚本来微调Pix2Struct模型。例如，微调Screen2Words任务的命令如下:

python -m t5x.train \
  --gin_search_paths="pix2struct/configs" \
  --gin_file="models/pix2struct.gin" \
  --gin_file="runs/train.gin" \
  --gin_file="sizes/base.gin" \
  --gin_file="optimizers/adafactor.gin" \
  --gin_file="schedules/screen2words.gin" \
  --gin_file="init/pix2struct_base_init.gin" \
  --gin.MIXTURE_OR_TASK_NAME='"screen2words"' \
  --gin.MODEL_DIR='"$PIX2STRUCT_DIR/experiments/screen2words_base"' \
  --gin.TASK_FEATURE_LENGTHS='{'inputs': 4096, 'targets': 128}' \
  --gin.BATCH_SIZE=32

推理:

Pix2Struct提供了命令行和Web demo两种推理方式。以下是使用命令行进行单张图像推理的示例:

python -m pix2struct.example_inference \
  --gin_search_paths="pix2struct/configs" \
  --gin_file=models/pix2struct.gin \
  --gin_file=runs/inference.gin \
  --gin_file=sizes/base.gin \
  --gin.MIXTURE_OR_TASK_NAME='"placeholder_pix2struct"' \
  --gin.TASK_FEATURE_LENGTHS='{'inputs': 2048, 'targets': 128}' \
  --gin.BATCH_SIZE=1 \
  --gin.CHECKPOINT_PATH='"gs://pix2struct-data/textcaps_base/checkpoint_280400"' \
  --image=$HOME/test_image.jpg

Pix2Struct的未来发展

作为一种新型的视觉语言预训练模型，Pix2Struct展现了巨大的潜力。未来可能的研究方向包括:

扩大模型规模：目前发布的最大模型为Large版本，未来可能会推出更大规模的模型。
改进预训练数据：探索更多样化、更高质量的预训练数据源。
应用到更多领域：除了已验证的9个任务，Pix2Struct可能在更多视觉语言任务中发挥作用。
与其他技术结合：如将Pix2Struct与大型语言模型或其他视觉模型相结合。
提高推理效率：优化模型结构和推理过程，以支持更快速的实时应用。

总结

Pix2Struct作为一种创新的视觉语言预训练模型，通过学习解析网页截图这一独特任务，实现了对多种视觉语言理解任务的出色泛化能力。它的开源不仅为研究人员提供了宝贵的资源，也为各种实际应用场景中的视觉语言理解任务带来了新的可能性。随着进一步的研究和优化，Pix2Struct有望在视觉语言理解领域发挥更大的作用，推动人工智能在理解和处理多模态信息方面取得新的突破。