{"id":"ed8d3a76-977d-4874-8ba1-0076187ce0eb","arxiv_id":"2505.18313","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":4.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":4,"one_line_summary":"PLUMAGE, an unbiased minimum-variance low-rank gradient estimator with optimizer-moment realignment, reduces the loss gap to full-rank training by about a third compared with GaLORE.","lead":"PLUMAGE is a new way to compress the gradient information used when training large language models, so training needs less GPU memory. Test runs on models from 130 million to 1 billion parameters show it lands closer to full-quality training than the prior GaLORE method, without extra memory cost.","discovery_kind":"new_application","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The reuse-unbiasedness claim in §3.3.2 fails in the optimization sense: with P fixed for τ steps, the conditional expectation of the update given the past is P D^{-1} P^T G_t, not G_t, and '#steps >> τ' does not remove this per-step bias.","rationale":"The reader and I converge on Section 3.3.2, but for different reasons. The simple claim E[P D^{-1} P^T] = I is actually true for any complete orthonormal basis, so a stale basis does not by itself break the estimator's unbiasedness for a fixed gradient. The real issue is that Algorithm 5 fixes P for τ steps; at a later step, P is part of the past, and the conditional expectation of the update given the past is P D^{-1} P^T G_t, not G_t. Moreover, because the trajectory depends on P, even an expectation over P conditional on the observed current gradient is not G_t; the two-parameter quadratic example makes this concrete. The paper's condition '#total_training_steps >> τ' does not remove this per-step conditional bias. The MVUE derivation itself (Eqs. 7-18) and the sampling construction are internally consistent, and the empirical comparisons are plausible and honestly report Adam's residual bias. The main deficiency is therefore a theoretical overclaim rather than a broken method; the requested revisions should either prove conditional unbiasedness under an explicit independence/refresh assumption or reclassify the optimizer as a biased but low-variance heuristic. This leaves the reader's conditional verdict unchanged.","tokens_in":16129,"tokens_out":34369,"duration_ms":297856,"concrete_test":"Run PLUMAGE-SGDM on the two-parameter quadratic L(w)=0.5||w-(1,1)||^2 with k=1, U=I, p=[0.5,0.5], τ=1, η=0.1, full-batch. Enumerate the two trajectories: if P=e1, the step-1 gradient is (-0.2,-1) and the update is 2e1 e1^T G_1=(-0.4,0); if P=e2, the gradient is (-1,-0.2) and the update is (0,-0.4). In each realized trajectory, the conditional mean update equals 2P P^T G_1, which is not G_1. This direct enumeration falsifies the claim that the reused-projection update is unbiased conditional on the optimization state; report the ratio ||E[\\hat G_1 | trajectory] - G_1||_F / ||G_1||_F, which is nonzero.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The load-bearing problem is the transfer of the §3.1 unbiasedness theorem to the deployed loop, but the mechanism is not primarily subspace stability. In Algorithm 5, P_kt is sampled once and then fixed for τ steps. At any later step t, P_kt is part of the conditioning information (the past), so the relevant conditional expectation is E[P D^{-1} P^T G_t | F_{t-1}] = P D^{-1} P^T G_t, not G_t, even when U_kt is a complete orthonormal basis. The identity E[P D^{-1} P^T] = I only justifies unbiasedness if the expectation is over an independent redraw of P at the same G_t; with reuse, P is part of the state, and the trajectory itself depends on P. A two-parameter quadratic example makes this concrete: with k=1, p=[0.5,0.5], U=I, τ=1, η=0.1, the realized trajectory determines whether P=e1 or e2, and the conditional mean update is 2P P^T G_t, which differs from G_t. The statement 'so long as #total_training_steps >> τ' speaks only to how often projections are refreshed, not to the dependence between P and future gradients, so it does not repair the conditional bias. Moment realignment (Eqs. 28-29) rotates old moments into the new basis but cannot remove the bias of using a stale, state-dependent projection. Thus the unbiasedness component of the central claim is not established for the algorithm as deployed; the empirical results remain useful, but should be presented as a heuristic or paired with an explicit independence/refresh assumption.","agreement_with_reader":"partial"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes PLUMAGE, a low-rank gradient estimator for training large models. In §3.1 the authors derive inclusion probabilities for a k-sparse estimator of the form (7), solve the variance-minimization problem (13)-(18), and in §3.2 show that the one-sided projection form (20) is equivalent to the two-sided form. Section 3.3 integrates the estimator with Adam: Eqs. (22)-(24) define low-rank moments, Eqs. (25)-(27) introduce projection reuse over τ steps, and Eqs. (28)-(29) realign first and second moments when the projection changes. Experiments compare PLUMAGE with Adam, GaLORE, and FLORA on LLaMA pretraining (C4) and RoBERTa fine-tuning (GLUE), reporting smaller gaps to full-rank Adam.","tokens_in":16406,"tokens_out":12994,"duration_ms":136350,"significance":"The derivation of the inclusion probabilities is mostly sound, the wheel-of-fortune sampling is efficient, and the empirical comparisons are relevant and consistently reported. If the unbiasedness claim held, PLUMAGE would be a valuable contribution to memory-efficient LLM training. The main empirical result—PLUMAGE reduces the gap to full-rank Adam relative to GaLORE by roughly a third in perplexity and by about 28% in GLUE loss—is a useful practical finding. However, the theoretical claim that the deployed algorithm remains unbiased is not established; this directly affects the paper's title and abstract. With that claim reframed or proved under explicit assumptions, the work would be publishable.","major_comments":[{"comment":"The sentence 'the estimate remains unbiased so long as #total_training_steps >> τ' is asserted without proof and is not correct as stated. Once P_kt is sampled, it is fixed for τ steps; at any later step t the projection is measurable with respect to the past, so the conditional expectation of the update is E[P D^{-1} P^T G_t | F_{t-1}] = P D^{-1} P^T G_t, not G_t. The identity E[P D^{-1} P^T] = I used in §3.1 requires an independent redraw of P at the same G_t, which is not what Algorithm 5 does. The condition '#steps >> τ' only describes refresh frequency and does not remove the dependence between the stale projection and future gradients. The unbiasedness guarantee of §3.1 therefore does not transfer to the deployed loop. Please either derive a corrected statement (e.g., a bias bound under an explicit subspace-stability assumption) or present PLUMAGE as a heuristic and remove 'unbiased' from the title and abstract.","section":"§3.3.2, Eqs. (25)-(27)"},{"comment":"The pseudocode is internally inconsistent in its handling of moment realignment. In the branch where the projection is refreshed, M_t and V_t are overwritten by update_state(...) before the moving-average updates 'M_t <- β1 M_{t-1} + ...' and 'V_t <- β2 V_{t-1} + ...', so the old-subspace moments that should feed the EMA are no longer available. This makes the described algorithm non-reproducible. Please clarify the intended ordering (e.g., compute R_t, update the EMA in the new subspace, and realign the previous moments before combining, or explicitly store the old moments).","section":"Algorithm 5"},{"comment":"The minimum-variance claim is established only within the restricted family of estimators in Eq. (7) with fixed inclusion probabilities p_i. The title's 'Minimum-vAriance' and the term 'MVUE' in §3.2 are stronger than what is proven; no optimality over all unbiased estimators is shown. Please qualify the claim (e.g., 'minimum variance within the class of component-wise estimators of the form (7)').","section":"§3.1 / title"}],"minor_comments":[{"comment":"The orthogonality relation should read u_i^T u_j = δ_ij (or equivalently Tr(v_i v_j^T) = δ_ij); v_i^T u_j is not defined when m ≠ n.","section":"Eq. (11)"},{"comment":"The subscripts are garbled; the right-hand side should reference M^{kt}_{t-1} and V^{kt}_{t-1}, not M^{kt}_t and V^{kt}_t.","section":"Eqs. (25)-(26)"},{"comment":"The claim that no hyperparameters are introduced beyond r and τ is contradicted by κ in §3.3.2 and by γ_shrink, γ_expand, γ_reset in Appendix A; please clarify that these are optional and were set to fixed values in the main experiments.","section":"Abstract and §3.3.2"},{"comment":"The 'Mean' row reports the average terminal loss, but the caption does not define whether these are training or validation losses; the abstract's 28% claim should reference this explicitly.","section":"Table 2"},{"comment":"Notation is inconsistent with §3.1 (d vs n, r* = d-c unexplained); please align the pseudocode with the main text.","section":"Algorithm 2"},{"comment":"Inconsistent capitalization of GaLore/GALORE appears throughout; please choose one spelling.","section":"General"}],"recommendation":"major_revision","confidential_remarks":"For the editor: the core theoretical claim is currently overclaimed, but the deficiency is fixable by reframing the contribution as an empirical low-rank optimizer with moment realignment, or by proving a bias bound under explicit assumptions. The derivation of the inclusion probabilities is correct, and the empirical tables support a useful practical improvement over GaLORE. I would not recommend rejection on the current evidence; a major revision asking for an explicit statement of what is and is not unbiased seems appropriate."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Two things you should know.\n\nFirst, the empirical method works. Replacing GaLORE's top-k projection with one-sided probabilistic sampling consistently closes a third of the gap to full-rank Adam in pretraining and about a quarter on GLUE, at the same memory and compute. That is a concrete, useful result for anyone training small-to-mid LLMs on limited hardware.\n\nSecond, the paper's central theoretical claim—that the estimator stays unbiased when the projection is reused for tau steps—is not established and, as stated, I think it is false. In Section 3.3.2 they say it \"remains unbiased so long as #total_training_steps >> tau\". That reasoning does not survive contact with the algorithm. Once P is fixed for tau steps, P is part of the conditioning history, so the conditional expectation of the update given the past is P D^{-1} P^T G_t, not G_t. The identity E[P D^{-1} P^T] = I only holds if P is redrawn independently at the current gradient. \"#steps >> tau\" only says projections get refreshed; it does not say the stale projection is independent of the trajectory. The stress-test's two-parameter quadratic example makes the failure concrete. So the unbiasedness guarantee from Section 3.1 does not transfer to the deployed loop. The min-variance property is already correctly conceded to be lost under reuse; unbiasedness should be conceded too.\n\nWhat is genuinely new: the one-sided projection variant, which preserves the MVUE property using only U and D, and the moment realignment formulas (Eqs. 28-29), which are a reasonable heuristic and empirically do help. The core sampling-probability derivation is a re-derivation of ATOMO and Chmiel et al., which the paper cites, so no foul there. The package as a drop-in GaLORE replacement is new and practical.\n\nThe paper is honest about Adam's nonlinearity reintroducing bias, and it does not hide that the min-variance property is lost with reuse. Reproducibility is only partial: no code release, no seed-level tables, and the adaptive interval appendix is a side show. Those are normal preprint issues.\n\nWho is this for? Anyone working on low-rank or memory-efficient training. It deserves a serious referee: the empirical claim is concrete, and the theory is mostly right except for the one overreach. I would send it out with a request to either prove or soften the reuse-unbiasedness claim, and to release code. The method itself is likely to be used regardless.","headline":"Useful empirical low-rank gradient method; the reuse-unbiasedness claim in §3.3.2 is not established and should be corrected, but the paper deserves a serious referee.","tokens_in":17070,"tokens_out":2327,"would_cite":true,"duration_ms":25934,"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":"PLUMAGE is a fixed-rank, unbiased, minimum-variance gradient estimator for LLM training; it samples singular components without replacement and realigns Adam moments across projection updates.","keywords":["low-rank gradient estimation","minimum-variance unbiased estimator","LLM pre-training","optimizer state alignment","Adam moments","sampling without replacement","memory-efficient training","gradient compression"],"falsifier":"Measure the per-step bias $E[P_t D_t^{-1} P_t^\\top G_t] - G_t$ when $P_t$ is sampled at step $t_0 < t$ and the gradient's top subspace rotates; train a small model on data with a known drifting subspace and compare PLUMAGE's update direction with the full-batch gradient. If the angular error grows with subspace drift, then the claimed unbiasedness does not hold in the amortized training loop.","tokens_in":15817,"feed_emoji":"📉","tokens_out":5446,"duration_ms":52805,"temperature":0.7,"pith_summary":"This paper tries to establish that low-rank gradient training need not sacrifice unbiasedness or pay high variance: it constructs a k-sparse estimator that samples singular components without replacement with inclusion probabilities chosen to minimize variance, then proves the one-sided projection version preserves those properties. If correct, PLUMAGE is a drop-in replacement for GaLore-style low-rank optimizers that closes about a third of the gap between low-rank and full-rank pre-training loss at the same memory footprint, and about 28% of the GLUE training-loss gap. The paper also claims to fix a subtle failure mode: when the projection is periodically updated, Adam's first- and second-moment statistics live in different subspaces, and it provides a realignment rule that prevents spurious weight updates.","feed_headline":"Unbiased low-rank gradients close 33% of full-rank pretraining gap","feed_subtitle":"New sampling scheme also cuts GLUE training loss by 28% at GaLore memory cost.","key_machinery":"The load-bearing object is the probability vector $p$ and the without-replacement wheel-of-fortune sampler: top $r^*$ singular modes are always kept, and the remaining $k-r^*$ modes are drawn with probability proportional to their singular values, so that the estimator's expected value equals $G$ and its expected squared error is minimized under the k-sparse constraint. The one-sided projection $P D^{-1} P^\\top G$ makes the estimator usable with Adam by storing a single $m\\times k$ projection per weight, and the realignment matrix $B=P_2^\\top P_1$ with $V_2 \\approx B^{\\circ 2} V_1$ is the mechanism that transfers first- and second-moment statistics between consecutive projection subspaces.","core_discovery":"On the paper's own terms, PLUMAGE's central claim is that a fixed-rank estimator of the form $\\hat{G}=\\sum_{i=1}^{n} \\frac{1}{p_i} I_i \\sigma_i u_i v_i^\\top$, with inclusion probabilities $p_i=1$ for the top $r^*$ modes and $p_i=(k-r^*)\\sigma_i / \\sum_{j>r^*} \\sigma_j$ for the remaining modes, and with exactly $k$ indices sampled without replacement, is unbiased and has minimum variance among k-sparse estimators, and that the one-sided estimator $\\hat{G}=P D^{-1} P^\\top G$ inherits these properties because it reduces to the same form. It further claims that realigning Adam's moments via $M_2 \\approx P_2^\\top P_1 M_1$ and $V_2 \\approx (P_2^\\top P_1)^{\\circ 2} V_1$ removes the spurious updates caused by changing projections, and reports that this shrinks the full-rank optimization gap by 33% on average in pre-training and 28% on GLUE while keeping GaLore's computational and memory footprint.","pith_inferences":["Inference: the unbiasedness proof covers a freshly sampled projection; the amortized claim that staleness does not matter relies on subspace stability, so the method's advantage should be tested under explicit subspace drift.","Inference: the same sampling scheme could be applied right-sided to activations, trading forward-pass projection cost for activation-memory savings; the paper mentions this direction but does not test it.","Inference: the principal-angle adaptive interval controller depends on manually chosen thresholds, so the reported gains may be sensitive to those settings.","Inference: the second-moment realignment approximates the off-diagonal entries of the gradient's second-moment matrix as zero; when gradients are strongly correlated across coordinates, this approximation could degrade alignment quality."],"forward_implications":["At the same memory and compute footprint as GaLore, PLUMAGE lowers pre-training validation perplexity from 30.18, 24.08, and 17.03 to 28.73, 21.81, and 16.29 for the 130M, 350M, and 1B Llama variants.","On GLUE fine-tuning, PLUMAGE's mean terminal accuracy is 85.66 versus GaLore's 84.89 and full-rank Adam's 85.91, with lower mean loss than GaLore.","PLUMAGE introduces no hyperparameters beyond the rank $r$ and update interval $\\tau$, and it reuses the full-rank learning rate without retuning.","Realigning first- and second-moment statistics across projection updates removes spurious weight updates and makes training stable across seeds.","The one-sided projection halves the communication cost of two-sided ATOMO-style estimators in data-parallel settings."],"supporting_citations":[{"why":"Supplies the GaLore baseline, the low-rank projection training setup, and the hyperparameters reused in the experiments.","marker":"[39]"},{"why":"FLORA is the main unbiased random-projection baseline; its moment-projection idea is adapted and compared.","marker":"[16]"},{"why":"ATOMO is the prior min-variance unbiased low-rank gradient estimator; PLUMAGE builds on its design while using one-sided projections.","marker":"[36]"},{"why":"Establishes the minimum-variance unbiased sparse gradient estimation methodology that PLUMAGE extends to low-rank projections.","marker":"[7]"},{"why":"Provides the importance-sampling variance-reduction perspective behind choosing inclusion probabilities.","marker":"[2]"},{"why":"Gives the SVD and truncation optimality context for low-rank gradient projections.","marker":"[11]"},{"why":"Supplies the wheel-of-fortune algorithm that samples exactly k indices without replacement.","marker":"[12]"},{"why":"Defines Adam, whose first- and second-moment statistics require the realignment procedure.","marker":"[20]"},{"why":"Provides principal angles used to measure subspace overlap in the adaptive interval controller.","marker":"[41]"}],"fun_headline_variants":["Unbiased low-rank gradients close 33% of full-rank gap","PLUMAGE: unbiased low-rank gradients, 33% closer to full-rank","Sampling trick yields unbiased min-variance low-rank gradients","Low-rank estimator hits 33% of full-rank gap with unbiased updates","Unbiased low-rank gradients: 33% gap reduction, same memory"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The projection sampled from an old gradient remains a fair sample of the current gradient because the dominant gradient subspace is stable over the reuse interval.","fun_headline_variants_meta":{"raw":{"variants":["Unbiased low-rank gradients close 33% of full-rank gap","PLUMAGE: unbiased low-rank gradients, 33% closer to full-rank","Sampling trick yields unbiased min-variance low-rank gradients","Low-rank estimator hits 33% of full-rank gap with unbiased updates","Unbiased low-rank gradients: 33% gap reduction, same memory"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000612,"raw_usage":{"total_tokens":2881,"prompt_tokens":1016,"completion_tokens":1865,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":632,"completion_tokens_details":{"reasoning_tokens":1771}},"tokens_in":632,"tokens_out":1865,"duration_ms":8955,"temperature":1.0,"reasoning_tokens":1771,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-07T14:33:57.774392+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Measure the per-step bias $E[P_t D_t^{-1} P_t^\\top G_t] - G_t$ when $P_t$ is sampled at step $t_0 < t$ and the gradient's top subspace rotates; train a small model on data with a known drifting subspace and compare PLUMAGE's update direction with the full-batch gradient. If the angular error grows with subspace drift, then the claimed unbiasedness does not hold in the amortized training loop.","supporting_citations":[{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Supplies the GaLore baseline, the low-rank projection training setup, and the hyperparameters reused in the experiments."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"FLORA is the main unbiased random-projection baseline; its moment-projection idea is adapted and compared."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"ATOMO is the prior min-variance unbiased low-rank gradient estimator; PLUMAGE builds on its design while using one-sided projections."},{"cited_title":"Chmiel, I","cited_arxiv_id":null,"evidence_quote":"Establishes the minimum-variance unbiased sparse gradient estimation methodology that PLUMAGE extends to low-rank projections."},{"cited_title":"Alain, A","cited_arxiv_id":null,"evidence_quote":"Provides the importance-sampling variance-reduction perspective behind choosing inclusion probabilities."},{"cited_title":"Eckart and G","cited_arxiv_id":null,"evidence_quote":"Gives the SVD and truncation optimality context for low-rank gradient projections."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Supplies the wheel-of-fortune algorithm that samples exactly k indices without replacement."}],"review_version":1}