Optimizing NOTEARS Objectives via Topological Swaps

TL;DR

This paper introduces a novel bi-level algorithm for optimizing NOTEARS objectives in DAG learning, using topological swaps to effectively handle non-convex constraints and improve solution quality.

Contribution

The work presents a new topological swap-based bi-level optimization method that guarantees convergence to a local minimum and enhances DAG learning performance.

Findings

Outperforms state-of-the-art methods in score optimization

Can significantly improve scores of existing algorithms as a post-processing step

Guarantees convergence to a local minimum under weaker conditions

Abstract

Recently, an intriguing class of non-convex optimization problems has emerged in the context of learning directed acyclic graphs (DAGs). These problems involve minimizing a given loss or score function, subject to a non-convex continuous constraint that penalizes the presence of cycles in a graph. In this work, we delve into the optimization challenges associated with this class of non-convex programs. To address these challenges, we propose a bi-level algorithm that leverages the non-convex constraint in a novel way. The outer level of the algorithm optimizes over topological orders by iteratively swapping pairs of nodes within the topological order of a DAG. A key innovation of our approach is the development of an effective method for generating a set of candidate swapping pairs for each iteration. At the inner level, given a topological order, we utilize off-the-shelf solvers that can handle linear constraints. The key advantage of our proposed algorithm is that it is guaranteed to find a local minimum or a KKT point under weaker conditions compared to previous work and finds solutions with lower scores. Extensive experiments demonstrate that our method outperforms state-of-the-art approaches in terms of achieving a better score. Additionally, our method can also be used as a post-processing algorithm to significantly improve the score of other algorithms. Code implementing the proposed method is available at https://github.com/duntrain/topo.