The Gradient-Causal Gap: Why Gradient Importance Fails on Complex Tasks

Removing ‘‘important’’ high-gradient components from a neural network can improve generalization, while removing unimportant’’ low-gradient components can destroy it. We demonstrate this paradox by formalizing the \textit{Gradient-Causal Gap} in Transformers trained on algorithmic tasks. While gradient magnitude and causal importance align on simple tasks ($ρ=0.73$ for reversal), this relationship collapses as task complexity increases ($ρ=0.32$ for sorting), sometimes becoming inverted ($ρ=-0.11$). Pruning experiments reveal that gradient magnitude is not merely inaccurate but \textit{unpredictably} so. Removing low-gradient ‘‘Hidden Heroes’’ consistently devastates OOD accuracy ($-32%$). Removing high-gradient ‘‘Gradient Bloats’’ is a coin flip: harmless in most seeds (indicating optimization noise), catastrophic in others (indicating overfitting circuits). This unpredictability means gradient-based pruning cannot reliably preserve model capabilities.

💡 Research Summary

The paper introduces the “Gradient‑Causal Gap” (GCG) to quantify the mismatch between gradient‑based importance scores and true causal importance of model components, focusing on out‑of‑distribution (OOD) generalization. The authors train small decoder‑only Transformers (4 layers, 4 heads per layer, 20 components total) on two algorithmic benchmarks: Sequence Reversal (a simple task) and Sequence Sorting (a more complex task). For each component they compute (1) Gradient Magnitude (G) as the average Frobenius norm of the weight gradients over 50 OOD batches, and (2) Causal Importance (C) via mean ablation, i.e., replacing the component’s output with its mean activation and measuring the drop in exact‑match accuracy.

The core metric is the Spearman correlation ρ between G and C. On Reversal, ρ ≈ 0.73 ± 0.12, indicating strong alignment; on Sorting, ρ drops to ≈ 0.32 ± 0.24, with one seed even showing a negative correlation (ρ ≈ ‑0.11). To capture individual mismatches, the authors define Δ = Rank(G) − Rank(C). Components with Δ ≤ ‑6 are labeled “Hidden Heroes” (low gradient but high causal impact), while those with Δ ≥ 6 are “Gradient Bloats” (high gradient but low causal impact).

Analysis of the component distribution reveals a clear layer‑wise pattern: Hidden Heroes cluster in later layers (especially layer 3 heads) for the Sorting task, whereas Gradient Bloats concentrate in early layers (layers 0‑1). This suggests early layers receive large back‑propagation signals because they contribute heavily to training loss reduction, yet they are not essential for the algorithmic logic that enables OOD generalization. Conversely, later layers implement sparse, task‑specific reasoning that yields weak gradients but is crucial for correct behavior.

Pruning experiments validate these observations. Removing the top two Hidden Heroes in each seed consistently degrades OOD accuracy by an average of 32 percentage points, confirming their indispensable role despite low gradients. Pruning Gradient Bloats yields a bimodal outcome: in the majority of seeds the removal has negligible effect on both in‑distribution (ID) and OOD performance, indicating these components are “Optimization Noise.” In a minority of seeds, however, pruning causes severe ID collapse (up to a 39 % drop), revealing that some high‑gradient components form overfitting circuits. Thus, a high gradient norm is an ambiguous signal that cannot reliably distinguish useful from redundant or harmful components.

The discussion emphasizes that gradient‑based importance measures, widely used for pruning (Han et al., 2015; Molchanov et al., 2017) and attribution (Integrated Gradients, Sundararajan et al., 2017), are fundamentally tied to loss reduction rather than to the functional circuits that support robust generalization. Consequently, relying solely on gradients for model compression or interpretability can inadvertently prune the very components that enable OOD performance, especially in models that solve tasks without explicit chain‑of‑thought reasoning.

Limitations are acknowledged: the study uses small Transformers and synthetic algorithmic tasks; it remains an open question whether the Gradient‑Causal Gap persists at the scale of modern language models or on natural language datasets. Moreover, causal importance is measured via mean ablation; more sophisticated techniques such as activation patching could yield different rankings.

In conclusion, the Gradient‑Causal Gap widens as task complexity grows, giving rise to Hidden Heroes and Gradient Bloats. The findings caution against uncritical use of gradient magnitude as a proxy for importance and call for causal, intervention‑based methods to identify safe pruning targets and to achieve trustworthy interpretability of neural networks. Future work should explore scaling this analysis, develop gradient‑agnostic importance metrics, and test the phenomenon on real‑world NLP benchmarks.