LLM-enhanced Score Function Evolution for Causal Structure Learning

1 Introduction

Causal structure learning (CSL) is a fundamental approach for understanding causality [Pearl, 2009; Spirtes et al., 2001]. It uncovers causal relationships from observational or interventional data and represents them using graphical models such as Directed Acyclic Graphs (DAGs) et al. [Glymour et al., 2019; Vowels et al., 2022]. Learning an exact DAG from data is NP-hard [Chickering, 1996], and the mainstream algorithms can be broadly categorized into two types: constraintbased methods and score-based methods. Constraint-based methods reconstruct the causal graph from a statistical perspective. They first identify the skeleton by performing independence tests and then infer edge directions while adhering the acyclicity and other rules [Koller, 2009]. However, the significance of independence tests cannot be accurately assessed, as it is influenced by variable dimensionality and the number of conditioning variables, among others. In contrast, score-based methods have gained popularity in recent research [Huang et al., 2018]. They treat DAG learning as a constrained optimization problem, and define a score function to evaluate the fitness between candidate graphs and the observed data. Optimization techniques, either combinatorial or continuous, are then applied to identify the optimal DAG in the space of possible causal structures.

Although much of the current research focuses on improving the effectiveness and efficiency of search methods in more complex spaces or under weaker causal assumptions [Cheng et al., 2024], these studies generally adopt consistent score functions, and when the sample size approaches infinity, CSL can reveal the true causal structure, represented by the minimal I-map that accurately captures the data distribution.

However, is the score function can be applied across various scenarios? To explore this, we still take BIC [Schwarz, 1978] as example and consider the limited sample size case, which frequently arises in medical field. As illustrated in Figure 1 [Sachs et al., 2005], BIC exhibits varying preferences in model selection depending on the sample size. In fact, this gap is universal across score functions used for model selection. For instance, abnormal noise can negatively impact the BIC-Gaussian score, and the choice of priors can influence the BDeu score. Although some improvements have been proposed under various conditions [Silander et al., 2008; Huang et al., 2018; Andrews et al., 2018]. it is impractical to manually design new scoring functions for every emerging scenario. Given these challenges, our aim is to improve the score functions in a automatic and systematic manner. In fact, this task can be also framed as an optimization problem, requiring to explore the landscape of possible scoring functions.

Previously, such optimization tasks have been hindered by the lack of a comprehensive understanding of the score function space. Fortunately, recent advancements in large language model (LLM) [Achiam et al., 2023; Zhao et al., 2023] offer promising tools. LLM encodes algorithms through code implementation and contextual description, enabling more effective search methods for algorithm design [Liu et al., 2024b; Romera-Paredes et al., 2024]. Inspired by it, we propose the L-SFE (LLM-enhanced Score Function Evolution) framework to discovery optimal score functions. Note that L-SFE still achieves the MEC-level identifiability under the causal faithfulness and causal suffiency assumptions [Pearl, 2009]. The contribution of this paper includes:

1. We frame L-SFE as a supervised learning process and formulate it using a bi-level optimization approach. At the lower level, a greedy local search (GLS) is employed to identify the optimal DAGs based on the specific score function. At the upper level, Evolutionary algorithm (EA) is utilized to optimize the scoring function with maximal fitness, which is quantified by the structural difference between the aforementioned DAGs and the ground-truth graphs.
2. For LLM, we take the BIC as standard and design the prompt integrating both conceptual and code-level information to initialize and evolve the scoring function. Additionally, we refine several operators, including mutation, crossover, and injection, to effectively balance convergence and diversity of EA.
3. We evaluate the generated scoring functions from both theoretical and experimental analysis. The former focuses on equivalence and consistency, while latter is performed across synthetic (including discrete and continuous) and real-world data. The results demonstrate the effectiveness, stability, and generality of L-SFE.

2 RelatedWorks

Causal Structure Learning. The score-based CSL algorithms can be categorized into combination-based and continuous-based types. The former execute approximate or exact search in: (1) DAG space G (e.g., HC [Heckerman et al., 1995] and MAHC [Constantinou et al., 2022]), which explore plausible DAGs by heuristically manipulating single or multiple directed edges; (2) EC (Equivalent space) space E (e.g., GES [Chickering, 2002] and fGES [Ramsey et al., 2017]), which greedy search the completely partially directed acyclic graph (CPDAG) via forward equivalent search and backward equivalent search. (3) Permutation space O (e.g., OBS [Teyssier and Koller, 2012], GRaSP [Lam et al., 2022], and BOSS [Andrews et al., 2023]), which identify the best causal structure by discovering the topological ordering that maximizes ancestral information. Continuous-based algorithms make assumptions about the data distribution and learn causal graphs from a structural equation model (SEM), such as LiNGAM [Shimizu et al., 2011], NOTEARS [Zheng et al., 2018], DAG-RL [Zhu et al., 2019] and DAGMA [Bello et al., 2022] et al.

LLM for CSL. Given the exceptional text comprehension capabilities, numerous studies have leveraged LLM to research the causality [Kıcıman et al., 2023; Takayama et al., 2024]. These approaches regard LLM as sole determiners of pairwise causal relationships [Zhiheng et al., 2022; Wan et al., 2024] or domain experts. The latter harness metainformation embedded in the LLM's training data to enhance the causal discovery, including initialization [Ban et al., 2023; Li et al., 2024], post-door adjustments [Khatibi et al., 2024], or structural constraints fusion [Ban et al., 2023; Zhou et al., 2024; Jiralerspong et al., 2024; Zhang et al., 2024].

LLM for Algorithm Design. Finally, we provide a brief overview of using LLM for auto algorithm design. Most existing frameworks employ evolutionary approaches and leverage operators defined within the algorithmic space for algorithm generation, such as FunSearch [Romera-Paredes et al., 2024], EoH [Liu et al., 2024b], and ReEov [Ye et al., 2024]. Additionally, these frameworks have been applied to a variety of scenarios, including the capacitated vehicle routing problem [Liu et al., 2024a], critical node discovery [Mao et al., 2024], and tensor network search [Zeng et al., 2024].

3 Background

DAG can be described as a tuple G = (V ,E), where V = {X1,X2, · · · ,Xn} represents the collection of variables, and E = {Xi → Xj|Xi,Xj ∈ V } denotes the directed edges between the variables. CSL can be typically formulated as the optimization problem

Ψ(.) is a score function that evaluates the fitness of DAG G on i.i.d dataset D = {D1, · · · ,Dm}, and it is usually decomposable and equivalence [Koller, 2009].

Definition 1. (Decomposable) If the score functionΨ(G|D) can be written as Ψ(G|D) = ∑Xi∈V Ψ(Xi, PaGi |D), then Ψ is decomposable.

Definition 2. (Equivalence) ∀G1, G2 ∈ G, if G1, G2 is Iequivalent, and Ψ(G1|D) = Ψ(G2|D) holds, then Ψ is equivalence.

where PaGi represent the parents of Xi in G. For discrete datasets, the most commonly used information theoretic score functions include BIC and AIC [Kitson et al., 2023], which evaluate model based on the multi-information content and entropy of the variables. Other Bayesian scores, such as BDeu, BDs, K2, and BDe[Kitson et al., 2023], place priors on the parameters of the tabular conditional probability distribution. For continuous datasets, BIC remains applicable when using conditional covariance matrix. Thus, we take BIC as benchmark for evolution [Andrews et al., 2023]

where ℓ(.) is a log-likelihood function, and ˆ θmle represents the maximum likelihood estimation (MLE) for the parameters of subgraph Gi : PaGi → Xi. Furthermore, BIC is also consistent if the underlying distribution P behind D belongs to a curved exponential family [Haughton, 1988].

Definition 3. (Consistent) Suppose G ∗ is a P-map under the distribution P ∗ . A score function Ψ is said to be consistent if as the sample size m→∞, the following conditions hold:

• The structure G ∗ maximum the Ψ(G ∗ |D).
• ∀G ∈ G, if G is not I-equivalent to G ∗ , then Ψ(G|D) < Ψ(G ∗ |D).

4 Framework

The search in score function space can be formulated as a bi-level optimization problem

where G ∗ denotes the arbitrary DAG within G, and DG∗ represents the dataset sampled from G ∗ under the gaussian or multinational distribution. Dis(Ga, Gb) is a structural distance measure for Ga, Gb. Note that in equation 3, the upper level optimization identify the optimal Ψ by minimizing the expectation of structural loss over all possible G ∗ assigned with DG∗ . This objective can be achieved through lower level optimization, which finds the best G †Ψ,DG∗ for a given score function Ψ. Furthermore, since the score function space is abstract, L-SFE employes EA to solve equation 3, as it does not require a clear characterization for the mathematical properties of the solution. As shown in figure 2, the workflow of the L-SFE can be divided into two stages. Firstly, EA searches the score function space to identify an improved Ψ; then, Ψ is combined with the GLS to solve the lower level optimization problem, ultimately leading to a best DAG for evaluating the fitness of Ψ.

4.1 Score Function Generation

In this subsection, we provide a detailed explanation of how the score function is discovered using LLM.

Prompts Design. Inspired by EoH [Liu et al., 2024b], all prompts pt are designed using a three-level hierarchical structure as shown in figure 3. Task Description informs the LLM of its intended role and the objective of CSL. Code Snippets presents few-shot examples to heuristic the LLM designing analogous score functions, and it contains two key components: Idea, which provides a textual explanation of the core thought behind the score function; Code, which presents the detailed Python implementation. In addition, there is an optional item that demonstrates the fitness of the score function. Task Assignment applies the chain-of-thought (CoT) reasoning [Wei et al., 2022] to derive better scoring functions from the given examples, and it includes four steps:

1. First, the LLM carefully reads and interprets the core idea provided in the code comments;
2. Second, if the input score function is labeled with a fitness, the LLM analyzes the reasons behind its effectiveness or shortcomings;

1. Third, based on the selected operators, different prompts are generated to inspire new ideas that diverge from the given examples;
2. Finally, the LLM implements the new idea in code, ensuring that the function name, input, and output match those in the example code.

Fitness Evaluation inCSLSpace. Based on equation 3, LSFE employs the Monte Carlo method [Metropolis and Ulam, 1949] to randomly sample np ground-truth graphs G ∗ and corresponding training datasets DG∗ for training. For convenience, the weight matrix W k ∈ {0, 1}n×n is used to represent the k-th G ∗ , where W k ij = 1 implies the directed edge Xi → Xj, and the sampled dataset is denoted as Dk . Then, the Normalized Hamming Distance [Kıcıman et al., 2023] is used to quantify the Dis(Ga, Gb). Thus, the fitness function can be formulated as

where dW k denotes the number of variables in W k . W †Ψ,Dk represents the best DAG learned by the pre-selected GLS under Ψ and Dk . Obviously, a lower L(Ψ) indicates a higher fitness. Note that GLS can source from different CSL spaces, such as HC in DAG space, GES in EC space and BOSS in permutation space.

EA in Score Function Space. The initial algorithm population Ψ = {Ψ1, · · · ,Ψna } is generated with standard BIC, and mutation, crossover, and injection are employed for evolution. More detailed prompt design can be referred in Supplementary material 1.1 1 .

1 https://github.com/wzd2502/L-SFE

The crossover operator facilitates the random combination of two score functions, generating a new one that exhibits different structural and conceptual features. Specifically, Ψ1 and Ψ2 are randomly selected from the Ψ. Based on their code implementations and the associated fitness L(Ψ1) and L(Ψ2), LLM analyses and assimilates the ideas of parents to creates a new score function Ψ′ .

The mutation operator modifies an existing score function to generate a variant with potentially altered characteristics. Specifically, Ψ1 is randomly selected from the Ψ with its corresponding fitness L(Ψ1). Based on the extent of the modifications, mutation can be classified into two types: (1) structural mutation, which involves significant alterations to the score function's underlying structure, and (2) parameter mutation, which only modifies the parameters of the score function without altering its fundamental structure.

The injection operator is designed to maintain diversity within the algorithm population, preventing EA convergence to a local optima. Specifically, a new score function Ψ′ is generated, which is entirely distinct from all previous Ψ ∈ Ψ. L(Ψ′ ) is not evaluated at this stage, meaning Ψ′ is directly incorporated into the mutation and crossover operations. Ψ′ will only be discarded at the end of the iteration if it does not contribute positively to the evolutionary.

In each iteration, all of the aforementioned operators are applied repeatedly, and the generated score functions are tested for legality, ensuring that they do not introduce cyclic graphs and can produce results within a specified time limit. After completing an iteration, the population Ψ would be updated. Only the top na algorithms are retained for the next iteration. The evolution continues until a predefined stopping condition is met. To provide a clear understanding of L-SFE, we outline the pseudo-code in Alg. 1 in Supplementary material 1.2.

4.2 Score Function Analysis

L-SFE is repeated nl times on different training datasets. To avoid ambiguity, we introduce the notation L-SFEabk , a ∈ {D,C}, b ∈ {1, · · · , nl}, k ∈ {1, · · · , na}. This represents the k-th score function learned by L-SFE over the last iteration in b-th repetitions under the a-type dataset (D for Discrete and C for Continuous). Here we take the L-SFED1 as example and analyze its evolution in figure 4, and more detailed results can be found in Supplementary material 1.3.

At the first iteration, L-SFE enhances the CSL by incorporating extra penalty terms that account for the length of the parent sets

Equation 5 modifies only the structure of the penalty, so it remains a consistent and equivalent score. At the end of the evolution, L-SFE identifies an optimal Ψ by incorporating priors into the BIC score, which can effectively mitigate overfitting under limited data. Furthermore, a quadratic penalty is introduced to enhance the model's preference for simplicity, and the penalty term related to the number of data is modified from log(m) to 1/m. This formulation is expressed as

where ˆ θMAP represents the maximum a posterior (MAP) estimate, which is based on Dirichlet prior. Notably, the score function in Equation 6 is decomposable and possesses several advantageous properties

Theorem 1. Score function Equation 6 is still consistent but not equivalence.

The proof can be found in supplementary material 1.3. Actually, all LLM learned scores provide similar insights: they leverage the prior to enhance robustness under limited data and use 1/m rather than log(m) to penalize the graph's complexity in a less restrictive manner.

5 Experiment

5.1 Settings

Datasets. GPT-4o mini is utilized for score function discovery in L-SFE, and synthetic datasets generated from pytetrad are employed for training and testing 2 . For the discrete dataset, L-SFE is trained on ten RandomGraphs with n = 30, where the variables follow a multinomial distribution, and tested on ErdosRenyi (ER) and ScaleFree (SF) graphs. For the continuous datasets, L-SFE is trained on ten RandomGraphs under the linear Gaussian SEM assumption with n = 30, and evaluation is performed on the linear exponential SEM and the linear gumbel SEM. The GLS used for training is HC with a tabu search. Each test is repeated 10 times with m = 5000. Additional details of experimental setup are provided in Supplementary material 2.1.

2 Code is avaliable on https://github.com/wzd2502/L-SFE

Metrics. We evaluate the structural accuracy of L-SFE learned graph against the ground-truth graphs using the Structural Hamming Distance (SHD), the F1 score for adjacency accuracy (F1adj ), and the F1 score for arrow accuracy (F1arr ).

Baselines. The score function generated by L-SFE is compared with commonly used scores, including BIC, BDeu, and BDs for discrete data [Kitson et al., 2023], and BIC-Gaussian for continuous data [Kitson et al., 2023]. Additionally, these methods are integrated with GLS across various search spaces and compared to baseline algorithms, including:

• Discrete. HC [Heckerman et al., 1995], PC [Colombo et al., 2014], BOSS [Andrews et al., 2023], fGES [Ramsey et al., 2017], GRaSP [Lam et al., 2022].
• Continuous. LinGAM [Shimizu et al., 2011], DAGMA [Bello et al., 2022], BOSS, fGES, HC.

5.2 Overall Analysis

In this subsection, we want to answer the following questions: Stability: Are the all score functions identified by L-SFE consistently superior to human-designed scores? In this topic, we conduct nl = 3 repetitions of L-SFE and record the all na algorithms identified in the last iteration of each repetition. HC is employed as the search method. Figure 5

presents the comparison with three human-designed scores on ER graphs under the discrete dataset, and detailed test results on other settings can be found in Supplementary material 2.2. From figure 5, most of LLM learned score functions are either comparable to or outperform the human-designed scores in each repetition. For instance, L-SFED15 performs best in Repetition 1, while L-SFED21 and L-SFED25 show superior performance in Repetition 2, and L-SFED31 and L-SFED32 lead in Repetition 3. A similar trend is also observed in the SF graphs. In the case of continuous datasets, all score functions learned by the LLM outperform the human-designed BICgaussian score, highlighting the stability of the L-SFE.

Therefore, we select the best score function from i-th repetition, denoted as L-SFEDi∗ and evaluate their generality.

Generality: Can the score functions discovered by L-SFE be applied across different search methods? In this topic, we explore whether the score function, trained using HC in the DAG (G) space, remains effective in the EC (E) and permutation (O) spaces. To investigate this, we employ the GES and BOSS, along with L-SFEDi∗ for testing. The results on ER graphs are displayed in figure 6, with additional results for other settings provided in Supplementary material 2.2. LLM learned score functions still demonstrate strong performance in the EC space as well compared with BIC and BDeu. However, in the permutation space, this superiority narrows significantly. This is likely due to the smaller scope of the permutation space (O(n2n-1 )) compared to the graphical space (O(n!2C2n)), where GLS are more likely to find the optimal solution, with less sensitive to the choice of score function. Notably, the computation cost of the GrowShrink Tree (GST) for BDeu is prohibitively high, preventing BOSS from identifying a DAG in test datasets when n ≥ 40. For the continuous dataset, L-SFE maintains excellent generality in the EC space. However, in the permutation space, due to higher capability of BOSS, L-SFED1∗ and L-SFED2∗ slightly under-perform compared to BIC-gaussian, with only L-SFED3∗ achieving a marginal advantage.

Effectiveness: Does L-SFE improve the performance of CSL compared to state-of-the-art algorithms? Finally, we investigate whether the LLM generated score function, when combined with GLS in different search spaces, can outperform SOTA causal learning methods in tables 1 and 2. Overall, L-SFE achieves the lowest SHD in 8/10 settings on discrete datasets and in 7/10 settings on continuous datasets. Counterexamples are observed in some small-scale networks, such as ER-10, ER-20, and Exp-10. These occurrences are mainly because 5000 instances provide a sufficiently sample size for the BIC to yield optimal results. However, as the problem scale increases, the advantages of L-SFE become more pronounced. For discrete datasets, L-SFE in E space identifies the most optimal DAGs among the eight algorithms, and achieves similar results on continuous datasets in O space. This observation aligns with our understanding that graphical-based GLS are more suited for discrete datasets, whereas permutation-based GLS can explore causal structures in larger steps, making them more effective for SEM causal models. For further results, including F1arr and F1adj scores, please refer to Supplementary material 2.2.

5.3 Other Experiments

Case Study. We present a case study using data from the real world COVID-19 pandemic in the UK 3 . The data set comprises 866 samples that encompass eight categories, including viral tests, infections, hospitalizations, and related factors. The expert-designed benchmark network includes 17 variables, 37 directed edges, and a maximum node degree of 10. The dataset is discretized using k-means and quartiles discretization, as introduced by [Constantinou et al., 2023]. A graphical accuracy analysis is provided in Table 3. Due to the absence of a predefined distribution, the expert-designed benchmark network is not a P-map for dataset. As a result, most existing algorithms fail to produce high F1 score graphs. Even then, the L-SFE-guided GLS approach can consistently identify more correct edges and accurately orient more arrows across different settings.

Ablation Study. We investigate howdifferent types of standard score function hints impact the performance of L-SFE under three configurations: 1) Seed + Prompt: The BIC is included in both the population seeds and the code snippets within the prompts; 2) Seed: The initialization code snippets provided only the basic inputs and outputs information, with the BIC placed in the seed; 3) Random Walk: BIC is absent from both the seed and the prompts. Table 3 evaluates their test performance on ER-20 and SF-20. As expected, L-SFE(Prompt + Seed) yields the best F1arr and F1adj , indicating that evolution helps refine the BIC to generate even better score functions, and adding the seed within the population also achieves this goal. However, allowing the LLM to randomly search the score function space leads to suboptimal results. Interestingly, when we extended the number of iterations for L-SFE(RandomWalk) to 10, L-SFE still learn the similar thought of BIC. This observation further supports that BIC represents an optimal solution within the score function space. Furthermore, we also analyze the performance of LSFE across varying sample sizes in Supplementary material 2.2.

3 https://bayesian-ai.eecs.qmul.ac.uk/bayesys/

6 Conclusion

In this paper, we introduce the L-SFE, which leverages LLM and EA to explore the score function space through a bi-level optimization framework. Three key insights from this work are concluded as follows:

• LLM can effectively learn valuable components, such as Dirichlet priors and relaxed penalties. These components are theoretically well-founded and align with human-like reasoning.
• EA are crucial for finding optimal solutions. Even in the absence of explicit guidance, LLM can discover similar principles through iterative refinement.
• The score function space is highly abstract and potentially multimodal, making best-of-n search a valuable approach for yielding improved results.

Although L-SFE focuses primarily on exploring the structure of score functions, additional work is required to optimize their parameters using LLM. Furthermore, future research could focus on guiding LLM in exploring more efficient causal search strategies or expanding the search space beyond conventional frameworks.