Technical Summary: CauTion – Knowing When to Trust LLMs for Ensemble Causal Discovery

1. Problem Statement

Causal discovery from observational data faces fundamental limitations when relying solely on statistical methods. Constraint-based approaches (e.g., PC, FCI) are sensitive to finite sample sizes and can only recover a Markov Equivalence Class (MEC), leaving many edge directions unresolved. Score-based and continuous optimization methods (e.g., GES, NOTEARS) often suffer from local optima, sensitivity to hyperparameters, and distributional assumptions that may not hold in practice.

To address these gaps, Large Language Models (LLMs) have been introduced to provide domain knowledge and semantic priors. However, existing LLM-augmented methods suffer from three critical flaws:

1. High Token Costs and Error Exposure: Global methods query LLMs over all $O(n^2)$ variable pairs, leading to prohibitive costs and exposing the entire graph to potential LLM hallucinations.
2. Error Amplification: Undirected-only methods (e.g., LLM-MEC) restrict queries to unresolved edges but suffer from error accumulation; a single LLM error can trigger a cascade of incorrect orientations via Meek rules.
3. Algorithmic Bias and Unverified Reliability: Most methods rely on a single data-centric algorithm, making results vulnerable to specific algorithmic biases. Furthermore, they often incorporate LLM knowledge without verifying its reliability against the specific dataset, treating the LLM as an infallible oracle.

2. Methodology: The CauTion Framework

The authors propose CauTion, a framework that integrates LLM domain knowledge into an ensemble of statistical causal discovery algorithms through consensus filtering and trust-calibrated arbitration. The pipeline operates in three distinct stages:

Stage 1: Algorithm Ensemble and Consensus Filtering

The framework aggregates outputs from multiple causal discovery algorithms (specifically PC, GES, and CAMML in the experiments). For each variable pair, votes are accumulated to classify edges into two sets:

• Consensus Set ($\mathcal{C}$): Pairs where all algorithms agree (e.g., all vote for $A \to B$ or all vote for no edge). These edges are resolved directly without LLM intervention. On the Win95pts dataset ($n=76$), this step resolves approximately 96.6% of edges with 99.6% accuracy, serving as high-confidence pseudo-ground truth.
• Disputed Set ($\mathcal{D}$): Pairs where algorithms disagree or are uncertain. These are passed to the arbitration stage.

Stage 2: Trust-Calibrated Arbitration

This stage determines when to trust the LLM by estimating the relative reliability of the LLM and the algorithm ensemble without external annotations.

• Reliability Estimation:
  • LLM Calibration: The LLM is queried on a stratified sample of the Consensus Set (acting as proxy labels) to estimate its accuracy on edge existence and direction tasks.
  • Algorithm Calibration: A Leave-One-Out (LOO) evaluation is performed where each algorithm's prediction is compared against the majority vote of the remaining algorithms to estimate its accuracy.
• Trust Weight Calculation: Based on the calibrated accuracies, trust weights ($\alpha_{exist}$ and $\alpha_{dir}$) are computed. These weights quantify the relative reliability of the LLM versus the algorithm ensemble. If both sources perform near chance levels, the system defaults to equal weighting.
• Calibrated Voting: For edges in the Disputed Set:
  • The system first checks if the algorithm ensemble's vote margin exceeds a threshold derived from the trust weights. If the algorithms are confident, the LLM is skipped.
  • If the algorithm margin is low, the LLM is queried. The final decision is a trust-weighted aggregation of the algorithm vote and the LLM response.
  • Direction queries employ a double-framing technique (querying $A \to B$ and $B \to A$) to mitigate position bias, retaining only consistent answers.

Stage 3: Cycle Repair

The voting process produces a directed graph that may contain cycles. A cycle repair step ensures the final output is a valid Directed Acyclic Graph (DAG).

• Cycles are detected, and edges within them are evaluated based on their algorithmic uncertainty margins.
• Uncertain edges are presented to the LLM with the cycle path and variable interpretations. The LLM is asked to either Flip the edge direction or Remove the edge.
• If the LLM cannot resolve the cycle or no actions remain, edges with the lowest existence probability are removed iteratively until the graph is acyclic.

3. Key Contributions

The paper outlines three primary contributions:

1. Conceptual: It reframes LLM-augmented causal discovery as a trust-calibrated arbitration problem. This perspective treats LLMs and statistical algorithms as complementary but imperfect evidence sources, emphasizing that LLM knowledge should only be integrated after its reliability is calibrated against an ensemble of statistical methods.
2. Methodological: The proposal of CauTion, a framework that combines cross-algorithm consensus, reliability-aware LLM arbitration, and DAG-preserving cycle repair. It uniquely restricts LLM queries to disputed edges and downweights LLM influence when its estimated domain accuracy is low.
3. Empirical: Demonstration that CauTion consistently outperforms both data-centric and LLM-augmented baselines across six benchmarks. The performance gains are shown to scale with graph size, and the method exhibits strong robustness to variations in LLM backend selection.

4. Experimental Results

Experiments were conducted on six datasets from the bnlearn repository (Cancer, Insurance, Water, Alarm, Barley, Win95pts) with variable counts ranging from 5 to 76.

• Performance: CauTion achieved the lowest Structural Hamming Distance (SHD) and highest F1 scores across all datasets. On the largest dataset, Win95pts ($n=76$), CauTion achieved an SHD of 27, significantly outperforming the second-best LLM-augmented baseline (SCP) which had an SHD of 63.
• Scalability: The performance gap between CauTion and baselines widened as the graph scale increased, validating the framework's design for handling larger, more complex structures.
• Robustness: CauTion demonstrated superior stability across six different LLM backends (ranging from weak models like Llama-3.1-8B to strong models like Claude-Sonnet-4.6). In contrast, fully LLM-driven methods (e.g., LLM-BFS, CORR-LLM) showed severe performance degradation and high variance when weaker models were used.
• Efficiency: By filtering out 90%+ of edges via consensus, CauTion drastically reduced token usage compared to global LLM methods. For Win95pts, CauTion used ~26k tokens, whereas global methods like SCP used over 3 million tokens.

5. Significance and Claims

The paper claims that CauTion addresses the fundamental trade-off between the high error exposure of global LLM methods and the bias of single-algorithm statistical methods. By introducing a trust-calibrated mechanism, the framework ensures that LLM domain knowledge is leveraged only when statistical evidence is insufficient and when the LLM is empirically calibrated to be reliable for the specific task.

The authors state that this work advances causal discovery methodology, offering a path toward more reliable structure recovery in domains such as medicine, epidemiology, and economics. They acknowledge limitations, noting that the current ensemble uses a fixed set of three algorithms and that experiments were conducted on discrete-variable benchmarks, leaving generalization to continuous or mixed-type data for future validation.