Masked Gradient-Based Causal Structure Learning

概
主要内容
- 最终的目标
代码

Ng I., Fang Z., Zhu S., Chen Z. and Wang J. Masked Gradient-Based Causal Structure Learning. arXiv preprint arXiv:1911.10500, 2019.

概

非线性, 自动地学习因果图.

主要内容

NOTEARS将有向无环图凝练成了易处理的条件, 本文将这种思想扩展至非线性的情况:

X_{i} = f_{i} (X_{p a (i)}) + ϵ_{i},

$X_i = f_i(X_{\mathrm{pa}(i)}) + \epsilon_i,$

其中 $X_i$ 是因果图的结点, $X_{\mathrm{pa}(i)}$ 是其父结点, $\epsilon$ 是无关的噪声.

上述等式等价于

X_{i} = f_{i} (A_{i} \circ X) + ϵ_{i},

$X_i = f_i(A_i \circ X) + \epsilon_i,$

$A_i$ 是邻接矩阵 $A=[A_1|A_2|\cdots|A_d] \in \{0, 1\}^{d\times d}$ 的第i列, $A_{ij}=1$ 表示结点 $X_i$ 直接作用于 $X_j$ .

所以本文的目标就可以转换为如何估计 $A$ (实际上有了 $A$ 也就知道了因果图了). $A$ 应当满足的条件:

$A$ 能够表示有向无环图;
$X_i$ 和 $f(A_i \circ X)$ 必须接近, 比如用常见的

‖ X_{i} - f (A_{i} \circ X) ‖_{2}^{2}

$\|X_i - f(A_i \circ X)\|_2^2$

来度量.

直接处理非常麻烦, 首先对上面的问题进行放松, 等价于

X_{i} = f_{i} (W_{i} \circ X) + ϵ_{i},

$X_i = f_i(W_i \circ X) + \epsilon_i,$

此时 $A = \mathcal{A}(W)$ , 即

W_{i j} \neq 0 \to A_{i j} = 1; W_{i j} = 0 \to A_{i j} = 0.

$W_{ij} \not = 0 \rightarrow A_{ij} = 1; W_{ij} = 0 \rightarrow A_{ij} = 0.$

本文更进一步, 令

W = g_{τ} (U), U \in R^{d \times d} .

$W = g_{\tau}(U), \quad U \in \mathbb{R}^{d \times d}.$

[g_{τ} (U)]_{i j} = σ ((u_{i j} + g) / τ) = \frac{1}{1 + \exp (- (u_{i j} + (g_{1} - g_{0})) / τ)},

$[g_{\tau}(U)]_{ij} = \sigma((u_{ij} + g) / \tau) = \frac{1}{1 + \exp(-(u_{ij}+ (g_1 - g_0)) / \tau)},$

其中

g = g_{1} - g_{0}, g_{i} \overset{i . i . d .}{\sim} G u m b e l (0, 1) .

$g = g_1 - g_0, \: g_i \mathop{\sim}\limits^{i.i.d.} \mathrm{Gumbel}(0, 1).$

注: Gumbel.

此类操作能保证 $g_{\tau}(U) \in (0, 1)^{d\times d}$ , 此时能够把 $[g_{\tau}(U)]_{ij}$ 看成是 $X_i$ , $X_j$ 的关系的紧密型的度量, 在这种情况下

[g_{τ} (U)]_{i j} \leq ω \Rightarrow A_{i j} = 0.

$[g_{\tau}(U)]_{ij} \le \omega \Rightarrow A_{ij} = 0.$

或许会问, 为什么不用sigmoid而用一个这么麻烦的东西, 原因是当 $\tau$ 足够小的时候(如本文取的0.2), $[g_{\tau}(U)]_{ij}$ 非常接近 $0$ 或者 $1$ , 而用sigmoid, 作者发现这些值都接近0, 不能很好的模拟有向无环图, 故采用了这个方案.

接下来, 只需要满足

E [t r (e^{g_{τ} (U)}) - d] = 0,

$\mathbb{E}[\mathrm{tr}(e^{g_{\tau}(U)}) - d] = 0,$

即可保证 $g_{\tau}(U)$ 能够代表有效无环图. 在实际中, 只需

E [t r (e^{g_{τ} (U)}) - d] \leq ξ .

$\mathbb{E}[\mathrm{tr}(e^{g_{\tau}(U)}) - d] \le \xi.$

注: 期望是关于 $g$ 的.

最终的目标

总结下来,

min_{U, θ} E_{g} [\frac{1}{2 n} \sum_{k = 1}^{n} L (x^{(k)}, f (g_{τ}, x^{(k)}; θ))] s . t . E_{g} [t r (e^{g_{τ} (U)}) - d] \leq ξ .

$\min_{U, \theta} \quad \mathbb{E}_g[\frac{1}{2n} \sum_{k=1}^n \mathcal{L}(x^{(k)}, f(g_{\tau}, x^{(k)}; \theta))] \\ \mathrm{s.t.} \quad \mathbb{E}_g[\mathrm{tr}(e^{g_{\tau}(U)}) - d] \le \xi.$

注: $\mathbb{E}$ 是关于 $g$ 的, $n$ 的观测数据的总数.

进一步地, 我们希望 $g_{\tau}$ 是稀疏的, 故加上正则化项:

min_{U, θ} E_{g} [\frac{1}{2 n} \sum_{k = 1}^{n} L (x^{(k)}, f (g_{τ}, x^{(k)}; θ)) + λ ‖ g_{τ} (U) ‖_{1}] s . t . E_{g} [t r (e^{g_{τ} (U)}) - d] \leq ξ .

$\min_{U, \theta} \quad \mathbb{E}_g[\frac{1}{2n} \sum_{k=1}^n \mathcal{L}(x^{(k)}, f(g_{\tau}, x^{(k)}; \theta)) + \lambda \|g_{\tau}(U)\|_1] \\ \mathrm{s.t.} \quad \mathbb{E}_g[\mathrm{tr}(e^{g_{\tau}(U)}) - d] \le \xi.$

利用augmented Lagrange multiplier, 可得

L_{p} (U, ϕ, α) = E_{g} [\frac{1}{2 n} \sum_{k = 1}^{n} L (x^{(k)}, f (g_{τ}, x^{(k)}; θ)) + λ ‖ g_{τ} (U) ‖_{1} + α h (U)] + \frac{ρ}{2} (E [h (U)])^{2},

$L_p(U, \phi, \alpha) = \mathbb{E}_g[\frac{1}{2n} \sum_{k=1}^n \mathcal{L}(x^{(k)}, f(g_{\tau}, x^{(k)}; \theta)) + \lambda \|g_{\tau}(U)\|_1 + \alpha h(U)] + \frac{\rho}{2} (\mathbb{E}[h(U)])^2,$

其中 $h(U):= \mathrm{tr}(e^{g_{\tau}(U)}) - d$ .

采用分布更新:

U^{t + 1}, θ^{t + 1} = \arg min_{U, θ} L_{ρ^{t}} (U, ϕ, α^{t}); α^{t + 1} = α^{t} + ρ^{t} E [h (U^{t + 1})]; ρ^{t + 1} = {\begin{cases} β ρ^{t}, & i f E [h (U^{t + 1})] \geq γ E [h (U^{t})], \\ ρ^{t}, & o t h e r w i s e . \end{cases}

$U^{t+1}, \theta^{t+1} = \arg \min_{U, \theta} L_{\rho^t}(U, \phi, \alpha^t); \\ \alpha^{t+1} = \alpha^t + \rho^t \mathbb{E}[h(U^{t+1})]; \\ \rho^{t+1} = \left \{ \begin{array}{ll} \beta \rho^t, & \mathrm{if} \: \mathbb{E}[h(U^{t+1})] \ge \gamma \mathbb{E}[h(U^t)], \\ \rho^t, & \mathrm{otherwise}. \end{array} \right .$

其中第一步使用Adam执行1000次迭代计算的.

文中还讨论了后处理的一些方法, 和 $A$ 是否唯一.