最近我在整理结构学习这块的资料时，又把 DAG-GNN 从头到尾认真看了一遍。很多人对它的第一印象就是“给结构学习套了个图神经网络的壳”，这个说法不算全错，但会让人忽略它真正的价值。DAG-GNN 解决的核心问题是：给定一批观测数据，如何推断出变量之间隐藏的有向无环依赖关系，也就是我们常说的 DAG 结构学习。它把这个问题放进变分自编码器的框架里，用图神经网络同时承担编码、解码和结构表达，最后输出一张带权重的邻接矩阵，作为候选因果图。
如果你做因果发现、贝叶斯网络结构学习，或者想从横截面数据里挖出特征之间的依赖方向，这篇工作都值得认真对待。本文我会从问题本身讲起，再把 DAG-GNN 的建模思路、核心公式、训练逻辑、复现细节和落地经验逐段拆开，尽量把那些论文里不会写的坑也一起说清楚。
1. 结构学习的卡点：组合搜索和不可微约束
1.1 我们要学的“图”到底是什么
假设手里有一批观测数据，总共 n 个样本，每个样本有 d 个变量，通常表示成一个矩阵 X。DAG 结构学习的任务，就是学一个加权邻接矩阵 A，其中 A[i][j] 非零，表示变量 j 对变量 i 存在直接依赖关系，而且所有这样的依赖关系放在一起不能形成有向环。
要注意，这里的“依赖”不一定等于严格意义上的因果，但它是因果发现的重要前置。我习惯把它叫“候选因果图”，因为它给出了变量之间可能的方向关系，后续是否要用“因果”这个词，还取决于你手里有没有干预数据、领域假设是否成立等问题。
这类问题在很多场景都会出现：基因调控网络构建、供应链故障根因定位、风控特征关系梳理、推荐系统里的用户意图路径建模。说白了，只要你想知道“是谁在影响谁”，而不是单纯地预测“谁是什么”，结构学习就有用武之地。
1.2 为什么这个任务长期以来很难做
DAG 结构学习难在哪里？我拆成两个维度。
第一个是组合爆炸。d 个节点上所有可能的 DAG 数量是超指数增长的，大概的规模是 O(2^(d(d-1)/2)) 量级。当 d 只有 20 的时候，这个数字已经大到天文级别，不可能靠枚举搜索。传统的贪婪搜索、爬山法、模拟退火，虽然有效，但本质上都是在一个巨大的离散空间里碰运气，很容易落入局部最优。
第二个是方向不可辨识。在纯观测数据下，很多统计方法只能告诉你两个变量相关，无法告诉你谁影响谁。举个最简单的例子，线性高斯模型里，X = aY + ε_X，Y = ε_Y，这里的 ε_X 和 ε_Y 都是高斯噪声，那么反向写成 Y = bX + ε_Y' 也能得到一模一样的联合分布。只靠观测数据，模型会把两个方向都当成可行方案，这时候就必须靠额外假设来打破对称性。
这两个问题叠加在一起，导致传统结构学习算法的链条特别脆弱：离散搜索不稳定，评分函数也不平滑，无法直接做梯度下降。所以 2018 年前后，最令人兴奋的方向就是“把离散约束变成连续优化约束”。
1.3 连续的 DAG 约束带来了范式转变
如果你接触过结构学习，大概率听过那个里程碑式的做法：在加权邻接矩阵 W 上定义一个平滑函数 h(W) = tr(exp(W ⊙ W)) − d。这个函数有个非常好的性质：h(W) = 0 当且仅当 W 对应的图是有向无环图。于是 DAG 约束就从一个离散的组合性质，变成了一个可以放进优化器里的可微约束。
这个约束的原理并不复杂，我可以拿“路径环”来类比。矩阵 (W ⊙ W) 的 k 次幂的迹，统计的是长度为 k 的闭环的加权总和；对幂级数求和之后取迹，等价于把所有可能的闭环都加在一起。如果 DAG 存在环，那沿着环走一圈终会回到起点，这种情况下闭包总和必然为正；反过来，如果没有环，闭合路径的权重总和就是 0。所以 h(W) 可以敏锐地嗅出图中的环。
但 NOTEARS 那一类方法的生成模型假设是 X ≈ W^T X + ε，属于线性结构方程模型。真实数据往往是非线性的，线性假设一旦不成立，学出来的结构就容易出现漏边、错边。DAG-GNN 的思路，就是保留这个连续的 DAG 惩罚，同时把生成模型换成一个表达能力更强的非线性框架：图神经网络。
2. DAG-GNN 的建模核心：用 GNN 把非线性结构方程装进 VAE
2.1 从线性 SEM 到非线性生成方程
DAG-GNN 把结构学习重新理解成了一个生成过程。它假设观测变量 x 并不是直接由另一个观测变量线性生成，而是有一个潜变量 z，在结构矩阵 A 的指引下，通过一个非线性函数映射到观测空间：
x = g(A^T z) + ε
这里 g 是逐元素作用的非线性激活函数，比如 tanh 或带平滑的线性单元；A 是 d×d 的加权邻接矩阵；ε 是噪声项。这个设定比线性 SEM 强在哪儿？关键就在于 g 可以描述非线性的数据生成机制。比如基因调控里，一个转录因子对下游基因的作用很多时候是饱和型的、开关型的，用线性系数很难刻画，非线性激活函数则有更大的拟合空间。
注意，这里 A 同时承担了结构表达和生成传播两个角色。它既是我们最终要学出来的 DAG，又是生成数据过程中的传播骨架。这个设计很有意思，它让你在重构数据的同时被迫把结构 A 学好，因为如果 A 错了，重构误差就会大，梯度会逼着优化器去修正 A。
2.2 编码器和解码器中的图神经网络角色
有了生成过程还不够，还要有对应的推理过程。DAG-GNN 使用了一个编码器，把观测 x 和当前的结构矩阵 A 共同映射到潜变量 z 的分布参数上：
z ~ q(z | x, A) = N(μ(x, A), diag(σ²(x, A)))
这里的编码器不是普通的 MLP，而是图神经网络。为什么必须用图结构？因为在潜变量空间里，变量之间不是独立的，节点状态需要沿着 A 的关系互相传播信息。你可以把每个变量想象成房间里的人，A 就是他们之间的通话线路；想知道某个节点当前大概处于什么状态，光看它自己是不够的，还要把和它相连的邻居信息汇总进来。
解码器那侧也是同样的道理。潜变量 z 沿着 A 的转置方向做消息传递，再经过非线性变换，才能重构出 x。每一步传播都对应结构方程中的一个依赖路径，这让模型可以表达多跳的、非线性的依赖关系，而不再是一次性乘一个系数矩阵就完事。
2.3 为什么要套一个 VAE 框架
如果只是把线性 SEM 换成 GNN，理论上也可以直接优化最小二乘重构损失。但实际做起来会发现两个问题：一是变量之间的随机噪声很复杂，单纯的点估计容易过拟合；二是潜变量 z 本身就带有不确定性，我们需要一个机制来合理处理这种不确定性。
VAE 框架恰好能同时解决这两个问题。它通过重构似然项来拟合数据，通过 KL 散度项约束潜变量分布不过度偏离先验。这相当于给模型加了一个“概率缓冲”：即使某些依赖关系在当前样本里体现得不明显，模型也知道不要过度相信某一个确定的 z 值，从而降低结构误判的概率。
从这个角度看，DAG-GNN 不只是一个换壳的 NOTEARS，它把结构学习从“最小二乘加惩罚”升级成了“变分推断加结构约束”。这个变化带来的直接好处是，在样本量有限、信噪比不高的场景下，结构学习的稳定性会明显变好。我后面在复现时也确认了这一点。
2.4 和传统连续优化方法的对比
我把 DAG-GNN 与纯线性连续优化路线放在一起比较一下。
| 维度 | 线性连续优化路线 | DAG-GNN |
|---|---|---|
| 基础模型 | 线性结构方程模型 | 非线性生成模型，带 GNN 传播 |
| 学习方式 | 似然最小二乘 | 变分自编码器，重构 + KL |
| 非线性能力 | 几乎无 | 强，可学习非线性激活和深层传播 |
| 对潜变量的处理 | 不显式建模潜变量 | 显式潜变量，带不确定性估计 |
| 高维稳定性 | 依赖稀疏正则和调参 | 依赖 VAE 噪声缓冲，相对更好 |
| 训练代价 | 较低 | 较高，需要迭代 E-step 和结构约束交替优化 |
如果你手里的数据基本是线性关系，那我诚心建议线性基线就够用了。但现实世界的数据很少这么听话，非线性、交互作用、隐藏变量到处都是，这种情况 DAG-GNN 的优势就出来了。
3. 从 ELBO 到 DAG 惩罚：训练过程中真正在算的东西
3.1 三个损失组件拼出一个目标函数
DAG-GNN 的训练目标可以拆成三块。
第一块是重构损失，也就是负对数似然项。对于连续数据，通常用高斯分布假设，等价于最小化 x 和 x_hat 之间的均方误差；如果是离散数据，可以换成交叉熵形式。这一项决定了模型能不能把观测数据解释好。
第二块是 KL 散度项。它度量编码器得到的潜变量分布与标准高斯先验之间的差异，对潜变量做正则化，防止模型在潜变量空间里走极端。这两块合起来就是变分自编码器里最常说的 ELBO。
第三块是 DAG 惩罚项。A 通过 h(A) 进入目标函数，用增广拉格朗日法处理等式约束 h(A) = 0。整体目标可以写成：
L = ELBO + λ h(A) + (ρ/2) h(A)²
其中 λ 是拉格朗日乘子，ρ 是惩罚系数。训练过程中这两个参数会不断更新，让模型先把数据重构好，再把结构逐步拉向 DAG。
3.2 增广拉格朗日法的调参节奏
直接上完整代码前，先把这个优化过程说透。它的工作方式有点像拧螺丝：你不能一次拧到底，要一次拧一圈，然后看效果再继续。
通常会这么设计：把训练分成若干轮。在每一轮里，先固定当前的 λ 和 ρ，用 Adam 优化器去最小化 L 中的变量部分（包括 A、编码器参数、解码器参数）；每优化一定步数之后，更新 λ 和 ρ，更新规则是 λ ← λ + ρ h(A)，如果 h(A) 还很大，就把 ρ 乘以一个放大因子，比如 10 倍。
我建议把惩罚系数放大节奏放慢一点，比如每 50 到 100 步再更新一次，而不是每个 batch 都调整。原因很简单：模型需要时间去适应新的惩罚力度，你更新得太频繁，重构损失和 DAG 惩罚会互相打架，训练曲线会一路发抖。
收敛条件也不是简单地把 h(A) 降到 0 就收工。我在实战中通常同时看三个指标：重构损失是否平稳、h(A) 是否降到 1e-8 量级、A 的非零边比例是否合理。如果重构损失还在明显波动，即使 h(A) 很小，说明模型还没真正稳定，建议继续多跑若干轮。
3.3 一个可以跑通的最小 PyTorch 骨架
这里我给一个极度简化的骨架代码，目的是帮助理解，不是直接复现论文里的完整结构。核心模块就三个：一个可学习的邻接矩阵参数 A、一个编码器 GNN、一个解码器 GNN。
import torch import torch.nn as nn import torch.nn.functional as F def dag_penalty(A): # A 是带权重邻接矩阵，先保证对角线为 0 A = A * (1 - torch.eye(A.size(0), device=A.device)) M = A * A # 用矩阵指数近似 h(A) = tr(exp(M)) - d return torch.trace(torch.linalg.matrix_exp(M)) - A.size(0) class DAGGNN(nn.Module): def __init__(self, dim, hidden_dim=64): super().__init__() self.A = nn.Parameter(torch.zeros(dim, dim)) self.enc1 = nn.Linear(dim + dim, hidden_dim) self.enc_mean = nn.Linear(hidden_dim, dim) self.enc_logvar = nn.Linear(hidden_dim, dim) self.dec1 = nn.Linear(dim + dim, hidden_dim) self.dec_out = nn.Linear(hidden_dim, dim) def reparameterize(self, mu, logvar): std = torch.exp(0.5 * logvar) eps = torch.randn_like(std) return mu + eps * std def forward(self, x): # 掩码对角线，避免节点自己影响自己 A = self.A * (1 - torch.eye(self.A.size(0), device=self.A.device)) # 编码器：把 x 和 A 拼起来做消息传递的初始输入 inp = torch.cat([x, x @ A], dim=-1) h = torch.tanh(self.enc1(inp)) mu = self.enc_mean(h) logvar = self.enc_logvar(h) z = self.reparameterize(mu, logvar) # 解码器：从潜变量和结构关系重构 x dec_inp = torch.cat([z, z @ A], dim=-1) h = torch.tanh(self.dec1(dec_inp)) x_hat = self.dec_out(h) return x_hat, mu, logvar, A这个骨架里，编码器和解码器都只用了一层线性层，实际复现时通常会再叠一层 GNN 层，并加入残差连接。注意我这里是直接把 x 和 x @ A 拼起来当作一种粗糙的消息传递，你也可以换成图卷积、图注意力等更复杂的算子。
训练循环里，要把 ELBO 和 DAG 惩罚分开算，再合到一起。
def elbo_loss(x, x_hat, mu, logvar): recon = F.mse_loss(x_hat, x, reduction='sum') kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) return recon + kl def train_step(model, x, lam, rho): x_hat, mu, logvar, A = model(x) loss = elbo_loss(x, x_hat, mu, logvar) + lam * dag_penalty(A) + (rho / 2) * dag_penalty(A) ** 2 optimizer.zero_grad() loss.backward() optimizer.step() return dag_penalty(A).item(), loss.item()3.4 判断收敛时最容易踩的视觉陷阱
很多人会看 h(A) 降到 0 就认为成功了，结果发现学出来的图全是一片空白或者一团乱麻。为什么？因为 h(A) 只是无环约束，它保证无环，不保证边都正确。优化器完全可以把 A 里的边大多压到极小值，从而在满足无环约束的同时牺牲重构质量。
所以我的习惯是把 h(A)、重构损失、有效边数量三个量画在同一张图上，观察它们是否同时稳定。如果 h(A) 一直在降，但重构损失同时在上升，说明惩罚和重构的对抗已经失衡，你需要调小 λ 或者 ρ 的更新节奏。如果 A 的有效边比例突然骤降，往往意味着 DAG 惩罚过强，把边都给压死了。
4. 复现 DAG-GNN 最容易翻车的五个细节
4.1 数据标准化没有做够，损失直接失衡
这几乎是新手复现 DAG-GNN 的第一道坎。X 的各个变量如果量纲不同，有的在 0 到 1，有的在 0 到 10000，那么重构损失会被大量纲变量主导，模型会优先拟合这些变量，结构学习部分基本没法正常运营。
我建议在数据进入模型之前，先做标准化，让每个变量均值为 0、方差为 1。虽然看起来是很基础的一步，但很多人会觉得“反正模型能学非线性变换，标准化无所谓”，实际上这个无所谓会直接毁掉重构和 KL 之间的平衡。用标准化之后，重构损失各维度的贡献才相对均匀，A 的学习才不会被个别变量带偏。
4.2 把 h(A)=0 当成训练成功，会得到一张空图
前面说过这个问题，这里详细展开一下。增广拉格朗日法在后期会把 λ 变得非常大，这时候模型面临的惩罚极强。如果重构项和 DAG 项没有搭配好，模型会发现最简单省事的方法是把 A 的所有边都压到接近 0，这样 h(A) 自然满足，DAG 惩罚就没了，但结构什么都没学到。
对付这个问题，我最常用的办法是给 A 加一个小幅度的先验，比如初始化成带噪声的小值，并且在训练过程中对 A 的绝对值做 soft threshold 或 hard threshold。更直接的方法是盯住重构损失，一旦发现重构损失异常上升，立刻降低 ρ 的放大倍数，或者把 λ 重置一下，给重构项一点喘息空间。
4.3 ρ 和 λ 更新太激进，训练曲线全程发抖
我把常见症状和解决方法整理成一张表，方便你对照着排查。
| 症状 | 可能原因 | 处理办法 |
|---|---|---|
| A 的边不断翻转、训练后期不稳定 | ρ 放大倍数太大，惩罚更新太频繁 | 把 ρ 放大因子调小，比如从 10 改成 2，拉长更新间隔 |
| h(A) 一直波动，降不下去 | λ 更新幅度过大，优化器跟不上 | 对 λ 做 clip，限制单次最大增量 |
| 重构损失先降后升 | 增广拉格朗日惩罚过重，重建被压制 | 降低初期 λ，让模型先充分拟合数据再逐步加结构约束 |
| A 几乎所有边都趋近 0 | DAG 惩罚压过了重构 | 减小 ρ 的峰值上限，或引入最小边阈值 |
这组参数本质上没有标准答案，跟数据维度、样本量、噪声水平都有关系。我的经验是先跑一个小数据集，把 ρ 和 λ 的曲线打出来，观察它们与 h(A) 的变化关系，找到一个节奏，再放到更大数据集上。这个调参过程很繁琐，但非常值得做。
4.4 图神经网络层数过深，反而丢失结构信息
结构学习里的图规模通常都不大，多数实验在 10 到 50 个节点之间。这个规模下，GNN 层数不是越多越好。我试过堆到 5 层以上，结果所有节点的表征开始趋同，邻接矩阵 A 里的边也变得模模糊糊，这就是常说的过度平滑。
节点数量少时，两层到三层的图传播已经足够覆盖绝大多数依赖路径了。如果你想加深网络来提升非线性表达能力，更好的方式是加深每层内部的 MLP，而不是增加消息传递的跳数。另外，给每层 GNN 加上残差连接，也能明显缓解深层的表征退化问题。
4.5 评估时阈值选太硬，结果波动巨大
DAG-GNN 学出来的 A 是连续权重，不是二值图。评估时通常要把权重大于某个阈值的边视为存在，小于阈值的边视为不存在。这个阈值怎么选，直接决定最终 F1 和 SHD 的数值。
如果你设成 0.01，那么几乎所有的弱边都会被当成真边，F1 的精确率急剧下降；如果设成 0.5，很多真实的弱依赖会被丢掉，召回率断崖式下跌。正确的姿势是先看看 A 对训练集重构的影响，或者直接在验证集上扫一遍阈值，选一个让结构误差最小的点。我通常会把 0.05、0.1、0.2、0.3 都试一遍，然后报告一个范围，而不是拍脑袋定死一个数。
5. 怎么设计一套能说明问题的实验
5.1 先在自己生成的合成数据上考试
结构学习这件事，没有真实结构做对照，就谈不上好不好。所以在跑真实数据之前，一定要先在合成数据上验证。
我的做法是随机生成一个 DAG，节点数一般在 5 到 20，边的期望数量可以控制，比如让平均出度约等于 2。然后给非零边随机赋权重，通常在 [-2, -0.5] 和 [0.5, 2] 里均匀采样。接下来分两套数据：一套从线性 SEM 采样，另一套从非线性 SEM 采样，比如把 W^T X 先过一个 tanh 再加噪声。
评估指标我用三个：
| 指标 | 含义 | 如何解读 |
|---|---|---|
| SHD | 结构汉明距离，计算漏掉、多出、方向错误的边数总和 | 越小越好，0 表示完全正确 |
| F1 | 边预测的精确率和召回率的调和平均 | 越大越好，反映边存在性判断 |
| 方向准确率 | 在预测存在的边中，方向正确的比例 | 越大越好，反映方向识别能力 |
这里有个细节要注意：SHD 把你“预测了一条边但方向反了”也计为一次错误，而 F1 通常只看边存在性，方向不算。所以我会同时报告这两个指标，防止只报 F1 而掩盖方向错误的问题。
5.2 真实数据没有金标准时的打榜策略
真实数据通常没有完整的 DAG 作为参考答案，尤其像蛋白质表达数据、股票收益率数据这类场景，你只知道部分公认的依赖关系，不可能拿到全图。
这种情况下，我建议不要纠结全局指标，转而做稳定性分析。把原始数据用 bootstrap 重采样很多次，每次都跑一遍 DAG-GNN，统计每条边出现的频率。出现频率超过 80% 的边，可以看成稳定边；在 50% 到 80% 之间的，属于候选边；低于 50% 的，基本可以认为是噪声。这样得到的候选结构，比单次运行的结果可信得多。
另一个做法是把少量已知的依赖关系作为先验边固定住，让模型只去学未知部分，然后看学出来的边是否与你所在领域的常识冲突。冲突越多，说明模型可能出了问题，或者你的先验本身需要调整。
5.3 我惯用的一个实验套餐
给想快速上手的同学一个参考套餐。首先做数据生成，固定节点数 d=10，样本量取 500、2000 两档；设定低噪声和高噪声两档。每种组合生成 10 个随机 DAG，并分别在线性和非线性 SEM 下采样。然后把 DAG-GNN 和线性基线同时跑一遍，比较 SHD 和 F1。最后把结果画成箱线图，看看分布，而不是只看均值。
这个套餐的意义在于：它能让你快速感受 DAG-GNN 在不同条件下的适用边界。在我的经验里，线性数据上 DAG-GNN 和线性基线差距不大，有时候甚至会因为模型复杂而稍弱；但非线性数据上，DAG-GNN 的 SHD 往往能下降 20% 到 40%，这就证明了引入 GNN 非线性表达是有回报的。如果非线性数据上你的结果和线性基线完全一样，那多半是代码或者参数设置出了问题。
6. 从“学会一张图”到“用好一张图”的落地经验
6.1 先验知识与业务约束怎么塞进去
真实业务场景里，你不可能完全不知道结构。比如某个流程中，A 步骤一定发生在 B 步骤之前；或者某两个环节之间明确不允许有依赖关系。这些先验如果不利用，等于把模型架在一个完全没约束的环境里瞎猜。
DAG-GNN 里有很自然的接口做这件事：在邻接矩阵 A 上加 mask。已知存在的边可以直接固定住，不让梯度去改它；已知不存在的边直接置为 0，不参与训练；甚至可以限制某些节点的出度或入度为 0，比如已知某个变量只能是原因，不可能是结果。
我强烈建议在实际项目里至少塞一层先验 mask。数据量通常不够支持模型从零学出一张可信的全局图，但加上领域约束后，搜索空间小很多，学出来的结构也会稳很多。
6.2 把邻接矩阵转化为下游决策
结构学习本身不是终点，让人头疼的是后面怎么用它做决策。
如果目的是特征选择，那可以看 A 的可达性。比如你想知道某个目标变量受哪些上游变量的影响，可以顺着 A 的方向反向追溯，把所有能到达目标变量的节点都作为候选特征。这比单纯用相关系数筛选更能避免间接依赖造成的假候选。
如果目的是根因定位，那就找 A 中的“根节点”，也就是没有上游依赖、但能通过路径影响很多下游节点的变量。这类似于在一个故障因果图里找源头节点，对运维和客服场景非常有用。
如果是做反事实解释，那就用学到的图作为生成模型骨架，干预某个上游节点，然后看下游变量分布会发生什么变化。注意，这需要更强的因果假设，且最好结合领域验证，不能只靠模型输出的图就下结论。
6.3 我对新入行者的使用建议
我的经验是，别一上来就把 DAG-GNN 当成万能的因果发现工具。它更适合的场景是：变量数量几十个以内，样本量在数千到几万之间，数据基本上是连续或平滑变换的，而且你有一些领域先验可以用来筛选结果。如果节点数上千，或者数据里有大量的复杂离散类别，DAG-GNN 的计算代价和结果稳定性都不太理想，这时候更建议把问题拆小，或者改用其他专门面向高维稀疏场景的方法。
还要牢记一点：结构学习给出的是“结构候选项”，不是“终审判决”。我在实际项目里的使用流程通常是，先用 DAG-GNN 学一个关系矩阵，再结合 bootstrap 稳定边筛选，最后把少量关键边交给业务方评估。经过这几轮过滤之后，剩下的边才敢进入后续的建模或决策环节。
最后分享一个我的个人习惯：每次跑完 DAG-GNN，不要急着拿一次随机初始化的结果去下结论。我会至少换三到五个随机种子，统计边的出现频率和权重波动范围。结构学习本质上是一个缩小怀疑范围的过程，而不是替你拍板的工具。理解这一点，你在使用 DAG-GNN 时的很多困惑，都会自动解开。
