[Paper Review] Truncated Matrix Power Iteration for Differentiable DAG Learning

Zhen Zhang, Ignavier Ng+6 more|arXiv (Cornell University)|Aug 30, 2022

Remote-Sensing Image Classification4 citations

TL;DR

This paper proposes a novel differentiable DAG learning method using truncated matrix power iteration (TMPI) to approximate geometric series-based DAG constraints, enabling larger coefficients on higher-order polynomial terms without numerical instability. The method outperforms state-of-the-art approaches by up to 3× in structural Hamming distance (SHD), especially in sparse graphs, by mitigating gradient vanishing while maintaining computational efficiency.

ABSTRACT

Recovering underlying Directed Acyclic Graph (DAG) structures from observational data is highly challenging due to the combinatorial nature of the DAG-constrained optimization problem. Recently, DAG learning has been cast as a continuous optimization problem by characterizing the DAG constraint as a smooth equality one, generally based on polynomials over adjacency matrices. Existing methods place very small coefficients on high-order polynomial terms for stabilization, since they argue that large coefficients on the higher-order terms are harmful due to numeric exploding. On the contrary, we discover that large coefficients on higher-order terms are beneficial for DAG learning, when the spectral radiuses of the adjacency matrices are small, and that larger coefficients for higher-order terms can approximate the DAG constraints much better than the small counterparts. Based on this, we propose a novel DAG learning method with efficient truncated matrix power iteration to approximate geometric series based DAG constraints. Empirically, our DAG learning method outperforms the previous state-of-the-arts in various settings, often by a factor of $3$ or more in terms of structural Hamming distance.

Motivation & Objective

• To address the gradient vanishing problem in differentiable DAG learning caused by small coefficients on high-order polynomial terms.
• To demonstrate that larger coefficients on higher-order terms are safe and beneficial when the spectral radius of the adjacency matrix is small.
• To develop an efficient algorithm for approximating geometric series-based DAG constraints with bounded error and low computational cost.
• To improve the accuracy and robustness of DAG learning in both synthetic and real-world settings.
• To replace existing DAG constraints in state-of-the-art models with the proposed TMPI-based constraint for performance gains.

Proposed method

• Proposes a geometric series-based DAG constraint as an order-d polynomial over the adjacency matrix, using larger coefficients on higher-order terms to improve approximation of the nilpotency condition.
• Introduces Truncated Matrix Power Iteration (TMPI), an efficient algorithm to compute the geometric series approximation in O(log k) time, where k is the effective order of the series.
• The TMPI algorithm maintains a theoretical error bound, ensuring that the approximation remains within a controlled tolerance of the true geometric series.
• The method leverages the fact that adjacency matrices of DAGs are nilpotent, implying small spectral radii, which prevents numeric explosion even with large coefficients.
• The constraint is integrated into existing differentiable DAG frameworks such as NOTEARS, DAG-GNN, and GRAN-DAG by replacing their original polynomial-based acyclicity constraints.
• Employs a heuristic to identify a reduced order k ≤ d such that the feasible set of the constraint remains unchanged, enabling further computational savings.

Experimental results

Research questions

• RQ1 Can larger coefficients on higher-order polynomial terms improve the accuracy of DAG constraints without causing numerical instability?
• RQ2 Is it possible to efficiently approximate the geometric series of an adjacency matrix with bounded error using a truncated iterative method?
• RQ3 Does replacing existing DAG constraints with a geometric series-based constraint lead to significant improvements in SHD across diverse graph structures?
• RQ4 How does the proposed TMPI algorithm compare in speed and accuracy to naive and existing implementations of geometric series constraints?
• RQ5 Can the method effectively mitigate gradient vanishing in sparse DAGs, where high-order terms are critical for enforcing acyclicity?

Key findings

• The proposed TMPI-based DAG constraint reduces structural Hamming distance (SHD) by up to a factor of 3 compared to state-of-the-art methods in various settings.
• On a 50-node ER1 nonlinear SEM, the method achieved an SHD of 22.2±4.2, outperforming DAG-GNN's 25.2±4.5.
• In a 50-node nonlinear MLP dataset, NOTEARS-MLP with the TMPI constraint achieved SHD 14.9±1.3, compared to 16.9±1.5 with the original NOTEARS-MLP.
• On the Sachs protein signaling dataset, replacing DAG constraints with TMPI reduced SHD from 16 to 16 (DAG-GNN) and from 13 to 12 (Gran-DAG), with corresponding SHDC improvements from 21 to 17 and 11 to 9.
• The fast TMPI implementation runs significantly faster than naive implementations, especially for larger graphs, while maintaining similar SHD performance.
• The method effectively mitigates gradient vanishing in sparse graphs by enabling informative high-order terms through larger coefficients, without inducing gradient or numeric explosion.