CauTion: LLM-Augmented Causal Discovery

• CauTion is a framework that integrates large language model domain knowledge with an ensemble of statistical algorithms to enable robust causal discovery from observational data.
• It employs a three-stage process—consensus filtering, trust-calibrated arbitration, and cycle repair—to resolve ambiguities and ensure the output is a valid directed acyclic graph.
• The framework achieves high accuracy and token efficiency by selectively incorporating LLM judgments only when algorithmic evidence is unreliable.

Searching arXiv for the cited paper and closely related causal-discovery work. CauTion is a framework for causal discovery from observational data that integrates LLM domain knowledge into an ensemble of statistical causal discovery algorithms through consensus filtering and LLM reliability estimation. It is designed for a setting in which purely statistical methods are limited by Markov equivalence classes and finite-sample sensitivity, while naïve LLM augmentation is vulnerable to LLM errors, algorithm-specific bias, and high token costs. The framework proceeds in three stages: consensus filtering by an algorithm ensemble, trust-calibrated arbitration between algorithmic and LLM evidence on disputed edges, and cycle repair to ensure that the final output is a valid directed acyclic graph (DAG) (Peng et al., 2 Jun 2026).

1. Problem setting and rationale

Causal discovery from purely observational data is constrained by two structural difficulties. First, statistical methods such as constraint-based, score-based, and continuous-optimization approaches can recover only a Markov equivalence class rather than a unique DAG. Any two DAGs in the same Markov equivalence class induce identical distributions over the data, so edge orientations cannot be fully identified without external background knowledge. Second, finite-sample errors in conditional-independence tests or local score evaluations can propagate through the search process, producing cascading mistakes in skeleton recovery and orientation decisions (Peng et al., 2 Jun 2026).

LLMs offer a complementary source of information because they encode broad semantic and causal priors. Prompted with variable names and brief domain descriptions, they can often suggest whether variables are plausibly linked by a direct causal mechanism and may also infer directionality. However, direct reliance on LLM judgments for all variable pairs introduces two problems emphasized in the framework description: the pairwise search space scales as $O(n^2)$, increasing both token cost and decision burden on larger graphs, and LLM judgments remain imperfect, so accepting them without calibration makes the final graph vulnerable to isolated LLM errors (Peng et al., 2 Jun 2026).

CauTion is therefore organized around a selective-use principle: resolve as many edges as possible using cross-algorithm agreement, estimate the relative reliability of algorithms and LLMs without human annotations, and consult the LLM only where algorithmic evidence is unreliable. This design is intended to exploit LLM knowledge without allowing the LLM to dominate the entire discovery process (Peng et al., 2 Jun 2026).

2. Three-stage architecture

The framework is explicitly divided into three stages, each targeting a different failure mode in LLM-augmented causal discovery (Peng et al., 2 Jun 2026).

Stage	Mechanism	Purpose
1	Consensus filtering	Resolve edges agreed on by the statistical ensemble
2	Trust-calibrated arbitration	Use LLM judgments only on disputed edges with unreliable algorithmic evidence
3	Cycle repair	Guarantee that the final graph is acyclic

In the first stage, a small ensemble of statistical algorithms is run on the data. In the second stage, the framework estimates the relative reliability of the ensemble and the LLM for edge existence and edge direction, then uses trust-weighted voting only on unresolved cases. In the third stage, any cycles introduced by aggregation are repaired through a constrained post-processing loop (Peng et al., 2 Jun 2026).

This decomposition also gives CauTion a specific operational interpretation. Rather than treating the LLM as a universal oracle or reducing the graph search to a single data-centric procedure, it uses agreement structure to partition the problem into easy cases and hard cases. A plausible implication is that the framework treats LLM input as a scarce and fallible source of domain knowledge rather than as a replacement for statistical evidence.

3. Consensus filtering by a statistical ensemble

CauTion begins by running the ensemble

$$\mathcal A=\{\mathrm{PC},\mathrm{GES},\mathrm{CAMML}\}$$

on the observational dataset (Peng et al., 2 Jun 2026). For each unordered pair $e=\{i,j\}$, each algorithm $A$ casts one of four votes,

$$v_A(e)\in\{\text{fwd }(i\to j),\ \text{rev }(j\to i),\ \text{undir }(i-j),\ \text{none}\}.$$

These votes are aggregated into weighted counts. For example, the forward and reverse weights are defined as

$$w_{\rm fwd}(e)=\sum_{A\in\mathcal A}\Bigl[\mathbf 1(v_A(e)=\text{fwd})+0.5\,\mathbf 1(v_A(e)=\text{undir})\Bigr],$$

$$w_{\rm rev}(e)=\sum_A\Bigl[\mathbf 1(v_A(e)=\text{rev})+0.5\,\mathbf 1(v_A(e)=\text{undir})\Bigr].$$

A pair is placed in the consensus set when one of the weighted counts reaches unanimity: $$E_{\rm consensus}= \bigl\{e:\max\{w_{\rm fwd}(e),w_{\rm rev}(e),w_{\rm none}(e)\}=|\mathcal A|\bigr\}.$$ The use of a half-weight for undirected votes allows partially oriented evidence to contribute without being treated as a full directional claim (Peng et al., 2 Jun 2026).

On six bnlearn benchmarks—Cancer ($n=5$), Insurance ($27$), Water ($$\mathcal A=\{\mathrm{PC},\mathrm{GES},\mathrm{CAMML}\}$$0), Alarm ($$\mathcal A=\{\mathrm{PC},\mathrm{GES},\mathrm{CAMML}\}$$1), Barley ($$\mathcal A=\{\mathrm{PC},\mathrm{GES},\mathrm{CAMML}\}$$2), and Win95pts ($$\mathcal A=\{\mathrm{PC},\mathrm{GES},\mathrm{CAMML}\}$$3)—this consensus step resolves over $$\mathcal A=\{\mathrm{PC},\mathrm{GES},\mathrm{CAMML}\}$$4 of edges with near-perfect accuracy, leaving only a small disputed set even for the largest graph. The abstract characterizes this stage as resolving up to $$\mathcal A=\{\mathrm{PC},\mathrm{GES},\mathrm{CAMML}\}$$5 of edges on which algorithms agree. For Win95pts, the disputed set is approximately $$\mathcal A=\{\mathrm{PC},\mathrm{GES},\mathrm{CAMML}\}$$6 of all $$\mathcal A=\{\mathrm{PC},\mathrm{GES},\mathrm{CAMML}\}$$7 pairs (Peng et al., 2 Jun 2026).

The methodological significance of this stage is twofold. First, it sharply reduces the number of LLM queries. Second, it converts cross-algorithm agreement into a high-confidence pseudoground-truth that can later be used for calibration.

4. Trust-calibrated arbitration and reliability estimation

The second stage addresses disputed edges by estimating, directly from the data and without human labels, the relative reliability of the algorithm ensemble and the LLM on two separate subproblems: edge existence and edge direction (Peng et al., 2 Jun 2026).

For LLM calibration, the framework samples $$\mathcal A=\{\mathrm{PC},\mathrm{GES},\mathrm{CAMML}\}$$8 pairs from the consensus set, stratified by the soft edge-existence probability inferred from ensemble votes. The LLM is prompted twice per pair, once in each variable ordering, to answer whether a direct causal link exists and, if so, which variable causes which. Comparing these responses against unanimous-consensus pseudoground-truth yields two accuracy estimates, $$\mathcal A=\{\mathrm{PC},\mathrm{GES},\mathrm{CAMML}\}$$9 and $e=\{i,j\}$0. Existence accuracy is then reweighted to the disputed-edge prior: $e=\{i,j\}$1 where $e=\{i,j\}$2 is the fraction of disputed pairs that contain any edge (Peng et al., 2 Jun 2026).

For algorithm calibration, CauTion applies a leave-one-out procedure. Each algorithm’s vote on each consensus edge is compared against the majority vote of the remaining algorithms, yielding per-algorithm accuracies for existence and direction. These are then aggregated to ensemble-level estimates $e=\{i,j\}$3 and $e=\{i,j\}$4 (Peng et al., 2 Jun 2026).

Trust weights are computed by offsetting accuracies relative to the $e=\{i,j\}$5 random baseline and sharpening with exponent $e=\{i,j\}$6: $e=\{i,j\}$7 with an analogous definition for $e=\{i,j\}$8 (Peng et al., 2 Jun 2026).

These trust weights do not merely interpolate between algorithmic and LLM votes. They also define whether the LLM is consulted at all. The framework first constructs an algorithmic existence score for a disputed pair and examines its margin. Only when the margin falls below a threshold determined by $e=\{i,j\}$9 is the LLM invoked. The final existence score is then a trust-weighted blend of algorithmic evidence and the LLM’s binary judgment: $A$0 with a parallel construction for directionality (Peng et al., 2 Jun 2026).

This arbitration mechanism embodies the framework’s name: it aims to know when to trust the LLM rather than assuming either universal reliability or universal unreliability. A plausible implication is that the framework treats reliability estimation as an adaptive resource-allocation problem, not just a voting heuristic.

5. Cycle repair and acyclicity guarantees

Because consensus filtering and trust-weighted arbitration can still yield a graph containing directed cycles, CauTion adds a third stage that guarantees a valid DAG (Peng et al., 2 Jun 2026).

Whenever a cycle is detected, the framework marks as repair candidates those edges on the cycle whose algorithmic vote margins remain below the same thresholds used for arbitration. For each such candidate, the LLM is prompted once to choose between two repair operations: FLIP or REMOVE. The chosen edit is applied immediately, and the cycle-detection loop repeats until the graph becomes acyclic (Peng et al., 2 Jun 2026).

If LLM-guided repair opportunities are exhausted and cycles still remain, the framework falls back to a deterministic rule: delete the edge with the lowest algorithmic existence-score margin. This fallback ensures termination with a valid DAG even when LLM-assisted repair is insufficient (Peng et al., 2 Jun 2026).

The cycle-repair stage is methodologically important because it separates local arbitration from global graph validity. Consensus and arbitration operate edgewise; acyclicity is a graph-level constraint. CauTion therefore treats global consistency as a dedicated post-processing problem rather than assuming that local decisions will automatically yield a valid causal graph.

6. Benchmarks, performance, and robustness

CauTion was evaluated on six bnlearn Bayesian-network benchmarks spanning graph sizes from $A$1 to $A$2 and edge counts from $A$3 to $A$4. It was compared against eight data-driven methods—PC, FCI, GES, BOSS, GRaSP, CAMML, NOTEARS-MLP, and DAGMA—and five LLM-augmented approaches—LLM-BFS, Causal-LLM, SCP, ET-MCMC, and LLM-MEC—plus a simple correlation-based LLM baseline. Evaluation used Structural Hamming Distance (SHD), directed-edge $A$5, and Structural Intervention Distance (SID) (Peng et al., 2 Jun 2026).

Across these datasets, CauTion achieved the lowest SHD and the highest $A$6 on every dataset. On Win95pts ($A$7), it obtained $A$8, compared with $A$9–$$v_A(e)\in\{\text{fwd }(i\to j),\ \text{rev }(j\to i),\ \text{undir }(i-j),\ \text{none}\}.$$0 for the best existing LLM-augmented methods. The reported gains were largest on larger graphs, matching the framework’s emphasis on selective querying and consensus-based pruning (Peng et al., 2 Jun 2026).

The framework was also evaluated across six different LLM backends, including Claude Sonnet 4.6, GPT-5.2, and Llama-3.3-70B. Its performance varied little across backends, whereas global LLM methods exhibited SHD swings of hundreds of edges. This is presented as evidence of robustness to LLM errors and backend variation (Peng et al., 2 Jun 2026).

Token efficiency is another explicit feature. Because CauTion queries the LLM only for the disputed set, plus a fixed calibration budget of $$v_A(e)\in\{\text{fwd }(i\to j),\ \text{rev }(j\to i),\ \text{undir }(i-j),\ \text{none}\}.$$1 edges and occasional cycle-repair prompts, its total usage on Win95pts is under $$v_A(e)\in\{\text{fwd }(i\to j),\ \text{rev }(j\to i),\ \text{undir }(i-j),\ \text{none}\}.$$2 tokens. The description contrasts this with global LLM methods that routinely exceed one million tokens (Peng et al., 2 Jun 2026).

Taken together, these results position CauTion as a selective-integration framework rather than a monolithic LLM-discovery method. Its core empirical claim is not simply that LLMs help causal discovery, but that calibrated and restricted use of LLM judgments can improve graph quality while limiting both token cost and sensitivity to LLM failure.