I Think, Therefore I Diffuse: Enabling Multimodal In-Context Reasoning in Diffusion Models

📄 arXiv: 2502.10458v1 📥 PDF

作者: Zhenxing Mi, Kuan-Chieh Wang, Guocheng Qian, Hanrong Ye, Runtao Liu, Sergey Tulyakov, Kfir Aberman, Dan Xu

分类: cs.LG, cs.AI

发布日期: 2025-02-12

备注: Project page: https://mizhenxing.github.io/ThinkDiff, 19 pages, 14 figures

🔗 代码/项目: PROJECT_PAGE

💡 一句话要点

ThinkDiff：通过对齐视觉-语言模型，赋予扩散模型多模态上下文推理能力

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 扩散模型 多模态学习 上下文推理 视觉-语言模型 模型对齐

📋 核心要点

  1. 现有方法侧重像素重建，缺乏多模态上下文推理能力，且推理数据集复杂性高、数量有限。
  2. ThinkDiff利用视觉-语言训练作为代理任务，将视觉-语言模型与大型语言模型的解码器对齐。
  3. 实验表明，ThinkDiff在多模态推理生成任务中显著提升了准确率，且训练成本较低。

📝 摘要（中文）

本文提出了一种名为ThinkDiff的新型对齐范式，它通过整合视觉-语言模型(VLM)的优势，使文本到图像的扩散模型具备多模态上下文理解和推理能力。现有的多模态扩散微调方法主要侧重于像素级别的重建，而非上下文推理，并且受到基于推理的数据集的复杂性和有限性的制约。ThinkDiff通过利用视觉-语言训练作为代理任务来解决这些挑战，将VLM与编码器-解码器大型语言模型(LLM)的解码器对齐，而不是与扩散解码器对齐。这种代理任务建立在这样的观察之上：$ extbf{LLM解码器}$与使用相应$ extbf{LLM编码器}$进行提示嵌入的$ extbf{扩散解码器}$共享相同的输入特征空间。因此，通过与LLM解码器对齐，可以简化VLM与扩散解码器的对齐。无需复杂的训练和数据集，ThinkDiff即可有效地释放扩散模型中的理解、推理和组合能力。实验表明，ThinkDiff在具有挑战性的CoBSAT基准测试中，多模态上下文推理生成的准确率从19.2%显著提高到46.3%，仅在4个A100 GPU上训练5小时。此外，ThinkDiff在将多个图像和文本组合成逻辑上连贯的图像方面表现出卓越的性能。

🔬 方法详解

问题定义：现有文本到图像的扩散模型在处理需要多模态上下文推理的任务时表现不佳。它们主要关注像素级别的重建，而忽略了图像和文本之间的深层语义关系。此外，训练这些模型需要大量的、专门设计的、基于推理的数据集，而这些数据集的构建成本很高且数量有限。

核心思路：ThinkDiff的核心思路是将视觉-语言模型（VLM）与扩散模型的解码器对齐，从而赋予扩散模型多模态上下文推理能力。关键在于，它不是直接将VLM与扩散解码器对齐，而是通过一个代理任务，即与大型语言模型（LLM）的解码器对齐。这是因为LLM解码器与扩散解码器共享相同的输入特征空间，从而简化了对齐过程。

技术框架：ThinkDiff的技术框架主要包含以下几个模块：1) 一个预训练的视觉-语言模型（VLM），用于提取图像和文本的特征；2) 一个编码器-解码器大型语言模型（LLM），其编码器用于生成提示嵌入，解码器作为VLM对齐的目标；3) 一个扩散模型，其解码器负责生成最终的图像。训练过程包括将VLM的输出与LLM解码器的输出对齐，从而使VLM能够理解和推理多模态上下文信息。

关键创新：ThinkDiff最重要的创新点在于它提出了一种新的对齐范式，即通过与LLM解码器对齐，间接地将VLM与扩散解码器对齐。这种方法避免了直接训练VLM和扩散模型，从而降低了训练成本和数据需求。此外，它还利用了LLM强大的语言理解和推理能力，从而提高了扩散模型的多模态上下文推理能力。与现有方法的本质区别在于，ThinkDiff关注的是模型之间的对齐，而不是像素级别的重建。

关键设计：ThinkDiff的关键设计包括：1) 选择合适的VLM和LLM，以确保它们具有良好的特征提取和语言理解能力；2) 设计合适的损失函数，以衡量VLM和LLM解码器之间的对齐程度；3) 调整训练参数，以优化模型的性能。具体的损失函数和训练参数的选择取决于具体的应用场景和数据集。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

ThinkDiff在CoBSAT基准测试中取得了显著的性能提升，准确率从19.2%提高到46.3%，提升幅度超过一倍。此外，ThinkDiff仅使用4个A100 GPU训练5小时，表明其具有较高的训练效率。实验结果表明，ThinkDiff能够有效地赋予扩散模型多模态上下文推理能力，并生成具有逻辑一致性的图像。

🎯 应用场景

ThinkDiff具有广泛的应用前景，例如智能图像编辑、创意内容生成、视觉故事讲述等。它可以用于生成具有复杂语义和逻辑关系的图像，从而为用户提供更加丰富和个性化的体验。此外，ThinkDiff还可以应用于机器人导航、自动驾驶等领域，帮助机器人更好地理解和推理周围环境。

📄 摘要（原文）

This paper presents ThinkDiff, a novel alignment paradigm that empowers text-to-image diffusion models with multimodal in-context understanding and reasoning capabilities by integrating the strengths of vision-language models (VLMs). Existing multimodal diffusion finetuning methods largely focus on pixel-level reconstruction rather than in-context reasoning, and are constrained by the complexity and limited availability of reasoning-based datasets. ThinkDiff addresses these challenges by leveraging vision-language training as a proxy task, aligning VLMs with the decoder of an encoder-decoder large language model (LLM) instead of a diffusion decoder. This proxy task builds on the observation that the $\textbf{LLM decoder}$ shares the same input feature space with $\textbf{diffusion decoders}$ that use the corresponding $\textbf{LLM encoder}$ for prompt embedding. As a result, aligning VLMs with diffusion decoders can be simplified through alignment with the LLM decoder. Without complex training and datasets, ThinkDiff effectively unleashes understanding, reasoning, and composing capabilities in diffusion models. Experiments demonstrate that ThinkDiff significantly improves accuracy from 19.2% to 46.3% on the challenging CoBSAT benchmark for multimodal in-context reasoning generation, with only 5 hours of training on 4 A100 GPUs. Additionally, ThinkDiff demonstrates exceptional performance in composing multiple images and texts into logically coherent images. Project page: https://mizhenxing.github.io/ThinkDiff.