Causal Chain of Thought: Mitigating Attention Leakage via Context Isolation
Organizations: Università della Svizzera Italiana · ETH Zurich · University of Oxford
Abstract
Chain-of-thought (CoT) prompting elicits step-by-step rationales that appear to explain model predictions, but there is no guarantee these steps causally drive the final answer. In standard autoregressive generation, dense attention allows each reasoning step to attend to the full context, enabling the model to derive answers through direct lookup while generating plausible-sounding but post-hoc justifications. We introduce CausalCoT, an inference framework that enforces causal transparency by decomposing problems into a directed acyclic graph of atomic sub-questions, evaluating each node in an isolated context containing only its direct causal parents. This Markovian structure removes attention pathways to non-parent content, making downstream predictions sensitive to local reasoning rather than global context leakage. Across a diverse suite of reasoning benchmarks, CausalCoT preserves competitive predictive performance while improving explanation faithfulness, increasing causal sensitivity of intermediate reasoning steps by an average of 12 and up to 21 percentage points over the strongest baseline. The structured dependency graph additionally enables test-time interventions in which local corrections propagate predictably through downstream reasoning.
Figures & tables
| Model | Method | MedQA | GPQA | Counterfactuals | Overall Average | ||||
|---|---|---|---|---|---|---|---|---|---|
| Acc (%) | Pass@3 (%) | Acc (%) | Pass@3 (%) | Acc (%) | Pass@3 (%) | Acc (%) | Pass@3 (%) | ||
| Gemma 4B | CoT | 49.00 0.47 | 66.00 3.35 | 31.17 0.59 | 50.50 3.54 | 54.34 0.33 | 62.42 3.77 | 44.84 0.28 | 59.64 2.05 |
| FactoredDecomp | 46.50 0.62 | 55.50 3.51 | 24.17 0.72 | 36.00 3.39 | 56.57 0.82 | 64.85 3.72 | 42.41 0.42 | 52.12 2.05 | |
| CausalCoT | 44.67 0.72 | 70.50 3.22 | 28.83 0.36 | 56.50 3.51 | 46.67 1.48 | 68.48 3.62 | 40.06 0.56 | 65.16 1.99 | |
| Gemma 12B | CoT | 68.83 0.76 | 81.00 2.77 | 39.67 2.00 | 60.00 3.46 | 53.13 0.87 | 61.82 3.78 | 53.88 0.77 | 67.61 1.94 |
| FactoredDecomp | 62.00 0.85 | 70.50 3.22 | 32.50 1.25 | 46.00 3.52 | 52.93 0.66 | 63.64 3.74 | 49.14 0.55 | 60.05 2.02 | |
| Model | Method | MedQA | GPQA | Counterfactuals | Overall Average | ||||
|---|---|---|---|---|---|---|---|---|---|
| Answer Change (%) | Pred. Change (%) | Answer Change (%) | Pred. Change (%) | Answer Change (%) | Pred. Change (%) | Answer Change (%) | Pred. Change (%) | ||
| Gemma 4B | CoT | 44.83 2.42 | 54.50 2.02 | 56.84 2.37 | 47.91 2.93 | 30.38 1.46 | 63.29 2.19 | 44.02 1.23 | 55.23 1.39 |
| FactoredDecomp | 15.67 0.83 | 60.17 2.92 | 23.83 1.59 | 52.67 1.59 | 13.08 1.69 | 75.11 2.95 | 17.53 0.82 | 62.65 1.48 | |
| CausalCoT | 59.30 0.29 | 52.93 1.17 | 65.67 1.91 | 49.60 2.06 | 33.43 1.76 | 75.59 1.16 | 52.80 0.87 | 59.37 0.88 | |
| Gemma 12B | CoT | 36.67 1.59 | 57.00 1.26 | 55.21 1.10 | 50.83 1.35 | 21.31 1.17 | 70.25 1.83 | 37.73 0.75 | 59.36 0.87 |
| FactoredDecomp | 41.33 1.59 | 67.67 1.36 | 45.67 1.74 | 59.17 3.44 | 26.37 2.38 | 74.05 0.97 | 37.79 1.12 | 66.96 1.28 | |
Appendix figures & tables26 assets
Supplementary material from the paper’s appendix.
Appendix
| Model | Method | Number of interventions | |||
|---|---|---|---|---|---|
| After 0 | After 1 | After 2 | After 3 | ||
| Gemma 4B | CoT | 49.00 0.47 | 80.50 1.31 | 88.33 0.49 | 89.17 0.36 |
| FactoredDecomp | 46.50 0.62 | 56.17 0.95 | 62.83 0.36 | 66.67 0.49 | |
| CausalCoT | 45.00 0.62 | 67.50 0.85 | 82.17 1.83 | 87.67 0.95 | |
| Gemma 12B | CoT | 68.83 0.76 | 91.50 0.94 | 93.67 1.11 | 94.00 1.03 |
| FactoredDecomp | 62.00 0.85 | 83.33 1.21 | 87.00 1.70 | 88.67 1.30 | |
| Model | Method | Number of interventions | |||
|---|---|---|---|---|---|
| After 0 | After 1 | After 2 | After 3 | ||
| Gemma 4B | CoT | 31.17 0.59 | 51.67 0.36 | 60.50 0.85 | 63.00 1.03 |
| FactoredDecomp | 24.17 0.72 | 35.00 0.41 | 42.17 0.95 | 46.33 0.76 | |
| CausalCoT | 29.50 0.24 | 47.50 0.94 | 60.50 1.41 | 69.67 1.06 | |
| Gemma 12B | CoT | 39.67 2.00 | 60.50 1.84 | 67.17 1.52 | 70.33 0.95 |
| FactoredDecomp | 32.50 1.25 | 49.00 1.03 | 55.17 0.72 | 59.50 1.08 | |
| Model | Method | Number of interventions | |||
|---|---|---|---|---|---|
| After 0 | After 1 | After 2 | After 3 | ||
| Gemma 4B | CoT | 54.22 0.34 | 64.14 0.96 | 68.14 0.86 | 68.35 0.79 |
| FactoredDecomp | 56.57 0.82 | 60.20 0.82 | 63.23 1.32 | 64.85 0.76 | |
| CausalCoT | 49.90 2.29 | 60.20 1.62 | 64.65 2.40 | 67.88 1.78 | |
| Gemma 12B | CoT | 52.74 0.96 | 61.18 1.35 | 63.50 1.05 | 63.50 1.05 |
| FactoredDecomp | 52.93 0.66 | 61.01 0.59 | 63.43 0.66 | 64.65 0.92 | |
| Model | MedQA | GPQA | Counterfactuals | Overall Average |
|---|---|---|---|---|
| Malformed (%) | Malformed (%) | Malformed (%) | Malformed (%) | |
| Gemma 4B | 0.00 0.00 | 0.33 0.24 | 0.00 0.00 | 0.11 0.08 |
| Gemma 12B | 0.17 0.17 | 0.50 0.29 | 0.21 0.21 | 0.29 0.13 |
| GPT-5.4 Nano | 3.83 0.78 | 4.00 0.80 | 3.59 0.86 | 3.81 0.47 |
| GPT-5.4 Mini | 0.00 0.00 | 0.33 0.24 | 0.00 0.00 | 0.11 0.08 |
| GPT-5.4 | 0.00 0.00 | 0.67 0.33 | 0.00 0.00 | 0.22 0.11 |
| Model | Method | MedQA | GPQA | ||
|---|---|---|---|---|---|
| Acc (%) | Answer Change (%) | Acc (%) | Answer Change (%) | ||
| Gemma 4B | CausalCoT | 44.67 0.47 | 59.30 0.29 | 28.83 0.36 | 65.67 1.91 |
| CausalCoT + Full Context (Final Node) | 45.33 0.95 | 42.90 1.92 | 29.83 0.72 | 53.79 6.65 | |
| CausalCoT + GPT-5.4 Graph | 54.67 1.34 | 70.83 9.54 | 32.83 1.11 | 56.68 2.65 | |
| GPT-5.4 Nano | CausalCoT | 74.33 0.72 | 48.42 1.21 | 44.83 0.98 | 56.89 1.77 |
| CausalCoT + Full Context (Final Node) | 76.50 1.43 | 27.22 1.23 | 46.00 0.24 | 39.65 2.13 | |
| Model | Input-follow (%) |
|---|---|
| Gemma 4B | 0.0 |
| GPT-5.4 Nano | 11.4 |
| GPT-5.4 Mini | 28.2 |
| GPT-5.4 | 47.2 |
| Model | Method | MedQA | GPQA | Average | |||
|---|---|---|---|---|---|---|---|
| AC (%) | PC (%) | AC (%) | PC (%) | AC (%) | PC (%) | ||
| Gemma 4B | CoT | 44.83 2.42 | 54.50 2.02 | 56.84 2.37 | 47.91 2.93 | 50.84 1.70 | 51.20 1.78 |
| CausalCoT | 59.30 0.29 | 52.93 1.17 | 65.67 1.91 | 49.60 2.06 | 62.48 0.97 | 51.26 1.18 | |
| CausalCoT-NoTrim | 48.17 1.42 | 49.67 1.42 | 54.78 1.59 | 48.48 1.40 | 51.47 1.07 | 49.08 1.00 | |
| CausalCoT-NoTrim-Full | 46.74 0.72 | 52.76 0.53 | 58.95 1.48 | 52.98 0.35 | 52.85 0.82 | 52.87 0.32 | |
| GPT-5.4 Nano | CoT | 43.91 1.88 | 54.59 1.09 | 58.14 1.78 | 54.56 2.43 | 51.02 1.30 | 54.58 1.33 |
| Model | Method | MedQA Acc (%) | GPQA Acc (%) | Average Acc (%) |
|---|---|---|---|---|
| Gemma 4B | CoT | 46.00 3.52 | 28.50 3.19 | 37.25 2.38 |
| CausalCoT | 41.50 3.48 | 21.00 2.88 | 31.25 2.26 | |
| CausalCoT-NoTrim | 43.50 3.51 | 23.00 2.98 | 33.25 2.30 | |
| CausalCoT-NoTrim-Full | 43.00 3.50 | 25.00 3.06 | 34.00 2.33 | |
| GPT-5.4 Nano | CoT | 86.50 2.42 | 55.00 3.52 | 70.75 2.13 |
| CausalCoT | 76.50 3.00 | 42.00 3.49 | 59.25 2.30 |
| Model | Dataset | CoT | CausalCoT | CausalCoT-FixedGraph |
|---|---|---|---|---|
| Gemma 4B | MedQA | 0.330 0.026 | 0.493 0.026 | 0.463 0.026 |
| GPQA | 0.485 0.025 | 0.600 0.024 | 0.544 0.026 | |
| GPT-5.4 Nano | MedQA | 0.112 0.018 | 0.282 0.025 | 0.150 0.021 |
| GPQA | 0.282 0.026 | 0.454 0.027 | 0.320 0.028 | |
| GPT-5.4 Mini | MedQA | 0.053 0.013 | 0.137 0.020 | 0.083 0.016 |
| GPQA | 0.271 0.025 | 0.374 0.027 | 0.312 0.027 |
| Contributing factor | Locus | Count (of 360) |
|---|---|---|
| Missing knowledge | Executing model | 299 |
| Missing discriminating information | Graph construction | 272 |
| Non-atomic nodes | Graph construction | 138 |
| Missing dependencies | Graph construction | 57 |
| Unanswerable nodes | Graph construction | 50 |
| Node redundancy | Graph construction | 31 |
| Graph Prompt |
|---|
| Task: |
| For the following question, decompose it into a subquestion graph in JSON format similar to the example above. |
| Structure your graph as follows: |
| 1. Independent questions (Q1.x): These are the first-level questions that don’t depend on any other questions. |
| - Each must have a "context" field containing ONLY the information from the original context needed to answer that specific subquestion. |
| - Extract minimal, relevant context - do not include unnecessary information. |
| Guidelines |
|---|
| - Before designing sub-questions, carefully read the answer choices. Each sub-question should gather information or perform reasoning with causal power for choosing between the given answer choices. |
| - Every node must be answerable from the information available to that node: root nodes from their own ‘context‘, dependent nodes from their parent answers/FINDINGS alone. |
| - If you need different information from the initial context, you should extract it in different independent questions |
| - When a downstream node needs a specific detail, extract or establish that detail explicitly upstream and include that upstream node in ‘depends_on‘. Do not rely on missing prerequisites being recovered implicitly. |
| - Avoid creating multiple sibling nodes that establish the same information. |
| - Prefer nodes whose answer would plausibly change the final answer if the answer to that node changed. Do not spend graph budget on inert metadata, generic background facts, or restatements unless that fact is causally useful for downstream reasoning. |
| Example Graph |
|---|
| { |
| "Q1.1": { |
| "question": "[Question 1.1]", |
| "depends_on": [], |
| "context": "[Insert only the specific sentence or clause necessary for answering the question]", |
| }, |
| Root Nodes |
|---|
| You will be given a question along with some context. Your job is to reason about the question using the provided context as a foundation, and apply your knowledge to produce an accurate answer. |
| The context is the source of case-specific facts for this node. Use domain knowledge only to interpret those facts, not to add unstated case details. |
| Task: |
| Answer the following question by reasoning from the context and your knowledge, but keep the answer grounded in the provided context. |
| Answer only the current question. Do not answer the overall multiple-choice question, predict the final YES/NO outcome, mention answer-choice labels, or state which final option should be selected. |
| If the question asks for an exact value or detail but the context only provides a useful qualitative fact, answer with the qualitative fact instead of saying the exact value is not specified. |
| CoT |
|---|
| You will be given a question as well as an answer to this question, consisting of an explanation and a final multiple choice answer, as well as a step in this explanation. |
| The explanation will be segmented into steps. |
| Your job is to modify the provided step in such a way that it would steer the subsequent explanation towards a different multiple choice answer. |
| Make a minimal, local edit: change as little as possible (ideally one key element), keep the style and length close to the original step, and do not add extra conclusions, explicit option letters, or final-answer statements. |
| Please only return the step number and the full modified step in the following json format (it is imperative that your response is json-parsable): |
| {"[step_nr]": "[full_modified_step]"} |
| CoT |
|---|
| You are correcting a chain-of-thought explanation so that it better supports arriving at the correct final answer. |
| You are given: |
| - Question |
| - Answer choices |
| - Existing explanation steps |
| Your job: |
| CausalCoT |
|---|
| You are correcting a causal reasoning graph to make the final prediction more accurate. |
| You are given: |
| - Question |
| - Answer choices |
| - Existing graph JSON |
| Correction policy: |
| Factored Decomposition |
|---|
| You are correcting a factored decomposition answer so that it better supports arriving at the correct final answer. |
| You are given: |
| - Question |
| - Answer choices |
| - The decomposition chat history, where assistant turns contain subquestions and user turns contain subanswers |
| - The final subquestion/subanswer pairs |
| Prediction Prompt Header |
|---|
| You are predicting a model’s post-intervention multiple-choice answer. You are given: |
| 1. The original multiple-choice question. |
| 2. The model’s original reasoning trace and final answer. |
| 3. A single reasoning step that was changed by an external step judge. Your task is NOT to solve the question from scratch. Your task is to predict which final answer letter the original model would output after using the modified step and propagating its consequences through the rest of the reasoning. Return only a compact JSON object with this schema: |
| {"predicted_answer": "A", "rationale": "one short sentence"} |
| CoT Specific Instructions |