CausalGate: Causal Importance Distillation for Transformer Module Pruning
Kiran Nair, Smriti Regmi, Rodrigue Rizk
cs.LG, cs.CL, cs.CV, stat.ML
Submitted: 2026-07-21
Comments: In-Review at a conference
License: http://creativecommons.org/licenses/by/4.0/
The gist: Existing adaptive inference methods for Large Language Models rely on observational heuristics, such as hidden-state similarity or activation magnitudes, to drop redundant modules.
Terminology
Abstract
Existing adaptive inference methods for Large Language Models rely on observational heuristics, such as hidden-state similarity or activation magnitudes, to drop redundant modules. However, these correlation-based metrics often fail to capture subtle, non-linear structural computations vital for semantic accuracy. We introduce CausalGate, an intervention-guided framework for compute-efficient transformer inference. During a calibration phase, CausalGate isolates individual Attention and MLP sub-layers, zeros out their respective outputs, and measures the exact semantic damage via the Kullback-Leibler divergence of the final logit distribution. To eliminate runtime routing overhead, this structural importance hierarchy is distilled into a global set of static, lightweight scalar gates using an Exponential Moving Average smoothing objective paired with a differentiable pairwise ranking loss. Evaluated on TinyLlama-1.1B, Qwen2.5-3B, and Llama-3.1-8B across language modeling and commonsense reasoning benchmarks, CausalGate consistently outperforms prominent dynamic routing and layer-skipping baselines, translating theoretical compute savings into concrete hardware latency reductions with zero operational overhead.
Sources
- Batch-Shaping for Learning Conditional Channel Gated Networks
- PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling
- GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers
- Dynamic Channel Pruning: Feature Boosting and Suppression
- What Layers When: Learning to Skip Compute in LLMs with Residual Gates
- Accelerating Inference in Large Language Models with a Unified Layer Skipping Strategy
- Adaptive Layer-skipping in Pre-trained LLMs
- Mixture-of-Depths: Dynamically allocating compute in transformer-based language models
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks