{"id":"f9eaee66-2d75-46e2-8b0e-aa7fb1d0039f","arxiv_id":"2501.18824","paper_version":1,"verdict":"CONDITIONAL","confidence":"HIGH","novelty_score":5.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":2,"one_line_summary":"TokenTune reduces the memory needed for fine-tuning transformers by computing gradients only for a random subset of tokens, cutting activation cache size with little accuracy loss.","lead":"This paper introduces TokenTune, a fine-tuning method that saves GPU memory by backpropagating through only a random subset of input tokens. It reports accuracy comparable to full fine-tuning and LoRA, and it can be combined with LoRA and QLoRA for larger memory savings.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Attention backward through Eq. 10 requires key and value activations for unselected tokens, so the abstract's claim that only a subset of intermediate activations are cached is not supported as written and the per-layer memory model is missing a 2(n-k)d term.","rationale":"The paper has genuine strengths: released code, extensive experiments on BERT and Llama2-7B, measured peak-memory figures, and a clean, combinable idea (token selection plus LoRA/QLoRA). The empirical claims are mostly measured, not asserted. The weakest link is the analytic core of the contribution: the derivation of the memory reduction. The reader's weakest_assumption identifies the same spot; my analysis goes slightly further and shows it is not merely under-specified. Given Algorithm 1 as written, autograd must retain K_Gbar and V_Gbar for the backward of Eq. 10, because the selected tokens attend to all positions. That is an O((n-k)d) per-layer cache that is absent from the paper's narrative, and it makes the abstract's blanket statement inaccurate for attention layers. If the released code in fact caches full K/V, the measured figures likely remain valid—the missing term is only a few GiB for Llama2-7B at 2,048 tokens—so this is a corrigible presentation/derivation error rather than a fatal flaw; but the paper cannot be verified from its own equations, which is precisely the condition for a conditional acceptance. Secondary but real: the performance-parity claim on Llama2 few-shot shows TokenTune alone at 61.23 versus QLoRA at 62.55 with no error bars or significance testing; GLUE shows 82.1 versus 82.8 for full fine-tuning. These gaps are small but unquantified. The full-fine-tuning memory baseline of 91.4 GiB is estimated, not measured, as the Figure 4 caption admits. None of these overturn the verdict: conditional acceptance with the attention-memory accounting and the K/V caching behavior as the explicit conditions is appropriate.","tokens_in":20551,"tokens_out":19746,"duration_ms":196641,"concrete_test":"Use the released code at github.com/facebookresearch/tokentune: wrap the forward pass of the Llama2-7B attention layer in torch.autograd.graph.saved_tensors_hooks to log the shapes of every tensor the autograd engine retains for the backward of the Eq. 10 operation, with n=2048 and k=256 (12.5% ratio). If tensors of shape (1792, 4096) for K_Gbar and V_Gbar appear among the saved tensors, the abstract's 'only a subset of intermediate activations are cached' is contradicted for attention layers. Then recompute the Table 3 token-tune peak-memory entries with the corrected per-layer accounting 3kd (selected Q/K/V) + 2(n-k)d (unselected K/V) + k*n attention logits, and check whether the reported 64-74 GiB figures are consistent with this model; if they are, the memory-reduction magnitude claim survives but the Section 3 analysis and abstract must be revised.","verdict_should_be":"UNCHANGED","load_bearing_attack":"Section 3.2's equations 10-11 make the backward pass through the attention layer require the full key and value matrices, yet the paper's central claim—'only a subset of intermediate activations are cached during the forward pass' (Abstract)—ignores this. In Eq. 10, h_G = softmax(Q_G[K_Gbar, K_G]^T/sqrt(d))[V_Gbar, V_G], so the attention backward must save the k x n softmax probabilities and the full K and V matrices: dL/dV needs P, dL/dP needs V, dL/dK needs Q_G, and dL/dQ needs K. The (n-k) unselected columns K_Gbar and V_Gbar are produced under torch.no_grad() in Algorithm 1 lines 6-7, but they are inputs to the grad-enabled matmul in line 5 and are therefore retained until the backward pass finishes. The paper never states how these tensors are handled (cached, freed, or recomputed), so the claimed memory reduction for attention layers is not derivable from the equations as written; the per-layer cache must include an O((n-k)d) term for keys and values. This does not necessarily invalidate the measured peak-memory numbers in Table 3 and Figure 4, which are empirical, but it does mean the method's stated mechanism is inaccurate or incomplete, and a reader cannot reproduce the claimed savings from Sections 3.1-3.2. Section 7's limitations do not mention this gap, and the estimated full-fine-tuning baseline (91.4 GiB, Figure 4 caption) is a further unverified quantity, although it is not the primary issue.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper introduces TokenTune, a method for reducing the memory used to store intermediate activations when fine-tuning transformer models. TokenTune selects k input tokens per sequence and backpropagates only through those selected tokens, while the forward pass still uses all tokens. The method is presented for dense, normalization, and attention layers, and the authors report empirical results on BERT-large for GLUE and on Llama2-7B for instruction tuning followed by few-shot evaluation, both with TokenTune alone and in combination with LoRA and QLoRA. The central claims are that only a subset of intermediate activations need to be cached and that TokenTune achieves accuracy comparable to full fine-tuning and to representative memory-efficient baselines while substantially reducing GPU memory.","tokens_in":20874,"tokens_out":5072,"duration_ms":58001,"significance":"If the memory accounting is clarified, TokenTune is a useful and simple contribution to memory-efficient fine-tuning. Its token-selection mechanism is orthogonal to parameter-efficient methods, and the reported combinations with LoRA and QLoRA plausibly yield cumulative memory savings. The empirical evaluation is extensive, covering both medium-size encoders and a 7B decoder, and the code is publicly available. The paper does not rely on circular reasoning: it evaluates on standard external benchmarks and the method is defined by an explicit approximation of the gradient. The main weakness is that the attention-layer treatment is underspecified, which directly affects the paper's central mechanism for a large portion of the model.","major_comments":[{"comment":"The paper does not specify what happens to the unselected tokens' key and value activations K_Gbar and V_Gbar during the backward pass. Equation (10) shows that the selected tokens' attention output h_G depends on [K_Gbar, K_G] and [V_Gbar, V_G], so the backward pass through this matmul requires the full key and value matrices, including the unselected columns. These quantities are produced inside the torch.no_grad() block in Algorithm 1 (lines 6-7) but are then used in the grad-enabled computation in line 5; under standard autograd semantics they are retained as constants for the backward pass unless the implementation explicitly frees or recomputes them. The abstract's claim that 'only a subset of intermediate activations are cached during the forward pass' is therefore not derivable from the equations as written, and the per-layer memory model appears to be missing a term of order 2(n-k)d for keys and values. The measured peak-memory numbers in Table 3 and Figure 4 may still be correct, but the authors should state whether K_Gbar and V_Gbar are cached, freed and recomputed, or handled otherwise, and if they are cached, they should update the memory model and the claims accordingly.","section":"Section 3.2, Eqs. (8)-(11), and Algorithm 1"},{"comment":"The attention equations omit the causal mask that is required for the language-modeling experiments on Llama2-7B. As written, Eqs. (10) and (11) describe bidirectional full attention, whereas instruction tuning of a decoder-only model uses causal masking. The paper reorganizes tokens into h_G and h_Gbar groups, which means the mask must be permuted consistently with the concatenation [K_Gbar, K_G]; this is not described anywhere. Without this detail, the method for the main large-model experiments is not fully reproducible from the equations. The authors should specify how the causal mask is applied after token reordering.","section":"Section 3.2, Eqs. (10)-(11)"},{"comment":"Equation (5) is written as if the gradient of the loss with respect to W is a concatenation of a selected-token term and a zero term, but the gradient for a dense layer is a sum (or matrix product) over the selected positions: dL/dW = sum_{i in G} (dL/da_i) sigma'(z_i) h_i^T. The notation in Eq. (5) is therefore incorrect as a mathematical statement and should be replaced with the summed form or a clear outer-product expression.","section":"Section 3.1, Eq. (5)"}],"minor_comments":[{"comment":"The full fine-tuning memory value of 91.4 GiB is an estimate, not a measured value, and the caption does not give the estimation formula. Because several reported percentage reductions in Figure 1 are relative to this estimate, the authors should provide the exact formula or state clearly that the full-fine-tuning baseline is extrapolated from the other measurements.","section":"Section 5.3 / Figure 4"},{"comment":"Several entries in Table 3 are missing spaces between numbers (e.g., '65.0152.6578.37' in the 10% row of part (a)), which makes the table hard to read and should be fixed.","section":"Table 3"},{"comment":"There are occasional typos and formatting issues, such as 'LLama' for 'Llama', 'TOKEN TUNE' vs. 'TokenTune' inconsistency, and broken line breaks in the abstract and references. These do not affect the technical content but should be cleaned up.","section":"Throughout"},{"comment":"The claim of being 'the first method that reduces GPU memory usage for fine-tuning via token selection' should acknowledge the preliminary workshop version more prominently and clearly state what is new in this submission relative to that version.","section":"Section 1 / Contributions"}],"recommendation":"major_revision","confidential_remarks":"The paper has a solid empirical core, but the attention-memory issue is central to the method's stated mechanism and needs to be resolved before publication. If the authors can show that K_Gbar and V_Gbar are indeed not retained (e.g., by recomputation) and update the memory model, or alternatively revise the claims to match the actual caching behavior, the paper would be close to acceptance. The causal-mask specification is also important for reproducibility."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Colleague,\n\nYou should know two things about TokenTune (arXiv:2501.18824). First, the core idea is as simple as it sounds: pick k random token positions per sequence, backprop only through those, and freeze the rest. It works. On GLUE with BERT-large, it lands within a point of full fine-tuning on average, and on Llama2-7B instruction tuning it roughly matches LoRA and QLoRA while shaving real memory. The paper also shows clean composition with LoRA/QLoRA, which is useful. Second, the marketing is ahead of the mechanism. The abstract says 'only a subset of intermediate activations are cached,' but the attention equations don't support that as written. In Eq. 10, h_G uses K_Gbar and V_Gbar from all unselected positions. To backprop through that matmul, autograd has to retain those full K and V tensors. So per layer you still cache O((n-k)d) for keys and values, plus the k-by-n attention matrix. The paper never says whether those tensors are cached, recomputed, or freed. That missing term is exactly what the stress-test note flags, and it's a real gap. It does not invalidate the empirical memory numbers in Table 4 and Figure 4—those are measured—but it means the claimed savings cannot be derived from Section 3, and the abstract is misleading.\n\nElsewhere, the paper is decent. The writing is clear, the GLUE evaluation uses the official server, and the few-shot harness is standard. The main weaknesses beyond the attention accounting are that Table 2 has no error bars or significance tests, and the full-fine-tuning baseline (91.4 GiB) is estimated, not measured. Neither is fatal; the accuracy differences are small but consistent, and the baseline estimate is clearly labeled.\n\nWho should read this: practitioners who want a cheap knob to trade activation memory for a bit of gradient fidelity, especially on encoder models or combined with LoRA. It is an incremental but genuinely useful empirical contribution.\n\nVerdict: send it to review. A referee should ask the authors to fix the memory accounting and add variance bars; the method itself is sound and the results are plausible. I would cite it if the authors clarify the attention memory story, because as it stands the headline claim overstates what is actually cached.","headline":"TokenTune's token-subbing trick is simple, useful, and mostly works, but the attention memory accounting has a real gap that makes the paper's headline claim overstate what is actually cached.","tokens_in":21428,"tokens_out":3708,"would_cite":true,"duration_ms":41046,"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":"Fine-tuning transformers by backpropagating through only a randomly selected subset of input tokens cuts activation memory while keeping accuracy within a few tenths of a point.","keywords":["memory-efficient fine-tuning","token selection","intermediate activations","transformer fine-tuning","large language models","LoRA","QLoRA","few-shot learning"],"falsifier":"Run TokenTune on an attention-heavy transformer while instrumenting peak activation memory layer by layer and compare with full caching: if the unselected tokens' keys and values still occupy memory for the backward pass, the measured savings will fall short of the paper's formula; separately, fine-tune Llama2-7B at a 10% selection ratio and check whether few-shot accuracy drops more than the roughly one-point spread reported in Table 3, which would indicate the approximation is not reliably on par.","tokens_in":20322,"feed_emoji":"🧠","tokens_out":9008,"duration_ms":81137,"temperature":0.7,"pith_summary":"TokenTune proposes that fine-tuning a transformer does not need gradients from every input token. It selects a random subset of positions, freezes the rest during the backward pass, and caches only the selected tokens' intermediate activations, so the forward pass still sees the whole sequence while activation memory shrinks. The paper shows on BERT-large and Llama2-7B that this approximation matches full fine-tuning or LoRA/QLoRA within a few tenths of a point on GLUE and few-shot benchmarks, while cutting peak GPU memory to as little as 21% of full fine-tuning when combined with QLoRA and by up to 28% when used alone. The method is offered as a simple, combinable knob: the selection ratio trades memory for gradient fidelity.","feed_headline":"Fine-tuning on a token subset cuts memory and keeps accuracy","feed_subtitle":"TokenTune caches only selected tokens' activations, using as little as 21% of full fine-tuning memory on Llama2-7B.","key_machinery":"The central mechanism is the split of each layer's hidden states into a selected group $h_G$ of $k$ positions and an unselected group $h_{\\bar G}$, with gradient computation disabled for the unselected group. During the forward pass the full sequence is processed, but only the selected positions' activations are retained for the backward pass; the gradient for unselected positions is set to zero, so the weight update uses only $h_G$. The paper spells this out for dense layers, where the gradient of the weight becomes a function of $h_G$ alone, and for attention, where selected queries attend to both selected and unselected keys and values while unselected branches do not contribute gradients. This selection-and-mask pattern turns a memory cost that scales with the full sequence length into one that scales with the chosen ratio $k/N$.","core_discovery":"TokenTune's central claim is that intermediate-activation memory during fine-tuning is largely redundant across token positions, so it suffices to backpropagate through a subset of k randomly chosen tokens and cache only their activations. The paper derives the resulting gradient approximation for dense, normalization, and attention layers: the backward error for unselected positions is set to zero, which makes the weight gradient depend only on the selected hidden states $h_G$. In attention, the selected tokens' output still attends to both selected and unselected keys and values, but the unselected branches run with gradient computation disabled, so they need not contribute gradient terms. On GLUE with BERT-large, TokenTune averages 82.1 versus 82.8 for full fine-tuning, and on Llama2-7B few-shot benchmarks it improves the base model by about half a point, with a 30% selection ratio matching LoRA and QLoRA within noise. The paper concludes that token selection is a valid standalone memory-saving technique and a drop-in complement to parameter-efficient and quantized methods.","pith_inferences":["The paper leaves open whether adaptive token selection, such as picking positions by attention weight or gradient norm, would outperform the uniform random sampling it uses; the method's design would allow such a swap without changing the memory argument.","Because the mechanism is architecture-agnostic, the same token-selection trick could transfer to vision or multimodal transformers, a setting the paper names as future work.","If a user wants to predict the memory savings for a new model, the paper's breakdown suggests the selection ratio helps most when intermediate activations dominate the memory budget, as with medium encoders at large batch sizes; for models where parameters and optimizer states dominate, TokenTune alone gives little relief and should be paired with LoRA or QLoRA."],"forward_implications":["TokenTune can be combined with LoRA and QLoRA, and the savings add: TokenTune+QLoRA on Llama2-7B uses about 11.7 GiB at a 12.5% selection ratio, roughly one quarter of QLoRA alone.","A selection ratio of 20–30% is enough in the paper's experiments; accuracy is relatively flat across 10–50%, while memory grows steadily with the ratio.","For medium encoders with large batches, where activations dominate memory, TokenTune alone cuts BERT-base activation memory from 23,196 MiB to 9,952 MiB at batch size 512.","Because the forward pass still uses all tokens, the method does not require re-training or architectural changes and drops into existing transformer fine-tuning pipelines."],"supporting_citations":[{"why":"It is the quantized fine-tuning baseline (QLoRA) that TokenTune is compared against and combined with, providing the memory figures for the combination.","marker":"Dettmers et al., 2023"},{"why":"It is the low-rank fine-tuning baseline (LoRA) used for comparison and combination, and the main PEFT method TokenTune builds on.","marker":"Hu et al., 2022"},{"why":"It defines the Llama2-7B model used in the large-decoder instruction-tuning and few-shot experiments.","marker":"Touvron et al., 2023"},{"why":"It defines BERT, the encoder model used for the GLUE experiments, and its evaluation protocol.","marker":"Devlin et al., 2019"},{"why":"It defines the GLUE benchmark whose tasks and test-server scores measure downstream classification performance.","marker":"Wang et al., 2018"},{"why":"It supplies the Open-Platypus instruction-tuning dataset used to fine-tune Llama2-7B.","marker":"Lee et al., 2023"},{"why":"It supplies the evaluation harness and prompt formatting used for the few-shot benchmark scores in Table 2.","marker":"Gao et al., 2021"}],"fun_headline_variants":["TokenTune: Fine-tune with fewer tokens, lower memory","Backprop through token subset slashes fine-tuning memory","Cache only some activations to fine-tune large transformers","Token selection cuts activation memory in LLM fine-tuning","Equal quality, far less memory: backprop on token subset"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The load-bearing premise is that fine-tuning can succeed with gradients computed from only a randomly chosen subset of the input tokens; the claimed attention-layer savings also depend on an unspecified treatment of unselected tokens' key and value activations.","fun_headline_variants_meta":{"raw":{"variants":["TokenTune: Fine-tune with fewer tokens, lower memory","Backprop through token subset slashes fine-tuning memory","Cache only some activations to fine-tune large transformers","Token selection cuts activation memory in LLM fine-tuning","Equal quality, far less memory: backprop on token subset"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000573,"raw_usage":{"total_tokens":2756,"prompt_tokens":1045,"completion_tokens":1711,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":661,"completion_tokens_details":{"reasoning_tokens":1629}},"tokens_in":661,"tokens_out":1711,"duration_ms":14540,"temperature":1.0,"reasoning_tokens":1629,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-09T22:20:43.381412+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Run TokenTune on an attention-heavy transformer while instrumenting peak activation memory layer by layer and compare with full caching: if the unselected tokens' keys and values still occupy memory for the backward pass, the measured savings will fall short of the paper's formula; separately, fine-tune Llama2-7B at a 10% selection ratio and check whether few-shot accuracy drops more than the roughly one-point spread reported in Table 3, which would indicate the approximation is not reliably on par.","supporting_citations":[],"review_version":1}