[ACL 2026 findings] C²DLM: Causal Concept-Guided Diffusion Large Language Model

June 23, 2026 · View on GitHub

This repository contains the official implementation for the paper "C²DLM: Causal Concept-Guided Diffusion Large Language Models".

Abstract

Autoregressive (AR) language models and Diffusion Language Models (DLMs) constitute the two principal paradigms of large language models. However, both paradigms suffer from hallucinations and unfaithful reasoning in reasoning-intensive tasks. We hypothesize that these limitations stem from a misalignment between the attention mechanism’s modeling priors on natural language and the causal priors underlying natural language. To address this issue, we propose the Causal Concept-Guided Diffusion Language Model (C²DLM). Starting from DLM's fully connected attention, C²DLM obtains concept-level causal supervision signals through a teacher model and explicitly guides attention to learn the causal relationships between concepts, which better aligns with the underlying causal prior of natural language. Simulation experiments demonstrate the significance of incorporating language priors. Compared with direct DLM fine-tuning, C²DLM improves the COT-OrderPerturb task by 12% with a about 3.2x training speedup, enhances STG tasks by 7.43% on average with about 2x speedup, and achieves an average gain of 1.31% across six downstream reasoning tasks. Alt text

Key Contributions

  • C²DLM Framework: A novel framework that injects causal priors into Diffusion Language Models to improve reasoning and reduce hallucinations.
  • V-aware Re-attention: A new mechanism to align the model's attention map with the causal structure of natural language, weighted by the norms of the value vectors.
  • COT-OrderPerturb Dataset: A new synthetic dataset designed to quantify the impact of causal order perturbations on model reasoning robustness.
  • Improved Performance & Efficiency: C²DLM demonstrates significant gains in performance and training speed across a variety of reasoning-intensive tasks.

Method

Alt text

File Structure

d1/
├── eval/                  # Scripts for evaluating on downstream tasks (which is the offical repository of "d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning")
│   ├── math500.py
│   ├── stg.py
│   ├── calculate_gpqa_arc_acc.py
│   ├── parse_and_get_acc.py
│   └── ...
├── eval_results/          # Stores raw model generation outputs for evaluation
│   ├── gsm8k_512_BLOCK_LENGTH32_C2DLM/
│   ├── math_512_BLOCK_LENGTH32_C2DLM/
│   └── ...
├── draw/                  # Utility scripts for plotting and analysis
│   ├── bar_viz.py
│   ├── causal_llada_util.py
│   └── ...
├── train_COT_OrderPerturb/  # Implementation of our proposed COT-OrderPerturb task
│   └── experiment/
│       ├── cot_dataset_test.jsonl
│       └── multi_gpu_train.py
└── train_downstream/      # Scripts for training on general downstream tasks
    └── ...

Training

The training scripts leverage torchrun for multi-GPU execution.

Training on COT-OrderPerturb

To train a model on our proposed COT-OrderPerturb dataset:

cd train_COT_OrderPerturb/experiment
torchrun --nproc_per_node=[NUM_GPUS] multi_gpu_train.py [YOUR_ARGUMENTS]

Training on Downstream Datasets

To train a model on other downstream reasoning tasks:

cd train_downstream
torchrun --nproc_per_node=[NUM_GPUS] multi_gpu_train.py [YOUR_ARGUMENTS]

Evaluation

The evaluation is a two-step process: first, generate model outputs, and second, calculate the accuracy.

1. Generate Model Outputs

Use the eval.py script to produce raw outputs from your trained model. The results will be saved in the eval_results/ directory.

cd eval
python eval.py [YOUR_ARGUMENTS]

2. Calculate Accuracy

Once generations are complete, use the provided parsing scripts to calculate the final accuracy. The scripts read the output folders generated in the previous step.

Each folder in eval_results contains the model's raw output with the following fields: question, prompt_input, generations, and ground_truth.

cd eval

# For GSM8K and MATH500 tasks, use the general accuracy script
python parse_and_get_acc.py --file ../eval_results/[MY_OUTPUT_FOLDER]

# For specific tasks like GPQA or ARC and MMLU_STEM, use the dedicated calculator
python calculate_gpqa_arc_mmlu_acc.py ../eval_results/[MY_GPQA_OUTPUT_FOLDER]
python calculate_sat_acc.py ../eval_results/[MY_SAT_OUTPUT_FOLDER]

Visualization

The draw/ directory contains scripts used to generate the attention visualizations and performance plots from our paper.

Citiation

@inproceedings{han-etal-2026-c2dlm,
    title = "{C}$^2${DLM}: Causal Concept-Guided Diffusion Large Language Models",
    author = "Han, Kairong  and
      Shan, Nuanqiao  and
      Zhao, Ziyu  and
      Hu, Zijing  and
      Dong, Xinpeng  and
      Jian, Ye Jun  and
      Pan, Lujia  and
      Wu, Fei  and
      Kuang, Kun",
    editor = "Liakata, Maria  and
      Moreira, Viviane P.  and
      Zhang, Jiajun  and
      Jurgens, David",
    booktitle = "Findings of the {A}ssociation for {C}omputational {L}inguistics: {ACL} 2026",
    month = jul,
    year = "2026",
    address = "San Diego, California, United States",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2026.findings-acl.40/",
    pages = "821--838",
    ISBN = "979-8-89176-395-1",
    abstract = "Autoregressive (AR) language models and Diffusion Language Models (DLMs) constitute the two principal paradigms of large language models. However, both paradigms suffer from insufficient reasoning capabilities. Human reasoning inherently relies on causal knowledge and thought, which are reflected in natural language. But in the AR paradigm, language is modeled as next token prediction (a strictly left-to-right, token-by-token order), whereas natural language itself exhibits more flexible causal structures. In the DLM paradigm, the attention mechanism is fully connected, which entirely disregards causal order. To fill this gap, we propose the Causal Concept-Guided Diffusion Language Model (C$^2$DLM). Starting from DLM{'}s fully connected attention, C$^2$DLM first obtains a concept-level causal graph from the teacher model, and then explicitly guides attention to learn causal relationships between concepts. By focusing on causal relationships and avoiding interference from difficult subgoals involving causal inversion, C$^2$DLM achieves a 12{\%} improvement and a 3.2{\texttimes} training speedup on the COT-OrderPerturb task, along with an average gain of 1.31{\%} across six downstream reasoning tasks. Code and data are available {~}here."
}