SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Bryan Andrews
Department of Psychiatry & Behavioral Sciences
Minneapolis, MN 55454
andr1017@umn.edu
Joseph Ramsey
jdramsey@andrew.cmu.edu
Rubén Sánchez-Romero
Center for Molecular and Behavioral Neuroscience
Newark, NJ 07102
ruben.saro@rutgers.edu
Jazmin Camchong
camch002@umn.edu
Erich Kummerfeld
Institute for Health Informatics
erichk@umn.edu
Learning graphical conditional independence structures is an important machine learning problem and a cornerstone of causal discovery. However, the accuracy and execution time of learning algorithms generally struggle to scale to problems with hundreds of highly connected variables—for instance, recovering brain networks from fMRI data. We introduce the best order score search (BOSS) and grow-shrink trees (GSTs) for learning directed acyclic graphs (DAGs) in this paradigm. BOSS greedily searches over permutations of variables, using GSTs to construct and score DAGs from permutations. GSTs efficiently cache scores to eliminate redundant calculations. BOSS achieves state-of-the-art performance in accuracy and execution time, comparing favorably to a variety of combinatorial and gradient-based learning algorithms under a broad range of conditions. To demonstrate its practicality, we apply BOSS to two sets of resting-state fMRI data: simulated data with pseudo-empirical noise distributions derived from randomized empirical fMRI cortical signals and clinical data from 3T fMRI scans processed into cortical parcels. BOSS is available for use within the TETRAD project which includes Python and R wrappers.
We present a permutation-based algorithm, Best Order Score Search (BOSS), for learning a Markov equivalence class of directed acyclic graphs (DAGs). A novel method for caching intermediate calculations for permutation-based algorithms, Grow-Shrink Trees (GSTs), is also developed and presented in this paper. Our implementation of BOSS using GSTs scales well in both accuracy and time to higher numbers of variables and graph densities. We demonstrate that in particular, this method scales at least to the complexity of dense cortical parcellations of fMRI data.
In real world systems, it is not unusual to have hundreds or thousands of variables, each of which may be causally connected to dozens of other variables. Models of functional brain imaging (fMRI), functional genomics, electronic health records, and financial systems are just a few examples. In order for structure learning methods to have an impact on these real world problems, they need to be not only highly accurate but also highly scalable in both dimensions.
This work follows earlier research on permutation-based structure learning [12, 22, 24, 26]. A previous permutation-based method, Greedy Relaxations of the Sparsest Permutation (GRaSP) [12], demonstrated strong performance on highly connected graphs, but struggled to scale to sufficiently large numbers of variables. By using BOSS and GSTs, we are able to overcome this challenge, while maintaining nearly identical performance to GRaSP. As our simulations show, other popular methods fall short on accuracy, scalability, or both, for such large, highly connected models.
In the remainder of this paper, we provide background for our approach in Section 2, followed by a discussion of GSTs in Section 3. We then introduce BOSS in Section 4 and validate it by comparing it to several alternatives. We show that BOSS using GSTs compares favorably to a number of best-performing combinatorial and gradient-based learning algorithms under a broad range of conditions, up to 1000 variables with an average degree of 20 in Section 5. Finally, to demonstrate its practicality, we apply BOSS to two sets of resting-state fMRI data: simulated data with pseudo-empirical noise distributions derived from randomized empirical fMRI cortical signals, and clinical data from 3T fMRI scans processed into 379 cortical parcels in Section 6. Section 7 provides a brief discussion.
Our novel contributions to the field include the following:
1. We present a new data structure, Grow-Shrink Trees (GSTs), for efficiently caching results of permutation-based structure learning algorithms. Using GSTs dramatically speeds up many existing permutation-based methods such as GRaSP.
2. We present a new structure learning algorithm, the Best Order Score Search (BOSS), which has similar performance as GRaSP (and thus superior performance to other DAG-learning methods), but is more convenient to use than GRaSP because it has fewer tuning parameters and is faster and more scalable.
3. We prove that BOSS is asymptotically correct and provide an implementation within the TETRAD project[17].
4. We present validations of BOSS's finite sample performance via standard graphical model simulations and on both real and simulated fMRI data.
The following conventions are used throughout this paper:  $V$  denotes a non-empty finite set of variables which double as vertices in the graphical context, lowercase letters  $a, b, c, \ldots \in V$  denote variables or singletons, uppercase letters  $A, B, C, \ldots \subseteq V$  denote sets, and  $X$  denotes a collection of random variables indexed by  $V$  whose joint probability distribution is denoted by  $P$ .
Probabilistic conditional independence between the members of  $X$  corresponding to disjoint sets  $A, B, C \subseteq V$  is denoted  $A \perp B \mid C [P]$  and reads  $X_A$  and  $X_B$  are independent given  $X_C$ .
2.1 Permutations
A permutation is a sequence of variables  $\pi = \langle a, b, c, \ldots \rangle$ . Let  $v \in V$  and  $i \in \mathbb{N}$  ( $i \leq |V|$ ). If  $\pi$  is a permutation of  $V$ , then  $\pi$  is equipped with two methods:  $\pi . \text{index}(v)$  which returns the index of  $v$  in  $\pi$ , and  $\pi . \text{move}(v, i)$  which returns the permutation obtained by removing  $v$  from  $\pi$  and reinserted at position  $i$ . Furthermore,  $\text{pre}_{\pi}(v) \equiv \{w \in V : \pi . \text{index}(w) < \pi . \text{index}(v)\}$  is the prefix of  $v$ .
2.2 DAG Models
A directed acyclic graph (DAG) model is a set probabilistic models whose conditional independence relations are described graphically by a DAG. In general, the vertices and edges of a DAG represent
variables and conditional dependence relations, respectively. We use the notation  $\mathcal{G} = (V,E)$  where  $\mathcal{G}$  is a graph,  $V$  is a vertex set, and  $E$  is an edge set. In a DAG, the edge set is comprised of ordered pairs that represent directed edges $^1$  and contains no directed cycles.
Let  $\mathcal{G} = (V,E)$  be a DAG and  $v\in V$  ..
$$
\operatorname {p a} _ {\mathcal {G}} (v) \equiv \{w \in V: w \to v \text {i n} \mathcal {G} \}
$$
$$
\operatorname {c h} _ {\mathcal {G}} (v) \equiv \{w \in V: v \to w \text {i n} \mathcal {G} \}
$$
are the parents and children of  $v$ , respectively.
Graphical conditional independence between disjoint subsets  $A, B, C \subseteq V$  can be read off of a DAG using  $d$ -separation and is denoted  $A \perp B \mid C$  [9]. This criterion admits the concept of a Markov equivalence class (MEC) which (in the context of this paper) is a collection of DAGs that represent the same conditional independence relations.
2.3 Causal Discovery
DAGs are connected to causality by the causal Markov and causal faithfulness assumptions. A DAG is causal if it describes the true underlying causal process by placing a directed edge between a pair of variables if and only if the former directly causes the latter.
Causal Markov: If  $\mathcal{G} = (V, E)$  is causal for a collection of random variables  $X$  indexed by  $V$  with probability distribution  $P$ , then for disjoint subsets  $A, B, C \subseteq V$ :
$$
A \perp B \mid C [ \mathcal {G} ] \Rightarrow A \perp B \mid C [ P ].
$$
More generally, this implication is called the Markov property and we say DAGs satisfying this property contains the distribution. Moreover, a DAG  $\mathcal{G}$  is subgraph minimal if no subgraph of  $\mathcal{G}$  contains  $P$ .
Causal faithfulness: If  $\mathcal{G} = (V,E)$  is causal for a collection of random variables  $X$  indexed by  $V$  with probability distribution  $P$ , then for disjoint subsets  $A, B, C \subseteq V$ :
$$
A \perp B \mid C [ P ] \Rightarrow A \perp B \mid C [ \mathfrak {G} ].
$$
Accordingly, under the causal Markov and causal faithfulness assumptions, finding the MEC of the causal DAG is equivalent to identifying the simplest model that contains the data generating distribution. From this point of view, causal discovery is a standard model selection problem which we address using the Bayesian information criterion (BIC) [7, 21].
Let  $\mathcal{G}$  be a DAG and  $X \stackrel{\mathrm{id}}{\sim} P$  be a dataset drawn from a member of a curved exponential family:
$$
\begin{array}{l} \operatorname {B I C} (\mathcal {G}, \boldsymbol {X}) = \sum_ {v \in V} \operatorname {B I C} \left(\boldsymbol {X} _ {v}, \boldsymbol {X} _ {\mathrm {p a} _ {\mathcal {G}} (v)}\right) \\ = \sum_ {v \in V} \ell_ {v | \mathrm {p a} _ {\mathcal {G}} (v)} (\hat {\theta} _ {\mathrm {m l e}} \mid X) - \frac {\lambda}{2} | \hat {\theta} _ {\mathrm {m l e}} | \log (n) \\ \end{array}
$$
where  $\lambda > 0$  is a penalty discount,  $\ell$  is the log-likelihood function, and  $\theta_{\mathrm{mle}}$  is the maximum likelihood estimate of the parameters; for details on curved exponential families, we refer reader to [9]. Importantly, the BIC is consistent.
Proposition 1. Haughton [7] If  $P$  is a member of a curved exponential family and  $X \stackrel{\mathrm{iid}}{\sim} P$ , then the BIC is maximized in the large sample limit by DAG models containing  $P$  that minimize the dimension of the parameters space.
Notably, Gaussian and multinomial DAG models form curved exponential families in which the causal DAG minimizes the dimension of the parameters space, and thereby maximize the BIC [11]. For these distributions, DAG models belonging to the same MEC correspond to the same DAG model,
<sup>1</sup>The edge  $v \gets w$  corresponds to the order pair  $(v, w)$ .
2This concept is also known as SGS-minimality [28].
so it is common for algorithms to return a graphical representation of the MEC called a CPDAG rather than a DAG. Choosing between DAG models in a MEC requires additional information.
We greedily search over permutations of variables using grow-shrink (GS) [13] to project and score DAGs from permutations. Algorithm 1 (grow) and Algorithm 2 (shrink) give the details of GS while Algorithm 3 (project) details the process of projecting a permutation to a DAG.
Algorithm 1: grow(X, v, Z)
Input: data:  $X$  var:v prefix:Z   
Output: parents:W   
 $W\gets \emptyset$    
repeat  $\begin{array}{rl} & w\leftarrow \mathrm{argmax}_{z\in Z}\mathrm{BIC}(\mathbf{X}_v,\mathbf{X}_{W\cup z})\\ & \mathrm{if~}w\neq \emptyset \mathrm{then}\\ & \mathrm{\lfloor W\leftarrow W\cup w} \end{array}$    
until  $w = \emptyset$    
Algorithm 2: shrink(X,v,W)   
Input: data:X var:v parents:W   
Output: parents:W   
repeat  $\begin{array}{rl} & w\leftarrow \mathrm{argmax}_{w\in W}\mathrm{BIC}(\mathbf{X}_v,\mathbf{X}_{W\backslash w})\\ & \mathrm{if~}w\neq \emptyset \mathrm{then}\\ & \mathrm{\lfloor W\leftarrow W\backslash w} \end{array}$    
until  $w = \emptyset$
Algorithm 3: project  $(X,\pi)$
Input: data : X perm : π  
Output: DAG : G  
E ← ∅  
foreach v ∈ π do  
Z ← preπ(v)  
W ← grow(X, v, Z)  
W ← shrink(X, v, W)  
foreach w ∈ W do  
E ← E ∪ (v, w)  
G ← (V, E)
However, how do we know that the causal DAG will even be constructed by project? This is guaranteed by the follow results. Let  $P$  is a member of a curved exponential family satisfying causal Markov and causal faithfulness:
- If  $\mathcal{G}$  is the causal DAG, then  $\mathcal{G}$  is subgraph minimal [11, 28].
- If  $\mathcal{G}$  is a DAG and  $X \stackrel{\mathrm{iid}}{\sim} P$  then  $\mathcal{G}$  is subgraph minimal if and only if there exists a permutation  $\pi$  such that  $\mathcal{G} = \operatorname{project}(X, \pi)$  in the large sample limit [12].
These results admit the following strategy: (1) Search over permutations while using GS and the BIC to construct and score subgraph minimal DAG from these permutations. (2) Return the MEC of the subgraph minimal DAG that maximizes the BIC. The question is how to do so in a way that is both comprehensive and efficient.
3 Grow-Shrink Trees
Grow-shrink trees (GSTs) are tree data structures for caching the results of the grow and shrink subroutines of GS [13]. This data structure is compatible with many permutation-based structure learning algorithms, including BOSS and GRaSP. As an initialization step, a GST is constructed for each variable in the dataset which are then queried during search rather than running GS. Each tree has a root node representing the empty parent set. Parents are accumulated by traversing edges of the tree, with every node in the tree corresponding to a "grown" parent set. Each grown parent set corresponds to running the grow subroutine for some prefix. The shrink subroutine is also run and cached at each node of the tree.
The benefit of using GSTs is in how efficiently they store information needed for running GS. In the grow subroutine, nodes are added to the GST one at a time and scored. A new child node is added to the tree for each possible addition, and the child nodes are sorted in descending order relative to their scores. After all candidate parents have been added to the tree, the first child in the sorted list that is also in the prefix is chosen and the corresponding edge traversed. The shrink subroutine can then be run and the removed parents and scores can be cached.
Figure 1: Minimal subgraphs.
(f)
Figure 2: A Grow-shrink Tree
3.1 An Example
This section walks through the usage of a GST applied to a toy example. Minimal subgraphs are depicted in Figure 1 of which (1a) is the true DAG. Figure 2 outlines the process of growing a GST for  $a$  while scoring permutations. In the following, we identify nodes in the GST by their path from the root. All ambiguities in sorting the branches of a node are resolved in alphabetical order. We proceed in sequence, but in theory some of these steps could be performed in parallel.
In what follows we run GS on a series of permutations while caching the results in a GST. We describe the nodes of the tree by the sequence of variables traversed to reach them from the root. We use colors to track our progress: light gray parts of the tree denote unexplored paths with no cached scores, dark gray parts of the tree denote explored paths with cached scores. Colored edges denote the paths considered during the execution of GS: red for rejected and blue for accepted. Colored nodes denote the same, but are only colored if they do not have a cached value and require a new score calculation. Numbered edges directed out of a parent node record the order of preference for the children from high to low. To simplify this example, we do not cache the results of shrink. In (2a) we initialize the tree by adding the root node  $a$ .
Our first permutation to evaluate is  $\langle b, d, a, c \rangle$  which is shown in (2b) and whose minimal subgraph is depicted in (1b). Node  $a$  has not been expanded, so we score and sort nodes  $ab$ ,  $ac$ , and  $ad$ . Node  $ab$  scored the highest and is contained in the prefix, so we travel to node  $ab$ . Node  $ab$  has not been expanded, so we score and sort nodes  $abc$  and  $abd$ . Node  $abd$  scored the highest and is contained in the prefix, so we travel to node  $abd$ . At this point no more variables are available in the prefix so we have completed growing. We run shrink at node  $abd$  which removes nothing so the GS result is  $\{b, d\}$ .
Our next permutation to evaluate is  $\langle c, d, a, b \rangle$  which is shown in (2c) and whose minimal subgraph is depicted in (1c). Node  $a$  has already been expanded so we check its outgoing edges. Node  $ab$
has the highest score but is not contained in the prefix. Node ad has the second highest score and is contained in the prefix, so we travel to node ad. Node ad has not been expanded, so we score and sort nodes adc. Node  $c$  has the highest score and is contained in the prefix, so we travel to node adc. At this point no more variables are available in the prefix so we have completed growing. We run shrink at node adc which removes nothing so the GS result is  $\{c, d\}$ .
Our next permutation to evaluate is  $\langle c, a, b, d \rangle$  which is shown in (2d) whose minimal subgraph is depicted in (1d). Node  $a$  has already been expanded so we check its outgoing edges. Node  $ab$  has the highest score but is not contained in the prefix. Node  $ad$  has the second highest score but is not contained in the prefix. Node  $ac$  has the third highest score and is contained in the prefix, so we travel to node  $ac$ . At this point no more variables are available in the prefix so we have completed growing. We run shrink at node  $ac$  which removes nothing so the GS result is  $\{c\}$ .
Our next permutation to evaluate is  $\langle b, c, a, d \rangle$  which is shown in (2e) whose minimal subgraph is depicted in (1e). Node  $a$  has already been expanded so we check its outgoing edges. Node  $ab$  has the highest score and is contained in the prefix, so we travel to node  $ab$ .  $ab$  has already been expanded so we check its outgoing edges. Node  $abd$  has the highest score but is not contained in the prefix. Node  $abc$  has the second highest score and is contained in the prefix, so we travel to node  $abc$ . At this point no more variables are available in the prefix so we have completed growing. We run shrink at node  $abc$  which removes nothing so the GS result is  $\{b, c\}$ .
Our last permutation to evaluate is  $\langle c, b, d, a \rangle$  which is shown in (2f) and whose minimal subgraph is depicted in (1a). Node  $a$  has already been expanded so we check its outgoing edges. Node  $ab$  has the highest score and is contained in the prefix, so we travel to node  $ab$ . Node  $ab$  has already been expanded so we check its outgoing edges. Node  $d$  has the highest score and is contained in the prefix, so we travel to node  $abd$ . Node  $abd$  has not been expanded, so we score node  $abcd$ . No scores result in an improvement to the overall score so we have completed growing. We run shrink at node  $abd$  which removes nothing so the GS result is  $\{b, d\}$ .
4 Best Order Score Search
Similar to GES, BOSS uses a two phase search procedure [5]. The first phase uses the best-move method, which takes a variable as input and greedily moves it to the position in the current permutation that maximizes the score. In this phase, best-move is repeatedly applied to each variable, one at a time, until there are no more moves that increase the score. This phase concludes with the find-compelled procedure of [4] which converts the DAG into a CPDAG. The second phase of BOSS is BES which is exactly second phase of GES. The BES step is optional but guarantees asymptotic correctness if executed.
Algorithm 4: BOSS(X,π,δ)
Input: data:  $X$  perm:  $\pi$  flag:  $\delta$   
Output: graph:  $\mathcal{G}$ $\mathcal{T} \leftarrow \mathrm{GST}(\mathbf{X})$   
repeat  
best  $\leftarrow \mathcal{T}.\mathrm{score}(\pi)$   
foreach  $v \in \pi$  do  
 $\begin{array}{r}\lfloor \pi \leftarrow \mathrm{best - move}(\mathcal{T},\pi ,v) \end{array}$   
until best = T.score(π)  
 $\mathcal{G} \leftarrow \mathcal{T}.\mathrm{project}(\pi)$ $\mathcal{G} \leftarrow \mathrm{find - compelled}(\mathcal{G})$   
if  $\delta =$  true then  
 $\begin{array}{r}\lfloor \mathcal{G} \leftarrow \mathrm{BES}(\mathcal{G},\mathbf{X}) \end{array}$
Algorithm 5: best-move  $(\mathcal{T},\pi ,v)$
Input: GSTs: T perm: π var: v  
Output: perm: π  
best ← T.score(π)  
for i ← 1 to |π| do  
    j ← π.index(v)  
    π ← π.move(v, i)  
if best < T.score(π) then  
    best ← T.score(π)  
else  
    π ← π.move(v, j)
Proposition 2. Let  $P$  be a member of a curved exponential family satisfying causal Markov and causal faithfulness. If  $X \stackrel{\mathrm{iid}}{\sim} P$  then  $\mathsf{BOSS}(X, \pi, \mathsf{true})$  returns the MEC of the causal DAG for all initial permutation  $\pi$  in the large sample limit.
Proof. Since project returns a subgraph minimal DAG, it contains  $P$ . Accordingly, asymptotic correctness follows from the correctness of BES.
In Algorithm 4 (B0SS) and Algorithm 5, GST constructs a collection of GSTs, denoted  $\mathfrak{T}$ , which contains one GST for each variable. The collection  $\mathfrak{T}$  is equipped with project and score methods which runs the GS algorithm to project and score a permutation, respectively.
5 Simulations
We evaluated the speed and performance of BOSS on simulated data compared to other algorithms: GRaSP, fGES, PC, DAGMA, and LiNGAM. Our evaluation used the performance metrics tabulated in Table 1 which were originally proposed by [10]. All algorithms were run on an Apple M1 Pro processor with 16G of RAM. The results reported in the main text are abridged but the complete results are available in the Supplement.
Table 1: Metrics
True	Estimated	Adjacency	Orientation
a←b	a←b	tp	tp, tn
a→b	tp	fp, fn
a-b	tp	fn
a...b	fn	fn
a...b	a←b	fp	fp
a→b	fp	fp
a-b	fp
a...b	tn
$$
\text {P r e c i s i o n} = \frac {\mathrm {t p}}{\mathrm {t p} + \mathrm {f p}}
$$
$$
\text {R e c a l l} = \frac {\mathrm {t p}}{\mathrm {t p} + \mathrm {f n}}
$$
We evaluated our implementation of BOSS using a BIC score with  $\lambda = 2$  and no BES step as it did not appear to improve performance. For GRaSP, we modified (to use GSTs) and used the TETRAD implementation with the same parameters as the authors and a BIC score with  $\lambda = 2$ . [12, 17]. For fGES we used the TETRAD implementation with default parameters and a BIC score with  $\lambda = 2$  [5, 15, 17]. We also used the implementation of PC in TETRAD with default parameters using a BIC score with  $\lambda = 2$  as a conditional independence oracle [17, 25]. For DAGMA we used the authors' Python implementation with the parameters reported in their paper but changed the threshold to 0.1 [2] and used the technique described in [14] to resolve cycles. Moreover, the output of DAGMA was converted to a CPDAG for linear Gaussian simulations. For LiNGAM we used the authors' Python implementation with default parameters [8, 22].
We generated Erdős-Rényi DAGs by applying an arbitrary order to vertices of an Erdős-Rényi graph. For scale-free networks, first we generated an Erdős-Rényi DAG and then redrew the parents of each vertex according to the Barabási-Albert model [1]. This resamples the out-degrees distribution to be scale-free while keeping in-degree distribution constant. This procedure was motivated by the results in Figure 5. Edge weights were sampled from a uniform distribution in the interval [-1.0, 1.0] and error distributions were generated with standard deviations sampled from a uniform distribution in the interval [1.0, 2.0]. The complete simulation details, including figures plotting the simulated scale-free in/out degree distributions, are include in the Supplement. Additionally, the data are available for download at: https://github.com/cmu-phil/boss.
Figure 3 compares BOSS against GRaSP, fGES, PC, and DAGMA on linear Gaussian data generated from (3a) Erdős-Rényi and (3b) scale-free networks. The output of DAGMA was converted to a CPDAG for these simulations. Lam et al. [12] attribute the excellent performance of GRaSP to it being robust against the ubiquity of almost-violations of causal faithfulness [29, 27]. Due to the algorithmic and performance parallels between GRaSP and BOSS, we conjecture that a similar argument could be made for BOSS.
Figure 4a compares BOSS against DAGMA and LiNGAM on linear exponential and linear Gumbel data from scale-free networks. Interestingly, LiNGAM, which takes advantage of non-Gaussian signal in data, does not perform appreciably better than BOSS.
Figure 4b compares BOSS against GRaSP on linear Gaussian data generated from scale-free networks with a focus on scalability; runs exceeding two hours were cancelled and are not reported. Here we see that BOSS maintains a high level of accuracy while scaling much better than GRaSP. It is also
(a) Erdős-Rényi
(b) Scale-free