TAG-DLM：面向文本属性图学习的扩散语言模型

# 文本属性图学习的扩散语言模型
来源：https://arxiv.org/html/2606.31166
凌洁 Chen UIUC lingjie7@illinois\.edu &袁晨贝∗ UIUC bei4@illinois\.edu &浩博徐∗ UIUC haoboxu@illinois\.edu 燕君赵 UIUC yanjunzh@illinois\.edu &宇中陈 VISA yuzchen@visa\.edu &航航童 UIUC htong@illinois\.edu

###### 摘要

文本属性图（TAGs）中每个节点都带有自然语言描述，要求模型能够同时对文本和图拓扑进行推理。现有方法通常将这两种模态分开处理：图神经网络在浅层文本特征上操作，而 LLM 与图的混合方法则将语言模型主要用作文本编码器，并将结构学习委托给一个单独的图模块。我们提出 TAG-DLM，它将文本推理和图消息传递统一在一个掩码扩散语言模型（一种具有双向注意力和生成式解码的语言模型）中。对于每个图实例，TAG-DLM 将采样的局部邻域线性化为一个 token 序列，并通过一个拓扑注意力掩码注入图结构，该掩码实现了图上的消息传递。由于扩散语言模型既能解释文本也能生成文本，TAG-DLM 仅通过改变提示就能适应不同任务，支持节点分类、链接预测和跨数据集迁移，无需针对目标进行微调。实验表明，TAG-DLM 在所有三个 TAG 基准测试的两个任务上均优于图神经网络、图变换器和基于 LLM 的基线方法，相比最强基线性能提升高达 3.9 个百分点。

TAG-DLM：用于文本属性图学习的扩散语言模型

凌洁 Chen††感谢平等贡献。UIUC lingjie7@illinois\.edu 袁晨贝∗ UIUC bei4@illinois\.edu 浩博徐∗ UIUC haoboxu@illinois\.edu

燕君赵 UIUC yanjunzh@illinois\.edu 宇中陈 VISA yuzchen@visa\.edu 航航童 UIUC htong@illinois\.edu

## 1 引言

节点携带自然语言描述的图出现在各种现实世界系统中，例如每篇论文由其标题和摘要表示的引文网络、节点描述研究主题或文章的生物医学文献图，以及实体由自由形式段落描述的知识库（Yan et al., 2023 (https://arxiv.org/html/2606.31166#bib.bib23); Wang et al., 2025 (https://arxiv.org/html/2606.31166#bib.bib24)）。这种图称为文本属性图（TAG）：每个节点都附有一片文本，边表示节点间的关系。TAG 学习要求模型能够从节点文本和图拓扑两方面对这样的图进行预测（Wang et al., 2024 (https://arxiv.org/html/2606.31166#bib.bib25)）。我们研究两个标准任务：在节点分类（NC）中，模型为目标节点分配一个标签；在链接预测（LP）中，模型判断候选对之间是否存在边。在两种情况下，答案由提示指定，并且必须从两个互补的证据来源——局部邻域的文本内容及其连通性——推断出来。因此，解决这些任务需要一个能够同时理解自然语言和推理图结构的模型。

参见标题图 1：TAG-DLM 与现有文本属性图学习流程的比较。现有方法通常将语言理解和图推理放在不同的组件中，这限制了两者交互的方式。图神经网络（GNN），如 GCN（Kipf and Welling, 2017 (https://arxiv.org/html/2606.31166#bib.bib3)）和 GAT（Veličković et al., 2018 (https://arxiv.org/html/2606.31166#bib.bib5)）有效建模图结构，但将节点文本表示为浅层的词袋或截断的嵌入，丢弃了大部分语义内容。最近的大语言模型（LLM）混合方法（Chen et al., 2024 (https://arxiv.org/html/2606.31166#bib.bib12); Tang et al., 2024 (https://arxiv.org/html/2606.31166#bib.bib26)）将语言模型视为冻结的文本编码器：节点文本被独立嵌入，然后附加一个单独的 GNN 来在图上传播结果嵌入。这种流程将两种模态结构性地分开。语言模型从不以图拓扑为条件，而图模块则在这些预计算的文本嵌入上运行，而非原始节点文本。因此，LLM 的上下文理解无法与图级推理交互：每一方都只看到对方的退化视图。另一个限制是任务特异性：由于任务由一个固定的输出头读出，这些方法在任务改变时必须重新训练。

首先，MDLM 使用完全双向的自注意力，因此在单次前向传播中信息可以在所有方向上流动。这适用于图提示，其中邻居节点的线性化顺序是任意的，目标及其邻居应该交换证据而不考虑顺序，而不是沿着固定的从左到右方向。其次，MDLM 是生成式的：它们填充序列中的掩码跨度，因此答案空间可以由提示指定，而不是由任务特定的输出头指定。这两个特性引出了我们工作的核心见解：*图邻域关系可以直接编码为注意力掩码*。通过限制哪些 token 可以关注哪些，我们将图结构嵌入到去噪过程中，而无需单独的图模块或对主干进行任何修改。并且由于答案仅仅是一个填充的提示位置，相同的公式通过提示变化即可表达 NC 和 LP，无需任务特定的图头。

基于此，我们引入 TAG-DLM（Text-Attributed Graph Diffusion Language Model），一个将 TAG 学习视为掩码填充问题的框架。对于目标图实例（例如要分类的节点或要评分的候选边），TAG-DLM 将采样的局部上下文线性化为一个 token 序列，用提示定义的选项填充掩码的答案位置，并附加一个二进制拓扑注意力掩码。该掩码在局部图上下文上实施结构化注意力，因此通过 MDLM 主干的单次前向传播同时读取文本证据并沿图传播信息，无需单独的 GNN 组件。

我们的贡献总结如下：

- • **拓扑注意力掩码作为图算子。** 图结构纯粹通过每个样本的二进制注意力掩码注入到 MDLM 中，消除了对单独 GNN 编码器或主干任何架构修改的需要。
- • **用于 TAG 学习的 MDLM。** TAG-DLM 将掩码扩散范式适应于图结构数据，在单一模型中结合了双向语言理解和图拓扑。
- • **理论基础。** 我们证明，使用拓扑注意力掩码的单次去噪前向传播在掩码诱导的图上执行基于注意力的消息传递。
- • **统一的图学习评估。** 由于 TAG-DLM 是生成式的，节点分类、链接预测和跨数据集迁移可以通过提示定义的答案选项来表达，而不是任务特定的输出头。

## 2 预备知识

### 2.1 文本属性图

*文本属性图*（TAG）是一个三元组 G = (V, E, X)，其中 V 是节点集，E ⊆ V × V 是边集，X = {x_v}_{v∈V} 为每个节点 v 分配一片自然语言文本 x_v。我们考虑在此类图上由提示定义的图预测任务。在*节点分类*（NC）中，每个节点带有一个离散标签 y_v ∈ Y，取自一个有限的、可用语言表达的标签集 Y（例如“神经网络”、“cs.CV”），子集 V_train ⊆ V 的标签是观察到的，目标是对未见的节点预测 y_v。在*链接预测*（LP）中，模型接收一个候选对 (u, v) 并在二元答案集下预测该边是否属于目标图。两个任务都被视为从提示指定的候选集中选择一个答案。目标节点 v 在 k 跳内的自我图表示为 N_k(v) ⊆ V，包含 v 以及从 v 出发最多 k 步边可达的节点；对于成对任务，我们使用候选端点周围的采样局部上下文。

### 2.2 掩码扩散语言模型

TAG-DLM 使用*掩码扩散语言模型*（MDLM）作为其主干（Sahoo et al., 2024 (https://arxiv.org/html/2606.31166#bib.bib13); Nie et al., 2025 (https://arxiv.org/html/2606.31166#bib.bib14)）。令 V 表示模型词汇表，附加一个特殊的掩码 token [M]，并令 x = (x_1, ..., x_L) ∈ V^L 是一个干净的 token 序列。*前向*过程通过连续时间 t ∈ [0,1] 以概率 β(t) 独立地将每个 token 替换为 [M] 来破坏 x，其中 β 是一个单调调度，β(0)=0 且 β(1)=1。记 x̃_t 为时间 t 的序列，前向核按位置分解：

q(x̃_t | x) = ∏_{i=1}^L q(x̃_{t,i} | x_i).   (1)

一个去噪网络 f_θ 从被破坏的序列中使用*双向*自注意力在所有位置上预测每个干净 token。训练最大化一个 Rao-Blackwellized 证据下界（Liu et al., 2019a (https://arxiv.org/html/2606.31166#bib.bib27)），该下界在依赖于时间 t 的重加权 w(t) 下，简化为被破坏位置上的掩码交叉熵：

L_MDLM(θ) = E_{t, x, x̃_t} [ w(t) ∑_{i: x̃_{t,i}=[M]} -log f_θ(x_i | x̃_t, t) ].   (2)

在推理时，生成通过一个迭代去噪循环进行，从完全或部分掩码的序列开始，通过从 f_θ 采样逐步解掩码 token。关键在于，对于我们在第4节中的分析，f_θ 使用的注意力不是因果的，因此层内信息流的唯一结构限制是外部提供的注意力掩码。

### 2.3 基于注意力的消息传递

许多图层可以写为在节点邻域上使用学习注意力的消息传递。给定第 l 层的节点表示 h_u^l，基于注意力的消息传递层具有以下形式：

m_u^l = ∑_{w∈N(u)} α_{uw}^l φ_l(h_w^l),   (3)
α_{uw}^l = exp s_l(h_u^l, h_w^l) / ∑_{r∈N(u)} exp s_l(h_u^l, h_r^l),   (4)
h_u^{l+1} = U_l(h_u^l, m_u^l),   (5)

其中 s_l 是注意力分数，φ_l 转换邻居特征，U_l 更新节点状态。GAT（Veličković et al., 2018 (https://arxiv.org/html/2606.31166#bib.bib5)）是这个家族的一个标准实例：它使用加性注意力分数 s_l(h_u, h_w) = LeakyReLU(a^⊤[W h_u ∥ W h_w]) 和一个共享的线性值变换。第4节将展示拓扑掩码去噪实现了这种更广泛的基于注意力的消息传递形式，并且 GAT 可以作为一个特例恢复。

## 3 方法

### 3.1 总体框架

TAG-DLM 将图预测视为在包含目标实例及其采样局部上下文的线性化序列上的掩码标签填充。图结构通过拓扑注意力掩码注入到自注意力中，而不是通过单独的 GNN 编码器。现有 LLM 与图的混合方法将语言模型视为冻结的文本编码器并附加一个图模块，而 TAG-DLM 则使语言模型本身成为图算子：唯一的图感知输入是每个样本的二进制注意力掩码，并且唯一训练的参数是线性投影上的 LoRA 适配器。这种设计提供了两个有用的性质。首先，底层 MDLM 权重被重用，无需架构更改或额外的图模块。其次，掩码自注意力在掩码诱导的图上执行消息传递（第4节），因此通过 MDLM 的单次双向前向传播读取文本证据并沿图传播信息。如图2所示，TAG-DLM 包含两个核心组件：从采样的局部图上下文中构建提示（§3.2）以及一个限制 MDLM 内部信息流的拓扑注意力掩码（§3.3）。

参见标题图 2：TAG-DLM 概述。目标节点 v 在 k 跳内的自我图被线性化为一个 token 序列 S_v，答案位置被替换为 [M]。一个二进制拓扑注意力掩码 M_v 强制实行星形注意力模式：v 关注所有 token；每个邻居只关注自身和 v。两者都被输入带有 LoRA 调优的 LLaDA-8B，它从答案位置的 logits 预测由提示定义的答案。

### 3.2 提示构建

##### 子图线性化。

对于 TAG G = (V, E, X) 中的每个目标节点 v ∈ V，我们采样一个 k 跳内的自我图，N_k(v) = {v} ∪ {u_1, ..., u_{n_v}}，方法是从每跳中均匀选择最多 K 个邻居（默认 k=2, K=10）。线性化后的序列是

S_v = T(v) ∥ T(u_1) ∥ ⋯ ∥ T(u_{n_v}),   (6)

其中 T(⋅) 表示每个节点的模板，∥ 是序列拼接。目标节点始终占据第一个片段，以便其 token 跨度可以直接恢复。对于成对预测，我们在候选端点周围使用相同的构造，并将候选关系编码到提示中。

##### 多项选择提示。

我们对节点分类和链接预测均使用*多项选择数字*提示格式。答案位置是一个单独的 [M] token，每个候选答案由一个数字 token 表示。对于节点分类，数字选项索引数据集标签名称。对于链接预测，使用相同的评分规则，二元选项指示候选边是否存在。邻居模板可以包含方括号内的观测训练标签，而没有观测标签的邻居则去掉方括号前缀。我们在附录C中提供了一个具体的提示示例。

### 3.3 拓扑注意力掩码

TAG-DLM 的一个核心设计选择是仅通过注意力暴露图邻接关系。基础 MDLM 权重保持冻结，图结构以输入依赖的注意力约束方式进入，而不是通过单独的训练模块。