PaperlayerView exact revision on arXiv

Source-linked AI summary

Task-Aware Spectral Pruning: A Mixture-of-Masks Framework for Efficient LLM Inference

Ibne Farabi Shihab, Fariya Afrin, Sanjeda Akter, Anuj Sharma

arXiv:2609.29499v1cs.LG

TL;DR

Static pruning applies one task-agnostic structure despite task-dependent module needs. TASP calibrates spectral features against task-specific ablation effects, builds dependency-closed masks, and routes each turn to a compiled plan; at a 43% active decoding FLOP reduction, it retains 97.7±0.2% in the BF16 benchmark harness.

  • Problem

    Static pruning commits to one task-agnostic compressed structure even though different tasks may depend on different model modules.

    Source: §1 · p. 1

  • Method

    TASP tests spectral-feature predictions against task-specific ablation effects, constructs dependency-closed masks, and routes each turn to a compiled task-conditioned plan.

    Source: §3 · p. 3, §3.1 · p. 4, §6 · p. 8

  • Results

    At a 43% active decoding FLOP reduction, TASP retains 97.7±0.2% in the BF16 benchmark harness; the deployment runtime retains 97.3±0.2% relative to dense BF16 with a 1.44× decode speedup.

    Source: §5.2 · p. 6, §6 · p. 8

  • Takeaways & Limitations

    Where spectral–task separability holds, calibrated spectral descriptors can guide task-conditioned hardware-valid masks that retain more capability than a universal mask at the same active-FLOP budget.

    Source: §6 · p. 8

    TASP requires task definitions stable enough to amortize calibration, which costs 136 GPU-hours on 70B; current candidates target decoder-only Transformers with grouped-query attention and SwiGLU.

    Source: Limitations · p. 9

Abstract

from arXiv · show

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

1 Introduction

Static pruning uses a task-agnostic structure even though different tasks may rely on different model modules. TASP tests whether spectral features predict task-specific ablation sensitivity, then uses them to build and route among task-conditioned masks.

1 Introduction

  • A single task-agnostic compressed structure may incur a “versatility tax” because tasks can depend on different modules.The paper asks whether a frozen model’s task-dependent structure can support a small family of executable subnetworks instead of one compromise mask.Source: §1 · p. 1
  • TASP couples an out-of-sample pilot, task-specific sensitivity calibration, dependency-closed mask construction, and per-turn routing to a compiled plan.The selected plan stays unchanged through prefill and decoding.Source: §1 · pp. 1–2
  • Removing spectral features costs 4.1 retention points, using one global mask instead of five costs 4.4 points, and layer-matched random masks cost 12.1 points at the same FLOP target.These ablations factor the contributions of spectral scoring, mask multiplicity, and learned mask structure.Source: §1 · p. 2
  • The method’s demonstrated scope includes Llama-3-8B and Llama-3-70B, while Qwen2.5-1.5B is near chance; the experiments do not isolate scale from model family or architecture.The stated operating rule is retrospective.Source: §1 · p. 2

2 Related Work

Prior pruning methods vary in their importance signals and structural choices, while task-aware methods use downstream calibration data. TASP instead predicts task-family sensitivity from spectral features and dispatches among precompiled masks.

2 Related Work

  • Structured pruning removes blocks, widths, layers, heads, or coupled channels, whereas unstructured masks may reduce parameters without producing smaller regular matrix operations.TASP’s candidate units are attention heads and coupled feed-forward groups.Source: §2 · pp. 2–3
  • Task-aware pruning uses downstream calibration data, while TASP builds static task-conditioned launch plans that share one packed backing store rather than separate one-task checkpoints.Its low-capacity predictor tests whether spectral features contain an out-of-sample task-sensitivity signal.Source: §2 · p. 2
  • ShadowLLM predicts prompt-specific importance from activation features, whereas TASP predicts task-family sensitivity from weight-space spectral descriptors and keeps one compiled plan fixed per turn.TASP’s out-of-fold pilot can decline to prune a checkpoint; ShadowLLM applies a fresh pattern per input.Source: §2 · p. 2
  • A task-specialized score can sacrifice general-purpose breadth when used globally, and multiple masks alone do not ensure task-relevant structure.TASP tests whether features predict intervention effects and whether routing selects among masks without a larger error than the masks remove.Source: §2 · p. 3

3 Task-Aware Spectral Pruning

TASP tests whether module-level spectral features predict task-specific structural-ablation effects, then uses those predictions to construct dependency-valid masks and route each turn to a compiled plan. A predeployment gate screens checkpoints before full labeling and calibration.

3 Task-Aware Spectral Pruning

  • TASP partitions a frozen decoder-only Transformer into hardware-valid candidate units, measures task effects of removing them, and learns scores for a predeployment applicability gate.The method targets a small set of deployment task families.Source: §3 · p. 3

3.1 Candidate Units and Spectral Features

  • Candidate-unit features include spectral descriptors computed from weight matrices and a calibration-time salience statistic; salience is predictive, not a mechanistic attribution.Feed-forward groups use mean absolute activation on the same token set for the corresponding statistic.Source: §3.1 · pp. 3–4

3.2 Disjoint Sensitivity Calibration and Applicability Gate

  • A module-disjoint pilot uses out-of-fold AUROC to decide whether full calibration proceeds, with a retrospective rule requiring mean AUROC of at least 0.62 and specified diagnostic and confidence-bound thresholds.Generation labels are collected but excluded from the gate decision because isolated-module evaluations are noisier.Source: §3.2 · p. 4

3.3 Dependency-Closed Masks and Cache-Consistent Routing

  • TASP ranks candidates by predicted task sensitivity per active FLOP and closes dependencies before accepting mask additions.Grouped-query attention retains a KV head whenever an associated query head is kept; SwiGLU channels are kept or removed jointly across Wgate, Wup, and matching Wdown columns.Source: §3.3 · p. 4
  • Task masks are compiled as static launch plans over a shared INT8 weight store, and a router selects a plan or falls back to dense INT8 when confidence is below threshold.The chosen plan stays fixed through prefill and decoding so the KV cache is not reused across incompatible masks.Source: §3.3 · pp. 4–5

4 Experimental Design

The experiments evaluate three checkpoints across five task families, using dense BF16 for benchmark claims and dense INT8 for deployment operating points. Structured baselines are compiled into a shared runtime for post-compilation evaluation.

4 Experimental Design

  • The evaluation covers Llama-3-8B, Qwen2.5-1.5B, and Llama-3-70B across reasoning, generation, retrieval, code, and translation.Llama-3-70B is the primary quality and systems reference, and heterogeneous benchmarks are reported with raw scores and macro-average relative retention.Source: §4 · p. 5
  • Dense BF16 is the scientific reference for benchmark claims, while dense INT8 is the systems reference for deployment operating points.This distinction avoids attributing quantization loss to pruning.Source: §4 · p. 5
  • The principal target retains 0.57 of active model FLOPs, a 43% reduction, and compares one global TASP mask, five routed masks, and an oracle task-family assignment.Structured systems are packed into the same groupwise INT8-weight/BF16-compute runtime and reevaluated after compilation.Source: §4 · p. 5
  • TASP variants report mean ± standard deviation over five independent runs, while deterministic baselines use one paired-bootstrap standard error and latency averages five timing blocks.Each latency block contains 200 runs after 50 warmups.Source: §4 · p. 5

5 Results

TASP performs well on the evaluated Llama-3 checkpoints but not Qwen2.5-1.5B, and its quality–latency results depend on the evaluation and serving regime. Ablations, routing tests, and calibration-cost estimates further bound its benefits and intended use.

5.1 Applicability Is Model-Dependent

  • The pilot passes on Llama-3-8B and Llama-3-70B but fails on Qwen2.5-1.5B; direct 8B deployment retains 96.4±0.3% across the eleven-benchmark suite.The reported evidence does not isolate scale from model family or architecture.Source: §5.1 · p. 5, Table 2 · p. 6
  • On untouched modules and disjoint prompts, mean AUROC falls from 0.78 to 0.75; removing least-important groups costs 2.6±0.3 points versus Wanda’s 4.8±0.3.The intervention results support out-of-sample prediction, but not mechanistic identification.Source: §5.1 · pp. 5–6

5.2 Quality and Realized Latency

  • Task-conditioned masks help most on GSM8K and HumanEval, with smaller gaps on NQ-Open and translation; TASP-MOM stays above the global mask across the practically useful 30–60% active-FLOP region.Retention is computed from all eleven benchmark-level ratios, not only the displayed subset.Source: §5.2 · p. 6
  • At 43% active-FLOP reduction, TASP-MOM retains 97.7±0.2% versus 93.3±0.3% for one global spectral mask and 94.0±0.3% for the strongest comparison baseline.In the BF16 scientific harness, the oracle reaches 98.7±0.2%, leaving a one-point routing gap.Source: §5.2 · p. 6, Table 3 · p. 7
  • In the deployment-matched INT8-weight/BF16-compute runtime, TASP-MOM retains 97.3±0.2% relative to dense BF16 and cuts decode latency from 45.2±0.4 to 31.3±0.4 ms/token, a 1.44× speedup.Short-GPT is faster at 1.63×, but TASP retains 5.8 points more BF16-relative quality.Source: §5.2 · p. 6, Table 4 · p. 7

5.3 Spectral Scoring and Mask Diversity Are Complementary

  • Removing spectral features costs 4.1 points, one global spectral mask costs 4.4 points, and five random masks lose 12.1 points.Across task families, mask overlap ranges from 0.24±0.03 for reasoning/generation to 0.81±0.02 for generation/translation, indicating a stable core with replaceable modules.Source: §5.3 · p. 7, Table 5 · p. 7
  • When top-1 route accuracy falls from 100% to 70%, no-fallback retention drops from 98.7±0.2 to 95.2±0.4; dense fallback restores 97.5±0.3.Using half the training module–prompt labels yields held-out AUROC 0.70±0.02 and retention 96.7±0.3, compared with 0.75±0.02 and 97.7±0.2 with all labels.Source: §5.3 · pp. 7–8

5.4 Calibration Cost and Serving Regime

  • Full calibration costs 136.0 A100 GPU-hours, while an unsuitable checkpoint can be rejected after about 15.8 hours.At τ=0.85, the selected policy breaks even after about 159,000 requests (40.8 million generated tokens), defining a sustained-serving regime with stable task families rather than one-off inference.Source: §5.2 · p. 6, §5.4 · p. 8

6 Discussion and Conclusion

Within its supported regime, TASP combines a pilot for spectral–task separability with dependency-valid task-conditioned masks. The study reports quality retention with lower decode latency, while noting limits on applicability and causal interpretation.

Discussion and Conclusion

  • TASP is supported when a frozen model exhibits out-of-sample spectral–task separability, and its applicability is not universal across the evaluated checkpoints.Qwen2.5-1.5B is near chance under the pilot, while Llama-3-8B and Llama-3-70B show useful separability.Source: §6 · p. 8
  • The study does not disentangle scale, architecture, training corpus, and model family, and its intervention analyses do not establish a unique reasoning circuit.Stronger causal claims would require activation patching or causal mediation.Source: §6 · p. 8
  • At a 43% active-FLOP reduction, TASP retains 97.7±0.2% in the BF16 benchmark harness and 97.3±0.2% relative to dense BF16 in the INT8-weight/BF16-compute runtime, with a 1.44× decode speedup.The runtime comparison is on a single A100; the paper distinguishes it from the BF16 benchmark harness.Source: §6 · p. 8
  • The applicability rule remains in need of prospective validation across broader model families, despite a six-checkpoint test reporting 100% precision and 75% recall.That test also reports Spearman ρ = 1.00 between pilot AUROC and downstream retention on six points.Source: §6 · p. 8

Limitations

TASP’s deployment regime is constrained by calibration amortization, model architecture, distribution shift, and the specific runtime used for latency measurement.

Limitations

  • TASP requires stable task definitions to amortize its 136-GPU-hour 70B calibration, and its retrospective pilot thresholds need prospective validation.Current candidates target decoder-only Transformers with grouped-query attention and SwiGLU; other architecture families require new structural units and compatibility rules.Source: Limitations · p. 9
  • Aggregate retention on the English-dominated suite does not establish safety under task or language shift, where the router can be overconfident.The measured latency is specific to one A100 runtime, sequence shape, batch size, quantization scheme, and compiler.Source: Limitations · p. 9

Ethics Statement

The supplied sections situate TASP among pruning and conditional-computation methods, describe its offline and online procedure, and delimit the reported evaluations and deployment scope.

Ethics Statement

  • Deployment should report task- and subgroup-specific quality, use conservative fallback thresholds for safety-critical traffic, and weigh calibration’s environmental burden against the stated break-even serving volume.The method introduces no human-subject data, and benchmark and model use should follow underlying resource licenses and intended-use restrictions.Source: Ethics Statement · p. 9

A Extended Related Work

  • Unlike methods that primarily produce one compressed architecture, TASP derives static task-conditioned subnetworks, routes once per turn, and keeps the selected mask fixed.The related-work discussion distinguishes this approach from expert routing, token-wise computation changes, and early exit.Source: A Extended Related Work · pp. 11–12

B End-to-End Procedure

  • For Llama-3-70B, the procedure evaluates 7,360 structured candidates and conducts 1,089,280 evaluations in the complete pipeline.The pilot is performed before full labeling, and compilation precedes router timing; only routing and dispatch are online.Source: B End-to-End Procedure · p. 12, Table 6 · p. 13

C Detailed Pipeline Views

  • Figures 2–4 trace sensitivity estimation, dependency-valid mask construction, and inference-time routing with cache-consistent execution.The diagrams expand the decision path and its feature, dependency, and execution interfaces.Source: C Detailed Pipeline Views · p. 12

D Full Benchmark Results

  • Table 7 reports Llama-3-70B quality at matched active-FLOP reduction, while the displayed task columns are representative of a broader benchmark set.The broader retention column includes MT-Bench, TriviaQA, MBPP, and both WMT14 directions; this is the BF16 scientific harness, not the quantized timing runtime.Source: D Full Benchmark Results · p. 12, Table 7 · p. 13

E Accuracy–FLOP Pareto Frontier

  • Across 30–60% active-FLOP reduction, TASP-MOM remains above the global mask, while the learned-to-oracle gap widens at the most aggressive sparsity.The paper retains the matched 43% point as the primary comparison because structured baselines were compiled and timed there under the same runtime.Source: E Accuracy–FLOP Pareto Frontier · p. 12

F Applicability and Disjoint Diagnostics

  • Across three evaluated checkpoints, Llama-3-8B falls between near-chance Qwen and stronger Llama-3-70B separation, but these checkpoints cannot distinguish scale, architecture, family, or training-data explanations.The 70B check reports 63% of measured reasoning sensitivity, compared with 28% for a layer-matched random control. Table 9 reports direct Llama-3-8B deployment at 43% active-FLOP reduction.Source: F Applicability and Disjoint Diagnostics · pp. 12–13, Table 9 · p. 13

G Exploratory Taskwise Ablation Diagnostic

  • The taskwise ablation diagnostic is exploratory and uses the original calibration/validation-module protocol; confirmatory claims instead rely on the module- and prompt-disjoint comparison.Figure 2 depicts calibrated sensitivity estimation, and Table 10 reports degradation on untouched modules and disjoint prompts.Source: G Exploratory Taskwise Ablation… · p. 14, Table 10 · p. 14, Figure 2 · p. 14

H Routing, Calibration, and Task-Switching Robustness

The section describes natural dense fallback, controlled routing corruption, limited-label calibration, and task switching across user turns.

Routing, Calibration, and Task-Switching Robustness

  • Below-threshold requests use dense INT8 with the hybrid path disabled, and the reported natural deployment curve uses mixture-implied decode values.Only the all-sparse endpoint is a directly timed sparse-kernel result; sparse-path usage is denoted by u(τ).Source: H Routing, Calibration, and… · p. 14, Table 12 · p. 15
  • The routing-corruption study replaces a fixed fraction of predictions with confusion-matrix-matched incorrect routes, so its 100% row is constructed rather than a natural confidence condition.Corrupted rows draw confidence from the empirical error-confidence distribution.Source: H Routing, Calibration, and… · p. 14, Table 13 · p. 15
  • The limited-label study samples complete training module–prompt pairs jointly while keeping validation and test modules and prompts fixed.Its 50% condition uses exactly half the labels, rather than independently subsampling modules and prompts.Source: H Routing, Calibration, and… · p. 14
  • Routing once per user turn allows task-family changes across a conversation; per-turn rerouting closes most of the gap to dense, while locking the first mask allows degradation to accumulate.The selected mask remains fixed within each turn’s prefill and decoding phases.Source: H Routing, Calibration, and… · pp. 14–15, Figure 3 · p. 15

I Cross-Task Structure and Within-Task Multiplicity

The results distinguish cross-task specialization from within-task nonuniqueness, and examine how those masks transfer, compare with low-rank compression, and amortize deployment costs.

I Cross-Task Structure and Within-Task Multiplicity

J Offline Cost, Memory, and Amortization

  • 121,000 requests is the workload-specific break-even point for the 136-GPU-hour calibration cost under the reported dense-versus-sparse request timings.With τ = 0.85 and 76% sparse usage, the estimate is about 159,000 requests; these are not universal efficiency claims.Source: J Offline Cost, Memory, and Amortization · p. 17

K Transfer and Compression Controls

  • At the same 43% active-FLOP reduction, a mask transferred from Llama-3-8B recovers 92% of the performance obtained by direct Llama-3-70B calibration.The largest gaps occur on GSM8K and HumanEval; the result supports partial reuse within a family, not a replacement for target-model calibration.Source: K Transfer and Compression Controls · p. 17
  • At matched active-FLOP budgets, structured pruning leads low-rank approximation on attention heads, while low-rank approximation is competitive for feed-forward matrices.The reported head advantages are 1.9 retention points on GSM8K and 1.4 on HumanEval; low-rank remains within 0.3 points on AlpacaEval and 0.5 on NQ-Open.Source: K Transfer and Compression Controls · p. 18

L Statistics, Seeds, and Reproducibility

  • TASP uncertainty is measured across five independent calibration, mask, and router runs, while deterministic baseline uncertainty comes from paired bootstrap over evaluation instances.Latency uses five timing blocks, each with 50 warmups and 200 measured requests; loading, packing, and compilation are outside steady-state timing.Source: L Statistics, Seeds, and Reproducibility · pp. 18–19

M Supplementary Result Visualizations

The supplementary visuals show complementary ablation contributions, matched-retention comparisons, run-to-run variability, and task-level differences across methods.

M Supplementary Result Visualizations

N Benchmark-Provenance Audit and a Genuine MMLU-Pro Evaluation

This block audits benchmark provenance and tests the label efficiency of spectral ranking, router generalization, matched baselines, and the applicability rule. It also specifies candidate construction and validates the shared-weight-store implementation.

N Benchmark-Provenance Audit and a Genuine MMLU-Pro Evaluation

  • The score of 79.2 is original MMLU, not MMLU-Pro; the genuine MMLU-Pro evaluation is reported separately.The fixed harness produces dense MMLU 79.2 and dense MMLU-Pro 55.9; the former differs from Meta’s instruct reference because of the prompt and harness, not a different checkpoint.Source: N Benchmark-Provenance Audit and a… · p. 24, Table 24 · p. 24

O Direct Measured-Ablation Ranking at Full and Matched Label Cost

  • At the 25% and 50% label budgets, the preregistered success criterion requires spectral ranking to beat direct-plus-stratum imputation by at least 0.7 retention points and finish within 1.5 points of the all-label oracle.At the method-full budget, TASP closes about 77% of the gap from the no-spectral variant to the all-label direct oracle, so it is close to, but not equal to, exhaustive ranking.Source: O Direct Measured-Ablation Ranking at… · pp. 24–25, Table 25 · p. 25

P Natural Router Accuracy, Calibration, and Source-Disjoint Generalization

  • The calibrated in-suite router still drops on source-disjoint data; the table identifies fallback, rather than raw routing, as the robustness mechanism.The evaluation includes prompt-, template-, dataset-, and mixed-domain splits and measures classification, calibration, coverage, and retention.Source: P Natural Router Accuracy,… · p. 25, Table 26 · p. 26

Q Matched Closest-Work Baselines

  • At 70B and matched 43% active-FLOP reduction, the comparison is a Pareto trade-off: depth methods remain faster while TASP retains more quality.At 8B, ShadowLLM’s native per-prompt head/neuron comparison retains slightly less than TASP-MOM and reports 1.51× speedup.Source: Table 29 · p. 26, Table 28 · p. 26

S Candidate Construction, FLOP Ledger, and One-Store Validation

  • Each candidate is an attention head or a 1,024-channel SwiGLU group, with attention-head descriptors built from the coupled Q, O, K, and V projections.SwiGLU descriptors likewise stack coupled gate, up, and down components; the alternatives test separate per-matrix descriptors or dropping K/V and down blocks.Source: S Candidate Construction, FLOP Ledger,… · pp. 26–27

R Prospective Validation of the Diagnostic

  • On six unseen checkpoints, the frozen applicability rule achieves 100% precision and 75% recall, with Spearman ρ = 1.00 between pilot AUROC and downstream retention.The authors characterize the rule as conservative and useful for avoiding poor fits, not as identifying every viable model; the association is based on six points.Source: Table 30 · p. 27

S Candidate Construction, FLOP Ledger, and One-Store Validation

  • The active-FLOP target is the retained per-token computation ratio against full width, with embeddings, final norm, and LM head remaining dense in both totals.Shared-KV projection cost is charged once per active GQA group, regardless of how many query heads survive.Source: S Candidate Construction, FLOP Ledger,… · p. 27
  • Across mask counts, the shared-store plan meets the stated equivalence target while memory rises with plan metadata rather than duplicated weights.The reported peak allocation is 77.3 GB; five duplicated 70B INT8 stores would require roughly 353 GB before scales and workspace.Source: S Candidate Construction, FLOP Ledger,… · p. 28, Table 32 · p. 29
Loading 2609.29499v1…