{"id":"f6905724-a959-4dec-94fe-866ac47c896e","arxiv_id":"2412.20677","paper_version":2,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":5.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":3,"one_line_summary":"Applying Procrustes-based orthogonal alignment to attention heads before L0-pruned merging produces standard GQA models with up to 87.5% KV-head compression on LLaMA2-7B.","lead":"Researchers propose aligning attention heads with orthogonal rotations before merging them, using Procrustes analysis and L0 pruning to convert multi-head attention into grouped-query attention. On LLaMA2-7B and Sheared-LLaMA-1.3B, the approach compresses up to 87.5% of KV heads with modest accuracy loss on eight commonsense benchmarks.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The Procrustes alignment's dependence on a single 128-sequence C4 calibration set is untested; a miscalibrated rotation would vitiate the core mechanism, and some ablations already show the transformation can hurt.","rationale":"The reader's weakest_assumption is exactly the representativeness of the 128-sequence C4 calibration set, and I agree that this is the most load-bearing condition for the central claim. The alignment step is the only difference between the method and the baseline; if the rotations are miscalibrated, the claimed benefit of 'aligning before merging' disappears or reverses. I verified that the paper provides no sensitivity analysis and that some ablations already show the transformation can hurt, which underscores the risk. The concrete test would settle the concern by measuring how much the final accuracy depends on the calibration distribution and whether the alignment gains hold out-of-distribution. I do not see a more fundamental flaw: the value-side invariance math is correct, the L0 pruning pipeline is plausible, and the released code allows verification. The reader's CONDITIONAL verdict is appropriate; my proposed test would either add a condition (calibration robustness) or confirm the current result, so no verdict change is needed.","tokens_in":13613,"tokens_out":7998,"duration_ms":78632,"concrete_test":"Re-run the LLaMA2-7B GQA-8 conversion in Table 1 (default grouping, dist) using three different calibration sets: (a) the original C4 sample, (b) WikiText-2, and (c) the SFT training data (the benchmark train splits). Keep all hyperparameters fixed and report average accuracy over the eight evaluation tasks. Also compute the average cosine-similarity gain between heads on a held-out C4 sample not used for calibration. If the best-configuration average accuracy varies by more than 1 point across calibration sets, or if the held-out similarity gain is less than half the calibration-set gain, the alignment is calibration-sensitive and the central claim is not robust.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The method's central advantage is the orthogonal alignment of KV caches before L0-pruned merging. The rotations are estimated from 128 C4 sequences of 2048 tokens (Section 4.1), then frozen before pruning training. The paper provides no sensitivity analysis over calibration data, and Figure 2 only reports cosine-similarity gains on the calibration set itself. If those rotations do not generalize to the deployed/fine-tuned distribution, the aligned heads are not actually mergeable at inference time, so the L0 pruning must compensate on its own and the reported gains over baseline would not be robust. The risk is not merely hypothetical: in Table 1, several transformed configurations underperform the no-transformation baseline (e.g., GQA-8 grouping-by-value cos, avg 80.32 vs baseline 81.65; GQA-8 grouping-by-key dist, 83.94 vs baseline 81.65 is better, but value cos is worse). This shows that the transformation is not universally beneficial and that its sign depends on the calibration-derived grouping and rotation. Because the central claim is that aligning attention heads improves conversion quality, the absence of any held-out validation of the alignment step is a load-bearing gap. The paper's stated limitation ('our method entirely relies on the statistical mathematical features of attention heads') does not address distribution shift or calibration robustness.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes a two-stage method for converting a pre-trained multi-head attention (MHA) model into a grouped-query attention (GQA) model. In the first stage, the authors collect KV caches from a small calibration set (128 C4 sequences of 2048 tokens), compute pairwise similarity scores between key heads and between value heads under an orthogonal Procrustes alignment, group heads by simulated-annealing search to maximize intra-group similarity, and then fuse the resulting orthogonal rotations into the query, key, and value projection matrices so that the model output is exactly preserved. In the second stage, L0 regularization with a hard-concrete mask is used to gradually transfer each original KV head to a new shared head initialized by mean-pooling within each group, with logits distillation from the fine-tuned MHA teacher. The final model is a standard GQA model. Experiments are reported on LLaMA2-7B and Sheared-LLaMA-1.3B for GQA-16, GQA-8, and GQA-4 configurations, evaluated on eight commonsense QA datasets. The authors claim the method compresses up to 87.5% of LLaMA2-7B KV heads and 75% of Sheared-LLaMA-1.3B KV heads with acceptable performance degradation.","tokens_in":13947,"tokens_out":10098,"duration_ms":97844,"significance":"If the central claim is correct, the paper offers a practical recipe for converting MHA checkpoints to GQA at arbitrary (divisor-based) compression ratios, with a computationally cheap alignment step and a training procedure based on L0 pruning and distillation. The output-invariance argument for the value-side transformation (Eqs. 12-13) and for the RoPE-compatible key-side transformation (Eqs. 14-17) is mathematically sound, and the release of code is a strength. The generalized Procrustes formulation and the use of calibration statistics to guide grouping are reasonable. However, the current evidence is insufficient to establish that the proposed alignment step is the cause of the reported gains: there are no comparisons against existing MHA-to-GQA conversion methods, no sensitivity analysis for the calibration set, and no multi-seed or error-bar information on a very narrow evaluation suite. The contribution is potentially valuable, but its robustness and incremental value over prior work remain unproven.","major_comments":[{"comment":"The entire alignment step depends on rotations and groupings computed from a single calibration sample of 128 C4 sequences (262K tokens), yet the paper provides no sensitivity analysis over calibration sets, no held-out validation of the alignment, and Figure 2 reports cosine-similarity gains only on the calibration set itself. If the learned rotations do not transfer to the deployment distribution (for example, after fine-tuning on the target tasks), the advantage over direct pruning could disappear. This risk is not merely hypothetical: in Table 1, GQA-8 with grouping-by-value cos reaches 80.32 average versus 81.65 for the no-transformation baseline, and Section 4.4 concedes that the transformation can hurt due to cumulative errors. The limitations paragraph stating that the method relies on statistical mathematical features of attention heads does not address calibration robustness. I request a sensitivity study with different calibration subsamples and a verification that the computed rotations improve mergeability on held-out data, not only on the calibration set.","section":"4.1, Figure 2, Tables 1-2"},{"comment":"The experiments compare only against the authors' own baseline of L0 pruning without transformation. No comparison is made to established MHA-to-GQA conversion methods, such as mean-pooling followed by uptraining (Ainslie et al., 2023), the principal-component-based method of Yu et al. (2024), or the grouped-query conversion methods cited in Section 2.2 (Chen et al., 2024; Chen et al., 2024a). As a result, the paper cannot support its central claim that the proposed alignment step is an effective way to convert MHA to GQA or that the overall pipeline is cost-effective; a reader cannot tell whether the reported gains are due to the Procrustes alignment, the L0-plus-distillation training budget, or other hyperparameters. I request at least one standard MHA-to-GQA baseline at matched token budgets and evaluation conditions.","section":"4 (Main results) and 2.2 (Related Works)"},{"comment":"The empirical evidence is too narrow to support the abstract's claims of acceptable performance degradation at 87.5% and 75% KV-head compression. All results are single runs on eight small commonsense QA datasets, with no error bars, no seeds, no language-modeling perplexity, and no generation or multitask benchmarks such as MMLU or GSM8K. Given the test-set sizes in Appendix B (e.g., ARC-C 299, OpenbookQA 500, WinoGrande 1267), average-accuracy differences of 0.5-1.5 points are within noise. The paper should report multiple seeds or confidence intervals and ideally include a broader evaluation, including held-out perplexity and a larger multitask benchmark, before claims about acceptable degradation can be accepted.","section":"Tables 1-2 and 4.1"},{"comment":"The key-side alignment is under-specified. The paper states that due to RoPE, the orthogonal matrix should be block-diagonal with 2D rotation blocks, and that one can apply Procrustes analysis in every two dimensions just like RoPE. It does not derive the constrained Procrustes solution or prove that solving independent per-subspace problems minimizes the Frobenius loss over block-diagonal rotations. The output-invariance derivation in Eq. (17) is correct once such an R is chosen, but it does not address optimality of the chosen R for the alignment objective. Please provide the constrained optimization (objective and SVD per 2D block) or a reference that justifies this step.","section":"Eq. (14)-(17), Section 3.3"}],"minor_comments":[{"comment":"The method as described assumes equal-size groups with D = H/G, which requires G to divide H; the paper tests only G = 16, 8, 4 for LLaMA2-7B. The abstract's claim of 'any compression ratio' should be qualified to avoid overclaiming.","section":"3.4, Eq. (23)-(24)"},{"comment":"The paper does not specify how the continuous hard-concrete masks are discretized after training, nor what happens if some masks remain nonzero when the training ends; a precise post-processing step is needed for reproducibility.","section":"3.5, Eq. (27)-(28)"},{"comment":"The target size T is described both as a target size and as something that equals zero after sparsity warm-up steps; the annealing schedule is only sketched in Appendix A. Please clarify the relationship between T, the number of retained heads, and the actual mask values during and after training.","section":"3.5, Eq. (29)"},{"comment":"The algorithm only accepts strict improvements and has no temperature-based acceptance probability, so it is random-restart hill climbing rather than simulated annealing; the nomenclature should be adjusted or the algorithm should be changed to match the name.","section":"Algorithm 2"},{"comment":"There are several typographical and formatting issues: 'theL0 method' in Section 2.1, 'SIQAAvg.' in the Table 1 header, 'Sheared-llama-1.3B' in Table 6, and inconsistent capitalization of 'Sheared-LLaMA' elsewhere. These should be corrected in a revision.","section":"Throughout"},{"comment":"The generalized Procrustes analysis is cited via a Wikipedia article; please cite the original methodology literature (e.g., Gower's work on generalized Procrustes analysis) for a citable and verifiable source.","section":"References"}],"recommendation":"major_revision","confidential_remarks":"The core idea is interesting and the invariance derivations are largely correct, but the experimental validation is not yet at the level expected for a conversion-method paper. The absence of any comparison with existing MHA-to-GQA baselines and the lack of calibration robustness analysis are the two most significant gaps. The paper would be substantially stronger if the authors add matched-budget comparisons, a sensitivity study of the calibration set, and multi-seed (or at least variance-aware) results. I recommend major revision rather than rejection because the proposed mechanism is plausible and the missing evidence is obtainable within the scope of the paper."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Here's my read. The paper does something real: it aligns attention heads with orthogonal Procrustes rotations before merging them via L0 regularization, and the invariance math checks out. Equations 12-17 are correct—the value-side rotation commutes through the output projection, and the key-side block-diagonal rotations commute with RoPE, which is the subtle part. That's not trivial, and the authors deserve credit for it.\n\nThe empirical story is also positive in its own terms. On LLaMA2-7B, the best transformed configs beat the no-transformation L0 baseline by a wide margin at high compression (GQA-4: 82.36 avg vs 77.03). Sheared-LLaMA-1.3B shows similar gains. The ablations over grouping criteria (adjacent, by-key, by-value) and similarity metrics (cosine vs distance) are useful.\n\nThe soft spots are real but not fatal. First, there is no comparison against Yu et al. (2024), which is the closest prior work—they also use orthogonal transformations on KV heads. Without that baseline, I can't tell how much of the gain is Procrustes + L0 versus just L0 + distillation. Second, single runs and no error bars: some of the per-task differences are within a few points, and the paper reports no variance. Third, the calibration set is a single 128-sequence C4 sample, and there's no sensitivity analysis. The stress-test note is right that some transformed configs underperform baseline (e.g., GQA-8 value-cos at 80.32 vs baseline 81.65), which tells me the transformation isn't uniformly beneficial and its effectiveness depends on grouping and criterion choice. That's a moderate concern, not a reason to reject—but it should be probed. The paper's stated limitation about relying on statistical features doesn't address distribution shift.\n\nThe citations look fine; the related work is honestly drawn, and the limitations section is candid. The L0 loss formulation (Eq. 29) with a shared budget across blocks is a nice detail, and Figure 3 showing block-wise pruning speeds is informative.\n\nBottom line: this deserves a serious referee. The method is new, the math is solid, and the results are promising, but the missing comparison and calibration robustness need to be addressed before I'd trust the size of the claimed advantage. I'd bring it to a reading group for the technique, though I'd pair it with Yu et al.","headline":"Procrustes alignment before L0 head merging is a real, mathematically sound trick with promising results, but the missing comparison to the closest prior work and untested calibration sensitivity keep me from fully trusting the effect size.","tokens_in":14448,"tokens_out":3432,"would_cite":true,"duration_ms":35148,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"The paper claims that aligning attention heads with Procrustes rotations before merging them lets multi-head attention be converted to grouped-query attention at up to 87.5 percent KV-head compression with acceptable accuracy loss.","keywords":["grouped-query attention","KV cache compression","Procrustes analysis","L0 regularization","attention head merging","multi-head attention","orthogonal transformation","LLM inference efficiency"],"falsifier":"Run the same MHA-to-GQA conversion twice, once with rotations computed from the paper's C4 calibration caches and once with rotations computed from an out-of-domain corpus such as code or mathematics. If the out-of-domain rotations erase or reverse the accuracy gain over the no-alignment baseline, the central claim that the alignment itself causes the improvement is falsified.","tokens_in":13450,"feed_emoji":"📉","tokens_out":9884,"duration_ms":84118,"temperature":0.7,"pith_summary":"The paper aims to establish that a pretrained multi-head attention (MHA) transformer can be converted into a grouped-query attention (GQA) transformer at almost any key-value (KV) head compression ratio through a two-stage procedure: a cheap, output-preserving alignment of attention heads followed by sparsifying fine-tuning. The alignment stage uses generalized Procrustes analysis to find orthogonal rotations that make the key and value caches of heads within the same group as similar as possible, and fuses those rotations into the projection matrices so the model's output is unchanged. The fine-tuning stage then uses L0-regularized masks to transfer the original heads into shared KV heads, producing a model that fits the standard GQA framework. On LLaMA2-7B the method reaches 87.5 percent KV-head compression and on Sheared-LLaMA-1.3B 75 percent, with what the authors describe as acceptable average accuracy loss relative to the MHA teacher. The practical stake is that KV-cache memory and inference cost can be cut sharply without the expensive uptraining that earlier MHA-to-GQA conversions required.","feed_headline":"Aligning heads before merging slashes key-value cache by 87.5%","feed_subtitle":"Procrustes rotations plus L0 pruning convert MHA to standard GQA while keeping most accuracy.","key_machinery":"The load-bearing mechanism is the orthogonal Procrustes rotation fused into projection matrices, paired with L0 masks for head transfer. Generalized Procrustes analysis rotates each head's KV cache toward the group mean; fusing the resulting orthogonal matrix into the value/output pair, or in RoPE-compatible block-diagonal form (2D rotations per coordinate pair) into the key/query pair, preserves the attention output exactly. This changes the geometry of the KV caches so that heads in a group are nearly parallel, turning what would otherwise be a large pruning step into a small perturbation. A simulated-annealing search over head groupings, scored by post-alignment cosine similarity or Euclidean distance, determines which heads should share a KV head.","core_discovery":"The central discovery is that the near-orthogonality of KV caches across attention heads is not fixed: because the caches are low-rank, a per-head orthogonal rotation can make them nearly parallel without changing the attention computation. For value heads, the paper computes the optimal rotation aligning one cache to another, or iteratively to a group mean, from the SVD of their cross-covariance, and fuses the rotation into the value and output projection matrices. For key and query heads, the rotation is restricted to a block-diagonal form of 2D rotations so that it commutes with rotary position embeddings, and it is fused into the key and query projections. The transformed model is mathematically identical to the original MHA, but its heads within each group are much more similar, so mean-pooling initialization plus L0 masks can transfer all original heads into shared KV heads with less damage. The paper reports that transformed and regrouped models consistently beat direct-pruning baselines on the tested benchmarks, and that the gap grows as the compression ratio increases.","pith_inferences":["If the alignment is doing the work the paper attributes to it, the same rotations should transfer across fine-tuned variants of a base model, since they are computed only from KV caches; that would make the conversion a one-time per-architecture transformation rather than a per-checkpoint procedure.","The comparable results for cosine-similarity and Euclidean-distance criteria suggest the exact objective matters less than the fact of rotating heads into a common frame; a cheaper heuristic such as aligning each head to the group centroid once might capture most of the gain.","A direct extension is to test alignment on out-of-domain calibration data; persistent gains would indicate robustness, while vanishing gains would reveal calibration distribution as a hidden hyperparameter.","Because the pre-pruning transformation is output-invariant, the same head-alignment trick could benefit other head-merging schemes, such as per-layer asymmetric key/value sharing, independent of L0 pruning."],"forward_implications":["The conversion is not locked to preset GQA group sizes: L0 masks can transfer original heads to shared heads at any target ratio, so practitioners can choose a KV-head budget freely.","The resulting model is a standard GQA model, meaning it can be served by existing GQA inference kernels and frameworks without custom attention code.","The alignment stage is cheap relative to training: calibration and transformation take at most about an hour on one A100 GPU for the tested models, according to the paper's appendix.","Because the rotation is output-invariant before pruning, the same alignment can be combined with other fine-tuning or distillation objectives, not only the L0-plus-BiLD setup used here.","At 87.5 percent KV-head compression on LLaMA2-7B and 75 percent on Sheared-LLaMA-1.3B, converted models retain average accuracy close to the MHA teacher on the eight commonsense benchmarks reported."],"supporting_citations":[{"why":"Supplies the orthogonal Procrustes solution (SVD-based optimal rotation) that the alignment stage applies to each pair of caches.","marker":"Schönemann, 1966"},{"why":"Provides the L0 regularized hard-concrete masks and training recipe used to transfer original heads into shared KV heads.","marker":"Louizos et al., 2017"},{"why":"Introduces computational invariance in transformers, the principle that orthogonal transformations can be fused into adjacent projection matrices without changing output.","marker":"Ashkboos et al., 2024"},{"why":"Defines GQA and the mean-pooling initialization of shared KV heads that the pruned model starts from.","marker":"Ainslie et al., 2023"},{"why":"Shows that KV caches are low-rank and uses orthogonal transformations on KV projections, the direct precursor of the paper's alignment idea.","marker":"Yu et al., 2024"},{"why":"Provides Sheared-LLaMA-1.3B, one of the two testbed models, and the target-structure pruning practice the method follows.","marker":"Xia et al., 2023"},{"why":"Gives the generalized Procrustes analysis algorithm used to align more than two caches within a group.","marker":"Wikipedia contributors, 2022"},{"why":"Supplies the BiLD distillation loss used alongside KL loss to keep student logits close to the teacher during pruning training.","marker":"Li et al., 2024"}],"fun_headline_variants":["Procrustes alignment enables 87.5% KV cache reduction","Pre-merge head rotation cuts KV cache by 87.5%","Rotate-then-merge converts MHA to GQA with minimal loss","Align-then-merge strategy trims KV cache 87.5%","Orthogonal head alignment reduces KV cache for GQA"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The whole alignment step rests on the 128 C4 sequences of 2048 tokens used for calibration being representative of the model's deployment distribution; if those caches are not representative, the computed rotations are miscalibrated and the advantage over direct pruning disappears.","fun_headline_variants_meta":{"raw":{"variants":["Procrustes alignment enables 87.5% KV cache reduction","Pre-merge head rotation cuts KV cache by 87.5%","Rotate-then-merge converts MHA to GQA with minimal loss","Align-then-merge strategy trims KV cache 87.5%","Orthogonal head alignment reduces KV cache for GQA"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000221,"raw_usage":{"total_tokens":1464,"prompt_tokens":973,"completion_tokens":491,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":589,"completion_tokens_details":{"reasoning_tokens":397}},"tokens_in":589,"tokens_out":491,"duration_ms":4844,"temperature":1.0,"reasoning_tokens":397,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-10T23:14:02.414559+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Run the same MHA-to-GQA conversion twice, once with rotations computed from the paper's C4 calibration caches and once with rotations computed from an out-of-domain corpus such as code or mathematics. If the out-of-domain rotations erase or reverse the accuracy gain over the no-alignment baseline, the central claim that the alignment itself causes the improvement is falsified.","supporting_citations":[],"review_version":1}