{"id":"cab4599c-224e-4079-b024-0152846b13bb","arxiv_id":"2506.05433","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":4.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":0,"one_line_summary":"Prefix Grouper reduces GRPO training cost by sharing the prompt's forward computation across a group of responses, claiming identical gradients to standard GRPO.","lead":"This paper proposes a way to speed up GRPO, a reinforcement learning method for language models, by encoding the shared prompt only once per group of answers instead of once per answer. The authors claim this shortcut keeps gradients identical to standard GRPO, so training results should stay the same while cost drops, especially with long prompts.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Lemma 2.1's gradient-equivalence proof is invalid: it cancels prefix-token terms using a final-layer-only fact, and the experiments never test equivalence.","rationale":"The reader identified the mask and RoPE construction as the weakest assumption; I agree that an incorrect mask or position assignment would break equivalence, and this is a subset of the broader concern. My stress-test focuses one level deeper: even granting that the mask is correct, the proof of Lemma 2.1 in Appendix A.1 is not a valid proof, because Eq. (12) is used beyond its domain of validity. The per-token gradient decomposition in Eqs. (10)-(11) is not justified for shared parameters in a network with cross-token attention, and no induction over layers is supplied. This matters because the paper's headline claim is exact equality of forward and backward passes, not approximate or asymptotic equivalence. The experiments in Section 3 are also purely computational; they do not test the equivalence that the abstract claims. A numerical difference check on a small model would settle the correctness question directly and is a cheap, decisive test. If it passes, the method is likely sound and the proof can be repaired; if it fails, the central claim is false. The released code makes this test feasible, which is a genuine positive. I therefore agree with the CONDITIONAL verdict; no verdict change is needed, but the paper should add the missing equivalence experiments and repair the proof before acceptance.","tokens_in":6923,"tokens_out":14939,"duration_ms":192945,"concrete_test":"Run the released implementation on a small two-layer, two-head transformer (embedding dimension 64) with a 128-token prefix and G=3 responses of lengths 17, 31, and 23, using identical initialization, random seeds, sampled responses, and fixed advantages. Compute a GRPO-style loss under Prefix Grouper and under standard Repeated-Prefix GRPO in float64, then check the maximum absolute and relative difference in every response-token logit and in every parameter gradient, including prefix QKV and MLP parameters. If any difference exceeds 1e-8 relative to the baseline gradient norm, the equivalence claim is falsified; if it passes, the mask and RoPE construction, and the gradient computation, are validated and the proof can be repaired.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The most load-bearing weakness is that Lemma 2.1 is not actually proved. In Appendix A.1, Eqs. (10)-(11) are asserted for \"any pointwise operation parameter theta\" as per-token sums, and Eq. (12) then states that dJ/dh_t = 0 for all prefix tokens. This equality holds only for the final hidden states that directly feed the loss; at every intermediate layer a prefix position has nonzero loss gradient because downstream suffix tokens attend to it in later layers. The proof uses Eq. (12) to drop the prefix-token terms and then substitutes the attention-gradient identity Eq. (13), but no inductive layer-by-layer argument is given to justify replacing those dropped terms. Since the central claim is exact equality of optimization trajectories, an invalid proof is a direct gap. The empirical section does not close it: Section 3 reports only FLOPs and memory (Figures 2-3); there are no forward-output comparisons, gradient-norm comparisons, loss curves, or final policy metrics, despite the abstract claiming \"consistent results.\" Without a correct proof or a numerical equivalence check, the drop-in-replacement claim is unverified. If forward equivalence were shown exactly, gradient equivalence would follow by the chain rule; the paper's separate gradient proof is the weak link.","agreement_with_reader":"partial"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes Prefix Grouper, a modification of GRPO training that concatenates the shared prompt prefix with all group responses into one sequence and computes self-attention in two parts: a single prefix-only attention pass and suffix attention over the full prefix-plus-suffix context. The authors claim that this scheme is exactly equivalent to standard GRPO in forward outputs and backward gradients, while reducing FLOPs and memory by roughly a factor of the group size in long-prefix regimes. The theoretical claims are stated as Lemma 2.1 (gradient equivalence) and Lemma 2.2 (computation reduction), with proofs in Appendix A. The experimental section reports only FLOPs and GPU-memory measurements; there are no training runs, loss curves, or policy-performance comparisons.","tokens_in":7118,"tokens_out":7478,"duration_ms":95114,"significance":"If the equivalence claim were rigorously established, Prefix Grouper would be a practically valuable and simple efficiency improvement for GRPO, particularly for long-context and multi-modal reasoning tasks. The computational savings analysis is clean and parameter-free, and the method is presented as a drop-in replacement with code released. However, the central theoretical claim is currently not proved rigorously, and the empirical section does not test equivalence at all. The idea is plausible, but the paper overclaims 'extensive experiments' and 'identical gradients' without the necessary evidence.","major_comments":[{"comment":"The statement '∀t∈P, ∂J/∂h_t = 0 (both algorithms)' is only valid for the final-layer hidden states that directly feed the loss. At every intermediate layer, a prefix position has nonzero gradient because suffix tokens attend to its key/value vectors in that layer and in subsequent layers. Equations (10)-(11) use Eq. (12) to drop prefix-token terms for all layers, and Eq. (13) is asserted without derivation; no inductive layer-by-layer argument shows that the gradient of the single shared prefix equals the sum of gradients of the G prefix copies in the baseline. As written, Lemma 2.1 is not proved. Please supply an inductive proof (e.g., by induction on depth, showing that the backpropagated gradient to the shared prefix equals the sum of gradients to the G baseline prefix copies) or add a numerical gradient-equivalence test comparing full parameter gradients on a small model.","section":"Appendix A.1, Eq. (12)"},{"comment":"The experiments do not validate training equivalence. Figures 2 and 3 report only FLOPs and GPU memory. There is no comparison of forward outputs, gradient norms, loss curves, or final policy performance. The abstract's claim of 'consistent results' from 'extensive experiments' is therefore unsupported. Please add at least (i) a numerical equivalence check (e.g., maximum absolute difference in logits and parameter gradients between the two implementations on a real model), and (ii) an end-to-end GRPO training comparison on a reasoning or instruction-following task showing matched reward curves and final performance.","section":"Section 3"},{"comment":"The suffix attention mask and RoPE position assignment are never specified. Equivalence depends on the mask blocking attention from one response to all other responses, and on each token's position id being identical to its position in the baseline [P;Ri] sequence. The paper says masks are 'pre-computed' but does not define them. Please give the exact mask construction and position-id mapping, and include a proof or an explicit test that the concatenated sequence reproduces the per-sample causal attention pattern of the baseline.","section":"Section 2.1 and Algorithm 1"},{"comment":"The paper does not address stochastic operations such as dropout. In the baseline, the G prefix copies are separate forward passes and can receive different dropout masks; in Prefix Grouper the prefix is computed once, so the dropout masks for prefix tokens differ from the baseline. If any stochastic regularization is used during training, exact gradient equivalence does not hold. Please state the assumptions (e.g., no dropout, or shared deterministic masks) under which Lemma 2.1 is claimed.","section":"Section 2.2"}],"minor_comments":[{"comment":"The heading contains a typo: 'Futher discussion' should be 'Further discussion'.","section":"Section 2.4"},{"comment":"The caption contains a typo: 'memeory' should be 'memory'.","section":"Figure 3 caption"},{"comment":"The captions do not define 'Ratio'. Please state explicitly that Ratio is prefix length divided by suffix length, and clarify whether the FLOPs and memory numbers are analytical estimates or measured on real hardware.","section":"Figures 2 and 3"},{"comment":"The experimental section provides no model architecture, dataset, training hyperparameters, or evaluation setup. Even for a computational-overhead study, these details are needed for reproducibility and to interpret the reported memory figures.","section":"Section 3"},{"comment":"The computation-reduction ratio is derived in the limit Lp/Lr → ∞, but Figure 2 covers finite ratios. It would be helpful to plot the exact ratio from Eq. (17) alongside the measured values to show where the asymptotic 1/G behavior becomes accurate.","section":"Appendix A.2, Eq. (17)"}],"recommendation":"major_revision","confidential_remarks":"The proposed method is plausible and the computational savings story is clear, but the manuscript overclaims exact equivalence without a rigorous proof or empirical verification. The proof gap in Lemma 2.1 is load-bearing: the central claim is exact gradient equality, and the current argument drops prefix-token terms using a final-layer fact that does not hold for intermediate layers. The experiments are also limited to FLOPs and memory, which do not test the equivalence claim. I see this as fixable within the scope of a revision: an inductive proof or a numerical gradient check, plus at least one end-to-end training comparison, would substantially strengthen the paper. I therefore recommend major revision rather than rejection."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Quick take: this is a plausible efficiency trick for GRPO training—share the prefix across a group by concatenating and splitting attention—but the paper's central claim of exact gradient equivalence is not actually proven, and the experiments never test it.\n\nWhat's genuinely good: the idea of moving shared-prefix computation into a single forward pass for training is a natural extension of inference-time prefix caching, and the paper correctly identifies a real bottleneck in long-context GRPO. The forward pass is straightforwardly equivalent if the masks and RoPE positions are built right, and the FLOPs/memory analysis in Section 3 is sensible. Code is released, which is a plus.\n\nNow the problem. Appendix A.1 is supposed to prove Lemma 2.1, but it doesn't. Eq. 9 inserts a 1/G factor on the suffix-gradient sum that appears without justification; the baseline gradient should accumulate over the G copies of the prefix, and the relationship to the shared-prefix gradient needs a careful accounting. More importantly, Eq. 12 uses ∂J/∂h_t = 0 for prefix tokens, which is only true at the top-layer final hidden states. At every intermediate layer, the prefix tokens receive nonzero gradient through the response tokens that attend to them, so you can't just drop those terms unless you give an inductive layer-by-layer argument. The proof doesn't. That's a load-bearing gap because the whole contribution is the exact-equivalence guarantee.\n\nThe empirical side doesn't close it. Figures 2 and 3 only report FLOPs and memory. There's no forward-output comparison, no gradient-norm check, no loss curves, no final policy metrics. So the abstract's 'consistent results' is not supported by any experiment.\n\nAm I convinced the method is wrong? No. If the forward pass is truly identical, gradient equivalence follows by the chain rule; the code likely does something reasonable. But the paper as written doesn't establish it. The right fix is either a proper inductive proof or a numerical equivalence check (e.g., compare gradients from both implementations on a small model). The related-work gap—no mention of inference-time prefix caching—is secondary but worth noting.\n\nWho should read it: anyone doing GRPO with long prompts will find the idea useful, and the paper is a good starting point for implementing it. But they should treat the equivalence claim as a hypothesis, not a proven fact.\n\nMy recommendation: send it to peer review. The method is useful, the core idea is sound, and a serious referee can push for a corrected proof and real equivalence experiments. It's not desk-reject material, but it needs revision before the claims can be trusted.","headline":"A useful but under-verified efficiency trick for GRPO: the forward idea is sound, but the gradient-equivalence proof is incomplete and the experiments don't test it.","tokens_in":7675,"tokens_out":5047,"would_cite":true,"duration_ms":58431,"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":"A training algorithm called Prefix Grouper claims to make GRPO encode the shared prompt once per group while producing gradients identical to standard GRPO.","keywords":["GRPO","Group Relative Policy Optimization","shared-prefix forward","grouped attention","training equivalence","long-context reinforcement learning","self-attention","efficient LLM training"],"falsifier":"Run the same GRPO training step on a small transformer with a long shared prompt and several responses, once with repeated-prefix forward and once with Prefix Grouper, using identical weights, inputs, and random seeds; if the maximum absolute difference in any parameter gradient exceeds floating-point tolerance, or if the attention outputs differ beyond numerical precision, the claimed equivalence is false.","tokens_in":6715,"feed_emoji":"⚡","tokens_out":7406,"duration_ms":76335,"temperature":0.7,"pith_summary":"GRPO trains policies by comparing several response candidates generated from the same prompt, and standard implementations re-encode that shared prompt separately for every candidate. Prefix Grouper claims to remove this redundancy by encoding the prompt once and letting each response suffix attend to the full prompt-plus-responses sequence through a two-part attention computation called Grouped Attention. The paper proves that this restructured forward pass produces exactly the same attention outputs and the same parameter gradients as the standard GRPO forward pass, so the optimization dynamics should be unchanged. In the long-prefix regime the method reduces attention and pointwise FLOPs to about $1/G$ of the baseline for a group of size $G$.","feed_headline":"Prefix sharing cuts GRPO compute to 1/G with identical gradients","feed_subtitle":"A two-part attention trick gives standard GRPO the same gradients while encoding the shared prompt only once.","key_machinery":"The load-bearing object is Grouped Attention, a decomposition of one self-attention block into two attention calls: self-attention over the shared prefix tokens only, and attention for the suffix tokens, where each suffix query attends to keys and values formed by repeating the prefix and appending all suffixes. Position IDs are assigned exactly as in the repeated-prefix baseline, and the attention masks are pre-computed so that each suffix token sees exactly the same causal context it would have seen in the baseline, without seeing other group members' responses. This two-kernel decomposition, implemented through a custom autograd function, lets the prefix embeddings be computed once and shared across the whole group while keeping every token's computational path equivalent to the baseline.","core_discovery":"The central claim is training equivalence: if the shared prefix and all response suffixes are concatenated as $[P; R_1; \\ldots; R_G]$, and self-attention is split into a prefix-only attention pass and a suffix pass whose queries attend to the concatenated prefix-plus-suffix keys and values, then the GRPO policy loss has the same forward output and the same gradients with respect to every model parameter as the repeated-prefix baseline. Lemma 2.1 states that the gradients are identical, $\\nabla_\\theta J_{\\text{ours}}(X_{\\text{ours}}, A) \\equiv \\nabla_\\theta J_{\\text{base}}(X_{\\text{base}}, A)$, and Lemma 2.2 states that when the prefix is much longer than the responses the method asymptotically costs $1/G$ of the baseline FLOPs. The paper also reports that the same sharing pattern can be applied to other shared-prefix workloads, such as judging many short questions against one long context.","pith_inferences":["A direct numerical check of gradient equivalence on a small causal transformer, comparing the two implementations under identical seeds, would make the proof's mask construction testable in isolation.","The proof assumes a per-token additive loss decomposition; GRPO variants with non-additive or sequence-level reward terms would need a separate equivalence argument.","The masking and position-ID construction is the part to port carefully if the technique is applied to sparse, sliding-window, or linear attention variants, where the causal pattern is expressed differently.","The same shared-prefix idea would plausibly extend to multi-turn RL settings, where one long conversation history is replayed against many sampled continuations."],"forward_implications":["If the equivalence holds, Prefix Grouper can replace repeated-prefix GRPO without retuning hyperparameters, because every gradient step is numerically the same.","Long-context reinforcement learning tasks whose prompts dominate the sequence will see the largest speedups, approaching a factor $1/G$ in FLOPs.","Memory for activations scales with the shared prefix once instead of once per response, so larger group sizes fit under the same memory budget.","The method is a drop-in change confined to input concatenation and attention computation, so existing GRPO pipelines can adopt it without architectural modification.","Shared-context workloads other than RL training, such as scoring many candidate outputs against one long prompt, can reuse the same two-part attention pattern."],"supporting_citations":[{"why":"Defines the GRPO loss and group-based comparisons that form the baseline which Prefix Grouper claims to match exactly.","marker":"[1]"}],"fun_headline_variants":["GRPO's redundant prefix compute reduced to 1/G, safe","Share prefix once, keep GRPO gradients exact, cost 1/G","Exact GRPO at 1/G compute via shared-prefix attention","Prefix Grouper: zero gradient change, 1/G FLOPs on long prefixes"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The equivalence stands or falls on the pre-computed attention masks and position IDs making the concatenated sequence exactly reproduce the per-sample causal attention of the baseline; any cross-response leakage or any mismatch in position embeddings changes both the forward outputs and the gradients.","fun_headline_variants_meta":{"raw":{"variants":["GRPO's redundant prefix compute reduced to 1/G, safe","Share prefix once, keep GRPO gradients exact, cost 1/G","Exact GRPO at 1/G compute via shared-prefix attention","Prefix Grouper: zero gradient change, 1/G FLOPs on long prefixes"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000644,"raw_usage":{"total_tokens":3002,"prompt_tokens":1027,"completion_tokens":1975,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":643,"completion_tokens_details":{"reasoning_tokens":1892}},"tokens_in":643,"tokens_out":1975,"duration_ms":16690,"temperature":1.0,"reasoning_tokens":1892,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-07T10:33:09.386238+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Run the same GRPO training step on a small transformer with a long shared prompt and several responses, once with repeated-prefix forward and once with Prefix Grouper, using identical weights, inputs, and random seeds; if the maximum absolute difference in any parameter gradient exceeds floating-point tolerance, or if the attention outputs differ beyond numerical precision, the claimed equivalence is false.","supporting_citations":[],"review_version":1}