Task-Aware Spectral Pruning: A Mixture-of-Masks Framework for Efficient LLM Inference
cs.LG
Submitted: 2026-08-24
Updated: 2026-08-24
Code: https://github.com/tatsu-lab/alpaca_eval
License: http://creativecommons.org/licenses/by/4.0/
The gist: Static pruning imposes one sparse structure on every prompt, even though reasoning, retrieval, generation, coding, and translation can depend on different parts of a language model.
Terminology
Abstract
Static pruning imposes one sparse structure on every prompt, even though reasoning, retrieval, generation, coding, and translation can depend on different parts of a language model. We introduce Task-Aware Spectral Pruning (TASP), a post-training framework that calibrates module-level spectral descriptors against measured task-specific ablation effects, closes grouped-query-attention and SwiGLU dependencies during sparse-mask construction, and routes each user turn to one compiled mask that remains fixed throughout prefill and decoding. A module-disjoint pilot first determines whether the spectral signal is informative before full calibration. Under the stated retrospective operating rule, the pilot passes on the evaluated Llama-3-8B and Llama-3-70B checkpoints but rejects Qwen2.5-1.5B, demonstrating that applicability is model-dependent rather than universal. At a 43% active-FLOP reduction, the Llama-3-70B benchmark harness retains 97.7 +/- 0.2% of the dense BF16 score. In the deployment-matched INT8-weight/BF16-compute runtime on a single A100 80GB, the compiled sparse path retains 97.3 +/- 0.2% relative to dense BF16 and reduces decode latency from 45.2 +/- 0.4 to 31.3 +/- 0.4 ms/token, yielding a 1.44x speedup. Factorized ablations, disjoint-module tests, compiled structured baselines, routing-corruption studies, and an explicit 136-GPU-hour calibration audit further delimit the source and operating regime of these gains
Sources
- Program Synthesis with Large Language Models
- Flextron: Many-in-One Flexible Large Language Model
- Evaluating Large Language Models Trained on Code
- Training Verifiers to Solve Math Word Problems
- DISP-LLM: Dimension-Independent Structural Pruning for Large Language Models
- Diet Your LLM: Dimension-wise Global Pruning of LLMs via Merging Task-specific Importance Score
- LLMs can Compress LLMs: Adaptive Pruning by Agents
- AlphaPruning: Using Heavy-Tailed Self Regularization Theory for Improved Layer-wise Pruning of Large Language Models
- Rethinking generalization requires revisiting old ideas: statistical mechanics approaches and complex learning behavior
- The Llama 3 Herd of Models
- Certificate-Guided Pruning for Stochastic Lipschitz Optimization
- Frustratingly Easy Task-aware Pruning for Large 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