[Paper Review] Efficient Causal Graph Discovery Using Large Language Models

Thomas Jiralerspong, Xiaoyin Chen+3 more|arXiv (Cornell University)|Feb 2, 2024

Biomedical Text Mining and Ontologies5 citations

TL;DR

Proposes a BFS-based framework that uses LLMs to perform full causal graph discovery with linear query complexity, optional observational data, and state-of-the-art results on multiple graphs.

ABSTRACT

We propose a novel framework that leverages LLMs for full causal graph discovery. While previous LLM-based methods have used a pairwise query approach, this requires a quadratic number of queries which quickly becomes impractical for larger causal graphs. In contrast, the proposed framework uses a breadth-first search (BFS) approach which allows it to use only a linear number of queries. We also show that the proposed method can easily incorporate observational data when available, to improve performance. In addition to being more time and data-efficient, the proposed framework achieves state-of-the-art results on real-world causal graphs of varying sizes. The results demonstrate the effectiveness and efficiency of the proposed method in discovering causal relationships, showcasing its potential for broad applicability in causal graph discovery tasks across different domains.

Motivation & Objective

• Motivate the use of LLMs for causal graph discovery beyond pairwise queries.
• Develop a BFS-based framework that guarantees DAG structure with O(n) queries.
• Demonstrate how observational data can be incorporated to improve performance.
• Evaluate on real-world graphs of varying sizes and compare to numerical baselines and pairwise LLM methods.

Proposed method

• Use a BFS-inspired prompting strategy to expand nodes and discover edges with linear query complexity O(n).
• Three-stage process: initialization to find independent variables, expansion to identify children of a node, and insertion with cycle checks to maintain DAG constraints.
• Cycle checking via DFS on the current predicted graph prior to edge insertion to ensure acyclicity.
• Incorporate observational statistics optionally by appending Pearson correlation coefficients to prompts to potentially improve guidance for edge discovery.
• Experiment with GPT-4 API, using topological-order BFS to traverse variables and build the DAG.
• Compare against classical causal discovery methods (PC, GES, NOTEARS, DAGMA) and previous LLM-based pairwise approaches, with additional experiments using observational data (1000 and 10000 samples).
• Provide code and details on experiment setup, including graphs Asia, Child, and Neuropathic Pain, and data regimes (100, 1000, 10000 samples).

Figure 1: Proposed framework for full graph discovery with LLMs

Experimental results

Research questions

• RQ1 Can LLMs perform full causal graph discovery without quadratic pairwise querying?
• RQ2 Does a BFS-based prompting strategy ensure DAG-consistent full graphs with linear query complexity?
• RQ3 How does incorporating observational data via prompt statistics affect LLM-based causal discovery?
• RQ4 How does the proposed method scale to large graphs compared to traditional statistical methods and prior LLM-based approaches?
• RQ5 What is the performance of the method across small, medium, and large real-world causal graphs?

Key findings

• The proposed BFS-based LLM framework achieves O(n) query complexity, versus O(n^2) for pairwise methods.
• On Asia (8 nodes, 8 edges), the method with no observational statistics achieves F score 0.93 and NHD ratio 0.067, outperforming baselines.
• On Child (20 nodes, 25 edges), observational statistics with 10000 samples yield best F score 0.63 and NHD ratio 0.37; without statistics it is competitive with GES/10000 samples.
• On Neuropathic Pain (221 nodes, 770 edges), the proposed method without observational data achieves F score 0.351 and NHD ratio 0.643, while other methods fail or are intractable.
• Pairwise LLM methods become impractical for large graphs (e.g., Neuropathic Pain) due to combinatorial query growth; numerical methods struggle with memory or scalability on large graphs.
• Observational data can improve performance, with notable gains on larger graphs, but effects vary by graph size and data regime.