CausalGate: Causal Importance Distillation for Transformer Module Pruning

2026-07-21

AI Summary

This work addresses the limitations of existing adaptive inference methods for large language models, which rely on heuristic metrics and struggle to accurately identify nonlinear computation modules critical to semantic fidelity. The authors propose CausalGate, a novel framework that introduces causal intervention into module pruning for the first time. During calibration, it quantifies the causal effect of attention and MLP sublayers on model outputs by masking them and measuring the resulting KL divergence. By integrating differentiable ranking loss with exponential moving average, CausalGate distills dynamic importance scores into static gating decisions without runtime overhead. Evaluated on TinyLlama-1.1B, Qwen2.5-3B, and Llama-3.1-8B, the method significantly outperforms dynamic routing and layer-skipping baselines, effectively translating theoretical computational savings into tangible hardware latency reductions.

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.

Problem

• adaptive inference
• module pruning
• causal importance
• transformer efficiency
• semantic accuracy

Innovation

• Causal Intervention
• Module Pruning
• Transformer Efficiency
• Static Gate Distillation
• KL Divergence-based Importance