KGdLLM：用离散扩散模型在知识图谱上学习逻辑推理

KGdLLM是由Tieumi221E创建的实验性研究框架，旨在探索离散掩码扩散语言模型（MDM/LLaDA风格）在知识图谱上的知识获取和逻辑推理能力。本文将解析其解耦架构、训练管线、评估方法等核心内容，探讨扩散模型在结构化知识推理领域的潜力。

自回归vs扩散生成范式

传统自回归模型（如GPT、Llama）存在误差累积、缺乏全局视角等局限；扩散模型通过前向加噪（逐步掩码token）和反向去噪（恢复原始token）的方式，具备双向上下文、可迭代修正、并行解码潜力等优势。

MDM与LLaDA

KGdLLM参考了MDM（2021年Austin等人提出，伯努利采样掩码）和LLaDA（改进掩码策略与训练目标）两种离散扩散模型，其diffusion_core模块实现核心算法。

核心引擎（diffusion_core/）

包含model.py（双向Transformer架构）、masking.py（LLaDA风格前向加噪）、loss.py（掩码交叉熵+1/p重要性采样）、inference.py（块级并行解码+置信度重新掩码），独立可复用。

实验脚本（scripts/）

• 数据管线：prepare_kg_dataset.py转换三元组为训练文本格式；
• 训练脚本：train_mdm.py（双向掩码预训练）、train_sft.py（监督微调）；
• 评估分析：eval_all_checkpoints.py（多维度推理评估）、plot_results.py（可视化）、plot_summary.py（对比分析）。

预训练阶段（知识获取）

将知识图谱三元组转为文本序列，通过动态掩码比例的前向加噪，模型预测掩码位置token，计算掩码交叉熵损失并结合重要性采样，学习结构化知识。

监督微调阶段（逻辑推理）

使用指令格式对话数据（如推理问题），以纯文本生成目标训练，让模型学会运用知识进行逻辑推理。

反向关系推理

测试模型对关系方向性的理解，如从“A是B的父亲”推断“B是A的孩子”。

多跳推理

测试模型通过中间关系推断间接关系的能力，如从“A是B父亲”“B是C父亲”推断“A是C祖父”。

传递关系推理

测试模型对传递性的理解，如从“A大于B”“B大于C”推断“A大于C”。

技术亮点

• 块级并行解码：每步可同时预测多个token，理论上提升推理速度；
• 置信度重新掩码：对低置信度预测位置重新掩码修正，类似“反复思考”；
• 双向Transformer：编码时关注所有token，利于多方向上下文推理。

研究方向

探索混合架构（结合自回归与扩散优势）、扩展到更多推理任务等。

局限性

• 以合成数据集为主，真实大规模KG（如Wikidata）表现未验证；
• 模型规模有限，扩展性未知；
• 迭代去噪过程仍比自回归模型慢。

未来方向

在大规模KG上验证、探索混合架构、扩展到更多逻辑推理任务。

KGdLLM为扩散模型应用于结构化知识推理提供了清晰的实验平台，其双向上下文感知和迭代修正能力为该领域带来新可能。尽管处于实验阶段，但对扩散语言模型与知识图谱推理研究者具有重要参考价值。项目地址：https://github.com/Tieumi221E/kg-diffusion-lm