论文状态:已完成

DiffTMR: Diffusion-based Hierarchical Alignment for Text-Molecule Retrieval

发表:2025/10/25
原文链接
价格:0.100000
已有 2 人读过
本分析由 AI 生成,可能不完全准确,请以原文为准。

TL;DR 精炼摘要

DiffTMR 是一种新颖的文本-分子检索框架,通过将检索视为逆向去噪过程,解决了传统方法在检测域外数据和保持多样性方面的局限性。结合分层扩散对齐和动态扰动嵌入机制,DiffTMR 提升了检索准确性和域外泛化能力,评估结果显示其在基准数据集上超越现有基线模型。

摘要

Molecular retrieval is critical in drug discovery and molecular design. Traditional discriminative methods often model the conditional probability distribution of retrieving candidates, treating the query text as a deterministic input. However, these approaches have notable limitations: (1) They often overlook the statistical properties of the original data distributions of queries and candidates, preventing the recognition of out-of-distribution data. (2) They struggle to balance retrieval accuracy and diversity when processing open-ended semantic queries. To address these challenges, we introduce DiffTMR, a novel framework that reformulates text-molecule retrieval as a reverse denoising process, progressively generating the joint distribution of candidates and queries from noises. DiffTMR uniquely integrates hierarchical diffusion alignment with dynamic perturbation embedding mechanisms. By employing text-anchored perturbations, it enhances the diversity of molecular representations, and through global-local progressive denoising, it achieves cross-modal hierarchical alignment. This leads to significant improvements in retrieval accuracy and out-of-domain generalization. Evaluations on benchmark datasets ChEBI-20 and PCdes demonstrate that DiffTMR surpasses current leading baselines by 4.2%–5.4% in Hits@1 metrics and exhibits superior performance in out-of-domain retrieval tasks.

思维导图

论文精读

中文精读

1. 论文基本信息

1.1. 标题

DiffTMR: Diffusion-based Hierarchical Alignment for Text-Molecule Retrieval (DiffTMR: 基于扩散的分层对齐文本-分子检索)

1.2. 作者

Chenxu Wang, Dong Zhou, Ting Liu, Jianghao Lin, Yongmei Zhou, Aimin Yang。 主要作者来自广东外语外贸大学 (Guangdong University of Foreign Studies),Aimin Yang 来自岭南师范学院 (Lingnan Normal University)。

1.3. 发表期刊/会议

ACM International Conference on Multimedia (MM '25),第33届。 ACM Multimedia (MM) 是多媒体领域顶级的国际会议,在计算机科学、人工智能和多媒体技术领域享有很高的声誉和影响力。在该会议上发表论文通常意味着该研究具有较高的创新性和学术价值。

1.4. 发表年份

2025年

1.5. 摘要

分子检索 (Molecular retrieval) 在药物发现和分子设计中至关重要。传统的判别式方法 (discriminative methods) 通常建模检索候选分子的条件概率分布 (conditional probability distribution),将查询文本视为确定性输入。然而,这些方法存在显著局限性:(1) 它们常常忽略查询和候选分子原始数据分布的统计特性 (statistical properties),从而无法识别域外数据 (out-of-distribution data)。(2) 在处理开放式语义查询时,它们难以平衡检索准确性 (retrieval accuracy) 和多样性 (diversity)。为解决这些挑战,本文引入了 DiffTMR,一个新颖的框架,将文本-分子检索重构为一种逆向去噪过程 (reverse denoising process),逐步从噪声中生成候选分子和查询的联合分布 (joint distribution)。DiffTMR 独特地集成了分层扩散对齐 (hierarchical diffusion alignment) 和动态扰动嵌入 (dynamic perturbation embedding) 机制。通过采用文本锚定扰动 (text-anchored perturbations),它增强了分子表示的多样性;通过全局-局部渐进去噪 (global-local progressive denoising),它实现了跨模态分层对齐 (cross-modal hierarchical alignment)。这带来了检索准确性和域外泛化 (out-of-domain generalization) 的显著提升。在基准数据集 ChEBI-20 和 PCdes 上的评估表明,DiffTMR 在 Hits@1 指标上超越了当前领先的基线模型 4.2%–5.4%,并在域外检索任务中表现出卓越性能。

1.6. 原文链接

/files/papers/694a002c3e1288a634f1be16/paper.pdf 发布状态:预印本,已提交 ACM International Conference on Multimedia (MM '25) 并被接收。

2. 整体概括

2.1. 研究背景与动机

在药物发现、分子设计和虚拟筛选等领域,文本到分子检索 (text-to-molecule retrieval) 是一项关键任务,近年来取得了显著进展。主流方法通常基于对比学习 (contrastive learning) 的判别式范式 (discriminative paradigm),将文本和分子表示投影到共享特征空间,并通过优化密集特征提取器来实现跨模态语义对齐 (cross-modal semantic alignment)。尽管这些方法在检索准确性上有所提升,但从概率论角度看,它们存在以下固有限制:

  1. 忽视原始数据分布的统计特性: 现有的判别式模型只关注在给定查询条件下候选分子的条件概率 (conditional probability),即 p(候选 | 查询),而忽略了查询和候选分子数据本身的统计特性,例如 p(查询)p(候选)。这导致模型在数据分布发生偏移(如跨数据集)时泛化能力不足,难以识别域外数据 (out-of-distribution data)。

  2. 准确性与多样性的权衡难题: 传统的固定点嵌入机制 (fixed-point embedding mechanism) 在处理开放式语义查询 (open-ended semantic queries) 时,难以平衡检索准确性 (precision) 和多样性 (diversity)。为了保持准确性,模型可能会过度压缩语义空间,导致响应僵化;而为了增强多样性,又可能引入不受约束的噪声,造成语义漂移 (semantic drift)。

    这些挑战促使研究者寻求新的方法,不仅能够实现精确的跨模态对齐,还能捕获数据内在的统计特性,并有效平衡检索的准确性和多样性。

2.2. 核心贡献/主要发现

本文提出了 DiffTMR (Diffusion-based Text-Molecule Retrieval),一个新颖的生成式框架,将文本-分子检索任务重构为一种渐进式去噪过程 (progressive denoising process),用于联合分布建模 (joint distribution modeling)。DiffTMR 的核心贡献和主要发现包括:

  1. 首个将生成式扩散模型应用于文本-分子检索任务: DiffTMR 是首次将扩散模型 (diffusion models) 引入跨模态文本-分子检索领域,通过逆向去噪过程逐步生成查询和候选分子的联合分布,从而更全面地捕获跨模态语义关联。

  2. 引入动态扰动采样机制 (dynamic perturbation sampling mechanism): 提出文本锚定扰动表示 (Text-Anchored Perturbation Representation, TAPR) 模块,该模块能够根据文本-分子间的语义耦合 (semantic coupling) 动态调节分子表示的扰动范围,在文本约束下增强分子表示的多样性,同时防止语义漂移。

  3. 结合判别式与生成式优化 (discriminative-generative collaborative optimization): DiffTMR 通过联合优化判别式对比学习 (contrastive learning) 和生成式扩散建模 (diffusion modeling),前者强化了显式语义对齐,后者捕获了隐式分布依赖,从而显著提高了模型的检索准确性和对未见数据 (unseen data) 的泛化能力。

  4. 卓越的性能表现: 在基准数据集 ChEBI-20 和 PCdes 上,DiffTMR 在 Hits@1 指标上分别超越了当前最先进的基线模型 5.4% 和 4.2%。更重要的是,它在域外检索 (out-of-domain retrieval) 任务中表现出优越的泛化能力和鲁棒性,特别是在处理模糊查询和长尾检索场景时。

    这些发现表明 DiffTMR 提供了一个强大且通用的解决方案,能够克服传统判别式方法在文本-分子检索中的局限性。

3. 预备知识与相关工作

3.1. 基础概念

3.1.1. 文本-分子检索 (Text-Molecule Retrieval)

文本-分子检索是指给定一个文本描述(查询),从一个分子数据库中找到最相关的分子;或者反之,给定一个分子,找到最相关的文本描述。这项任务在药物发现、材料科学和化学生物学等领域具有重要应用。

3.1.2. 判别式模型 (Discriminative Models) vs. 生成式模型 (Generative Models)

  • 判别式模型 (Discriminative Models): 主要学习输入特征 XX 到输出标签 YY 的映射关系,即直接建模条件概率分布 p(YX)p(Y|X)。它们侧重于区分不同类别,通常用于分类、回归和检索(通过学习相似性分数)等任务。判别式模型不需要理解数据的内在生成机制,通常在给定足够数据时表现出高准确性,但可能难以处理域外数据或生成多样性结果。
  • 生成式模型 (Generative Models): 旨在学习数据本身的分布 p(X) 或联合分布 p(X, Y)。通过理解数据的生成过程,它们不仅可以区分类别,还可以生成新的、与训练数据相似的样本。在检索任务中,建模联合分布 p(X, Y) 可以更全面地捕获不同模态之间的关联,提高模型的鲁棒性和泛化能力。

3.1.3. 扩散模型 (Diffusion Models)

扩散模型 (Diffusion Models),全称去噪扩散概率模型 (Denoising Diffusion Probabilistic Models, DDPMs),是一种新兴的生成式模型。其核心思想是通过一个前向扩散过程 (forward diffusion process) 逐步向数据中添加高斯噪声,直到数据完全变为随机噪声。然后,模型学习一个逆向去噪过程 (reverse denoising process),从纯噪声中逐步恢复出原始数据。

  • 前向扩散过程: 在每个时间步 tt,根据预设的噪声调度向数据 xt1x_{t-1} 添加少量噪声,得到 xtx_t。这个过程是马尔可夫链式的,可以通过高斯分布表示。 q(xtxt1)=N(xt;1βtxt1,βtI) q(x_t | x_{t-1}) = \mathcal{N}(x_t; \sqrt{1 - \beta_t} x_{t-1}, \beta_t \mathbf{I}) 其中,βt\beta_t 是在时间步 tt 添加的噪声方差,I\mathbf{I} 是单位矩阵。
  • 逆向去噪过程: 模型学习从 xtx_t 预测 xt1x_{t-1} 的后验分布,通常通过预测在 xtx_t 中添加的噪声来实现。通过迭代地去除噪声,最终可以从纯噪声 xTx_T 中生成出干净的数据 x0x_0。这个过程通常用神经网络来建模。 pθ(xt1xt)=N(xt1;μθ(xt,t),Σθ(xt,t)) p_\theta(x_{t-1} | x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \Sigma_\theta(x_t, t)) 其中,μθ\mu_\thetaΣθ\Sigma_\theta 是由神经网络学习的均值和方差。 扩散模型的“粗粒度到细粒度”的生成特性使其非常适合逐步揭示文本和分子之间的关联。

3.1.4. 对比学习 (Contrastive Learning)

对比学习是一种自监督学习范式,旨在学习使相似样本在嵌入空间中彼此靠近,而不相似样本彼此远离的表示。在跨模态检索中,这意味着将匹配的文本-分子对的嵌入距离拉近,而不匹配的对的嵌入距离推远。它通常通过最大化匹配对的相似性并最小化不匹配对的相似性(例如使用 InfoNCE Loss)来实现。

3.1.5. 图神经网络 (Graph Neural Networks - GNN / GCN)

图神经网络 (Graph Neural Networks, GNN) 是一类专门用于处理图结构数据的神经网络。图卷积网络 (Graph Convolutional Networks, GCN) 是GNN的一种,通过在图的节点及其邻居之间聚合信息来学习节点表示。对于分子而言,分子可以表示为图,其中原子是节点,化学键是边。GCN 能够有效地捕获分子的局部和全局结构信息,例如官能团 (functional groups) 和原子亚结构 (atom substructures)。

3.1.6. 重参数化技巧 (Reparameterization Trick)

重参数化技巧是一种在训练带有随机采样的模型时,允许梯度通过采样过程传播的技术。例如,从高斯分布 N(μ,σ2)\mathcal{N}(\mu, \sigma^2) 中采样一个变量 zz,可以表示为 z=μ+σϵz = \mu + \sigma \cdot \epsilon,其中 ϵN(0,1)\epsilon \sim \mathcal{N}(0, 1) 是一个标准正态分布的随机噪声。这样,对 μ\muσ\sigma 的梯度就可以通过 ϵ\epsilon 进行计算,而 ϵ\epsilon 本身是无梯度的,从而实现了反向传播。本文利用此技巧引入动态扰动。

3.1.7. 注意力机制 (Attention Mechanism)

注意力机制是一种允许模型在处理序列或集合数据时,对不同部分赋予不同重要性的技术。它通常通过计算查询 (Query, QQ) 与键 (Key, KK) 之间的相似度来生成注意力权重,然后将这些权重应用于值 (Value, VV) 的加权和。 Attention(Q,K,V)=softmax(QKTdk)V \mathrm{Attention}(Q, K, V) = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V 其中 Q, K, V 是查询、键和值矩阵,dkd_k 是键向量的维度,用于缩放,防止内积过大导致 softmax 饱和。注意力机制使得模型能够聚焦于输入中最相关的部分。

3.1.8. KL 散度 (Kullback-Leibler Divergence)

KL 散度 (Kullback-Leibler Divergence) 是一种衡量两个概率分布之间差异的非对称度量。如果两个分布越相似,KL 散度越接近零。 KL(PQ)=iP(i)logP(i)Q(i) \mathrm{KL}(P \| Q) = \sum_i P(i) \log \frac{P(i)}{Q(i)} 其中 PPQQ 是两个概率分布。在本文中,它用于衡量模型学习到的数据分布与真实数据分布之间的差异,作为生成式损失函数的一部分。

3.2. 前人工作

3.2.1. 文本-分子检索 (Text-Molecule Retrieval) 领域

现有的跨模态文本-分子检索方法主要集中在构建文本和分子的表示,并将它们对齐到共享嵌入空间。这些表示可以分为:

  • 一维序列表示: 通常使用 SMILES (Simplified Molecular Input Line Entry System) 字符串,并利用预训练模型如 KV-PLM [40]、Text+^+Chem T5 [7] 和 MolT5 [10] 进行序列建模。

  • 二维分子图表示: 将分子建模为图结构,原子为节点,化学键为边。代表性工作包括 MoMu [33] 和 MoleculeSTM [23] 使用对比学习构建共享嵌入空间,MolCA [24] 引入专门的跨模态投影器,AMAN [43] 探索对抗学习进行模态对齐。

  • 三维空间构象表示: 3D-MoLM [19] 扩展图神经网络处理 3D 坐标信息。

  • 多模态辅助信息: MolFM [28] 整合知识图谱,GIT-Mol [21] 利用分子图像特征。

  • 混合粒度对齐: Orma [29] 和 Atomas [41] 通过联合建模原子级、基元级(如官能团、原子亚结构)和分子级特征,增强语义一致性。

    局限性: 大多数现有方法仅建模给定查询条件下候选分子的条件概率,忽略了查询和候选分子数据的潜在联合分布。

3.2.2. 基于扩散模型的表示学习 (Representation Learning with Diffusion Models)

扩散模型在多个领域取得了突破:

  • 图像生成: UNIT-DDPM [31] 利用马尔可夫链推断近似联合分布,ILVR [6] 优化 DDPM 生成过程。

  • 逆问题求解: Kadkhodaie 等 [16] 提出随机梯度上升方法结合 CNN-based 隐式先验,用于去噪、去模糊和超分辨率。

  • 3D 点云生成: Luo 等 [27] 将点云合成建模为热力学粒子扩散过程。

  • 视频理解: DiffusionVMR [42] 集成去噪生成策略优化视频片段检索和边界检测;Luo 等 [26] 引入文本引导的扩散模型编辑视频。

    局限性: 尽管扩散模型在其他领域表现出色,但在文本-分子检索领域尚未被探索,尤其是在建模复杂分子结构语义和化学分析固有的双向推理方面。

3.3. 技术演进

文本-分子检索领域的技术演进经历了从早期基于规则或特征工程的方法,到近年来以深度学习为核心的嵌入和对齐方法。其中,对比学习范式在将不同模态的数据映射到共享嵌入空间方面取得了显著成功,催生了大量判别式模型。这些模型专注于优化一个单向或双向的条件概率(如 p(分子|文本)p(文本|分子)),旨在最大化匹配对的相似性。

然而,这种判别式范式在处理数据分布偏移和平衡检索多样性与准确性方面存在固有缺陷。生成式模型,特别是扩散模型,作为一种强大的工具,能够学习并生成数据的联合分布 p(文本, 分子),从而提供了解决这些局限性的新思路。

本文的工作正处于这一技术演进的交汇点,它将扩散模型的生成能力与判别式方法的对齐优势相结合。通过将文本-分子检索重构为联合分布的生成问题,DiffTMR 克服了传统判别式模型对数据分布统计特性忽视的弊端,并利用扩散模型的渐进去噪特性来建立更鲁棒和细粒度的跨模态对齐。

3.4. 差异化分析

DiffTMR 与相关工作中的主要方法相比,核心区别和创新点在于:

  1. 从条件概率到联合分布建模: 现有方法大多是判别式模型,专注于建模条件概率 p(候选|查询)。DiffTMR 则从生成式角度出发,通过扩散模型建模查询和候选分子的联合分布 p(查询, 候选)。这种方法能够更全面地捕获数据内在的统计特性和双向依赖关系,从而增强模型对数据分布变化的鲁棒性和泛化能力。
  2. 融合判别式与生成式优化: DiffTMR 不仅仅是纯粹的生成式模型,它巧妙地结合了判别式对比学习和生成式扩散建模。判别式损失强化了显式的语义对齐,而生成式损失通过学习联合分布来捕捉隐式的分布依赖。这种协同优化 (co-optimization) 机制使得模型同时具备了判别式方法的精准性和生成式方法的泛化性。
  3. 动态文本锚定扰动表示 (TAPR): 引入了 TAPR 模块,通过分析文本和分子特征之间的细粒度关联,动态生成分子的扰动范围。这种“文本锚定”的扰动方式使得分子表示在增强多样性 (diversity) 的同时,避免了语义漂移 (semantic drift),解决了传统固定点嵌入在准确性和多样性之间的权衡难题。
  4. 层级扩散对齐网络 (HDAN): 提出了 HDAN,通过全局语义关联和局部细粒度对应(文本词元与分子基元)的建模,实现了跨模态的层级对齐。这种多粒度建模结合扩散模型的渐进去噪特性,能够动态地弥合模态间隙,提供更稳健的匹配。
  5. 对域外数据的卓越泛化能力: 由于其联合分布建模的特性,DiffTMR 在域外 (out-of-domain) 检索任务中表现出显著优于基线模型的性能,验证了其更强的泛化和迁移能力,解决了现有判别式方法在跨数据集泛化能力受限的问题。

4. 方法论

4.1. 方法原理

DiffTMR (Diffusion-based Text-Molecule Retrieval) 的核心思想是将传统的文本-分子检索任务从判别式方法(建模条件概率 p(mt)p(m|t)p(tm)p(t|m))重构为一种生成式任务,即通过逆向去噪过程 (reverse denoising process) 建模文本和分子的联合概率分布 p(t, m)。这种方法能够更全面地捕获跨模态数据的统计特性,从而提升模型的鲁棒性和泛化能力。

DiffTMR 框架主要由两个关键组件构成:

  1. 文本锚定扰动表示 (Text-Anchored Perturbation Representation, TAPR) 模块: 负责生成更丰富、文本语义感知的分子表示。它通过分析文本和分子特征间的细粒度关联,动态地为分子嵌入定义一个扰动范围,并在此范围内进行采样,以增强分子表示的多样性,同时通过文本锚定防止语义漂移。

  2. 层级扩散对齐网络 (Hierarchical Diffusion Alignment Network, HDAN): 这是一个基于扩散模型的网络,用于在全局(句子-分子)和局部(词元-基元)两个层级上进行跨模态对齐。它通过渐进式去噪过程,从噪声中逐步生成文本和分子的联合分布,从而实现鲁棒且细粒度的对齐。

    这两个组件共同工作,并通过判别式对比学习和生成式扩散建模的协同优化进行训练,以同时提升显式语义对齐的精确性和隐式分布依赖的捕获能力。

4.2. 核心方法详解 (逐层深入)

DiffTMR 的整体架构如下图(原文 Figure 2)所示:

该图像是示意图,展示了DiffTMR框架在文本-分子检索中的过程,通过文本锚定的扰动表示和层级扩散对齐网络进行建模。图中显示了全局和局部查询、候选分子与噪声的关系与逐步去噪的过程,以及最终的检索结果。 该图像是示意图,展示了DiffTMR框架在文本-分子检索中的过程,通过文本锚定的扰动表示和层级扩散对齐网络进行建模。图中显示了全局和局部查询、候选分子与噪声的关系与逐步去噪的过程,以及最终的检索结果。

4.2.1. 文本锚定扰动表示 (Text-Anchored Perturbation Representation, TAPR)

TAPR 模块旨在解决传统分子表示可能存在的次优问题,并充分利用文本信息来增强分子表示。它通过动态调节分子表示的扰动范围,使其更好地与文本语义对齐。TAPR 的架构如下图(原文 Figure 3)所示。

Figure 3: TAPR architecture: A text-molecule encoder extracts hierarchical features: text has word-level \(\\pmb { h } ^ { w }\) and sentence-level global \(\\mathbf { \\Omega } _ { h } t\) ; molecules have motif-level \(\\pmb { h } ^ { m }\) and global \(h ^ { g }\) .Based on text-molecule semantic coupling, each molecule's perturbation range \(\\mathcal { P }\) is modeled. Within this range, Gaussian-sampled \(( \\epsilon \\sim \\mathcal { N } ( \\mathbf { 0 } , \\mathbf { 1 } ) )\) perturbations generate text-guided molecular perturbation representation \(h _ { p } ^ { g }\) . 该图像是示意图,展示了TAPR架构的文本-分子编码器。文本编码器提取了词级特征hw\pmb { h } ^ { w }和句子级全局特征Ωht\mathbf { \Omega } _ { h } t;分子编码器提取了基元级特征hm\pmb { h } ^ { m }和全局特征hgh ^ { g }。通过语义耦合感知模块,结合每个分子的扰动范围P\mathcal { P },生成文本引导的分子扰动表示hpgh _ { p } ^ { g },其中扰动依据高斯分布采样,即ϵN(0,1)\epsilon \sim \mathcal { N } ( \mathbf { 0 } , \mathbf { 1 } )

  1. 编码器 (Encoders):

    • 文本编码器 (ϕt\phi_t): 采用预训练的 SciBERT [1] 模型,用于编码文本描述。它能够捕获句子级和词元级 (word-level) 的表示。
      • 文本词元级表示:\pmb{h}_i^w = \phi_t(t),其中 i[1,...,Nw]i \in [1, ..., N_w]NwN_w 是文本描述的长度。
      • 全局文本表示:Ωht=[cls]w\boldsymbol{\Omega}_h^t = [\mathrm{cls}]_w,通常从 [CLS] 词元获得,代表句子级的语义。
    • 分子编码器 (ϕm\phi_m): 采用一个三层图卷积网络 (Graph Convolutional Network, GCN) [37] 来学习分子表示。GCN 通过聚合邻居信息,能够捕获分子中不同粒度(如官能团、原子簇或原子亚结构)的特征。
      • 分子基元级表示:\pmb{h}_j^m = \phi_m(m),其中 j[1,...,Nm]j \in [1, ..., N_m]NmN_m 是分子基元的数量。
      • 全局分子表示:hg=1Nmj=1Nmhjmh^g = \frac{1}{N_m} \sum_{j=1}^{N_m} \pmb{h}_j^m,通过对所有分子基元表示进行平均池化 (average pooling) 得到。
  2. 动态扰动嵌入 (Dynamic Perturbation Embedding): 为了增强分子表示的多样性,并使其更好地与文本语义对齐,TAPR 模块引入了动态扰动。给定一个全局分子嵌入 hgh^g,模型通过重参数化技巧 (reparameterization trick) 引入一个动态扰动范围 P\mathcal{P},并在该范围内进行随机采样,得到分子扰动嵌入 hPGh_{\mathcal{P}}^{\mathcal{G}}hPG=hG+Pϵ,ϵN(0,1) h _ { \mathcal { P } } ^ { \mathcal { G } } = h ^ { \mathcal { G } } + \mathcal { P } \cdot \epsilon , \quad \epsilon \sim N ( 0 , 1 )

    • 符号解释:
      • hPGh_{\mathcal{P}}^{\mathcal{G}}:经过文本锚定扰动后的全局分子嵌入。
      • hGh^{\mathcal{G}}:原始的全局分子嵌入,如方程 (1) 定义。
      • P\mathcal{P}:动态扰动范围,它控制着扰动的大小。
      • ϵ\epsilon:服从标准正态分布 N(0,1)\mathcal{N}(0, 1) 的噪声项,用于辅助采样过程。
  3. 语义耦合感知模块 (Semantic Coupling Perception Module): 确定合适的扰动范围 P\mathcal{P} 至关重要:过小的扰动会限制模型捕获文本语义多样性的能力;过大的扰动则可能引入偏离语义核心的噪声嵌入。为解决此问题,作者提出了语义耦合感知模块,它通过细粒度的跨模态交互来动态校准扰动强度。

    • 计算语义耦合强度 (cic_i): 该模块首先分析分子 hgh^g 与文本中的每个词元 hiwh_i^w 之间的语义对齐强度: ci=hghiwhghiw,i=1,...,Nw \pmb { c } _ { i } = \frac { \pmb { h } ^ { g } \cdot \pmb { h } _ { i } ^ { w } } { \lVert \pmb { h } ^ { g } \rVert \lVert \pmb { h } _ { i } ^ { w } \rVert } , \quad i = 1 , . . . , N _ { w }
      • 符号解释:
        • ci\pmb{c}_i:分子与第 ii 个文本词元之间的语义耦合强度,通过余弦相似度 (cosine similarity) 计算。
        • hg\pmb{h}^g:全局分子嵌入。
        • hiw\pmb{h}_i^w:第 ii 个文本词元嵌入。
        • \lVert \cdot \rVert:L2 范数。
    • 确定分子扰动范围 (P\mathcal{P}): 基于计算出的语义耦合强度,使用一个可学习的线性层 Θ\Theta 来确定分子扰动范围: P=exp(CΘ),C=[c1,c2,..,cNw] \mathcal { P } = \exp ( C \Theta ) , \quad C = [ \boldsymbol { c } _ { 1 } , \boldsymbol { c } _ { 2 } , . . , \boldsymbol { c } _ { N _ { w } } ]
      • 符号解释:
        • P\mathcal{P}:动态扰动范围,它被设计为非负值,因此使用 exp\exp 函数。
        • C=[c1,c2,...,cNw]C = [\boldsymbol{c}_1, \boldsymbol{c}_2, ..., \boldsymbol{c}_{N_w}]:语义耦合张量,捕获分子与文本描述中每个词元之间的交互。
        • Θ\Theta:一个可学习的线性层 (learnable linear layer),用于将语义耦合张量 CC 映射到合适的扰动范围。
  4. 判别式损失函数 (LDis\mathcal{L}_{\mathrm{Dis}}): 在训练阶段,模型计算全局文本嵌入 Ωht\boldsymbol{\Omega}_h^t 与分子扰动嵌入 hPGh_{\mathcal{P}}^{\mathcal{G}} 之间的相似度分数 st,ms_{t,m}。然后,根据 InfoNCE 损失的思路,构建对应的条件概率: p(mt;ϕt,ϕm)=exp(st,m/τ)mMexp(st,m/τ) p ( m | t ; \phi _ { t } , \phi _ { m } ) = \frac { \exp { \left( s _ { t , m } / \tau \right) } } { \sum _ { m ^ { \prime } \in M } { \exp { \left( s _ { t , m ^ { \prime } } / \tau \right) } } }

    • 符号解释:
      • p(mt;ϕt,ϕm)p(m|t; \phi_t, \phi_m):给定文本 tt 检索到分子 mm 的条件概率,由文本编码器 ϕt\phi_t 和分子编码器 ϕm\phi_m 决定。
      • st,ms_{t,m}:文本嵌入 Ωht\boldsymbol{\Omega}_h^t 与分子扰动嵌入 hPGh_{\mathcal{P}}^{\mathcal{G}} 之间的相似度分数(通常是余弦相似度)。
      • τ\tau:温度参数 (temperature parameter),用于调节相似度分布的平滑程度。
      • MM:当前批次 (batch) 中的候选分子集合,包括正样本和负样本。 相应的判别式损失函数 LDis\mathcal{L}_{\mathrm{Dis}} 定义为: LDis=12E(t,m)[logp(mt;ϕt,ϕm)+logp(tm;ϕt,ϕm)] \mathcal { L } _ { \mathrm { Dis } } = \frac { 1 } { 2 } \mathbb { E } _ { ( t , m ) } \big [ \log \mathnormal { p } ( m | t ; \phi _ { t } , \phi _ { m } ) + \log \mathnormal { p } ( t | m ; \phi _ { t } , \phi _ { m } ) \big ]
    • 符号解释:
      • E(t,m)\mathbb{E}_{(t,m)}:对文本-分子对 (t, m) 的期望。
      • logp(mt;ϕt,ϕm)\log p(m|t; \phi_t, \phi_m):给定文本 tt 检索到分子 mm 的对数概率。
      • logp(tm;ϕt,ϕm)\log p(t|m; \phi_t, \phi_m):给定分子 mm 检索到文本 tt 的对数概率。 该损失函数鼓励模型学习能够使匹配的文本-分子对之间的条件概率最大化的表示。

4.2.2. 层级扩散对齐网络 (Hierarchical Diffusion Alignment Network, HDAN)

为了弥补判别式范式忽略原始数据分布 p(t)p(m) 的局限性,并进一步提升对齐准确性,HDAN 从生成式角度建模文本和分子的联合概率分布 p(t,m)。HDAN 架构如下图(原文 Figure 4)所示。

Figure 4: HDAN architecture: Multi-granularity representations are batch-concatenated. Noise is added to \(H ^ { t }\) (sentencelevel text) and \(H _ { p } ^ { g }\) (molecular perturbation) to generate global alignment distribution \(\\varepsilon ^ { \\mathrm { g l o b a l } }\) via attention. Local distribution \(\\varepsilon ^ { \\mathrm { l o c a l } }\) is built from \(H ^ { w }\) (word-level) and \(H ^ { m }\) (motif-level). Their weighted fusion is fed into the denoising network, which after reverse denoising iterations predicts clean alignment distribution \({ \\hat { X } } _ { 0 }\) . 该图像是示意图,展示了HDAN架构的工作原理。图中显示了多层次表示的批量连接过程,噪声被添加到句子级文本HtH^t和分子扰动HpgH^g_p,通过注意力机制生成全局对齐分布εglobal\varepsilon^{\mathrm{global}}。局部分布εlocal\varepsilon^{\mathrm{local}}由单词级HwH^w和基元级HmH^m构成。两者加权融合后输入去噪网络,经过反向去噪迭代,预测出干净的对齐分布X^0\hat{X}_0

  1. 目标函数 (Objective Function): HDAN 的目标是通过扩散网络参数 ψ\psi 来近似文本-分子语料库 Z\mathcal{Z} 中所有文本-分子对 (t,m) 的联合概率分布的乘积。这个过程从标准高斯噪声 N(0,I)N(0, \mathrm{I}) 开始,逐步生成联合分布。 (t,m)Zp(m,t)ψ(m,t,N(0,I)) \prod _ { ( t , m ) \in \cal Z } p ( m , t ) \approx \psi \big ( m , t , N ( 0 , \mathrm { I } ) \big )

    • 符号解释:
      • (t,m)Zp(m,t)\prod_{(t,m) \in \mathcal{Z}} p(m,t):文本-分子语料库 Z\mathcal{Z} 中所有文本-分子对的联合概率的乘积。
      • ψ()\psi(\cdot):扩散网络的参数,用于建模逆向去噪过程。
      • N(0,I)N(0, \mathrm{I}):标准高斯噪声,作为扩散过程的起点。
  2. 扩散理论启发下的噪声注入与投影 (Noise Injection and Projection): 为增强跨模态表示对数据分布变化的鲁棒性,HDAN 在查询 (query) 和键值 (key-value) 投影中引入了基于扩散理论的时间步自适应噪声 (timestep adaptive noise)。 具体来说,我们将批次中的分子扰动嵌入 hpgh_p^g 拼接成一个候选集张量 HpgRN×DH_p^g \in \mathbb{R}^{N \times D},并将其投影为键 KmK_m 和值 VmV_m。同时,将全局文本嵌入 Ωht\boldsymbol{\Omega}_h^t 拼接成文本张量 HtRN×D\boldsymbol{H}^t \in \mathbb{R}^{N \times D},并将其投影为查询 QtQ_t。在投影过程中,根据时间步 kk 向查询和候选集引入不同级别的噪声: Qt=WQ(Ht+Proj(Noisek)),Km=WK(HPg+Proj(Noisek)),Vm=WV(HPg+Proj(Noisek)), \begin{array} { r } { Q _ { t } = W _ { Q } \big ( \boldsymbol { H } ^ { t } + \mathrm { P r o j } ( \mathrm { N o i s e } _ { k } ) \big ) , } \\ { K _ { m } = W _ { K } \big ( \boldsymbol { H } _ { \mathcal { P } } ^ { g } + \mathrm { P r o j } ( \mathrm { N o i s e } _ { k } ) \big ) , } \\ { V _ { m } = W _ { V } \big ( \boldsymbol { H } _ { \mathcal { P } } ^ { g } + \mathrm { P r o j } ( \mathrm { N o i s e } _ { k } ) \big ) , } \end{array}

    • 符号解释:
      • Qt,Km,VmQ_t, K_m, V_m:文本查询、分子键和分子值张量。
      • WQ,WK,WVW_Q, W_K, W_V:对应的可学习投影矩阵。
      • Ht\boldsymbol{H}^t:批次中的全局文本嵌入张量。
      • HPg\boldsymbol{H}_{\mathcal{P}}^g:批次中的分子扰动嵌入张量。
      • Proj(Noisek)\mathrm{Proj}(\mathrm{Noise}_k):一个投影函数,将时间步 kk 处的噪声映射到嵌入维度 DD。这种噪声模拟了扩散模型的核心原理,使模型在不同时间步关注不同复杂度的对齐。
  3. 全局语义对齐分布 (Eglobal\mathcal{E}^{\mathrm{global}}): 为了建模文本-分子对齐的数据分布,HDAN 采用注意力机制 (attention mechanism),并引入了时间步 kk 的数据分布 XkX_k 到注意力权重中,以捕获扩散模型的逆向去噪过程。 Eglobal=(Softmax(QtKmT)+Xk)Vm+DWC(Vm) \mathscr { E } ^ { \mathrm { g l o b a l } } = \left( \mathrm { S o f t m a x } ( Q _ { t } K _ { m } ^ { T } ) + X _ { k } \right) \cdot V _ { m } + \mathrm { D W C } ( V _ { m } )

    • 符号解释:
      • Eglobal\mathcal{E}^{\mathrm{global}}:全局语义对齐分布,代表文本-分子在句子/全局层面的关联。
      • Softmax(QtKmT)\mathrm{Softmax}(Q_t K_m^T):标准的注意力权重计算部分,衡量查询和键的相似性。
      • XkX_k:在时间步 kk 的数据分布,反映了前一噪声级别中每个候选分子的联合概率,高值表示高置信度。将其加到注意力权重中,动态调整模型对高置信度候选的响应。
      • VmV_m:分子值张量。
      • DWC(Vm)\mathrm{DWC}(V_m):深度可分离卷积模块 (Depthwise Convolution Module),用于保留特征多样性。
  4. 局部细粒度对齐分布 (Elocal\mathcal{E}^{\mathrm{local}}): 为了进一步建模细粒度的局部对齐,HDAN 计算文本词元与分子基元之间的相似度矩阵 S=[sij]Nw×NmS = [s_{ij}]^{N_w \times N_m}sij = (hiw)hjmhiwhjm s _ { i j } ~ = ~ \frac { ( \pmb { h } _ { i } ^ { w } ) ^ { \top } { \pmb { h } } _ { j } ^ { m } } { \| \pmb { h } _ { i } ^ { w } \| \| \pmb { h } _ { j } ^ { m } \| }

    • 符号解释:
      • sijs_{ij}:第 ii 个文本词元 hiw\pmb{h}_i^w 与第 jj 个分子基元 hjm\pmb{h}_j^m 之间的余弦相似度。 然后,通过选择每个词元的最大基元对齐分数 maxjsij\max_j s_{ij},并基于词元的重要性权重 {fiw}i=1Nw\{f_i^w\}_{i=1}^{N_w} 进行聚合,得到细粒度局部对齐分布: Elocal=i=1Nwfiwmaxjsij,fiwi=1Nw=Softmax(MLP({hiw}i=1Nw)), \begin{array} { r } { \mathcal { E } ^ { \mathrm { l o c a l } } = \displaystyle \sum _ { i = 1 } ^ { N _ { w } } f _ { i } ^ { w } \operatorname* { m a x } _ { j } s _ { i j } , } \\ { \displaystyle \left. \left. f _ { i } ^ { w } \right. _ { i = 1 } ^ { N _ { w } } = \mathrm { S o f t m a x } \big ( \mathrm { M L P } ( \{ h _ { i } ^ { w } \} _ { i = 1 } ^ { N _ { w } } ) \big ) , \right. } \end{array}
    • 符号解释:
      • Elocal\mathcal{E}^{\mathrm{local}}:局部细粒度对齐分布。
      • fiwf_i^w:第 ii 个词元的重要性权重,通过一个多层感知机 (Multi-Layer Perceptron, MLP) 对词元嵌入进行处理后,再通过 Softmax 函数计算得到。
      • maxjsij\max_j s_{ij}:每个词元与所有分子基元的最大相似度。
  5. 联合数据分布 (Ejoint\mathcal{E}^{\mathrm{joint}}) 与去噪解码器 (Denoising Decoder): 通过加权融合全局和局部对齐分布,得到最终的联合数据分布 Ejoint\mathcal{E}^{\mathrm{joint}}Ejoint=γElocal+(1γ)Eglobal \mathcal { E } ^ { \mathrm { j o i n t } } = \gamma \mathcal { E } ^ { \mathrm { l o c a l } } + ( 1 - \gamma ) \mathcal { E } ^ { \mathrm { g l o b a l } }

    • 符号解释:
      • γ\gamma:平衡系数 (balance coefficient),用于调节局部和全局对齐的贡献。 Ejoint\mathcal{E}^{\mathrm{joint}} 作为输入送入去噪解码器 (denoising decoder)。该解码器由一个多层感知机 (MLP) 构成,包含一个线性层进行特征编码(通过 ReLU [12] 激活),然后是另一个线性层计算输出分布。经过多轮去噪迭代,最终得到一个干净的对齐分布 X^0\hat{X}_0,并基于此进行检索排名。
  6. 反向去噪过程与生成式损失 (LGen\mathcal{L}_{\mathrm{Gen}}): 在反向去噪过程中,数据分布通过马尔可夫链逐步重建: x^t=αtX^t1+1αtϵ \hat { x } _ { t } = \sqrt { \alpha _ { t } } \hat { X } _ { t - 1 } + \sqrt { 1 - \alpha _ { t } } \epsilon

    • 符号解释:
      • x^t\hat{x}_t:在时间步 tt 处的噪声数据。
      • X^t1\hat{X}_{t-1}:在时间步 t-1 处去噪后的数据(即模型预测的更干净的数据)。
      • αt\alpha_t:噪声调度系数 (noise scheduling coefficient),控制在每个时间步去噪的程度。
      • ϵ\epsilon:代表需要优化的噪声预测网络 ϵθ\epsilon_\theta 预测的噪声。 生成式损失函数旨在优化噪声预测和数据分布建模,其定义为: LGen=E(t,m)Z[KL(X0ψ(m,t,Xk))]+E(m,t)Z[KL(X0ψ(t,m,Xk))] \mathcal { L } _ { \mathrm { Gen } } = \mathbb { E } _ { ( t , m ) \in Z } \Big [ \mathrm { K L } \big ( X _ { 0 } \| \psi ( m , t , X _ { k } ) \big ) \Big ] + \mathbb { E } _ { ( m , t ) \in Z } \Big [ \mathrm { K L } \big ( X _ { 0 } \| \psi ( t , m , X _ { k } ) \big ) \Big ]
    • 符号解释:
      • E(t,m)Z\mathbb{E}_{(t,m) \in \mathcal{Z}}:对语料库 Z\mathcal{Z} 中所有文本-分子对 (t,m) 的期望。
      • KL()\mathrm{KL}(\cdot \| \cdot):Kullback-Leibler (KL) 散度,衡量真实数据分布 X0X_0 与扩散网络 ψ\psi 在给定当前噪声数据 XkX_k 和输入 (m,t)(t,m) 时预测的数据分布之间的差异。 该损失函数促使模型学习文本-分子对的联合分布,从而更好地捕捉潜在的语义关联。
  7. 总训练损失 (Ltotal\mathcal{L}_{\mathrm{total}}): 最终的训练损失是生成式损失和判别式损失的总和: Ltotal=LGen+LDis \mathcal { L } _ { \mathrm { total } } = \mathcal { L } _ { \mathrm { Gen } } + \mathcal { L } _ { \mathrm { Dis } } 这种联合优化策略使得 DiffTMR 能够利用判别式学习的精确对齐能力和生成式扩散建模的鲁棒分布捕获能力,实现双向“表示-分布”增强,从而显著提高跨模态检索的准确性和泛化能力。

5. 实验设置

5.1. 数据集

实验在两个基准数据集上进行:

  1. ChEBI-20 [11]:

    • 来源: ChEBI (Chemical Entities of Biological Interest) 数据库的一个子集。
    • 规模: 包含 33,010 个分子及其对应的文本描述。
    • 划分: 数据集被划分为训练集、验证集和测试集,比例分别为 8:1:1。
    • 特点: 在推理阶段,测试样本从整个数据集中检索,通常包含丰富的生物活性小分子信息。
    • 样本示例: 文本描述如“一种具有镇静和催眠作用的哌啶衍生物,用于治疗失眠。”对应的分子结构(SMILES)可能为 CCN(CC)CCC1CCN(C1)CC(=O)c2ccccc2CCN(CC)CCC1CCN(C1)CC(=O)c2ccccc2
  2. PCdes [40]:

    • 来源: PubChem [17] 数据库。

    • 规模: 包含 15,000 对分子和文本描述。

    • 划分: 数据集被划分为训练集、验证集和测试集,比例分别为 7:1:2。

    • 特点: PubChem 是一个大型的化学信息数据库,PCdes 数据集包含了广泛的化学实体和描述,适用于评估模型在更通用化学领域的性能。

    • 样本示例: 文本描述如“一种有效的酪氨酸激酶抑制剂,用于癌症治疗。”对应的分子结构(SMILES)可能为 COc1cc(NC(=O)c2ccc(F)cc2)c(C)n1CCOc1cc(NC(=O)c2ccc(F)cc2)c(C)n1C

      选择这两个数据集是因为它们是文本-分子检索领域的标准基准,能够有效地验证模型在不同规模和化学领域下的性能。

5.2. 评估指标

论文使用了以下四个标准评估指标来衡量检索性能:

  1. Hits@k (Recall@k)

    • 概念定义 (Conceptual Definition): Hits@k(也称为 Recall@k)衡量的是在检索模型返回的前 kk 个结果中,至少包含一个相关项的查询所占的比例。它关注的是模型能否在有限的检索列表中找到正确的匹配项,值越高越好。
    • 数学公式 (Mathematical Formula): Hits@k=q=1NqI(rankqk)Nq \mathrm{Hits}@k = \frac{\sum_{q=1}^{N_q} \mathbb{I}(\text{rank}_q \le k)}{N_q}
    • 符号解释 (Symbol Explanation):
      • NqN_q:表示查询的总数量。
      • rankq\text{rank}_q:表示第 qq 个查询的正确答案在检索结果中的排名。
      • I()\mathbb{I}(\cdot):是指示函数,如果括号内的条件为真,则其值为 1;否则为 0。
  2. 平均倒数排名 (Mean Reciprocal Rank - MRR)

    • 概念定义 (Conceptual Definition): MRR 衡量的是查询的第一个正确答案排名倒数的平均值。如果第一个正确答案排名非常靠前(例如排名为 1),那么它的倒数就是 1;如果排名靠后,倒数会很小。因此,MRR 值越高,表示模型在找到第一个正确答案方面的平均表现越好。
    • 数学公式 (Mathematical Formula): MRR=1Qi=1Q1ranki \mathrm{MRR} = \frac{1}{|Q|} \sum_{i=1}^{|Q|} \frac{1}{\text{rank}_i}
    • 符号解释 (Symbol Explanation):
      • Q|Q|:表示查询的总数量。
      • ranki\text{rank}_i:表示第 ii 个查询的第一个相关结果在检索列表中的排名。
  3. 平均排名 (Mean Rank - MR)

    • 概念定义 (Conceptual Definition): MR 衡量的是所有查询中正确答案排名的平均值。该指标直接反映了正确答案在检索列表中平均出现的位置。MR 值越低,表示模型性能越好,因为正确答案的平均排名更靠前。
    • 数学公式 (Mathematical Formula): MR=1Qi=1Qranki \mathrm{MR} = \frac{1}{|Q|} \sum_{i=1}^{|Q|} \text{rank}_i
    • 符号解释 (Symbol Explanation):
      • Q|Q|:表示查询的总数量。
      • ranki\text{rank}_i:表示第 ii 个查询的正确答案在检索列表中的排名。

5.3. 对比基线

论文将 DiffTMR 与两类基线模型进行了比较:

  1. 任务特定模型 (Task-Specific Models):

    • MLP-Ensemble [11], GCN-Ensemble [11], All-Ensemble [11]: 这是早期针对文本-分子检索任务提出的集成方法,分别使用多层感知机 (MLP) 和图卷积网络 (GCN) 来学习表示。
    • MLP+Atten [11], MLP+FPG [14]: 同样是早期基于 MLP 的方法,加入了注意力机制 (Attention) 或指纹图 (FPG) 来增强表示。
    • AMAN [43]: 对抗模态对齐网络 (Adversarial Modality Alignment Network),利用对抗学习来增强跨模态对齐。
    • Memory Bank [32]: 通过记忆库来存储和利用更多的负样本,以提升对比学习的效果。
    • ORMA [29]: 基于最优传输 (optimal transport) 的多粒度对齐方法,在不同粒度级别(原子、基元、分子)进行对齐。
    • CLASS (ORMA) [38]: 在 ORMA 基础上增强了性能和训练效率。
  2. 大规模预训练多模态模型 (Large-Scale Pretrained Multimodal Models):

    • SciBERT [1]: 一个在科学文本语料库上预训练的语言模型,用于编码文本。此处作为基线可能表示仅用其提取文本特征后进行简单匹配。

    • KV-PLM [40]: 基于知识增强的语言模型,旨在通过知识图谱等外部知识来提升分子和文本的对齐。

    • MoMu [33]: 分子多模态基础模型,将分子图与自然语言关联起来。

    • MolFM [28]: 一个多模态分子基础模型,可能利用了更广泛的多模态信息。

    • MolCA [24]: 分子图语言建模,通过跨模态投影器和单模态适配器进行对齐。

    • MoleculeSTM [23]: 用于文本检索和编辑的多模态分子结构-文本模型。

    • Atomas-base [41], Atomas-large [41]: 基于层级自适应对齐的模型,通常在不同规模下进行预训练。

      这些基线模型涵盖了从早期方法到近年来最先进的预训练模型和特定任务优化方法,具有较强的代表性,能够全面评估 DiffTMR 的性能。

5.4. 实现细节

  • 文本编码器: 采用预训练的 SciBERT [1] 模型,最大序列长度设为 256。
  • 分子图编码器: 采用一个三层图卷积网络 (GCN) [37],输出维度为 300。
  • 训练策略: 采用判别式-生成式联合优化策略。
  • 学习率:
    • SciBERT [1] 的学习率为 3e-5。
    • 语义耦合层 (semantic coupling layer) 的学习率为 1e-5。
    • 其他组件的学习率为 1e-4。
  • 扩散去噪网络: 使用 Adam 优化器 [18],基础学习率为 1e-3,并采用余弦退火调度器 (cosine annealing scheduler) [25] 调整学习率。
  • 扩散过程步数: 固定为 50 步。
  • 平衡系数 (γ\gamma): 全局和局部对齐的平衡系数 γ\gamma 设为 0.4。
  • 训练周期: 模型训练 60 个 epoch。
  • 批次大小: 批次大小 (batch size) 设为 32。
  • 硬件: 所有实验均在 NVIDIA A40 GPU 上进行。
  • 候选集: 假设候选集是预定义的。

6. 实验结果与分析

6.1. 核心结果分析

论文在 ChEBI-20 和 PCdes 这两个基准数据集上对 DiffTMR 的检索性能和泛化能力进行了系统评估。

6.1.1. ChEBI-20 数据集性能

以下是原文 Table 1 的结果:

Models Text-Molecule Retrieval Molecule-Text Retrieval
Hits@1(↑) Hits@10(↑) MRR(↑) MR(↓) Hits@1(↑) Hits@10(↑) MRR(↑) MR(↓)
MLP-Ensemble [11] 29.4% 77.6% 0.452 20.78
GCN-Ensemble [11] 29.4% 77.1% 0.447 28.77
All-Ensemble [11] 34.4% 81.1% 0.499 20.21 25.2% 74.1% 0.408 21.77
MLP+Atten [11] 22.8% 68.7% 0.375 30.37
MLP+FPG [14] 22.6% 68.6% 0.374 30.37
AMAN [43] 49.4% 92.1% 0.647 16.01 46.6% 91.6% 0.625 16.50
Atomas-base [41] 50.1% 92.1% 0.653 14.49 45.6% 90.3% 0.614 15.12
Memory Bank [32] 56.5% 94.1% 0.702 12.66 52.3% 93.3% 0.673 12.29
ORMA [29] 66.4% 93.7% 0.775 18.63 61.2% 92.8% 0.738 10.21
CLASS (ORMA) [38] 67.4% 93.4% 0.774 17.82 62.0% 92.7% 0.738 14.59
DiffTMR (ours) 72.8% 96.5% 0.823 16.24 66.7% 96.3% 0.784 10.07

在 ChEBI-20 数据集上,DiffTMR 在文本到分子检索任务中取得了 72.8%Hits@1,相比此前最先进的基线模型 CLASS (ORMA)67.4% 提升了 5.4%。在分子到文本检索任务中,DiffTMR 达到了 66.7%Hits@1,也比 ORMA 提升了 4.7%。此外,DiffTMR 在 Hits@10 上也表现出色,达到 96.5%,这表明其在处理模糊查询时具有强大的语义建模和理解能力。MRRMR 指标也显示 DiffTMR 取得了最佳或接近最佳的性能。

6.1.2. PCdes 数据集性能

以下是原文 Table 2 的结果:

Models Text-Molecule Retrieval Molecule-Text Retrieval
Recall@1(↑) Recall@5(↑) Recall@10(↑) MRR(↑) Recall@1(↑) Recall@5(↑) Recall@10(↑) MRR(↑)
Pretrained Model + Finetuning
SciBERT [1] 16.3% 33.9% 42.6% 0.250 15.0% 34.1% 41.7% 0.239
KV-PLM [40] 20.6% 37.9% 45.7% 0.292 19.3% 37.3% 45.3% 0.281
MoMu [33] 24.5% 45.4% 53.8% 0.343 24.9% 44.9% 54.3% 0.345
MolFM [28] 29.8% 50.5% 58.6% 0.396 29.4% 50.3% 58.5% 0.393
Pretrained Model + Zero-shot
MolCA [24] 35.1% 62.1% 69.8% 0.473 38.0% 66.8% 74.5% 0.508
MoleculeSTM [23] 35.8% 39.5%
Atomas-base [41] 39.1% 59.7% 66.6% 0.473 37.9% 59.2% 65.6% 0.478
Atomas-large [41] 49.1% 68.3% 73.2% 0.578 46.2% 66.0% 72.3% 0.555
From-scratch
ORMA [29] 64.8% 82.3% 86.3% 0.727 62.1% 81.4% 86.3% 0.710
DiffTMR (ours) 69.0% 85.6% 89.2% 0.755 66.8% 84.3% 88.7% 0.738

在 PCdes 数据集上,DiffTMR 在文本到分子检索任务中获得了 69.0%Recall@1,比 ORMA 提升了 4.2%,并且显著优于其他大规模预训练模型。在分子到文本检索任务中,模型达到了 66.8%Recall@1,比 ORMA 提升了 4.7%。尤其是在长尾检索场景(即稀有或不常见分子类别)中,DiffTMR 依然保持了强劲性能,Recall@10 达到 89.2%,明显优于 AtomasKV-PLMMoMu 等模型,展现了卓越的跨域泛化能力和鲁棒性。

综合分析: DiffTMR 在两个数据集上的卓越性能主要归因于其提出的判别式-生成式协同优化框架。该框架利用对比学习来构建精确的跨模态对齐,同时通过扩散模型隐式捕获潜在空间中复杂的联合数据分布,从而增强了表示能力和泛化能力。这种动态联合优化有机地整合了判别式方法的语义对齐精度与生成式模型的泛化能力,为跨模态分子检索提供了一个强大而有效的解决方案。

6.2. 消融实验/参数分析

论文通过消融实验分析了 TAPR 模块中扰动范围计算方法、扰动采样频率以及扩散步数对模型性能的影响。以下是原文 Table 3 的结果。

(a) Calculation Method of Perturbation Range
Perturbation P Text-to-Video
Hits@1↑ Hits@ 10↑ MRR↑ MR↓
exp( ∑cic_i / NwN_w ) 63.4 90.1 0.792 17.3
exp( θ\thetacic_i / NwN_w ) 66.1 93.2 0.805 16.7
exp(CΘ\Theta) 72.8 96.5 0.823 16.24
(b) Effect of Perturbation Sampling Frequency.
Frequency (F) Hits@1↑ Hits@10↑ MRR↑ MR↓
5 63.5 90.2 0.795 17.9
10 68.4 93.6 0.811 17.1
15 72.8 96.5 0.823 16.24
20 71.5 95.8 0.820 16.8
(c) Effect of Diffusion Steps on Hits@1.
Train Eval
10 50 100 500
10 71.3 × X X
50 71.7 72.8 × ×
100 70.3 71.2 70.5 ×
1000 70.0 70.7 70.9 71.0

Table 3: Ablation Experiment on Text-Molecule Retrieval using the CHEBI-20 Dataset. "×" indicates that evaluation is not possible. The default setting is highlighted in blue

6.2.1. 扰动范围计算方法 (Perturbation Range Calculation)

原文 Table 3a 比较了三种扰动范围计算策略:

  • 固定参数 (exp(1Nwci)\exp(\frac{1}{N_w} \sum c_i)): 静态计算平均语义耦合强度,没有可训练参数。表现最差,Hits@163.4%
  • 可伸缩参数 (exp(θNwci)\exp(\frac{\theta}{N_w} \sum c_i)): 引入一个可训练的标量参数 θ\theta 进行动态调整。性能有所提升,Hits@1 达到 66.1%,表明扰动机制的灵活性和有效性。
  • 矩阵变换 (exp(CΘ)\exp(C\Theta)): 使用一个可学习的线性层 Θ\Theta 对语义耦合张量 CC 进行矩阵变换,从而建模更复杂的高阶交互。这种方法表现最佳,Hits@1 达到 72.8%。这说明通过矩阵变换实现更强的非线性建模能力,能更精准地根据文本语义动态校准扰动范围,从而显著提升性能。默认采用矩阵变换策略。

6.2.2. 扰动采样频率 (Perturbation Sampling Frequency)

原文 Table 3b 分析了推理阶段分子扰动采样频率 FF 的影响。

  • 当采样频率 FF 从 5 增加到 15 时,模型性能持续提升。这是因为在推理时,模型进行多次随机采样 HpgH_p^g,并选择与文本锚点最接近的扰动嵌入。增加采样频率使得模型能够更充分地探索分子语义空间,从而找到最佳的跨模态对齐。
  • F=15F=15 时,模型达到最佳性能 (Hits@172.8%)。
  • FF 增加到 20 时,性能略有下降 (Hits@171.5%),这可能是因为过高的采样频率引入了额外的计算成本,且边际效益递减,甚至可能引入更多不必要的噪声或过拟合。 为了平衡检索准确性和计算成本,最终选择 F=15F=15 作为最优采样频率。

6.2.3. 扩散步数 (Diffusion Step)

原文 Table 3c 探讨了扩散步数对模型性能的影响。

  • 在文本到分子跨模态检索任务中,仅需 50 个扩散步数即可达到最佳性能 (Hits@172.8%)。
  • 这与图像生成任务中通常需要 1,000 个甚至更多步数(如 DDPM [15]、Dhariwal 和 Nichol [8])形成了鲜明对比。
  • 这种差异可能源于不同数据分布的复杂性:图像生成需要建模复杂的高维像素分布,而文本-分子对齐任务更侧重于语义关联。相对不那么复杂的表示空间使得较少的扩散步数就能捕获关键的跨模态关系。
  • 过多的迭代步数反而可能引入额外的噪声,降低对齐质量。

6.3. 域外检索 (Out-of-domain Retrieval)

为了评估模型在未见数据上的泛化能力,论文采用了“域内训练,域外测试”的范式。模型首先在一个源数据集上进行训练(例如,在 ChEBI-20 上训练的模型),然后在另一个未见的(目标)数据集上进行评估(例如,在 PCdes 上测试)。

  • 判别式方法的局限性: 实验结果(具体数据在补充材料中,但论文摘要和结论中均提及)表明,传统的判别式方法在从源域迁移到未见的目标域时,泛化能力有限。例如,虽然 ORMA [29] 在域内检索中显著优于 AMAN [43],但两者在域外场景下表现均不佳。这验证了判别式模型在处理数据分布偏移时存在的固有挑战。

  • DiffTMR 的优越性: 相比之下,DiffTMR 在域内和域外检索任务中都持续展现出强大的性能。为了进一步说明生成式建模的优势,论文可视化了模型在两种场景下的相似性分布,如下图(原文 Figure 5)所示。

    该图像是三维散点图,显示了在ChEBI-20和PCdes数据集上,DiffTMR、ORMA和AMAN模型的正负配对情况。正配对用蓝色表示,负配对用紫色表示,展示了不同模型在分子检索任务中的表现差异。 该图像是三维散点图,显示了在ChEBI-20和PCdes数据集上,DiffTMR、ORMA和AMAN模型的正负配对情况。正配对用蓝色表示,负配对用紫色表示,展示了不同模型在分子检索任务中的表现差异。

Figure 5: VisizasiiveNegaivain-Dn OuRevlsks - iicaeheilariissinhereraiHEB-taodeheHEB-e ; -0 PCdes" represents the generalization results on the unseen PCdes test et using the pre-trained CHEBI-20 dataset model.

从 Figure 5 可以看出,与基线模型相比,DiffTMR 在目标域中保持了正样本对和负样本对之间更清晰的分离。这表明 DiffTMR 在处理未见数据时具有更强的泛化能力和可迁移性,进一步证实了其联合分布建模和判别式-生成式协同优化的有效性。

7. 总结与思考

7.1. 结论总结

本文提出了 DiffTMR,一个新颖的基于生成式扩散的文本-分子跨模态检索框架。与传统的依赖确定性条件概率建模的判别式方法不同,DiffTMR 通过渐进式去噪过程重构联合跨模态分布。这一创新使得模型能够有效捕捉文本查询和分子结构之间固有的多对多语义对应关系。

DiffTMR 的关键贡献在于:

  1. 联合分布建模: 将检索任务重新定义为从噪声中生成文本和分子联合分布的过程,克服了传统判别式方法忽视数据固有统计特性的局限性。

  2. 层级扩散对齐网络 (HDAN): 结合了全局语义对齐(句子级)和局部亚结构匹配(词元-基元级),通过渐进去噪实现细粒度且鲁棒的跨模态对齐。

  3. 文本锚定扰动表示 (TAPR): 引入动态扰动嵌入机制,根据语义耦合动态调整分子表示的扰动范围,在增强多样性的同时避免语义漂移。

  4. 判别式-生成式协同优化: 通过联合优化对比学习(强化显式对齐)和生成式扩散建模(捕获隐式分布依赖),显著提升了模型的对齐准确性和泛化能力。

    实验结果表明,DiffTMR 在 ChEBI-20 和 PCdes 等基准数据集上显著超越了最先进的判别式基线,尤其在域外检索场景中展现出强大的泛化和迁移能力。

7.2. 局限性与未来工作

论文虽然未在专门章节中详细阐述自身局限性,但在背景和相关工作中暗示了以下挑战:

  • 复杂分子结构语义的建模: 分子结构具有高度复杂性和多样性,其语义建模仍然是一个挑战。尽管 DiffTMR 引入了分层对齐和动态扰动,但如何更深层次地捕获和理解分子结构的复杂语义仍有待探索。

  • 双向推理的固有挑战: 化学分析本身可能涉及复杂的双向推理过程,即从文本到分子和从分子到文本的更深层次的逻辑关联,而非简单的相似度匹配。DiffTMR 通过双向损失来处理,但其在推理链条上的深度仍可加强。

  • 候选集预定义: 论文提到“候选集是预定义的”,这意味着模型在实际应用中可能需要一个预先筛选好的分子库,这在某些开放式或大规模的真实场景中可能是一个限制。

  • 计算效率: 扩散模型虽然在性能上表现出色,但其迭代去噪过程通常计算成本较高,尤其是在推理阶段。尽管论文指出文本-分子对齐所需步数少于图像生成,但与轻量级判别式模型相比,仍可能存在效率权衡。

    基于上述潜在局限性,未来研究方向可能包括:

  • 更高效的扩散采样策略: 探索更快的采样方法或更轻量级的扩散模型架构,以提高推理效率。

  • 结合领域知识图谱或先验信息: 将更丰富的化学领域知识或专家先验信息融入扩散过程或扰动机制中,以进一步增强模型的语义理解和推理能力。

  • 零样本 (zero-shot) 或少样本 (few-shot) 域外检索: 进一步提升模型在数据稀缺或完全未见域上的泛化能力,使其在真实世界新药发现中更具实用性。

  • 可解释性增强: 提高扩散模型在跨模态对齐过程中的可解释性,帮助化学家更好地理解模型做出的检索决策。

7.3. 个人启发与批判

7.3.1. 个人启发

  1. 判别式与生成式模型融合的强大潜力: DiffTMR 最重要的启发在于,它成功地将判别式模型的精确对齐能力(通过对比学习)与生成式模型(扩散模型)的鲁棒分布建模能力相结合。这种“最佳结合”的范式在解决传统单一范式的局限性方面显示出巨大潜力,尤其是在跨模态检索这类既要求准确性又要求泛化能力的任务中。它促使我们思考,在其他机器学习领域,是否也能通过融合不同范式的优势来构建更强大的模型。
  2. 动态扰动在表示学习中的创新应用: “文本锚定扰动”机制是一个非常巧妙的创新点。它解决了表示多样性与语义一致性之间的经典矛盾,通过文本语义动态引导分子表示的探索,使得模型在不牺牲准确性的前提下获得更丰富的表达。这种思想可以推广到其他模态(如图像-文本、视频-文本)的表示学习中,作为一种自适应的数据增强或表示正则化手段。
  3. 扩散模型在非生成任务中的新视角: 传统上,我们更多地将扩散模型与高保真图像生成联系起来。DiffTMR 证明了扩散模型“粗粒度到细粒度”的去噪特性同样可以被有效利用于非生成性的对齐和检索任务。这为扩散模型在更多领域(如推荐系统、数据融合、知识图谱对齐)的应用开辟了新的研究方向。
  4. 跨模态语义粒度对齐的重要性: 论文通过全局-局部两个层级的对齐来捕捉跨模态语义,这再次强调了在复杂模态对齐中,仅依赖单一粒度(如全局嵌入)可能不足以捕捉所有相关信息。多粒度甚至层级化的对齐策略是提升模型理解能力的关键。

7.3.2. 批判

  1. Proj(Noisek)Proj(Noise_k)DWC(Vm)DWC(V_m) 的具体机制缺乏细节: 论文提到了 Proj(Noisek)Proj(Noise_k) 将时间步 kk 的噪声映射到嵌入维度,以及 DWC(Vm)DWC(V_m)(深度可分离卷积模块)用于保留特征多样性。然而,这些组件的具体实现细节、设计动机和它们如何精确地与扩散过程结合的内在机制在方法论中相对简略。对于初学者而言,这些“黑盒”操作可能需要更深入的解释或图示。
  2. 计算成本与实际部署考量: 尽管论文指出文本-分子对齐所需扩散步数较少(50步),但与纯粹的判别式模型(只需一次前向传播)相比,扩散模型的迭代去噪过程在推理时仍可能带来更高的计算开销。在药物发现等对效率要求极高的场景中,这种额外的计算成本是否可接受,以及如何进一步优化推理速度,是实际部署中需要考虑的问题。
  3. “首次应用”声明的严谨性: 论文声称“我们是第一个将基于生成式的模型应用于文本-分子检索任务的”。虽然在扩散模型方面可能是首次,但如果考虑到其他类型的生成式模型(如变分自编码器 VAEs 或生成对抗网络 GANs)可能在文本-分子匹配或相关任务中有过探索,这个“首次”声明可能需要更精确的限定。
  4. 对负样本采样的依赖性: 判别式损失函数 LDis\mathcal{L}_{\mathrm{Dis}} 依赖于批次内的负样本进行对比学习。负样本的质量和数量对对比学习的效果至关重要。虽然这是常见做法,但论文没有详细说明负样本的采样策略,这可能会影响其在不同场景下的表现和鲁棒性。
  5. 对平衡系数 γ\gamma 的敏感性: 联合数据分布 Ejoint\mathcal{E}^{\mathrm{joint}} 通过平衡系数 γ\gamma 融合局部和全局对齐。0.4 是一个经验值,这个参数的敏感性如何?在不同数据集或任务中是否需要调整?论文未提供详细分析。

相似论文推荐

基于向量语义检索推荐的相关论文。

暂时没有找到相似论文。