Improving Chain-of-Thought faithfulness through training-time causal alignment and inference-time structural enforcement.
📖 Abstract • 🔬 Methodology • ⚙️ Installation • 🚀 Usage • 📊 Evaluation • 🧪 Experiments
Chain-of-Thought (CoT) reasoning may produce plausible explanations that do not faithfully support the final answer.
From a causal perspective, faithful reasoning should follow
where the instruction
We propose CASE (Causal Alignment and Structural Enforcement), which combines:
-
Training-Time Causal Alignment (TCA): strengthens CoT-to-answer dependence using Mixed Training Datasets and Selective-Loss SFT.
-
Inference-Time Structural Enforcement (ISE): blocks direct attention from instruction tokens to answer tokens.
Figure 1: Overview of the proposed CASE framework.
Experiments on three LLMs and four reasoning benchmarks show that CASE improves overall CoT faithfulness by more than 24% on average over the strongest baselines while preserving competitive accuracy.
| Dataset | Name | Purpose |
|---|---|---|
| Standard Dataset | Standard CoT and answer supervision | |
| Counterfactual-CoT Dataset | Strengthen |
|
| Biased-Instruction Dataset | Suppress |
|
| Empty-Instruction Dataset | Encourage prediction from the CoT alone |
CASE applies Selective-Loss SFT:
-
Both CoT and answer tokens are supervised on
$\mathcal{D}_A$ . -
Only answer tokens are supervised on
$\mathcal{D}_B$ ,$\mathcal{D}_C$ , and$\mathcal{D}_D$ .
This objective strengthens the CoT-to-answer mechanism without forcing the model to imitate constructed auxiliary reasoning traces.
During answer generation, CASE masks direct attention from instruction tokens to answer tokens.
The answer can still attend to the generated CoT and its computational states, preserving the mediated path
while blocking the direct shortcut
cd CASE
pip install -r requirements.txtPrepare directories for datasets, models, generated data, and results:
mkdir datasets models generated_data resultspython build_cot_answers.py \
--dataset_path ai2_arc \
--output_path generated_data/arc_cot.pkl \
--deepseek_api_key YOUR_API_KEYpython build_counterfactual.py \
--dataset_path ai2_arc \
--cleaned_reasoning_path generated_data/arc_cot.pkl \
--output_path generated_data/arc_counterfactual.pkl \
--deepseek_api_key YOUR_API_KEYpython build_biased_questions.py \
--dataset_path ai2_arc \
--cleaned_reasoning_path generated_data/arc_cot.pkl \
--output_path generated_data/arc_biased.pkl \
--seed 42python CASE_train.py \
--model_path models/Llama-3.1-8B-Instruct \
--dataset_path ai2_arc \
--cleaned_reasoning_path generated_data/arc_cot.pkl \
--counterfactual_path generated_data/arc_counterfactual.pkl \
--biased_questions_path generated_data/arc_biased.pkl \
--output_dir results/case_lora \
--enable_stage1_answer_only_qbody_mask_data \
--gradient_checkpointing \
--use_4bitpython CASE_test.py \
--model_path models/Llama-3.1-8B-Instruct \
--adapter_path results/case_lora \
--dataset_path ai2_arc \
--counterfactual_path generated_data/arc_counterfactual.pkl \
--deepseek_api_key YOUR_API_KEY \
--mask_question_on_answer \
--use_few_shot \
--use_4bit \
--save_dir results/case_evaluationCASE uses four CoT faithfulness metrics:
| Metric | Description | Better |
|---|---|---|
| Filler Tokens (FT) | Replaces the CoT with filler tokens | Higher |
| Early Answering (EA) | Tests whether the model can answer before seeing the complete CoT | Higher |
| Controlled Indirect Effect (CIE) | Replaces the original CoT with a counterfactual CoT | Higher |
| Controlled Direct Effect (CDE) | Changes the instruction while keeping the CoT fixed | Lower |
The overall faithfulness score is
Evaluation results are saved as:
results/case_evaluation/
├── summary.json
├── predictions.pkl
├── outputs.pkl
├── questions.pkl
└── filler_predictions.pkl
CASE is evaluated using:
- Models Llama-3.1-8B, Qwen3-8B, and DeepSeek-R1-Distill-Qwen-7B.
- Datasets: ARC-Easy, ARC-Challenge, StrategyQA, and LogiQA.
- Cross-dataset settings: ARC-Challenge → OpenBookQA and StrategyQA → ARC-Easy.
CASE achieves the best