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

TL;DR

This work exposes a Gradient-Causal Gap where gradient magnitude lacks reliable alignment with causal importance for out-of-distribution generalization in Transformers trained on algorithmic tasks. By comparing gradient-based importance with ablation-based causal importance across components and tasks, the authors identify two failure modes: Hidden Heroes (low-gradient but causally essential) and Gradient Bloats (high-gradient yet causally weak, sometimes harmful). The correlation between gradient signals and causal impact degrades with task complexity, and pruning reveals that removing Hidden Heroes consistently harms OOD performance while pruning Gradient Bloats yields highly seed-dependent effects, challenging the reliability of gradient-based pruning for preserving capabilities. The study highlights the need for causal verification in interpretability and model compression, as gradient-based methods may track training dynamics rather than the reasoning circuits responsible for generalization.

Abstract

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.

Figures (4)

• Figure 1: The Gradient-Causal Gap. Gradient magnitude aligns with causal importance in Reverse ($\rho=0.73$), but collapses in Sort ($\rho=0.32$). Scores are min-max normalized per seed; coordinates (1.0,$y$) and ($x$,1.0) denote the maximum gradient and causal components. See Figure \ref{['fig:all_seeds_scatter']} in Appendix \ref{['sec:appendix_components']}.
• Figure 2: Layer-wise distribution of Hidden Heroes and Gradient Bloats in Sorting. Heroes cluster in later layers (L2--L3), while Bloats concentrate in early layers (L0--L1).
• Figure 3: Pruning effects on OOD accuracy for Sorting (Seed 456). Removing low-gradient "Hidden Heroes" devastates generalization. In this seed, pruning high-gradient "Gradient Bloats" paradoxically improves OOD accuracy ($\Delta=+3.5\%$), illustrating the Optimization Noise regime where gradients highlight functionally redundant components.
• Figure 4: Spearman $\rho$ for 10 seeds of Reversal vs. 10 seeds of Sorting. The divergence confirms that as complexity increases, gradient magnitude becomes a stochastic rather than deterministic proxy for causal importance.