Gradient-Guided Confounder Pruning (LeaF)

• The paper introduces LeaF, a two-stage framework that detects and prunes confounding tokens to align student attention with the true causal structure.
• It employs teacher-student gradient discrepancies and causal interventions to systematically remove spurious dependencies, improving performance in reasoning and code generation tasks.
• Empirical evaluations demonstrate that LeaF enhances interpretability and robustness of distilled models, delivering significant gains in benchmarks like GSM8K and HumanEval+.

Gradient-Guided Confounder Detection and Pruning (LeaF) is a two-stage framework designed to systematically identify and neutralize spurious, non-causal dependencies in transformer-based models, with a focus on distilling LLMs for robust reasoning and program synthesis tasks. By leveraging teacher-student gradients and causal intervention principles, LeaF isolates “confounding tokens” and enforces their removal during knowledge distillation, thereby aligning student model attention with the authentic causal structure underlying instructional data (Guo et al., 9 Jun 2025).

1. Theoretical Foundations of Gradient-Guided Pruning

LeaF is grounded in Pearl’s structural causal model (SCM) framework, operationalizing back-door adjustment through token masking in the sequence-to-sequence transformer setting. Input tokens $X = [x_1,\ldots,x_n]$ may include a subset $A \subset X$ termed “confounding tokens,” which fulfill three concrete criteria: (a) they introduce a spurious dependency between prompt and label $Y$, (b) their presence increases student model loss, and (c) they receive little attention from a stronger teacher model that correctly predicts $Y$.

The observed conditional distribution $P(Y|X)$ intermixes the true causal path $X{\rightarrow}Y$ with spurious $X{\leftarrow}A{\rightarrow}Y$ interactions. LeaF leverages hard interventions ($do$-operator) by removing confounders:

$P(Y|do(X)) = P(Y|X, X \setminus A),$

thus blocking the confounding back-door. Pruning $A$, followed by student-teacher distillation under observational ($A \subset X$0) and interventional ($A \subset X$1) regimes, ensures exposure to both spurious and causally pruned contexts and drives the student toward robust, causal attention behavior.

2. Algorithmic Procedure for Gradient-Guided Confounder Detection

Stage 1 of LeaF identifies confounders by quantifying the gradient-based sensitivity of model loss to each token, contrasting a strong teacher ($A \subset X$2) with a smaller student ($A \subset X$3). For token $A \subset X$4, the importance scores are

$A \subset X$5

Scores are min–max normalized per instance:

$A \subset X$6

The discrepancy $A \subset X$7 detects student-overweighted tokens. A further normalization across tokens yields

$A \subset X$8

and token $A \subset X$9 is labeled a confounder if $Y$0 and pruning $Y$1 alone induces a correct prediction by the student. $Y$2 (typically $Y$3–$Y$4) is tuned via validation.

Pseudocode for confounder detection: $X{\rightarrow}Y$7 Detected confounders are grouped into spans for subsequent intervention (Guo et al., 9 Jun 2025).

3. Causal Intervention and Attention-Pruned Distillation

After detecting confounder spans $Y$5, Stage 2 executes hard-masking over tokens in $Y$6 during student distillation. The pruning is implemented via a mask $Y$7 applied to each attention head $Y$8 and position $Y$9:

$Y$0

This is equivalent to zeroing out key/value vectors for pruned tokens.

The distillation loss blends two KL-divergence branches:

$Y$1

across all detected spans, typically with $Y$2. An auxiliary attention-matching loss may optionally regularize

$Y$3

though the core LeaF approach relies on the output-alignment objective.

The student is alternately exposed to unmasked ($Y$4) and pruned ($Y$5) contexts, enforcing both observational and interventional consistency. Key hyperparameters: $Y$6, $Y$7, batch size (32–64), epochs (3), learning rate ($Y$81e–5).

4. Empirical Evaluation and Comparative Metrics

LeaF’s efficacy was demonstrated on several benchmarks:

• Mathematical reasoning: GSM8K, MATH, OlympiadBench (pass@1 accuracy).
• Code generation: HumanEval+, LeetCode, LiveCodeBench (pass@1, pass@10).

Quantitative improvements over standard knowledge distillation (KD without mask): | Model | Math Gain (%) | Code Gain (%) | |--------------|--------------|--------------| | LLaMA-1B | +1.4 | +2.4 | | LLaMA-3B | +1.6 | +1.5 | | Qwen-1.5B | +1.7 | +2.1 |

Ablation studies revealed:

• Span-pruning (masking contiguous spans) outperformed collective pruning on MATH-500 (+3.4% vs. +0.2%).
• Gradient-based masking surpassed both random masking (which harmed performance) and PPL-based masking (modest $Y$90.5% gain), yielding $P(Y|X)$02.4% gain.
• Segmenting and masking confounders in model-generated prefixes led to further gains (+0.3–0.8%).

5. Interpretability and Model Consistency

LeaF yields both quantitative and qualitative improvements in interpretable attention. In representative MATH cases, attention-difference heatmaps indicate that standard distillation frequently misallocates weight to irrelevant tokens (e.g., author names), while LeaF suppresses these and emphasizes semantically critical elements (e.g., “real number,” “discriminant $P(Y|X)$1”).

Evaluations using Jaccard similarity show that instruction-pruned students align more closely with teacher rationales, indicated by a right-shifted Jaccard distribution. Smaller models (1B) are observed to benefit from more aggressive pruning thresholds ($P(Y|X)$2–$P(Y|X)$3) versus larger models (3B, $P(Y|X)$4–$P(Y|X)$5), consistent with the view that lower-capacity students are more susceptible to spurious correlations (Guo et al., 9 Jun 2025).

6. Generalization to Non-Transformer Causal Discovery

Related causal confounder-pruning principles are evident in differentiable backdoor discovery for causal inference (Gultchin et al., 2020). There, rather than symbolic or combinatorial search for a minimal backdoor adjustment set, a differentiable “score” $P(Y|X)$6 is learned, subject to L1-regularization and tuning via auxiliary “witness” variables. The optimization target balances three criteria: (i) independence of $P(Y|X)$7 and $P(Y|X)$8 conditional on $P(Y|X)$9 and $X{\rightarrow}Y$0, (ii) residual dependence of $X{\rightarrow}Y$1 and $X{\rightarrow}Y$2 conditioning only on $X{\rightarrow}Y$3, and (iii) sparsity in $X{\rightarrow}Y$4.

Gradient-descent on objective

$X{\rightarrow}Y$5

efficiently prunes $X{\rightarrow}Y$6 to a valid adjustment set, circumventing combinatorial search found in, e.g., the Entner–Hoyer greedy approach. Empirical results show competitive or superior performance in simulation and real-world datasets, illustrating the general utility of differentiable, gradient-guided confounder selection (Gultchin et al., 2020).

7. Significance and Outlook

LeaF operationalizes causal attention pruning for LLM distillation using rigorous gradient-based detection, spanning intervention, and objective-driven distillation. By enforcing consistency with both the teacher’s observational and interventional behaviors, LeaF systematically prunes spurious shortcuts, resulting in models with improved reasoning, code generation, and interpretability. The integration of causal inference frameworks and gradient-guided backdoor discovery suggests broad applicability for robustifying neural models across domains characterized by latent confounding (Guo et al., 9 Jun 2025, Gultchin et al., 2020).