Task-Aware Spectral Pruning: A Mixture-of-Masks Framework for Efficient LLM Inference
Organizations: Department of Computer Science, Iowa State University · Department of Computer Science, Kalinga Institute of Industrial Technology · Department of Civil, Construction & Environmental Engineering, Iowa State University
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
Figures & tables
| Component | Configuration | Component | Configuration |
| 70B candidates | 80 layers; 64 query-head candidates and 28 coupled 1,024-channel SwiGLU groups per layer; 7,360 total | Module split | Stratified by layer and type: 70/15/15% train/validation/untouched test, or 5,152/1,104/1,104 modules |
| Pilot | Validation modules; first 16 sequences per module–task pair; stratified five-fold outer module CV | Ablation labels | 32 sequences per train/validation module–task pair; 16 per untouched-test pair; batch size 8 |
| Targets | Normalized gold-answer log-likelihood for reasoning, retrieval, translation; teacher-forced completion NLL for code, generation | Feature preprocessing | Per-task z-normalization fitted only on the modules available to that fitting split |
| Predictor | Ridge regression; coefficient grid | Ridge selection | Lowest validation MSE; ties resolved by validation AUROC |
| Mask objective | Greedy predicted sensitivity per active FLOP with dependency closure; 0.57 active-FLOP target | Dependencies | Retain a KV head with any associated query head; jointly retain/remove gate, up, and down channels in groups of 1,024 |
| Router | Three-layer MLP with hidden widths 512 and 256 | Router optimization | AdamW; learning rate ; weight decay 0.01; batch 64; at most 20 epochs; patience 3 |
| Model | Mean | Minimum | Mean LCB | Gate |
| Qwen2.5-1.5B | 0.53 | 0.51 | 0.50 | reject |
| Llama-3-8B | 0.66 | 0.63 | 0.62 | pass |
| Llama-3-70B | 0.78 | 0.74 | 0.75 | pass |
| Method | MMLU | GSM8K | MGSM | AlpacaEval | NQ-Open | HumanEval | WMT14 avg. | Full retention |
| Dense BF16 | 79.2 | 87.5 | 74.3 | 82.1 | 68.4 | 72.6 | 28.3 | 100.0 |
| Wanda | ||||||||
| SliceGPT | ||||||||
| LLaMaFlex | ||||||||
| TASP -Single | ||||||||
| TASP - MoM |
| Method | Selected density | Active FLOPs | Ret. vs. BF16 | Ret. vs. INT8 | Prefill (ms) | Decode (ms/tok.) | Speedup |
| Dense INT8 runtime | |||||||
| TASP - MoM mask-as-zero | |||||||
| ShortGPT, compiled | |||||||
| LLM-Pruner, compiled | |||||||
| SliceGPT, compiled | |||||||
| TASP -Single, compiled |
| Variant | Retention vs. dense BF16 |
| Five masks, oracle router | |
| Five masks, learned router | |
| Without salience | |
| Without spectral features | |
| One global spectral mask | |
| Five layer-matched random masks |
Appendix figures & tables41 assets
Supplementary material from the paper’s appendix.
Appendix
| Stage | Inputs and split | Operation and output |
| Candidate formation | Frozen checkpoint and architecture metadata | Slice every query head and each 1,024-channel SwiGLU group; record GQA and gate/up/down dependencies. |
| Streaming profile | One BF16 layer or matrix group at a time | Compute spectral descriptors without placing the full BF16 checkpoint on the GPU. |
| Pilot labels | Validation modules; 16 sequences per module–task pair | Measure structural-ablation targets and obtain out-of-fold predictions through nested, stratified module CV. |
| Applicability gate | Four diagnostic-family AUROCs and hierarchical bootstrap | Stop if mean , any family , or the mean 95% LCB ; otherwise finish label collection. |
| Final calibration | Train/validation modules and disjoint untouched test | Fit task-specific ridge predictors, select regularization on validation modules, and evaluate once on unseen modules and prompts. |
| Mask construction | Scores, active-FLOP budget, dependency graph | Greedily select by score per active FLOP, close GQA and SwiGLU dependencies, and emit one mask per task family. |
| Method | FLOP red. | MMLU | GSM8K | MGSM | AlpacaEval | NQ-Open | HumanEval | WMT14 avg. | Full retention |
| Dense BF16 | 0% | ||||||||
| Magnitude | 43% | ||||||||
| SparseGPT | 43% | ||||||||
| Wanda | 43% | ||||||||
| Gradient | 43% | ||||||||
| AlphaPruning | 43% |
| Task | Pilot AUROC | Test AUROC | Change |
| Reasoning | 0.82 | 0.80 | |
| Retrieval | 0.76 | 0.72 | |
| Code | 0.79 | 0.76 | |
| Translation | 0.74 | 0.73 | |
| Mean | 0.78 | 0.75 |
| Llama-3-8B task family | TASP - MoM retention |
| Reasoning | |
| Generation | |
| Retrieval | |
| Code | |
| Translation | |
| Full-suite mean |
| Selection rule | Mean degradation |
| TASP least-important groups | |
| Wanda least-important groups | |
| Layer-matched random groups |
| Rule | Reasoning | Retrieval | Code | Translation | Average |
| Random | |||||
| Magnitude | |||||
| Activation norm | |||||
| Wanda | |||||
| TASP |
| Ret. vs. INT8 | Sparse use | ms/token | Speedup | |
| 0.50 | 31.3 | |||
| 0.70 | 32.6 | |||
| 0.85 | 34.6 | |||
| 0.95 | 38.0 |
| Top-1 accuracy | No fallback | With | Fallback rate |
| 100% | 0% | ||
| 90% | 12% | ||
| 80% | 24% | ||
| 70% | 37% |
| Labels | Held-out AUROC | TASP - MoM ret. |
| 25% | ||
| 50% | ||
| 75% | ||
| 100% |
| Scenario | Dense | First-turn locked | Per-turn routing |
| Math creative | |||
| Code translation | |||
| Retrieval reasoning | |||
| Multi-task average |
| Task-mask pair | Jaccard overlap |
| Reasoning / Code | |
| Generation / Translation | |
| Reasoning / Generation | |
| Retrieval / Translation | |
| Code / Translation |
| Task | Across-run within-task Jaccard | Distinct masks within 0.5 points | Retention range |
| Reasoning | 13 | 97.0–97.3 | |
| Retrieval | 18 | 97.2–97.6 | |
| Code | 10 | 96.7–97.1 | |
| Translation | 16 | 97.9–98.2 |
| Offline component | Prompt–module evaluations | A100 GPU-hours | Peak allocated memory |
| Layer-streamed BF16 spectral profiling | 0 | 5.6 | 11.8 GB |
| 15% pilot separability stage | 88,320 | 10.2 | 77.1 GB |
| Remaining train/validation labeling | 912,640 | 102.6 | 77.3 GB |
| Module- and prompt-disjoint test labeling | 88,320 | 11.7 | 77.3 GB |
| Predictor fitting and cross-validation | — | 0.3 | 3.2 GB |
| FLOP-constrained mask construction | — | 0.2 | 1.7 GB |
| Method | A100 GPU-hours | Peak allocated |
| Wanda | 7.4 | 76.5 GB |
| SparseGPT | 13.1 | 77.0 GB |
| SliceGPT | 10.3 | 76.8 GB |
| LLM-Pruner | 19.6 | 77.1 GB |
| TASP | 136.0 | 77.3 GB |
| Live component | Allocation |
| Packed INT8 weights | 70.60 GB |
| Scales and quantization metadata | 1.10 GB |
| Five mask plans and routing metadata | 0.54 GB |
| KV cache, activations, and workspace | 5.06 GB |
| Peak allocated | 77.30 GB |
| Variant | S1 | S2 | S3 | S4 | S5 | Mean SD |
| TASP -Single | 92.9 | 93.1 | 93.7 | 93.5 | 93.3 | |
| TASP - MoM | 97.5 | 97.6 | 98.0 | 97.8 | 97.7 | |
| TASP -Oracle | 98.5 | 98.6 | 99.0 | 98.8 | 98.7 |
| Variant | S1 | S2 | S3 | S4 | S5 | Mean SD |
| Oracle router | 98.5 | 98.6 | 99.0 | 98.8 | 98.7 | |
| Learned router | 97.5 | 97.6 | 98.0 | 97.8 | 97.7 | |
| Without salience | 95.2 | 95.4 | 95.9 | 95.6 | 95.5 | |
| Without spectral features | 93.0 | 93.4 | 94.1 | 93.8 | 93.7 | |
| One global mask | 92.9 | 93.1 | 93.7 | 93.5 | 93.3 | |
| Layer-matched random | 85.0 | 85.4 | 86.2 | 86.0 | 85.5 |
| Metric | S1 | S2 | S3 | S4 | S5 | Mean SD |
| MMLU | 77.5 | 77.8 | 78.3 | 78.0 | 77.9 | |
| GSM8K | 84.7 | 85.0 | 85.7 | 85.4 | 85.2 | |
| MGSM | 72.4 | 72.6 | 73.1 | 72.9 | 72.8 | |
| AlpacaEval | 80.0 | 80.2 | 80.8 | 80.5 | 80.4 | |
| NQ-Open | 67.0 | 67.1 | 67.5 | 67.3 | 67.2 | |
| HumanEval | 69.1 | 69.5 | 70.5 | 70.0 | 69.8 |
| Method | Original MMLU | MMLU-Pro | MMLU-Pro ret. |
| Dense BF16 | 79.2 | 55.9 | 100.0 |
| TASP -Single | |||
| TASP - MoM | |||
| TASP -Oracle |
| Ranking rule | Labels used | Rel. ret. at 43% |
| No-spectral calibrated variant | 31,280 | |
| Direct + imputation, 25% | 7,820 | |
| Spectral predictor, 25% | 7,820 | |
| Direct + imputation, 50% | 15,640 | |
| Spectral predictor, 50% | 15,640 | |
| Direct + imputation, full | 31,280 |
| Split | Top-1 | Macro-F1 | Cal. ECE | No-fb ret. | Fb ret. | Sparse |
| Prompt-disjoint | ||||||
| Template-disjoint | ||||||
| Dataset-disjoint | ||||||
| OOD/mixed (consensus) |
| True Pred. | Reas. | Gen. | Retr. | Code | Transl. |
| Reasoning | 93 | 1 | 2 | 3 | 1 |
| Generation | 1 | 94 | 2 | 1 | 2 |
| Retrieval | 3 | 2 | 91 | 2 | 2 |
| Code | 5 | 1 | 1 | 92 | 1 |
| Translation | 1 | 2 | 2 | 1 | 94 |
| Method | Training mode | Ret. @43% | Speedup | A100-h |
| Dense | none | 100.0 | 0 | |
| TaBP | calibration only | 5.8 | ||
| PuDDing | router + masks | 7.6 | ||
| ShadowLLM | per-prompt predictor | 9.4 | ||
| IG-Pruning | masks + routing | 11.9 | ||
| Instr.-Following Pr. | joint training | 86.0 |
| Method | Ret. @43% | Decode speedup |
| TaBP | ||
| PuDDing | ||
| ShadowLLM | ||
| IG-Pruning | ||
| TASP - MoM |
| New checkpoint | Mean | Min | LCB | Decision | Ret. |
| Llama-3.1-8B-Instruct | 0.68 | 0.64 | 0.63 | pass | |
| Gemma-2-9B-It | 0.65 | 0.60 | 0.61 | pass | |
| Mistral-Nemo-12B-Inst. | 0.63 | 0.59 | 0.60 | pass | |
| Qwen2.5-7B-Instruct | 0.59 | 0.56 | 0.55 | reject | |
| Mistral-7B-Inst.-v0.3 | 0.60 | 0.57 | 0.56 | reject | |
| Qwen2.5-14B-Instruct | 0.61 | 0.58 | 0.58 | reject |
| Candidate feature construction | AUROC | Ret. |
| Canonical stacked Q/O/K/V + gate/up/down | ||
| Per-matrix descriptors, separate fits | ||
| Query/output and gate/up only | ||
| Independently normalized, then averaged |
| Validation quantity | 1 mask | 3 masks | 5 masks |
| Plan/coordinate metadata | 0.14 GB | 0.34 GB | 0.54 GB |
| Peak allocated memory | 76.9 GB | 77.1 GB | 77.3 GB |
| Analytical active-FLOP ratio | 0.570 | 0.570 | 0.570 |
| Ratio from launched tile shapes | 0.571 | 0.571 | 0.571 |
| Max abs. logit diff. (achieved) | |||
| Greedy-token agreement (achieved) |
LLM Pruning
Large language model pruning removes less useful weights or components to reduce memory and inference costs while preserving model capabilities.
See all 178 papers