PaperlayerView exact revision on arXiv

Source-linked AI summary

C$^2$DLM: Causal Concept-Guided Diffusion Large Language Models

Kairong Han, Nuanqiao Shan, Ziyu Zhao, Zijing Hu, Xinpeng Dong, Junjian Ye, Lujia Pan, Fei Wu, Kun Kuang

arXiv:2511.22146v1cs.CL

TL;DR

AR’s fixed token order and DLM’s fully connected attention do not capture the flexible causal structures implicated in reasoning. C2DLM extracts concept-level causal graphs and guides attention with them; it reports gains on COT-OrderPerturb and across six downstream reasoning datasets.

  • Problem

    AR and DLM attention priors do not match the flexible causal structure of natural language, while both paradigms have insufficient reasoning capabilities.

    Source: §1 · pp. 1–2

  • Method

    C2DLM uses teacher-extracted concept-level causal graphs and V-aware Re-attention to guide DLM attention toward causal relationships.

    Source: §1 · p. 2

  • Results

    Across six reasoning-related datasets, C2DLM averages a 1.31% gain, alongside 12% higher performance and 3.2× faster training on COT-OrderPerturb.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results support combining causal knowledge with language models to guide attention toward natural-language causal priors.

    Source: §5 · p. 9

    Experiments focus on LLaDA-8B-Instruct and SFT; larger DLMs and pretraining remain underexplored, and complex causal graphs or ultra-long CoT remain challenging.

    Source: Limitations · p. 9

Abstract

from arXiv · show

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 a \underline{\textbf{C}}ausal \underline{\textbf{C}}oncept-Guided \underline{\textbf{D}}iffusion \underline{\textbf{L}}anguage \underline{\textbf{M}}odel (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 improves 12\% with about 3.2 times training speedup in the COT-OrderPerturb task, and achieves an average gain of 1.31\% across six downstream reasoning tasks. More details in the repository ~\href{https://github.com/Kairong-Han/C-2-DLM}{here}.

1 Introduction

The introduction argues that AR and DLM attention priors do not match the flexible causal structure of natural language, limiting reasoning. C2DLM uses concept-level causal guidance and reports gains on COT-OrderPerturb and downstream reasoning tasks.

1 Introduction

  • AR’s strict left-to-right order and DLM’s fully connected attention both mismatch the flexible causal structures involved in natural-language reasoning.The introduction links these mismatches to local greediness, weak global planning, and DLMs producing answers before intermediate reasoning steps.Source: §1 · pp. 1–2
  • C2DLM extracts concept-level causal meta-knowledge and uses V-aware Re-attention to align attention with those causal priors.The teacher model extracts concept-level graphs; the V-aware mechanism aligns the attention map weighted by the Value matrix’s L2 norm.Source: §1 · p. 2
  • 12% higher performance and 3.2× faster training on COT-OrderPerturb; C2DLM also averages a 1.31% gain across six reasoning-related datasets.On explicit-prior tasks, reported gains are 7.43% on STG and 10.84% on Sudoku (training set size 200).Source: §1 · p. 2
  • The paper identifies misalignment between attention mechanisms and natural-language causal priors as a risk and motivates combining causality with language models.The proposed approach targets causal relationships between concepts rather than superficial correlations.Source: §1 · p. 2

2 Preliminaries and Related Works

The section reviews DLM training and attention, then frames causal alignment as a response to limitations in AR and fully connected DLM attention.

2.1 Diffusion Large Language Model

  • DLMs use bidirectional attention and a negative evidence lower bound in training; C2DLM focuses on causal alignment during supervised fine-tuning.The section describes LLaDA’s SFT loss and notes that C2DLM’s focus is the SFT stage.Source: §2.1 · p. 2
  • Related DLM work has pursued reinforcement learning and faster inference, including methods such as d1 and Fast-dllm.The passage also names BranchGRPO and SlowFast among these lines of work.Source: §2.1 · p. 2

2.2 Combining Causality and Attention Mechanism

  • Transformer multi-head attention models token dependencies, while prior work has explored causal debiasing and Re-attention to inject causal knowledge.The cited examples include causal backdoor mechanisms for text and vision and a student-model Re-attention mechanism.Source: §2.2 · p. 3
  • C2DLM investigates guiding DLM attention to align with data-generation priors.The section presents this as a new paradigm for examining attention mechanisms in DLMs.Source: §2.2 · p. 3

2.3 The Limitations of Attention in AR and DLM

  • AR’s lower-triangular attention cannot handle outcomes preceding causes, while unconstrained DLM attention can dilute CoT causal signals with redundant context.The passage says this makes stable and effective reasoning over CoT difficult.Source: §2.3 · p. 3

3 Method

C2DLM has two stages: a teacher extracts concept-level causal knowledge, then V-aware Re-attention aligns model attention with those causal priors. The method represents reasoning as concept relationships and applies loss-based guidance to attention.

3 Method

  • C2DLM’s two stages are concept-level causal meta-knowledge extraction and causal alignment through V-aware Re-attention.Figure 2 shows teacher-generated textual links mapped onto weighted attention maps, followed by loss-based intervention.Source: §3 · p. 3, Figure 2 · p. 4

3.1 Concept-level Causal Meta-knowledge Extraction

  • The teacher extracts semantically complete concepts from reasoning steps, treats the remaining text as context, and builds a graph of causal dependencies among concepts.The concepts are intended to capture information needed for reasoning and reflect flexible causal priors.Source: §3.1 · p. 3
  • The graph permits dependencies beyond pairwise causal forms, prunes information unnecessary for generating a concept, and can use coarser inter-rule steps to prune irrelevant earlier information.The rule-based, semi-autoregressive signal is introduced for certain tasks.Source: §3.1 · p. 3

3.2 Causal Alignment via the V-aware Re-attention Mechanism

  • V-aware Re-attention weights attention maps with Value-matrix L2 norms because attention sinks can distort raw attention values.The mechanism targets alignment between model decision dependencies and natural-language causal mechanisms while reducing attention-map outlier instability.Source: §3.2 · p. 4
  • The alignment loss encourages attention on supervised causal links relative to neutral entries and penalizes weighted attention on entries marked for suppression.The row loss combines ratio and negative-attention terms.Source: §3.2 · pp. 4–5
  • The constraint loss is combined with DLM SFT loss on valid supervised rows, while a warm-up and cool-down schedule varies its balancing coefficient to stabilize training.The coefficient rises from γmin to γmax, then returns to γmin; t, T1, and T2 specify the schedule.Source: §3.2 · p. 5

4 Experimental Results

C2DLM improves performance and training efficiency on COT-OrderPerturb, then evaluates causal-prior benefits on Sudoku and STG and broader reasoning tasks. Results include stronger STG performance and attention patterns that distinguish causal factors from spurious or unrelated ones.

4.1 Experimental Setup

  • The experiments use COT-OrderPerturb, Sudoku, STG, and six broader reasoning datasets, with LLaDA-8B-Instruct as the primary model and direct SFT as the DLM baseline.The setup also compares four AR models and uses LoRA fine-tuning.Source: §4.1 · p. 5

4.2 Quantifying the Impact of Priors

  • COT-OrderPerturb compares AR and DLM using causal-order CoT, shuffled concept sequences, and a No-CoT control.Shuffling strategies include DFS, local reverse, output first, reverse, and three random shuffles.Source: §4.2 · p. 5, Figure 3 · p. 6
  • AR consistency declines when outcomes precede causes, while DLM is more robust to shuffled CoT but weaker on Normal CoT and slower to converge.DLM typically requires nearly 80 epochs to converge, compared with 4 epochs for AR models.Source: §4.2 · pp. 5–6
  • Under Normal COT, aligning with causal priors yields C2DLM a 12% performance gain relative to direct SFT on LLaDA-8B-Instruct; on COT-OrderPerturb, its training efficiency is 3.2 times DLM’s.Source: Table 2 · p. 6, Figure 4 · p. 6, §4.2 · p. 6

4.3.1 Sudoku dataset

  • Sudoku experiments use a 4 × 4 grid, where each cell depends on its row, column, and subgrid; C2DLM performs better than LLaDA-8B-Instruct with 200 training examples.The reported advantage is especially apparent in this small-data setting, where LLaDA-8B-Instruct lacks clear supervisory guidance.Source: §4.3.1 · pp. 6–7

4.3.2 STG dataset

  • C2DLM improves over direct SFT by an average of 7.43% across STG IID and OOD settings, and surpasses the best AR baselines on OOD STG_S, STG_M, and STG_L.Figure 5 is described as showing performance change across training epochs on STG_H.Source: §4.3.2 · p. 7, Figure 5 · p. 7
  • STG attention visualizations show C2DLM distinguishing causal factors from spurious correlations and irrelevant factors, unlike direct LLaDA fine-tuning.Figure 6 uses purple for causal factors, green for spurious correlations, and yellow for unrelated factors.Source: §4.3.2 · p. 7, Figure 6 · p. 7

4.4 Evaluate on Broader Math and Reasoning Downstream Tasks

  • For broader tasks without explicit causal structure, the paper uses GLM-4.5 in its automated workflow to address causal-relation extraction from long CoT sequences.The passage also states that the workflow starts from the s1k dataset, but its description is incomplete.Source: §4.4 · p. 7

Setting STG_S STG_M STG_L STG_H

The reported STG_S evaluation gives accuracy among successfully decoded instances. Ablations examine causal-emphasis strength, the γ scheduler, and Value-matrix weighting.

Setting STG_S STG_M STG_L STG_H

4.5 Ablation Study

  • Performance first improves and then declines as α increases within a certain range, and removing the γ scheduler reduces performance on all datasets.The table identifies α = 3 as the default setting and labels the no-scheduler condition w/o Sγ.Source: §4.5 · p. 8, Table 6 · p. 8
  • Without V-aware weighting, direct manipulation of attention scores decreases performance and introduces substantial training instability.The authors interpret this ablation as evidence for the V-aware strategy.Source: §4.5 · p. 8

5 Conclusion

The conclusion presents C2DLM as a distinct approach combining automated causal meta-knowledge extraction with V-aware Re-attention, and reports validation across Sudoku, STG, and six downstream tasks. The included limitations, ethical guidance, appendix prompt, and hyperparameter details qualify the scope and implementation.

5 Conclusion

  • C2DLM combines automated causal meta-knowledge extraction with V-aware Re-attention and is evaluated on Sudoku, STG, and six broad downstream tasks.The conclusion says the method improves reasoning ability and training efficiency and identifies attention–causal-prior misalignment as a risk.Source: §5 · p. 9

Limitations

  • Experiments focus on LLaDA-8B-Instruct and the SFT phase; larger-scale DLMs, pretraining, and more intricate causal graphs or ultra-long CoT remain underexplored.The authors also report that C2DLM still lags behind SOTA AR models, citing training-resource and base-model constraints.Source: Limitations · p. 9

Ethical considerations

  • The method could be misused to inject illegal or unethical information, so the authors urge compliance with relevant laws and ethical guidelines for training data..Source: Ethical considerations · p. 9

A Details of Prompt

  • The teacher-LLM prompt extracts concepts and causal relations from questions and solutions to form reasoning graphs used as surrogates for natural-language causal priors.Extracted concepts should reflect essential reasoning, carry richer information, influence the final answer, be unique, and come from the source text; the output maps effects to causes.Source: A Details of Prompt · pp. 11–12

B Details of Hyperparameters

  • Training epochs and evaluation generation lengths are configured by task, with different settings for COT-OrderPerturb, Sudoku, STG, and downstream tasks.The appendix specifies, for example, 10 epochs for COT-OrderPerturb and Sudoku, and 10 epochs for STG_E versus 40 for STG_H.Source: B Details of Hyperparameters · p. 13

C Generation Details of COT-OrderPerturb Dataset

The data-generation process aims to create chain-of-thought trajectories with controlled order perturbations while preserving the underlying process.

C Generation Details of COT-OrderPerturb Dataset

  • The generation process aims to produce chain-of-thought trajectories with controlled order perturbations while preserving the underlying process..Source: C Generation Details of… · p. 13

C.1 Graph Construction and Chain-of-Thought Generation

COT-OrderPerturb data are generated from a DAG whose computed values define a reasoning trace, then tested under changes to the order of its intermediate steps. The perturbations retain the final answer, and the dataset includes distinct training and test samples.

C.1 Graph Construction and Chain-of-Thought Generation

  • A DAG assigns functional dependencies to variables: two source values are sampled uniformly from [0,100], other variables are computed in topological order, and Stardust is the target.Step-wise traces record each rule, its inputs, and its output; in the figure, yellow nodes are inputs and Stardust is the final answer.Source: §C.1 · p. 14, Figure 7 · p. 14

C.2 Order Perturbations

  • Order perturbations include complete and local reversal, output-first and depth-first orderings, three fixed random permutations, and omission of reasoning.These changes shuffle intermediate steps while preserving the final Stardust answer.Source: §C.2 · p. 14
  • The dataset contains 2,000 unique base samples and a separate test set of 500 different samples..Source: §C.2 · p. 14

D Cost Analysis of C2DLM

The cost analysis estimates GLM-4.5 causal annotation at about 0.46$ per million tokens, using token counts from a representative sample.

D Cost Analysis of C2DLM

  • Annotation costs are estimated from a random representative subset of 100 examples, with average input length 865.2 tokens plus 1,981 prompt-overhead tokens and average output length 295.3 tokens..Source: D Cost Analysis of C2DLM · p. 15
  • About 0.46$ per million tokens is the estimated cost of causal annotation with GLM-4.5.The estimate uses empirical input and output token statistics and the stated official GLM-4.5 pricing.Source: D Cost Analysis of C2DLM · p. 15

E Attention Visualization of C2DLM

The visualization analysis finds that C2DLM distinguishes causal factors, while direct fine-tuning of LLaDA does not. It presents value-weighted attention as a more faithful view of token influence.

E Attention Visualization of C2DLM

F Human Evaluation of Teacher Model Generated Causal Graphs

Human evaluation assessed the logical consistency of teacher-generated causal graphs, yielding 93.42% ± 1.41% accuracy on correctly decoded graphs. The section also reports using ChatGPT to check manuscript grammar and syntax.

F Human Evaluation of Teacher Model Generated Causal Graphs

G Use Of AI Assistants

Loading 2511.22146v1…