📊 statistics

Efficient Causal Graph Discovery Using Large Language Models

This paper proposes a novel, efficient framework that leverages Large Language Models with a breadth-first search approach to discover full causal graphs using only a linear number of queries, thereby overcoming the quadratic query limitations of previous pairwise methods while achieving state-of-the-art performance on real-world datasets.

Original authors: Thomas Jiralerspong, Xiaoyin Chen, Yash More, Vedant Shah, Yoshua Bengio

Published 2026-04-06

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Thomas Jiralerspong, Xiaoyin Chen, Yash More, Vedant Shah, Yoshua Bengio

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a detective trying to solve a massive mystery: Who caused what?

In the world of data science, this is called Causal Discovery. You have a list of variables (like "Smoking," "Lung Cancer," "Pollution," or "Coughing"), and you need to draw a map showing exactly which one leads to the other. This map is called a Causal Graph.

Traditionally, to draw this map, scientists had to crunch millions of numbers from observational data (like hospital records) using complex math. It was slow, expensive, and often required huge amounts of data.

Recently, researchers tried using Large Language Models (LLMs)—the same AI brains behind chatbots—to do this. They asked the AI, "Does Smoking cause Lung Cancer?" and "Does Pollution cause Coughing?" one pair at a time.

Here is the problem with that old approach:
If you have 10 variables, you have to ask about 45 pairs. If you have 100 variables, you have to ask about 4,950 pairs! If you have 200 variables (like in a real medical study), you'd need to ask 20,000 questions. That's like trying to map a whole city by asking a tourist about every single street intersection one by one. It takes forever and costs a fortune.

The New Solution: The "Breadth-First Search" Detective

The authors of this paper, Thomas Jiralerspong and his team, came up with a much smarter way to use the AI. Instead of asking about every possible pair, they treat the AI like a detective exploring a city using a Breadth-First Search (BFS).

Here is how their method works, using a simple analogy:

1. The Setup: Finding the "Roots"

Imagine you are exploring a forest. You need to know where the trees start.

The AI's Job: The researchers first ask the AI: "Which of these variables are the 'roots'? Which ones are not caused by anything else?"
The Result: The AI identifies the starting points (e.g., "Smoking" and "Genetics" might be roots because nothing in the list causes them). These go into a "To-Do List" (a queue).

2. The Expansion: "Who do you affect?"

Now, the detective picks the first item from the "To-Do List" (let's say "Smoking").

The AI's Job: Instead of asking "Does Smoking cause X?", "Does Smoking cause Y?", "Does Smoking cause Z?" separately, the AI is asked one single question: "Given that we are at 'Smoking', list ALL the variables that Smoking directly causes."
The Magic: The AI might say: "Smoking causes Bronchitis and Lung Cancer."
The Result: In just one question, the detective has found two connections!

3. The Safety Check: No Loops Allowed

Causal graphs must be Directed Acyclic Graphs (DAGs). In plain English: No time travel allowed. You can't have a situation where A causes B, and B causes A (that would be a paradox).

Before adding a new connection to the map, the system checks: "If we draw a line from Smoking to Lung Cancer, does it create a weird loop?" If yes, it's discarded. If no, it's added to the map.

4. Repeat Until Done

The newly discovered variables (Bronchitis, Lung Cancer) are added to the "To-Do List." The detective picks the next one, asks "Who do you cause?", and repeats the process until every variable has been visited.

Why is this a Big Deal?

1. It's Lightning Fast (Linear vs. Quadratic)

Old Way: If you have $N$ variables, you need to ask roughly $N \times N$ questions. (If you double the variables, you need four times the questions).
New Way: You only need to ask $N$ questions (one for each variable). (If you double the variables, you only need twice the questions).
Analogy: The old way is like checking every single door in a building to see if it's locked. The new way is like walking down the hallway and asking the person in each room, "Who else is in this building?"

2. It Doesn't Need Numbers (Mostly)
The AI uses its "common sense" and training data (like knowing that smoking is bad for lungs) to figure this out. It doesn't need a spreadsheet of 10,000 patient records. It just needs the names and descriptions of the variables.

Note: If you do have data (like correlation numbers), you can feed it to the AI to make it even smarter, but it's not strictly necessary.

3. It Works on Huge Problems
The researchers tested this on three maps:

Asia: A tiny map (8 variables). The AI crushed it.
Child: A medium map (20 variables). The AI did great.
Neuropathic Pain: A massive map with 221 variables and 770 connections.
- The old "pair-by-pair" method would have needed 24,310 questions. It was impossible to run.
- The new method only needed 221 questions. It successfully mapped the complex relationships between nerves and symptoms, outperforming all traditional math-based methods.

The Bottom Line

This paper introduces a way to use AI to draw "Cause-and-Effect" maps much faster and cheaper than before. Instead of asking the AI millions of tiny questions, they ask it to explore the map step-by-step, like a detective walking through a city, discovering connections as they go.

It's a shift from brute force (checking everything) to smart exploration (following the trail), making it possible to solve complex causal mysteries in fields like medicine, economics, and climate science that were previously too big to handle.

1. Problem Statement

Causal Graph Discovery aims to identify the causal relationships between variables, typically represented as a Directed Acyclic Graph (DAG). Traditional methods rely on numerical observational data and statistical optimization (e.g., PC algorithm, NOTEARS, GES), which can be computationally expensive and require large datasets.

Recent approaches have utilized Large Language Models (LLMs) to infer causal relationships using only metadata (variable names and descriptions) rather than raw data. However, existing LLM-based methods suffer from a critical scalability bottleneck:

Pairwise Query Approach: Previous works query the LLM for every possible pair of variables to determine if $A \to B$ , $B \to A$ , or no relationship exists.
Complexity: This requires $O(n^2)$ queries for $n$ variables. For large graphs (e.g., hundreds of nodes), this becomes computationally impractical and prohibitively expensive in terms of API costs and time.

The paper addresses the need for a scalable, data-efficient method to discover full causal graphs using LLMs without relying on exhaustive pairwise comparisons.

2. Methodology

The authors propose a novel framework that leverages the Breadth-First Search (BFS) algorithm to construct the causal graph, reducing query complexity from quadratic to linear.

Core Algorithm: BFS-Based Graph Construction

Instead of querying all pairs, the method iteratively expands the graph by querying the LLM about the effects of a specific node. The process consists of three stages:

Initialization Stage:
- The LLM is prompted with descriptions of all variables to identify root nodes (variables not caused by any other variable in the graph).
- These root nodes are added to a BFS queue.
- Prompt Example: "Identify the variable(s) that are unaffected by any other variables."
Expansion Stage:
- The algorithm dequeues a node (the "current node") and prompts the LLM to identify all variables caused by this current node.
- The prompt includes the current graph structure (edges already discovered) to maintain context.
- Prompt Example: "Select the variables that are caused by [Current Node]."
Insertion & Cycle Check Stage:
- For every edge predicted by the LLM, the system performs a cycle check (using Depth-First Search) before adding the edge to the graph.
- If adding an edge creates a cycle, it is discarded to strictly enforce the DAG constraint.
- Newly connected nodes are added to the BFS queue if they haven't been visited yet.

This process repeats until the queue is empty. The total query complexity is $O(n)$ , as each node is visited and expanded exactly once.

Incorporating Observational Data

The framework is flexible and can incorporate observational data when available to boost performance.

Mechanism: Pearson correlation coefficients between the current node and other variables are appended to the expansion prompt.
Hypothesis: While correlation does not imply causation, providing statistical context helps the LLM resolve ambiguities when its internal knowledge is insufficient.

3. Key Contributions

Linear Query Complexity: The primary contribution is the shift from $O(n^2)$ pairwise queries to an $O(n)$ BFS-based approach, making full graph discovery feasible for large-scale problems (e.g., 200+ nodes).
DAG Constraint Enforcement: The method explicitly includes a cycle-checking mechanism during the insertion stage, ensuring the output is a valid Directed Acyclic Graph, a common failure point in generative LLM tasks.
Data Efficiency: The method functions effectively without any observational data, relying solely on the LLM's internal knowledge base (metadata). It also demonstrates how to seamlessly integrate statistical data when available.
Scalability: It is the first LLM-based method demonstrated to handle large causal graphs (221 nodes) where previous pairwise methods fail due to query limits.

4. Experimental Results

The method was evaluated on three datasets of varying sizes:

Asia: Small graph (8 nodes, 8 edges).
Child: Medium graph (20 nodes, 25 edges).
Neuropathic Pain: Large graph (221 nodes, 770 edges).

Performance Metrics: Accuracy, F-Score, and Normalized Hamming Distance (NHD) Ratio.

Key Findings:

Small Graphs (Asia): The proposed method achieved State-of-the-Art (SOTA) results (F-Score: 0.93) without any observational data, outperforming statistical methods (GES, PC, NOTEARS) and the pairwise LLM baseline. Interestingly, adding correlation data did not improve performance here, suggesting the LLM's internal knowledge was sufficient.
Medium Graphs (Child): The method with 10,000 samples of observational data achieved the best performance (F-Score: 0.63). Notably, the method without data still performed competitively, matching the best statistical baseline (GES) while using zero observational data.
Large Graphs (Neuropathic Pain):
- Pairwise methods and GES were intractable (requiring ~24,310 queries or edges to check).
- Statistical methods (PC) ran out of memory.
- The proposed method was the only approach capable of running on this scale, achieving a reasonable F-Score (0.35) and NHD Ratio (0.64), significantly outperforming the "worst-case" baseline.

5. Significance and Implications

Scalability: This work breaks the barrier preventing LLMs from being used for large-scale causal discovery, moving beyond toy examples to real-world, complex systems.
Cost and Time Efficiency: By reducing queries from quadratic to linear, the method drastically lowers the cost and time required for discovery, making it accessible for domains with limited resources.
Hybrid Potential: The framework demonstrates that LLMs can serve as a powerful prior or standalone engine for causal reasoning, bridging the gap between knowledge-based reasoning and statistical inference.
Domain Applicability: The ability to discover causal graphs using only variable descriptions makes this approach highly applicable to fields where data is scarce, proprietary, or difficult to quantify, but expert knowledge (metadata) is available.

6. Limitations & Future Work

Limitations: The method relies on the quality of the LLM's training data; it may struggle with novel domains not well-represented in training corpora. It also lacks transparency (black-box nature).
Future Directions:
- Combining LLM priors with statistical validation (hybrid approaches).
- Utilizing advanced prompting strategies like "Tree of Thoughts" for complex reasoning.
- Evaluating performance across different LLM architectures and sizes.