Table of Contents
Fetching ...

Can Post-Training Transform LLMs into Causal Reasoners?

Junqi Chen, Sirui Chen, Chaochao Lu

TL;DR

This work investigates whether post-training can endow LLMs with reliable causal reasoning. It introduces CauGym, a synthetic SCM-based dataset with seven causal tasks and five post-training methods applied to a 14B-scale model, evaluated across nine test sets targeting generalization, internalization, and robustness. Online RL, especially GRPO, yields the strongest gains, achieving up to 93.5% accuracy on CaLM and outperforming larger models, while also generalizing to paraphrased prompts and resisting noise and incomplete information. The findings demonstrate that careful post-training can enable smaller LLMs to function as principled and robust causal reasoners, with CauGym providing a benchmark and GRPO a strong method for future work. The dataset and GRPO model are publicly available to advance research in LLM-based causal inference.

Abstract

Causal inference is essential for decision-making but remains challenging for non-experts. While large language models (LLMs) show promise in this domain, their precise causal estimation capabilities are still limited, and the impact of post-training on these abilities is insufficiently explored. This paper examines the extent to which post-training can enhance LLMs' capacity for causal inference. We introduce CauGym, a comprehensive dataset comprising seven core causal tasks for training and five diverse test sets. Using this dataset, we systematically evaluate five post-training approaches: SFT, DPO, KTO, PPO, and GRPO. Across five in-domain and four existing benchmarks, our experiments demonstrate that appropriate post-training enables smaller LLMs to perform causal inference competitively, often surpassing much larger models. Our 14B parameter model achieves 93.5% accuracy on the CaLM benchmark, compared to 55.4% by OpenAI o3. Furthermore, the post-trained LLMs exhibit strong generalization and robustness under real-world conditions such as distribution shifts and noisy data. Collectively, these findings provide the first systematic evidence that targeted post-training can produce reliable and robust LLM-based causal reasoners. Our data and GRPO-model are available at https://github.com/OpenCausaLab/CauGym.

Can Post-Training Transform LLMs into Causal Reasoners?

TL;DR

This work investigates whether post-training can endow LLMs with reliable causal reasoning. It introduces CauGym, a synthetic SCM-based dataset with seven causal tasks and five post-training methods applied to a 14B-scale model, evaluated across nine test sets targeting generalization, internalization, and robustness. Online RL, especially GRPO, yields the strongest gains, achieving up to 93.5% accuracy on CaLM and outperforming larger models, while also generalizing to paraphrased prompts and resisting noise and incomplete information. The findings demonstrate that careful post-training can enable smaller LLMs to function as principled and robust causal reasoners, with CauGym providing a benchmark and GRPO a strong method for future work. The dataset and GRPO model are publicly available to advance research in LLM-based causal inference.

Abstract

Causal inference is essential for decision-making but remains challenging for non-experts. While large language models (LLMs) show promise in this domain, their precise causal estimation capabilities are still limited, and the impact of post-training on these abilities is insufficiently explored. This paper examines the extent to which post-training can enhance LLMs' capacity for causal inference. We introduce CauGym, a comprehensive dataset comprising seven core causal tasks for training and five diverse test sets. Using this dataset, we systematically evaluate five post-training approaches: SFT, DPO, KTO, PPO, and GRPO. Across five in-domain and four existing benchmarks, our experiments demonstrate that appropriate post-training enables smaller LLMs to perform causal inference competitively, often surpassing much larger models. Our 14B parameter model achieves 93.5% accuracy on the CaLM benchmark, compared to 55.4% by OpenAI o3. Furthermore, the post-trained LLMs exhibit strong generalization and robustness under real-world conditions such as distribution shifts and noisy data. Collectively, these findings provide the first systematic evidence that targeted post-training can produce reliable and robust LLM-based causal reasoners. Our data and GRPO-model are available at https://github.com/OpenCausaLab/CauGym.
Paper Structure (37 sections, 7 equations, 11 figures, 9 tables)

This paper contains 37 sections, 7 equations, 11 figures, 9 tables.

Figures (11)

  • Figure 1: An example of the base dataset.
  • Figure 2: (a) Model performance on CauGym-rephrased. (b) Model performance difference between the dataset and CaLM. "R1" denotes DeepSeek-R1-0528-671B, "CS" denotes Cold Start Base.
  • Figure 3: (a) Model performance on CauGym-omitted. (b) Model performance difference between the dataset and CaLM. "R1" denotes DeepSeek-R1-0528-671B, "CS" denotes Cold Start Base.
  • Figure 4: Model performance on CauGym-deconfounding, "R1" denotes DeepSeek-R1-0528-671B, "CS" denotes Cold Start Base.
  • Figure 5: Model performance on CauGym-redundant dataset and CauGym-insufficient dataset.
  • ...and 6 more figures