Learning Arborescence with An Efficient Inference Algorithm

We consider a class of structured learning problems on arborescence (i.e., the directed spanning tree) from the input graph. The key step involved in this problem is predicting the minimal weight arborescence (MWA) from the learned model. In literature, there are two lines of research for predicting MWA: the Chu-Liu Edmonds (CLE) and the Lovasz methods. The CLE method is easy to implement while it takes $\mathcal{O}(n)$ cycle contractions. Here $n$ is the graph size. The Lovasz method reduces to the multi-pair shortest path (MPSP) problem and takes only $\mathcal{O}(\log n)$ contractions. Nevertheless, in the CPU setting, MPSP has the same time complexity as finding MWA. The Lovasz method only attains time efficiency under a sufficient GPU setting. Both the aforementioned methods are painfully slow for large-scale learning tasks. In this research, we find the general MPSP problem can be simplified when working with machine learning models. This is because the learning model predicts edge weights for all pairs of vertices and the graph we process is always complete. Therefore, we only need to handle those paths that directly enter every weakly connected component (WCC) while the classic Lovasz method needs to handle all possible paths. This allows us to propose LAzy LoVAz (Lava) method that enjoys $\mathcal{O}(\log n)$ contractions as well as efficient performance in both CPU and GPU settings. In experiments, we consider synthetic datasets and two real-world learning tasks, i.e., graph-based dependency parsing and unsupervised parsing on ListOps. The empirical results exhibit important gains of our Lava method to the classic CLE and Lovasz methods, that Lava boosts the training time for arborescence learning tasks.