QK-Wanda: Coupling Queries and Keys for Unstructured Pruning
QK-Wanda attacks a blind spot in activation-aware LLM pruning: Wanda scores query and key weights independently even though their dot products jointly determine attention. The paper derives closed-form single-deletion costs under a coupled QK reconstruction objective, letting queries and keys share a pruning budget without gradients or weight updates. Across 15 models up to 72B parameters, the method cuts local QK reconstruction error by roughly half, yet downstream gains are uneven—Llama 2 and Qwen2.5 improve while Llama 3 degrades—highlighting that local reconstruction fidelity is an imperfect proxy for end-to-end model quality.
The paper makes a clean, theoretically grounded contribution to unstructured pruning by coupling query and key scores through exact single-deletion costs. The shared-budget mechanism is elegant and inexpensive, adding only 1.3–3.1% runtime overhead versus Wanda. However, the downstream story is bifurcated: strong Llama 2 and Qwen2.5 results coexist with severe perplexity regression on Llama 3.1-70B, and the authors candidly note that lower QK error does not guarantee better PPL or task accuracy. The work is a solid incremental advance, but its practical reliability depends heavily on model family.
The mathematical derivation is the clearest strength. Equation 7 extends the base score to grouped-query attention by summing query-head energies for shared keys, and Equation 8 frames mask selection as minimizing an additive surrogate of exact individual costs. The empirical protocol is also praiseworthy: 15 checkpoints, controlled ablations separating scoring changes from budget allocation, and calibration-sensitivity studies. The authors' transparency about model-dependent outcomes is refreshing.
The central limitation is the disconnect between local reconstruction and downstream quality. The paper documents this explicitly—Llama 3.1-70B at 80% sparsity jumps from 4.98/9.89 to 18.89/32.59 WikiText-2/C4 perplexity even though QK error drops from 5.27% to 3.27%—yet offers no deeper diagnostic beyond labeling it a limit of local objectives.
Additionally, the base objective ignores causal masking, RoPE, and the softmax nonlinearity, all central to real attention; while Appendix P tests eight variants, none consistently improve downstream metrics. The narrow QK-only scope also means the paper does not validate end-to-end inference speedups, and the claim of minimal overhead applies only to the pruning script, not to sparse inference.
The evidence robustly supports the narrow claim that coupled scoring improves QK reconstruction, but it does not consistently generalize to downstream metrics. Comparisons to Wanda are fair—same calibration, same mask-only setting, same sequential pruning—but the authors deliberately exclude stronger baselines like SparseGPT and Thanos that update surviving weights. The ablation distinguishing scoring from allocation is well-executed, yet without statistical testing, sign counts like 'lower in 103 of 135' remain merely descriptive. The near-absence of Llama 3 improvements suggests the method is not family-agnostic.
Reproducibility is strong on paper: the authors publish code and report exact calibration hyperparameters, including 524,288 C4 tokens segmented into 8192 sequences of length 64, seed 0, and FP16/BF16 settings. Hardware and timing protocols are documented. The biggest barrier to independent reproduction is computational access to 70B/72B checkpoints and the GPUs needed to prune them. A minor gap is that the exact token indices for the C4 windows are not fully specified in the excerpt, though the seed and segmentation logic may suffice.
Wanda (Sun et al., 2024) prunes large language models by scoring weights independently within each linear projection, although queries and keys interact through dot products. We introduce QK-Wanda, which scores query and key weights by their individual deletion costs under an unmasked pre-RoPE reconstruction objective. It augments Wanda scores with information from the opposite projection (keys for query weights, and queries for key weights), allowing both projections to share a pruning budget. Its closed-form scores require no gradients or weight updates; full pruning takes 1.3% longer than Wanda on A100 and 3.1% longer on H200 with the calibration used in our main experiments. We evaluate QK-only pruning across 15 models from TinyLlama, Llama 2, Llama 3, and Qwen2.5, spanning 0.5B-72B parameters. Relative to Wanda, QK-Wanda reduces QK reconstruction error by an average of 60% at 50% sparsity and 45% at 80%. Downstream gains depend on the model. At 80% sparsity on Llama 2 70B, WikiText-2 and C4 perplexity decrease by 20.3% and 13.5%, while mean zero-shot accuracy rises by 5.94 percentage points. Qwen2.5-72B also improves, but Llama-3.1-70B has substantially higher perplexity despite lower reconstruction error. These results show both the promise of coupled pruning criteria and the limits of local reconstruction as a predictor of model quality.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.