{"id":"fae3a942-d72f-42d4-a68d-5b6fd3430d2f","arxiv_id":"2505.10518","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":3,"one_line_summary":"Interleaved register tokens with offset-encoded position ids add a training-only multi-token prediction objective that improves fine-tuning, PEFT, and image-generation pretraining over next-token baselines.","lead":"MuToR inserts small learnable register tokens into training sequences so a model also predicts tokens several steps ahead, while the original next-token task is unchanged, and removes the registers at inference. It reports consistent gains over standard fine-tuning and prior multi-token methods in math reasoning, summarization, and image generation.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Compute-matching controls understate MuToR's per-step FLOPs; gains may be from added compute, not the register objective.","rationale":"I reviewed the method and the evidence. The core idea is coherent: register tokens with a masked auxiliary loss should add supervision without changing inference, and the paper reports consistent gains plus useful ablations in Tables 3, 5, 7, and 8. The weakest point is not the mechanism but the control for training compute. Equation (4) interleaves a register after every regular token, so training sequences roughly double in length; causal self-attention FLOPs grow quadratically, meaning per-step cost is likely three to four times the Next-Token baseline, not roughly doubling as claimed in Section 4.1. The extra-epochs control in Tables 11-12 therefore leaves MuToR with more total FLOPs and also confounds compute with data repetition; the observed overfitting at 10 epochs is exactly why that control cannot certify equal-compute comparisons. The image result similarly compares 100K MuToR-2D steps to 200K Next-Token steps without accounting for the doubled sequence length. Because every headline gain is measured against less-compute baselines, this is the load-bearing concern. A FLOP-matched baseline, with checkpoints selected to avoid overfitting, would settle it. This matches the reader's weakest assumption, so I keep the CONDITIONAL verdict: no change from the reader's assessment.","tokens_in":14568,"tokens_out":11023,"duration_ms":114846,"concrete_test":"Train a FLOP-matched Next-Token baseline rather than an epoch-matched one. In the Gemma-2B/1M-GSM setup, measure total training FLOPs for MuToR (dmax=4, 5 epochs, sequence length about 2T) and for Next-Token (5 epochs, length T), including the quadratic causal-attention cost; then train Next-Token for the epoch count that exactly matches MuToR's total FLOPs, selecting checkpoints on the validation set to avoid overfitting. If that baseline reaches MuToR's MATH500/GSM8K numbers, the advantage is explained by compute; if not, the register objective is supported. The image comparison in Table 6 should be re-run with Next-Token at the FLOP-matched step count, likely above 200K, to settle the similar-compute claim.","verdict_should_be":"UNCHANGED","load_bearing_attack":"MuToR's central claim that gains come from the register objective rather than added training compute depends on the compute-matching controls in Section 4.1 and Table 6. Those controls are not actually matched. Equation (4) interleaves a register token after every regular token, doubling or nearly doubling sequence length; causal attention FLOPs scale quadratically with sequence length, so a MuToR training step on length 2T costs roughly 3-4x a Next-Token step on length T, not the 'roughly doubling' stated in Section 4.1. Yet the control trains baselines for only 2x epochs (Tables 11-12), leaving MuToR with substantially more total FLOPs; and it does so by repeating data, which can overfit, as the tables themselves show. The image comparison in Table 6 similarly calls 100K MuToR-2D steps 'similar training-time compute' to 200K Next-Token steps without accounting for the doubled sequence length. Therefore the gains in Tables 1, 2, 4, and 6 could be due to the larger per-step compute budget, not to the register multi-token objective. The absence of a FLOP-matched next-token baseline or a compute-only control is the load-bearing gap.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes MuToR, a training-only multi-token prediction objective that interleaves learnable register tokens into the input sequence. Each register token is assigned a random future offset d and is trained with a cross-entropy loss to predict the token d steps ahead, while regular tokens retain the ordinary next-token loss and are masked so they never attend to register tokens. At inference the register tokens are discarded. The authors fine-tune Gemma 2B and Llama 3 8B on mathematical reasoning and dialogue summarization, apply the method with LoRA, pretrain LlamaGen-B on ImageNet with 1D and 2D offsets, and evaluate on the star-graph path-finding task of Bachmann and Nagarajan. They report consistent improvements over next-token and multi-token baselines with negligible additional parameters, and claim that the improvements are not due to higher training compute because baselines trained for twice as many epochs do not improve.","tokens_in":14819,"tokens_out":10213,"duration_ms":102377,"significance":"If the empirical claims survive a properly FLOP-matched comparison, MuToR is a genuinely useful training-time technique: it requires no architectural change, adds only a small embedding, is compatible with off-the-shelf pretrained models and PEFT, and supports arbitrarily large prediction horizons without extra output heads. The paper is strong on breadth: it covers full fine-tuning, LoRA, image pretraining, and a synthetic planning task, and it ablates the maximum offset, register sparsity, and shared versus distinct embeddings. The central weakness is that the compute-matched controls in Section 4.1 and Appendix B.1 do not actually match compute, and many of the headline gains are small and reported without uncertainty. If the authors add a genuine FLOP-matched next-token baseline and provide error bars, the contribution would be convincing.","major_comments":[{"comment":"The compute-matched controls do not actually match compute. Eq. (4) inserts a register token after every regular token except the last, so a MuToR training sequence has length about 2T. Causal attention FLOPs grow quadratically with sequence length; even counting non-attention layers, a MuToR step is at least about twice as expensive as a Next-Token step on the original length T, and in attention-dominated regimes it is about four times as expensive. Training the baselines for only twice as many epochs, as in Tables 11 and 12, therefore leaves MuToR with substantially more total FLOPs. The extra-epoch control is also confounded by repeated data, which produces the overfitting visible in those tables. The same issue affects the image experiments: the text around Table 6 calls 100K MuToR-2D steps versus 200K Next-Token steps 'similar training-time compute' without accounting for the doubled sequence length. Consequently, the gains in Tables 1, 2, 4, and 6 could be attributable to additional training FLOPs rather than to the register objective. Please provide a true FLOP-matched or compute-matched baseline (e.g., Next-Token trained for enough steps to match MuToR's measured FLOPs, or with correspondingly increased sequence length and batch size), and report measured FLOPs or wall-clock time for all compared runs.","section":"Section 4.1, Eq. (4), Tables 11-12"},{"comment":"Many of the headline improvements are small and are reported without uncertainty. Gemma 2B numbers are stated to be averaged over three seeds but no standard deviations or significance tests are given; Llama 3 8B results are from a single seed, as acknowledged in Appendix A.1.1; and the image experiments in Table 6 appear to be single runs. In this regime, differences such as GSM8K 38.87 versus 42.10, or 1M-GSM 66.09 versus 68.33, may be within seed variance, and the Llama 3 8B gains are particularly vulnerable because there is only one run. Please report per-seed values, standard deviations, and significance tests, or add multiple seeds for the Llama 3 8B and image-generation runs.","section":"Tables 1-6 and Appendix A.1.1"},{"comment":"The multi-token baseline is described only as 'adapted for finetuning by adding dmax-1 auxiliary prediction heads and applying a loss-weighting strategy,' with no details on head initialization, whether the heads are trained from scratch, or the loss-weight schedule; the reported added parameters (110M-550M) are marked 'in approximation.' A baseline with poorly tuned or under-trained heads could make MuToR look better than it is. Please specify the baseline architecture and tuning procedure, and consider reporting an additional multi-token baseline with a comparable parameter budget or comparable per-step FLOPs.","section":"Section 4.1 and Tables 1-3"}],"minor_comments":[{"comment":"The phrase 'minimize the the expected negative log-likelihood loss' contains a duplicated article and should be corrected.","section":"Section 3.1"},{"comment":"The displayed augmented sequence ends with x_T and contains T-1 register tokens; clarify whether a register is ever inserted after the final token, and state how the register loss handles offsets with t+d > T, i.e., whether such targets are masked out.","section":"Section 3.2, Eq. (4)"},{"comment":"For d=1, a register after x_t receives position t, which is identical to x_t's position; please explain why this does not cause a RoPE collision or ambiguity, since both tokens are present in the same sequence.","section":"Section 3.2, Position Embeddings"},{"comment":"Specify whether all MuToR image runs use 256 registers per image and how the checkpoints at 100K, 200K, and 360K iterations are selected; the main text says 'at 100K steps' but it is unclear if these are final checkpoints or best-validation checkpoints.","section":"Section 4.2 and Table 6"},{"comment":"Please report the search ranges for dmax and alpha, not just the best configurations, and correct the typo 'uniformally' in the dataset description.","section":"Appendix A.1.1 and Tables 9-10"},{"comment":"Add error bars or report the number of runs for the star-graph solve rates, as this task is known to be sensitive to initialization.","section":"Figure 4, Section 4.3"}],"recommendation":"major_revision","confidential_remarks":"The paper is within the journal's scope and the method is simple enough to be reproducible. My main concern is the FLOP-matching control; if the authors provide a direct compute-matched baseline and error bars, I would support publication. The relation to pause tokens and lookahead tokens is cited and discussed, so I do not see a novelty-disclosure problem."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Colleague,\n\nKnow these two things about arXiv:2505.10518. First, the core idea is simple and worth knowing: interleave a learnable register token after each regular token, give it a position id that encodes a random future offset, mask it so it can only see past regular tokens, and train it to predict the token d steps ahead. At inference you drop the registers. It works in fine-tuning and LoRA, where head-based multi-token prediction is awkward, and it extends to 2D offsets for image generation. Second, the evidence is real but modest: consistent single-digit accuracy gains on math and summarization, a nice synthetic result on star-graph path-finding, and FID improvements in LlamaGen. The title slightly oversells it—it does not show that all multi-token prediction needs registers, only that this register trick works in several settings.\n\nWhat's genuinely new: the specific combination of interleaved registers, offset-encoded position ids, and the masked training-only objective. Prior lookahead-token work (Monea et al.) freezes the base model and uses the tokens for inference. Here the registers are a pure training-time supervision channel, with negligible parameters and no inference change. That's a clean idea, and the paper executes it honestly: ablations on dmax, shared vs per-offset embeddings, sparse register placement, and a compute-matching control.\n\nThe soft spots, in proportion. The gains are small—1-2 points on GSM8K, 0.3-0.9 ROUGE—and most numbers have no error bars; only Gemma 2B is run with 3 seeds and even then no variance is shown. The Llama 3 8B and the image runs are single-seed. The compute-matching control is the main weakness. The stress-test note claims MuToR's per-step FLOPs are 3-4x a next-token step; that's off. With the restricted mask, attention pairs scale roughly 2x, and other costs are linear in sequence length, so 'roughly doubling compute' is about right. But the control trains baselines for twice the epochs on the same data, which introduces overfitting as a confound. That the baselines fail to improve with more epochs is weak evidence, not proof, that compute isn't the driver. A cleaner control would use more data or a larger model to absorb extra compute without repetition. Also, the image pretraining section never compares against the head-based multi-token baseline, which weakens the title's general claim. These are addressable rather than fatal.\n\nWho's this for? Anyone working on training objectives for LLM fine-tuning or autoregressive image generation. It deserves a serious referee, and I'd engage with it. My recommendation: send it to review, but ask for error bars, a multi-token baseline in the image experiments, and a better compute-matched control.","headline":"A clean training-only register mechanism for multi-token prediction with consistent but modest gains; worth engaging, though the compute control and missing baselines need work.","tokens_in":15355,"tokens_out":3668,"would_cite":true,"duration_ms":34218,"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":"MuToR claims that interleaving trainable register tokens to predict tokens d steps ahead delivers multi-token prediction benefits during training with negligible parameter overhead and no inference change.","keywords":["multi-token prediction","register tokens","next-token prediction","supervised fine-tuning","parameter-efficient fine-tuning","autoregressive image generation","lookahead supervision","language model pretraining"],"falsifier":"Train MuToR and a next-token baseline with identical total regular-token processing (equal optimizer steps scaled to match sequence length) on a fixed fine-tuning set; if next-token then matches or beats MuToR, the claimed advantage collapses. A sharper control is to insert the same register tokens but remove their auxiliary loss: equal performance would show the extra positions, not the lookahead objective, are responsible.","tokens_in":14376,"feed_emoji":"🧠","tokens_out":9519,"duration_ms":81644,"temperature":0.7,"pith_summary":"MuToR is a training-time technique that slips learnable register tokens between regular tokens, each trained to predict a token d positions ahead while the usual next-token objective keeps running on all regular tokens. The paper argues this produces the denser supervision of multi-token prediction without the extra output heads or architectural changes that earlier multi-token methods require, and without touching inference. The authors report consistent gains over next-token and multi-token-head baselines in supervised fine-tuning, LoRA fine-tuning, and autoregressive image generation pretraining. If the claims hold, lookahead supervision becomes a cheap, portable addition to any off-the-shelf autoregressive transformer.","feed_headline":"Register tokens add lookahead supervision with zero inference cost","feed_subtitle":"The training-only trick lifts math, summarization, and image generation scores with no inference slowdown.","key_machinery":"The load-bearing object is the register token r_d: one learnable embedding inserted after x_t, given position id t+d-1, and trained to predict x_{t+d}. The offset d is sampled uniformly per sequence from {1,...,d_max}, so registers share a single embedding and the offset is encoded through positional bias such as rotary position embeddings. A custom attention mask ensures regular tokens cannot attend to registers and registers cannot attend to other registers, keeping the next-token loss and inference graph unchanged. Because there is one register per position regardless of d_max, the prediction horizon can be scaled almost for free; in the 2D image variant a sampled offset pair (d_h,d_w) maps to a rasterized offset and lets one register cover an entire spatial neighborhood.","core_discovery":"The central claim is that an auxiliary lookahead objective can be implemented by interleaving a single shared learnable register token after regular tokens, giving each register the task of predicting the token d steps ahead. The register is assigned position id t+d-1 so that its prediction looks like ordinary next-token prediction, and a custom attention mask lets registers attend only to preceding regular tokens while regular tokens never see registers. That mask preserves the next-token objective exactly and makes registers discardable at inference. Empirical results show MuToR outperforming next-token and multi-token-head baselines on mathematical reasoning and summarization after fine-tuning, improving LoRA fine-tuning to near full-fine-tuning level, and, in image pretraining, reaching better FID at 100K steps than the next-token model at 200K steps.","pith_inferences":["The attention-mask pattern is effectively an invisible auxiliary slot: the same mechanism could carry other training-only objectives, such as predicting latent features, correctness labels, or auxiliary modalities, without changing inference.","The sparse-register experiment suggests placement is a resource; learning where to insert registers per task could produce similar gains with less compute than uniform interleaving.","If the gain is caused by supervision density rather than architecture, the same register-token idea should transfer to encoder-decoder and non-autoregressive training, though the paper does not test this.","The 2D offset formulation suggests registers can encode spatial structure generically, which may open a route to structured lookahead in non-image domains such as graphs or molecules."],"forward_implications":["In fine-tuning, MuToR can be applied to an off-the-shelf pretrained LLM without architecture changes, so lookahead supervision becomes a training-only addition to the standard recipe.","LoRA fine-tuning with registers reaches or exceeds full-fine-tuning next-token performance, so parameter-efficient pipelines could get stronger models at lower cost.","In autoregressive image generation, one register embedding covers a whole 2D prediction neighborhood that would otherwise require many separate prediction heads.","Because the number of registers is fixed as d_max grows, extending the prediction horizon does not add training cost.","The star-graph results indicate the objective can recover useful supervision in settings where teacher forcing collapses into shortcut learning."],"supporting_citations":[{"why":"It supplies the parallel-head multi-token prediction baseline and the motivation that predicting several future tokens improves pretraining.","marker":"Gloeckle et al. [2024]"},{"why":"It provides the sequential-head variant of multi-token prediction, another approach MuToR is compared with.","marker":"Liu et al. [2024]"},{"why":"It contributes the register-token concept that MuToR adapts for lookahead supervision.","marker":"Darcet et al. [2024]"},{"why":"It defines the star-graph path-finding task where teacher forcing fails, the setting used to show MuToR recovers supervision.","marker":"Bachmann and Nagarajan [2024]"},{"why":"It supplies the autoregressive image generation model and tokenizer used for the vision pretraining experiments.","marker":"Sun et al. [2024]"},{"why":"It provides the rotary position embedding scheme that MuToR uses to encode prediction offsets.","marker":"Su et al. [2024]"},{"why":"It supplies LoRA, the parameter-efficient fine-tuning setup where MuToR is tested.","marker":"Hu et al. [2022]"}],"fun_headline_variants":["Register tokens give next-token models a lookahead boost","Cheap lookahead: one shared token predicts farther ahead","Zero-cost future prediction with discardable registers","Multi-token power without architecture changes"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The load-bearing premise is that running baselines for extra epochs faithfully matches the extra training compute that register tokens consume; if doubled-epoch baselines overfit or misrepresent the cost of longer sequences, the reported gains could come from training budget rather than from the register objective.","fun_headline_variants_meta":{"raw":{"variants":["Register tokens give next-token models a lookahead boost","Cheap lookahead: one shared token predicts farther ahead","Zero-cost future prediction with discardable registers","Multi-token power without architecture changes"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000176,"raw_usage":{"total_tokens":1257,"prompt_tokens":880,"completion_tokens":377,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":496,"completion_tokens_details":{"reasoning_tokens":315}},"tokens_in":496,"tokens_out":377,"duration_ms":3857,"temperature":1.0,"reasoning_tokens":315,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-15T21:08:39.852271+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Train MuToR and a next-token baseline with identical total regular-token processing (equal optimizer steps scaled to match sequence length) on a fixed fine-tuning set; if next-token then matches or beats MuToR, the claimed advantage collapses. A sharper control is to insert the same register tokens but remove their auxiliary loss: equal performance would show the extra positions, not the lookahead objective, are responsible.","supporting_citations":[],"review_version":1}