Corrector Sampling in Language Models
Authors: Itai Gat
Neta Shaul
Uriel Singer
Yaron Lipman
Venue: NeurIPS 2025 poster ()
Keywords: Language modeling
Sampling
Primary area: deep_learning
TL;DR: We introduce Resample-Previous-Tokens (RPT), a sampling method that allows models to revisit and replace previously generated tokens, leading to ~10% relative improvements in reasoning and coding after a short fine-tuning.

Abstract: Autoregressive language models accumulate errors due to their fixed, irrevocable left-to-right token generation. To address this, we propose a new sampling method called Resample-Previous-Tokens (RPT). RPT mitigates error accumulation by iteratively revisiting and potentially replacing tokens in a window of previously generated text. Fine-tuning a pretrained 8B parameter model with RPT for only 100B resulted in ~10% relative improvements on reasoning and coding benchmarks compared to the standard sampling.

Corrector Sampling in Language Models
Itai Gat
Neta Shaul
Uriel Singer
Yaron Lipman
FAIR, Meta
Abstract
Autoregressive language models accumulate errors due to their fixed, irrevocable left-to-right token generation. To address this, we propose new sampling method called Saumple-Previous-Toee iteratively revisiting and potentially replacing tokens window of previously generated text. Fine-tuning: pretrained 8B parameter model with RPT for only 100B resulted in 100% relative improvements on reasoning and coding benchmarks compared to the standard sampling.
(RPT). RPT mitigates error accumulation by
Introduction
Autoregressive (AR) language models represent pivotal advancement in sequence generation, deliv- ering state-of-the-art results in translation, code generation, text summarization, question answering. and many other text tasks. By factoring the probability p(r) of sequence x with the chain rule, AR training reduces to learning the next-token-prediction (NTP) conditional probabilities, making both training and inference highly efficient. At inference time, however, this left-to-right process is irrevocable: once a token drawn the model cannot revise it, resulting error propagating from the NTP sampling.
Recent progress LLMs has largely focused on improving training via: better data quality (Meta
2024 DeepSeck 2024 Groeneveld et al., [2024] Gemini 2025; Li et al. 2025 OLMo 2025 Lambert
al. 2025);
(ii) modifications the transformer architecture, such as, positional embeddings (Su
et al. 2021 Press et al. 2022), attention mechanism (Ainslie et al.) 2023] Beltagy et al.) (020);
and (ii) Remforcement learning fine-tuning (Ouyang et al.) [022] | Ramesh al. 2004] DeepSeee
2025, Muennighoff et al.) 2025 Kimil 2025). Nevertheless, the sampling itself remains relatively
underexplored, with most models still relying on the vanilla NTP sampling (Holtzman et al. 2020
In this work, we introduce Resample- Dexsumple-rrevous-Tiee Pre---rs--ooueeeen (RPT), a novel sampling method that iteratively revisits fixed-size window of previously sampled tokens. Unlike NTP sampled tokens, which are unchangeable once sampled, RPT allows local token replacements during generation, potentially reducing error accumulation. RPT training lightweight, allows fine-tuning pretrained AR models, preserves their NTP sampling quality and speed (i.e., with rey-value caching), and integrates seam- lessly into existing AR models and code. We fine-tuned pretrained 8B AR model for 10% of its final training iterations and compared it to the fully-trained pretrained model to find RPT pretraining and sampling provides 5%-10% relative improvements common benchmarks. Our contributions include: (1) Introducing RPT, simple yet effective sampling process from L.LMs; (2) Developing a training algorithm that incorporates seamlessly in AR training loop with minimal overhead and parameter cost; (3) Demonstrating empirically that RPT outperforms standard NTP sampling in both reasoning and coding tasks, as well as a controlled error analysis.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

2
Reaaple-PreviousTTTeen sampling
We denote by x: (01,
sequence of tokens, where each token z; is an element of vocabulary set V. i.e., E V, and V = V denotes the size of the vocabulary. Autoregressive (AR)
, Fn) ( ,an.
modeling uses the probability chain-rule model the joint probability of sequences æ,
p(z) Il p(z))a<)).
where the notation Fzt (I1, (. 44-1)) 1), and learn model (denoted with
(2)
that, in inference time, can be used for sampling from the joint via next-token-prediciinn (NTP) sampling:
Next Token Prediction
~ x(2)(
(3)
for i = 1, This procedure obviously introduces some errors into the NTP sampling process that originates from the approximation errors equation? 2
Our goal is to reduce the errors of the NTP sampling process by using Resample-Previone Prerpple-Previous-Tokesss (RPT) sampling. In its simplest form, consider pairs of adjacent tokens 2i, Ii initialized with NTP sampling, and perform iterations of the form
flas (141) Fi+1 (Miiil/rcit)
Resample Previous Tokens
(4)
where the number of iterations is a hyper- parameter set by threshold (c.g., confidence), see Figure [ for an illustration. To implement RPT sampling we require, in addition the learned NTP conditional p(x,/res) (used in the second equation in 4), to learn the previous- token- prediction (PTP) conditional (used in the first equation in i.e.
secured
she
secured
she secured some
she fetched some
[( I4+1) & p(z,/Z20T (5)
she fetched
same
supply
that predicts the i-th token 2; given the previ- ous tokens ee and future token _;-1. key observation of this paper is that since the error of predicting æ; using future token smaller than standard NTP error, the overall error in the RPT sampling
she fetched
fresh
water
Figure 1: Windows (in black) show 2 iterations of RPT sampling (see equation 4 green context t< blue are tokens introduced with NTP and red are tokens changed during RPT iterations.
4 can shown (under certain conditions) be smaller than NTP and often empirically leading to a better (yet more expen- sive) sampling procedure than NTP. We will theoretically compare and analyze the errors in NTP and RPT sampling in the next section.
More generally, we can consider longer sequences of w E N tokens, J..... -itw-1+ where w is called window size. Initialize with NTP, and iterate
p(x;
Feitw, Fi
Fi+1
M(1+1
Teitwy Fi+1)
T:42
p(r++2
Feitw, Li+2)
(6)
Fitw-1 1 1

---

Output
plasle.a))
plas)
plus Res))
...
pi (c)))
Model
Input
Rig
24
.ng
Figure 2: RPT training with window size 2: with some probability adjacent tokens are swapped (bottom row, 2y ++ IT), the model predicts the conditional probabilities as shown in the top row.
where z; means that the token v; is excluded from the condition, that is,
(w<i+w: 2.) (. -4 2i+1, ................................ -++w-1))
Rearranging the indices, we need learn conditionals of the form
ostsw- w-
(8)
The & = 0 case corresponds to the NTP conditional. As above, all the conditional with 0 < l w, enjoy lower error than the standard NTP prediction, and in fact, the more future tokens are used (i.e., the larger w E), the lower the error (see Figure Lastly, we that performing RPT with window size w = 0 collapse to standard NTP sampling.
2.1 Learning the conditionals required for RPT
For RPT sampling we need learn to sample from conditionals of the form p(r;:e/Fee+1, , 0 S f < w (see equation 8) To this end, we consider the following two sequences of indices: first, σ (01, ..... n.). permutation of {1....,7) n that encodes the order in which tokens are fed into the model, i.e., the i-th token that is fed into the model is T we denote the entire permuted sequence by H, = (Way Togs... ).a)) consider only a particular type permutations, namely, where a random index k E (I....,,) ..) is pushed w places to the right, that is
places to right
(1,
1 +
1 +
The second sequence T = (TI,..., Rn) prescribes the target token index Wr; for the model i-th output,
++ 1 i if for 22.................... o2,,,,,..,, = -1,,2...,.
(10)
k
o/w
That is, for the permutation in equation
we have
prodict pist k-th token
(11)
== (2,3....,
k + w
+ w w 1,...,,
Note that all the tokens predict the next token, as encoded in the respective т. except those that correspond to the part marked by over-brace that specifically predict the (past) k-th token, Zg. With a, T our network models
[p(a+11 (cit1)
if < k or o/w 0 = - k
p(Z (o
(12)
where we denote Fresil (Nirs Togy No.) The first case in equation 12 corresponds to standard NTP, while the second case corresponds exactly to the PTP conditionals required for our RPT sampling with window equal or smaller than w, as summarized in equation Figure 2 shows an example corresponding to the choice of a (1, 2, 4, 3, 6,...,
Training During training, we permute each given training data sequence with a probability of s using permutation a as defined in equation The permutation is performed with k sampled uniformly from the set f1.....M modeling, we perform the swap more than single token. In detail, we set some probability (0, 1) and then construct a permutation a as follows: (i) start with the identity permutation q -= (1, .......,.. 2..,,...,,? ..) and traverse from left to right with indices l: = 1, ..; (ii) for each index k; with probability y we move forward by w - places. (iii) if index is moved forward, continue from index & + (to avoid overlaps); otherwise, proceed to the next index A + 1.
w + 1). In fact, due the large sequence length typically used in text

---

Algorithm Resample-Preevius--Tore (RPT) training
1: Input: dataset D; pretrained or initialized model fo with params 80
2: Hyperparameters: Probabilities 9 E (0, 1); window size w 2 number of iterations
3: <- Pg
b Initialize parameters
4: for iteration i to do
5: Draw ~ D
> Draw a training sample
6: Set (1,2, 1)
The identity permutation
7: S (2,3. (2.3....,,)
b Next token index
8:
With probability permute using q and w
"training" section 2.1
9:
Compute т
Use equation TO
X (Ior
Set input to network
11:
Y It
Set the target
12:
C + (fo(fo (X)
Evaluate cross-entropy loss, equation
13
4- optimize(
Update 9 with optimization step
14: end for
Modeling the NTP together with the PTP conditionals (equation 12) requires the model to aware whether the input and predicted tokens were permuted, as it cannot infer the true order of the sequence otherwise. To address this, we use learned positional embedding layer that takes in the input (i.e., σ) and target (i.e., 7) positions, c.g., for NTP, the value is always one. The model is provided with the input token T its input position ag its target position тı (note that this is not a permutation in our case), and the output is compared to the target token using a standard cross-entropy loss:
((0) = -EE log po (2r, Toae 1),
(13)
where D is the training dataset. Algorithm summarizes this training procedure, which can be used to train from scratch or finetune a pretrained model.
3 Analysis of RPT versus NTP sampling
In this section develop the relevant theory to compare the error in wewppple-ppeviosttokkeeeeeee (RPT) sampling and the standard next-token-preedee -reekictprriiitionn (NTP) sampling. We provide the analysis for the simple 2-token (w 2). For notational conciseness we denote by T 2++1 p(æ,, di+1 (ي<²)) the ground-truth joint of the next two tokens (#,, given the context Tei. When there notational ambiguity we use more explicit notation for joints, conditionals, and marginal, i.e. (I, 2 = π [2,, xz)), Til [ (r,) = <xr]a,)) and Tı( = Tæx). We will perform an asymptotic error analysis. That is, we assume introduce asymptotically small errors to the approximated/lernnel conditionals used in the NTP (see equation 3) and RPT (see equation 4 sampling schemes.
*(x,) = π (xx) + E(xx), #(0+11)))) = x ( (r:) + e( A(r,(1141) = (xx; (m++) 1) <(xa(1)41)
(14a)
(14b)
(14c)
where c denotes the approximation errors and assumed to be very small (meaning that we ignore second and higher powers of c). The marginal æ, and conditional æz+1j; represent the next-token- prediction (also abbreviated NTP) while TiJ++1 represents the previoustoocen-predictio (PTP).
The two sampling methods NTP and RPT use the approximated conditionals in equation 14 and consequently introduce a certain error to their sampled joint 61,
A (x; 2441) = x (m; 11+1) ) + c(a,, 2:41)
(15)
Our goal: Quantify the errors e(2,, F441) for each of the sampling method: NTP and RPT.

---

To quantify the size of the errors we will use the total variation distance, defined between two probability distributions p. P over a state space A {a} by
Pp - p = sup (p()) - :(A)] = 1lp pp--pli -il1
(16)
ACA
where we denote the I-norm of a vector E R by
||l||
Σ |o(a)]
(17)
aEA
The RPT sampling, under certain conditions, introduces improved error compared to NTP. In more details, we define the RPT factor to be the ratio of RPT and NTP error bounds. takes the form
|esst22lox + |es42|sllo ||all1 + 11
RPT factor
(18)
where for matrices m E BxA we use the max-norm norm
|||| max
mmbbaa).
(19)
b
Theorem 1. If < then ressmmle-previions-toenn (RPT) sampling achieves a lower error bound compared a tight error bound of next-token-rediction (NPT) sampling.
The RPT factor is composed of two parts: the first is к which is a function of the ground truth conditionals Mij++1 and 1i+1/e which we cannot control (explained more later), and the ratio of the sum of PTP and NPT errors ||E ol + || and the NPT errors || + ||xles Our main observation is that, in practice, the PTP eπror Eji+1 is considerably lower compared to the NTP errors, 6 and 4441/-- Intuitively, this means that, given context T<is predicting the next token i; given a future token T4+1 is easier than without it. That is
Sp PIP emorr
NTP crror
en NTP eroor
(20) and consequently the second part (and the part controlled by the approximation method) of the factor in practice smaller than To justify 20] empirically we use the Pinsker inequality to bound the total variation distance of general probability distributions p, á with their KL-divergence,
[Glh
||
p1l - - (pllp))))
(21)
where DxL. (pl()) H(p where Hp. P) is the cross entropy and (p) is the entropy. Now, making the further as- sumption that the entropy of the ground truth conditionals Ti
NTP
Cross Hetrogy
T1a++ Te+1js near zero for non trivial context deei and large vocabulary sizes V, we focus on the cross entropies and com- pare those of the NTP, i.c., H HTT,, #.) and H (T 1j))
and the PТP, H (mqp=1 #1/+1). Figure 3 shows the cross en- tropies Hm; Rt) and H(T; 1 #2/(11 ) computed during the training of 1.5B parameter model on a validations set (training less than epoch). As be noticed, the PTP enjoys a considerably lower cross entropy loss and therefore under our assumptions also lower total variation norm. Next, we derive asymptotic bounds the sampling errors for both NTP and RPT, and compute the RPT factor. We start with the NTP emror..
50
150
Thousands
Figure 3: Cross entropy training curves of NTP and PTP.
Next-token- prediction eeretoken--rredctiion asymptotic error The error e(x,, 2+41) in the NTP sampling of the next two tokens can be calculated directly from the probability chain rule
A (x, x++1) (
(22)
(I((; + e(x:)) (π(;,11))) + +((ss41l2))
(23)
((0,,,241) + clx ( (ma) + (( o(c),
(24)
5

---

where by o(E) we denote terms that is asympiotically smaller than e, e.g., in this case
Therefore the asymptotic error in this case
e(zi, I++1) ((z))=((411 (m,) + ((2441l1))) d(e).
(25)
We can use this asymptotic error expansion to derive a simple tight bound on the NTP error (see Appendix A. for more details),
|ell] - ||e Ill + lel + o(c).
(26)
Resample-prewe previous-token asymptotic error Analyzing the error ((41) in RPT sampling requires more claborate computation it directly involves the stationary distribution of the RPT iterative sampling procedure in equation 4 The RPT iterations can be seen as a Markov chain where the states include all pairs of possible tokens (æz, 1;+1) V V. The RPT iterations can be interpreted as sampling from the following Markov probability transition kernel, which corresponds exactly to single iteration of equation 4
(27) (x)) is a standard nex--tooen--reeditionn the conditional π T (++1) predicts the i-th token given the context zzi and the next token 2i41 and therefore introduces a smaller error according to the assumption in equation 20 Before estimating the asymptotic error in the joint c(z,, 11+1) we first need to derive the asymptotic error in the kernel. We denote this error by e(2%, 4/41 12,, and it can be computed similar to before by expanding
p(a;; 2'4 [2 (1+1) n =2( m((44
Note that while T T
p(a x |z, 12+1) = (π(2|(0441) + e(x)) (1(4))111()) + c(<(+1la())
(28)
and get the kernel asymptotic error
c(a%, 2 [ Ti+1) 1) = c(2/|=41)) (m,) e( [2) + o(e).
(29)
Consequently bound on the kernel error is
||e| || + || ol + ooe),
(30)
where we use the max-norm again; see Appendix A.2 for the exact derivation. Next, to achieve
bound on the RPT sampling error 6,,++1 of the stationary distribution of the Markov chain 28 we can
use standard perturbation bounds (Cho and Meyer 2001 Seneta 2006
Kirkland et al. 2008
|||/|, S K ||ellso
(31)
where K is constant that depended on the ground truth Markov chain 27 known as its conditional number. There are different definitions of K. some of them are tighter than others but harder to compute/provide intuitive explanation. A relatively intuitive one is α (1 A(P))-1. Here, P the transition matrix of the Markov process 27 with entries p(a/, #/ [24, F1+1). and A(P) is its ergodicity coefficient, which is a scalar in T0. that quantifies how quickly the influence of the initial state diminishes over time, or equivalently, how quickly the chain converges to its stationary distribution. Plugging our kernel error (equation 30) into equation 31 we get the RPT error bound
||||| S (11 + [|<++)j(loo) + o(c).
(32)
Comparing this bound to the one in equation 26 we get the RPT factor in equation 18
In the next paragraph we provide a synthetic experiment that shows the benefit of the RPT over NTP
in controlled setting experiment.
Synthetic example We have conducted synthetic (toy) experiment to compare the NTP and RPT
sampling. To that end we considered a random ground truth joint π (1+1) with V/ = 20 vocabulary size. That is, for each entry ( 11+1) we sampled i.i.d. from 2/ 1) (uniform distribution) and normalized. To add noise to, e.g., Tj we random a noisy distribution (x;)) (as done before) and set
(x;) (1 - (z.) + e,n(2;)
(33)
where 6, >> is the noise level, and similarly for the conditionals T4+1/s (with noise level e,) and Ti (with noise level e,). We used three cases for the noise levels e, > cp: (i) Oracle where ep = 0 (ii) Medium where 6, 0.5c,; and (iii Low where ep = 0.75e, We take 6,
(we show other
choices of this base noise in Appendix A.33 We numerically computed the sampling error E1, of NTP and RPT, and repeated this experiment 1000 times. Figure 4sows
shows the histograms of the total variation distance ||r,++1/rrv (bottom row) and max-norm ||c,5+1|l (top row) per experiment, and their respective mean. Note that the smaller the ratio c//c, the larger the benefit in RPT sampling over NTP, and that RPT already exhibits non-trivial improvement for cases where e,/e, 0.75.

---

tp/tn
"p/ta
"p/"n
"p/ta
"p/ta ta
"p/"n
0.75
107 NTP
Denng
Donatt
Maximum errors
(b) Total variation
Figure 4: Synthetic sampling experiment comparing next-token -reeeictonniioon (NTP) and resample- previous-token (RPT). Maximum errors and total variation errors are shown for three different error levels E < e.. The lower the relative eπror s, the better RPT sampling over NTP, where RPT already shows benefit over NTP for cp/c. 0.75.
Table I: Results of RPT sampling on coding and reasoning task. We fine-tuned the model with an additional 100 billion tokens. Metrics are reported for both the fully autoregressively trained model (AR-F, 1 trillion tokens) and the initial checkpoint (AR-C, 900 billion tokens). Note, k = 0 is NTP sampling from our model. For all methods, we report best temperature in {0.0, 0.05, 0.1).
k
HvmanEval+
MBPP
GSM8K
C++
C#
PHP
Bash
Java
TypeScript
AR-F
25.6
39.0
35.2
28.5
17.7
28.5
6.9
37.9
35.8
AR-C
24.4-125
38.8
35.4
28.6
22.8
24.2 -33
8.8
32.9
33.9
0
2135
39.6
35.5 -0.15
31.0
22.8
26.1
9.5 -3.6
38.3
37.7
0.5
27.4 41.3%
39.6 .0.4%
37.3 +2.15
31.0 -2.9%
22.8 45.1%
26.1 -1.4%
9.5 42.6%
40.5 -2.6%
37.7 +1.9%
27.4 +1.8
40.6 -144%
37.5
-2.355
31.0 -2.9%
22.2 +1.5%
25.5 -300
9.5 -2.6%
41.1 -3.0
38.4 +2.00
28.6 +505
40.6 +1.45
31.0
22.2
25.5
.7.69
2
28.0
39.0
36.3 +1.5%
31.0 -2.5%
19.6 +1.0%
26.7
9.5 +2.6%
37,9
37.7 +1.9%
4 Related work
Sampling methods Autoregressive models typically employ NTP sampling, which is based on the chain rule of probabilities. However, NTP sampling has a limitation: it cannot correct previously predicted tokens, potentially propagating errors throughout the generated sequence. To address this issue, prior works have explored searching the model's probability space using techniques like Beam Search (Freitag and Al-Onaizan 2017). Despite their appeal. these methods are not widely adopted in practice due to issues like repetition and collapse (Holtzman et al. (020). In contrast, our work extends the traditional NTP sampling by introducing RPT, novel sampling approach that allows arguably improved sampling from the joint probability of tokens.
Predictor-corrector Discrete diffusion models often apply remasking of previously revealed to- kens (Lezama et al.] 2022 Campbell et al. 2024 Gat et al.) 2024), but these methods are not be directly applicable to autoregressive LLMs. Alternatively, other works suggest prompting the model to self-correct through carefully designed inputs (Chen er al., 2025 Muennighoff et al. 2025). Other approaches modify the model architecture or generation process to achieve similar goals: Li al.] (2024) introduce predictor-comeeeeee corrector mechanism by modifying the transformer architecture to accumulate states; (Stern et al. (2019); : Gu al. (2019) enable insertion, deletion, and replacement of tokens during generation, elfectively performing implicit corrector iterations. While previous works have made significant progress, they fall short of achieving tfat-oof-hhe-ar performance. In this work, we propose novel training and sampling procedures that enable corrector iterations for LLMs, improving the current state-of-the-art
Any-order architectures Our method augments standard next-token- with explicit infor- mation about permutations of the input and target tokens. Prior work has explored any-order language modeling. Pannatier et al. (2024) propose an architecture that uses double positional embeddings to encode the absolute positions of both the observed tokens and the targets to be predicted. Yang et al. (2019) achieve permutation awareness by permuting the attention bias. A complementary line of research employs discrete flow or diffusion models, which-br virtue of their bidirectional attention and path training objectives can naturally generate tokens in arbitrary orders (Lezama et al. 2022 Gat et al. 2024 Lipman al.) 20244 (Shaul al.) 2024) HHolerrieth al.) [2025] Nie al. 2025). In contrast to these global-order methods, our sampler requires only local permutation awareness within a window size w, and uses non permutation targets r. This locality requires only a small embedding table of dimension + that encodes the relative order of the most recent context tokens. This simple addition preserves the conventional sext-toke--predictin loss while requiring only minor modification to the training pipeline.
7

---

5 Experiments
We study the performance of our porrupl--ppeiioo-ttkkennn (RPT) in practice. First, in Section 5.2 we demonstrate that RPT sampling can be trained with minimal fine-tuning of pretrained next- voken-prediction (NTP) model, while maintaining its NTP performance. Second, in Section [5.3 we perform an error analysis, comparing RPT to traditional NTP sampling showing improvement in total variation distance of learned conditional probabilities p(z;) (جمم) measured over test data, as partially predicted by the theory in Section 3 Finally, in Section 5.4 we experiment with RPT sampling on code generation and reasoning benchmarks demonstrating improvements in both RPT pretraining and sampling.
5.1 Implementation details
Data, architectures, training, and baselines Our data consists of corpus of one trillion (IT) language tokens. Throughout all experiments, we use the same training dataset and maintain the same data order. We pretrained an autoregressive model with 8B parameters, using the standard cross entropy loss (i.e., equation 13 with o; = and 7; = + 1) and the same architectural design as in Meta (2024) on this IT token data for 240K iterations as our baseline, denoted AR-F. We denote its 224K iteration checkpoint (i.e., after 90% of the training tokens) by AR-C. We next finetuned AR-C to reach 240K iteration and the remaining 100B tokens (10% of total training tokens) with
16K = 0.5 and q = 0.02,
iterations in Algorithm| with window size w = 3, and which corresponds to 80 expected swaps in each sequences of n 4096 tokens (equation 9 We train on 256 H100 GPUs and a batch size of 4M tokens. We use A optimizer with a warmup of 2000 steps, peak learning rate of 1c-3 and a cosine scheduler.
Positional encoding Our method requires information on whether the predicted token is permuted.
see line 10 and 12 in Algorithm [ To incorporate this information, we introduce a learned positional embedding layer. One option to learn this layer applied to the source-target pair, (Ø; R.), however this incorporates global positional encoding. Inspired by relative positional encoding (Su et al. (021) we note that the source-target pair can be encoded relatively by using the difference т; o, as input to the learned positional embedding layer. For example, in Figure 2 the positional embeddings are {1, 11,1,-11 1, I,...,1}, 1}.
Practical sampling We focus on sampling with w 2 in the paper as we didn't 't see a practical benefit in sampling with w 3 window size, which is more computationally demanding (see Appendix [.J).. A benefit in RPT sampling algorithms is that it allows incorporating different sampling heuristics not usually available in standard AR sampling. Two useful heuristic we tested are: (i) Greedy decoding: using arg max in PTP sampling p(zg|=241, res); and (ii) Confidence: accepting token only if its probability is greater than threshold, i.e., 1) = 0.9. In the main paper we always use Greedy decoding; we show ablation on Greedy decoding in Appendix B.2
5.2 Evaluation
We first verify that our fine-tuning process does not de- grade the NTP sampling performance. F the autoregressive training loss for 900B tokens and the fine-tuning phase, where we train our proposed method for an additional 100B tokens. The NTP error continues the same trend where, convergence, higher NTP loss (~0.02), which is negligible in terms of cross-entropy (compared to 1.75 at convergence). At the same time the PTP losses (there are two such conditionals in w 3 training corresponding L = 1, 2 in equation rapidly converge to substantially lower error values com- pared to the NTP loss, as conjectured in our theoretical analysis in equation 20 in Section 3 Intuitively, the more the better.
Cross Entropy (CIE)
observe slightly
PSP-2
NTP
0
50
100
150
200
250
Thousands steps
Figure 5: NTP losses of the pretrain (in red) and fine-tune (in green) stages. We also show the loss curves of PTP-1 (in orange) and PTP-2 (in blue) that corre- spond to l and l = 2 in eq. 8
future tokens are used for sext-token-predittin Moreover, in T we show that sampling with k = O (NTP sampling with the fine-tuned model), often achieves better scores on the benchmarks compared to autore-

---

gressive bassline at convergence, AR-F. This suggests that our fine-tuning in fact could potentially improve the model's performance.
5.3 Error analysis
We compute the empirical total variation (TV) distance (equation 16) between the conditional p( (xi<c) com- puted by RPT, and the ground truth distribution p(z, (r<a) represented by the ground truth token r2 in validation set. In Table we report empirical TV distances computed with 128K validation set tokens from each dataset, all of them with context of at-least 20 tokens, and compute the (a; (x))) 0, 2, RPT
Table 2:
Empirical total variation dis- tance (lower is better) of RPT condition- als measured on validation data for an increasing number of RPT iterations k.
k
delm
github
wiki
arxiv
se
0
.34
.18
.26
.27
.27
RPT error e (x.) = (=)(()))) (z) where ((*)
1
.30
.15
.23
.23
.24
the predicted probability of x; after
2
.14
.23
.22
iterations (see 4. Iteration k = 0 corresponds to NTP
3 .30
.14
.22
.22
.23
sampling with our trained model, and the TV error decreases mostly after the first RPT iteration, and continue to decrease moderately afterwards; this is consistent across all tested validation datasets. discusses the error in the limit (mixed) case, which seems to happen after
The analysis in Section 3 only few iterations in practice (for our chosen sampling method).
Figure [ compares the probabilities assigned ground truth tokens I; by RPT (after k = iteration) and NTP sampling process, i... (x+)= (00) on the DCLM dataset. RPT improves the probability of the ground truth
Probability deastty
token in 64.5% cases which is another indication of the
improved sampling properties fo RPT compared to NTP.
5.4 Coding and reasoning tasks
0.50
1.00
((Ours) P(AR)
We evaluate our method on popular coding and reasoning benchmarks: HumanEval+ (Chen al. [021 Liu al.) 2023) is task where the model is required to complete a given function signature with docstring. MBPP (Austin et al. 2021) contains few-shot code generation tasks from problem descriptions. GSM8K (Cobbe al.) 2021) consists of lede-school--leel mathematical word problems. Finally, we report results on Mult a non-Pythonic version of HumanEval (Ben Allal
Figure 6: Difference of RPT (k: 1) and NPT probabilities on validation tokens. RPT improves probabilities of validation tokens.
al. 2022).
Table summarizes the results of our experiments, comparing the performance RPT sampling to NTP on several baselines. AR-F is our fully trained AR model using NTP sampling. We also report results using the initial fine-tuning checkpoint AR-C, providing a direct comparison to the initial fre-fiie-tuning state. Additionally,
include results from NTP sampling on our trained model (k: = 0). Our results demonstrate a consistent and significant improvement in performance when using RPT sampling compared to the NTP baseline. RPT improves both in pretraining (0 iterations) and RPT sampling (>0 iterations) rather consistently across all 9 benchmarks. All results are reported with confidence sampling. In Appendix B.2 we also provide the ablation results for non-confidence sampling, and non-greedy sampling. For a fair comparison between the methods, we follow Gloeckle al. (2024) and report oracle score over the temperatures {0.0, 0.05, 0.1).
6 Conclusions and future work
We introduce Resample- Preplleerrvvous---ee Prekeous-Tokens (RPT), sampling method that allows models to revisit and replace previously generated tokens. After a short fine-tuning of a pretrained model, RPT leads to approximately 10% relative improvements in reasoning and coding tasks. In the paper, we analyze the method in a controlled environment as well as large-scale validation datasets. We believe that RPT sampling has the potential to further enhance AR performance by incorporating larger window sizes and other pre-defined permutations, such as block biock-pernutatioon. As our work introduces an alternative sampling paradigm autoregressive models, does not seem to introduce significant socictal risks beyond those that already exist with large language models.

---

References
Joshua Ainslie, James Lee-Thorp: Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints, 2023. attps::/arxxv..ogg// 2200 13245
Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, and Charles Sutton. Program synthesis with large language models, 2021. https://arxiv.orr org/
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer, 2020. https://arxiv http://axxiv..ogabb/////as
05150
Loubna Ben Allal, Niklas Muennighoff, Logesh Kumar Umapathi, Ben Lipkin, and Leandro von Werra. A framework for the evaluation of code generation models. https://gi http://ggttttu.... thub. com/ bigcode-prrjeec//iggoddd eeeeeeealllllaiioe harness, 2022.
Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Generative flows on discrete spate-spaces: Enabling multimodal flows with applications to protein co-design. arXiv preprint wwXiv:2402.049999 2024.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert- Voss,, William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojeiech Zaremba. Evaluating large language models trained code, 2021.
Weizhe Chen, Sven Koenig, and Bistra Dilkina. Iterative deepening sampling for large language models. arXiv preprini arXiv:2502.05449, 2025.
Grace E Cho and Carl D Meyer. Comparison of perturbation bounds for the stationary distribution of markov chain. Linear Algebra and its Applications, 335(1-3):377-150, 2001.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian. Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021.
DeepSeck. Deepseek Ilm: Scaling open-souree source language models with longtermism. 2024. https:
//arxiv.org/// wttp.bab/240..225...
DeepSeck. Deepseek-rl: Incentivizing reasoning capability in llms via reinforcement learning, 2025.
https://axx..orrgaab////// 22t 12948
Markus Freitag and Yaser Al-Onaizan. Beam search strategies for neural machine translation. In Pro- ceedings of the First Workshop on Neural Machine Translation. Association for Computational Lin- guistics, 2017. doi: 10.18653/v1/w 17-3227. http::/dxx: http:/xxxxoi..orggo wtt../..rr/0...111///11111.11.1. 18653/v1/W17-3222
Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky T. Q. Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman. Discrete flow matching. arXiv preprint arXiv:2407. 15595, 2024.
Gemini. Gemini: A family of highly capable multimodal models, 2025. 2312.11805
Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, and Gabriel Synnaeve. Better & faster large language models via multi-token prediction, 2024. https://arxxiv..
abs/ 19737
10

---

Dirk Groeneveld, Iz Beltagy, Pete Walsh, Akshita Bhagia, Rodney Kinney, Oyvind Tafjord, Ananya Harsh Jha, Hamish Ivison, Ian Magmusson, Yizhong Wang, Shane Arora, David Atkinson, Russell Authur, Khyathi Raghavi Chandu, Arman Cohan, Jennifer Dumas, Yanai Elazar, Yuling Gu, Jack Hessel, Tushar Khot, William Merrill, Jacob Morrison, Niklas Muennighoff, Aakanksha Naik, Crystal Nam, Matthew E. Peters, Valentina Pyatkin, Abhilasha Ravichander, Dustin Schwenk, Saurabh Shah, Will Smith, Emma Strubell, Nishant Subramani, Mitchell Wortsman, Pradeep Dasigi, Nathan Lambert, Kyle Richardson, Luke Zettlemoyer, Jesse Dodge, Kyle Lo. Luca Sol- daini, Noah A. Smith, and Hannaneh Hajishirzi. Olmo: Accelerating the science of language models, 2024. https: htaps:///axiv... org/ abs/ 24000/03220008888888
Jiatao Gu, Changhan Wang, and Jake Zhao. Levenshtein transformer, 2019. https://arxiv.onr abs/1903 11006.
Peter Holderrieth, Marton Havasi, Jason Yim, Neta Shaul, Itai Gat, Tommi Jaakkola, Brian Karrer, Ricky T. Q. Chen, and Yaron Lipman. Generator matching: Generative modeling with arbitrary markov processes, 2025. hctps://arxiv.org//aa
20587
Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text 09751
degeneration, 2020.
Kimi. Kimi k1.5: Scaling reinforcement learning with Ilms, 2025. https://arxiv.or//bb// 2501. 12599.
Stephen J Kirkland, Michael Neumann, and Nung-Sing Sze. On optimal condition numbers for markov chains. Numerische Mathemati,, 110(4):521-537, 2008
Nathan Lambert, Jacob Morrison, Valentina Pyatkin, Shengyi Huang, Hamish Ivison, Faeze Brahman, Lester James V. Miranda, Alisa Liu, Nouha Dziri, Shane Lyu, Yuling Gu, Saumya Malik, Victoria Graf, Jena D. Hwang, Jiangjiang Yang, Ronan Le Bras, Oyvind Tafjord, Chris Wilhelm, Luca Soldaini, Noah A. Smith, Yizhong Wang, Pradeep Dasigi, and Hannaneh Hajishirzi. Tulu 3: Pushing frontiers in open language model post-training. 2025. https: www. hrrg/bbs/411 15124
Jose Lezama, Tim Salimans, Lu Jiang. Huiwen Chang, Jonathan Ho, and Irfan Essa Discrete predictor-corrector diffusion models for image synthesis. In The Eleventh International Conference on Learning Representations, 2022.
Bei Li, Tong Zheng, Rui Wang, Jiahao Liu, Junliang Guo. Xu Tan, Tong Xiao, JingBo Zhu, Jingang Wang, and Xunliang Cai. Predictor-corrector enhanced transformers with exponential moving average coefficient learning. Advances in Neural Information Processing Systems, 37:20358-20000 -23382, 2024.
Jeffrey Li, Alex Fang, Georgios Smyrnis, Maor Ivgi, Matt Jordan, Samir Gadre, Hritik Bansal, Ftash Guha, Sedrick Keh, Kushal Arora, Saurabh Garg, Rui Xin, Niklas Muennighoff, Reinhard Heckel, Jean Mercat, Mayee Chen, Suchin Gururangan, Mitchell Wortsman, Alon Albalak, Yonatan Bitton, Marianna Nezhurina, Amro Abbas, Cheng- Yu Hsieh, Dhruba Ghosh, Josh Gardner, Maciej Kilian, Hanlin Zhang, Rulin Shao, Sarah Pratt, Sunny Sanyal, Gabriel Ilharco, Giannis Daras, Kalyani Marathe, Aaron Gokaslan, Jieyu Zhang, Khyathi Chandu, Thao Nguyen, Igor Vasiljevic, Sham Kakade, Shuran Song, Sujay Sanghavi, Fartash Faghri, Sewoong Oh, Luke Zettlemoyer, Kyle Lo, Alaaeldin El-Nouby. Hadi Pouransari, Alexander Toshev, Stephanie Wang, Dirk Groeneveld, Luca Soldaini, Pang Wei Koh, Jenia Jitsev, Thomas Kollar, Alexandros G. Dimakis, Yair Carmon, Achal Dave, Ludwig Schmidt, and Vaishaal Shankar. Datacomp-lmn In search of the next generation of training sets for language models, 2025. https://arxiv wwtpsiivv.rrggaaa wbs./as///./ee/.............eeae 2406. 11794
Yaron Lipman. Marton Havasi, Peter Holderrieth, Neta Shaul, Matt Le, Brian Karrer, Ricky T. Q. Chen, David Lopez-Paz, Heli Han-H and Itai Gat. Flow matching guide and code, 2024.
ftttp:://axxv.orrg/aa 24122 06264
Jiawei Liu, Chunqiu Steven Xia, Yuyao Wang, and Lingming Zhang. Is your code generated by chatGPT really correct? rigorous evaluation of large language models for code generation. In Thirty- seventh Conference on Neural Information Processing Systems, 2023. /ttps://openrevien
net:/fooummiid-lqvxx1000C
11

---

Llama 3 Team Meta. The llama 3 herd of models, 2024.
Niklas Muennighoff, Zitong Yang, Weijia Shi, Xiang Lisa Li, Li Fei-Fei, Hannaneh Hajishirzi, Luke Zettlemoyer, Percy Liang, Emmanuel Candès, and Tatsunori Hashimoto. sl: Simple test-time scaling, 2025.
/tt////ww..ogg/////......... 19393
Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai Lin, Ji-Rong Wen, and Chongxuan Li. Large language diffusion models, 2025. https://arrx: org/ abs/2500 09992
Team OLMo. 2 olmo furious, 2025.
Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida. Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback, 2022. https://arxiv.org/ httpg//.a//22/....ee 2203.02152 02155
Arnaud Pannatier, Evann Courdier, and François Fleuret. o-gpts: A new approach to autoregressive models. In Joint European Conference Machine Learning and Knowledge Discovery in Databases, pages 143-159. Springer, 2024.
Ofir Press, Noah A. Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation, 2022. https://arriv..rrgaabb//////////
12409
Shyam Sundhar Ramesh, Yifan Hu, lason Chaimalas, Viraj Mehta, Pier Giuseppe Sessa, Haitham Bou Ammar, and Ilija Bogunovic. Group robust preference optimization in reward-free rlhf, 2024.
https://arxiv.org/ absaabbs52555 20304
Eugene Seneta. Non-negative matrices and Markov chains. Springer Science & Business Media, 2006.
Neta Shaul, Itai Gat, Marton Havasi, Daniel Severo, Anuroop Sriram, Peter Holderrieth, Brian Karrer, Yaron Lipman, and Ricky T. Q. Chen. Flow matching with general discrete paths: A kinetic-optimal perspective, 2024. happ:://aaxiv.org/bbb.... 034877
Mitchell Stern, William Chan, Jamie Kiros, and Jakob Uszkoreit. Insertion transformer: Flexible sequence generation via insertion operations. In International Conference on Machine Learning. pages 5976-5985 PMLR, 2019.
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: enhanced transformer with rotary position embedding. arxiv. arXiv preprint arXiv:2 102.000409990 2021.
Zhilin Yang, Zihang Dai, Yiming Yang. Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. Advances in neural information processing systems, 32, 2019.
12

---

A Analysis of RPT versus NTP sampling
A.1 Next-token-prediction asymptotic error
Let us first derive the NTP error bound
||l| S ||e:]l + 1|e4115loo + o(e).
(34)
consider equation 25 i.e.,
(aa, 1i+1)
+ c( (x;))((;)))))))))))))))))))))( + o(c).
(35)
and sum both sides after taking absolute values and using the triangle inequality on the r.h.s. to get
|c1,011/12
Σ
(1)1)))))))
+ |e( o(c²)
(36)
Σ le(xx)) + Σ (1s)
Σ
(37)
Mi
Zi+1
S ||l//a + ||
(38)
This bound will be tight for any scenario where
c(z;)e(<<41 (x)) 2 0.
(39)
We provide simple example when this holds. Consider V (vocabulary of size 2), and let
a
+
(40)
3
B _-
with a, 3 >> 0 and a and 0 < E min {a, 6). In this case have (T4++1/2;)) (x<+11=)) for all Tj, Ø++1 and therefore e(r,+1|0;; = 0 and equation holdsls trivially.
A.2 Resample-prerious-toeen asymptotic error
We derive the asymptotic enror in the RPT Markov chain transition kernel and consequently the bound equation 32 First, the asymptotic kernel error is
p(a,, 2/41 [1, 2i41)
# (1// ( 1 11/+)))
(m)+1) ) ((/l==+) (T(2/11)))) + c(<1411<<))
= =p(2%, 2/4 [z, #141) +
+ e(x (x) (2/1==+1) + o(e),
where o(e) are asymptotically smaller than terms. Therefore the asymptotic kernel error
e(ri, 2'41 ba, 1++1) = ( + c(al4 ( o(c).
(41)
Let us compute the max-norm the kernel errors,
||ell
max Σ le(,,, las, #+41)
(42)
2/,7/111
Σ (c(ol=1lat) ( moolleet)) + o(()) (43)
max
क alaglacte) } +
< max 21
|e(a|1=0+111 + Σ
Σ (((//11l)))
+ole) o(e)
(44)
x
< max sit1
e [af(1=+411 + max
Σ
|e(2/41 [xt) + o(e)
(45)
#.
of
= =||spetlle + [||4+t1 + +o(e).
(46)
13

---

A.3 Synthetic example- base noises
In section we present results with G, 1. Here we report results of our synthetic experiment with different base noises, E (0.1,0.01). The RPT improvements are invariant to absolute noise level.
cp/ Kn
ep/ 6o = 0.5
cp/es = 0.75
RPT
RPT
RPT
NTP
PIRARRY GOAEE
(a) Maximum errors, = 0.01.
c
c/ 0.5
9 = 0.75
proatt Beats
6.
0C
(b) Total variation errors, en 0.01.
<p/ 4m
Ep/ E,
0.5
cp Fn 0.75
26.
416
0.18
122
0.7
818
n3h
45
(c) Maximum errors, 6n = 0.1.
c/ = 0.5
5/G 0.75
72.
00
0692
0.06
0,00
912
6.54
607
600
012
60
0006.00
612
(d) Total variation errors, en
6p 6,
cp/G, 0.5
G/ 0.75
NTP
2429
(c) Maximum errors, En.
5 0
cp/C,
0.5
f = 0.75
RPT NTP
NTP
NT7
0.56
(f) Total variation errrrs, 6p
14

---

B
Ablations
B.1
Window size
Our trained model supports a sampling window of three tokens (u = 3). In the main paper, we present 2) sampling. A comparison between the two window sizes is
results for a two-token window presented in the table below, which shows that sampling with two and three tokens yields roughly comparable results. Given the similar performance and the increased efficiency of the two-token window during sampling, we opted to present results using the two-token window in the main paper.
k
HvaanEval+
MBPP
GSMSK
C++
C#
PHP
Bash
Java
TypeScript
AR-F AR-C
25.6
39.0
35,2
28.5
17.7
28.5
6.9
37.9
35,8
24.4
38.8
35.4
28.6
22.8
24.2
8.8
32.9
33.9
0
27.4
39.6
35.5
31.0
22.8
26.1
9.5
38.3
37.7
1
2
27.4
40.6
37.5
31.0
22.2
25.5
9.5
41.1
38.4
2
2
28.0
39.0
36.3
31.0
19.6
26.7
9.5
37.9
37.7
1
3
28.0
39.2
35,5
30.4
22.2
25.5
9.5
41.7
38.4
2
3
28.6
39.2
35.6
30.4
19.6
26.7
9.5
39.2
38.4
B.2
Sampling parameters
Table reports results on various task benchmarks using our default RPT sampling that includes greedy decoding and confidence (set to 0.9), the practical sampling paragraph in Section 5.1 Here we complete the picture by showing in the table below ablation results of samplings without greedy sampling and confidence. We find that the chosen setup is the most stable among the options considered. Note however, that naive sampling already introduces non-trivial improvements.
HumanEval+
MBPP
GSM8K
C++
C#
PHP
Bash
Java
TypeScript
AR-F
25.6
39.0
35.2
28.5
17.7
28.5
6.9
37.9
35.8
AR-C
24.4
38.8
35.4
28.6
22.8
24.2
8.8
32.9
33.9
0
27.4
39.6
35.5
31.0
22.8
26.1
9.5
38.3
37.7
28.6
39.0
36.0
32.3
20.9
27.9
7.0
34.8
37.7
2
29.9
39.0
35.5
30.4
22.2
27.9
8.9
41.1
35.2
Greedy decoding
29.9
38.8 38.8
35.0
31.7
19.0
27.9
7.0
37.6
35.2
2
28.6
35.6
32.3
20.9
27.9
7.0
37.6
35.2
Greedy decoding + confidence
1
27.4
40.6
37.5
31.0
22.2
25.5
9.5
41.1
38.4
28.0
39.0
36.3
31.0
19.6
26.7
9.5
37.9
37.7
15

---

C Temperature analysis
We compare RPT sampling (k 0) NTP sampling (k 0) using nucleus sampling (Holtzman et al. 2020) with top-p value of 0.95 and various sampling temperatures. The figures below suggest that our method is highly and consistently effective, particularly reasonably large temperatures, which are larger than the commonly used temperature values for pass@ evaluations. We demonstrate these results across various sampling parameter choices (see Appendix [B) and find that RPT sampling consistently improves performance compared to NTP sampling. Experiments here are performed on Hvaa.
32
32
Iteattions
Rerations
30
30
28
Paskel
3
24
24
22
22
0.0
0.1
0.2
0.3
0.4
0.5
0.6
0.0
0.1
0.2
0.3
0.4
D.5
0.6
Temperature
Temperature
(a) Naive sampling (no greedy decoding/confidence))
(b) Greedy decoding confidence
32
Iterations
Iterations
30
30
28
28
Pass@1
26
26
22
22
20
20
0.0
0.1
0.2
0.3 Temperature
0.4
0.5
0.6
0.0
0.1
0.2
0.3 Temperature
0.1
0.5
0.6
(c) Greedy decoding, no confidence.
(d) Greedy decoding and confidence 0.9.
16

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: Implementation straightforward and general to AR LLMs.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [NA]
Justification: We follow common practice in reporting standard benchmarks.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics httpp:///eanpp..cco/aaaa/ttccaca
[Yes]
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
11. Safeguards
17

---

Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [No]
Foundational research and not tied to particular applications, let alone deploy- ments.
Justification:
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
13. New assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMSs if it is important, original, non-standard component of the core methods this research? Note that the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [No]
18


Reviews:
[{"review_id": "5E2IfufHQ7", "reviewer": "Reviewer_o26f", "summary": "The paper introduces Resample-Previous-Tokens (RPT) for autoregressive (AR) LLMs that aims to mitigate the issue of error accumulation during token generation. Traditional AR models suffer from irrevocable errors in token generation, as once a token is produced, it cannot be revised. RPT iteratively revisits previously generated tokens within a fixed-size window and replacing them based on context. This method is designed to be easily integrated into existing AR models. The paper demonstrates that fine-tuning a pretrained 8B parameter model with RPT results in a 5-10% improvement in reasoning and coding tasks.", "questions": "1. What is the ratio of NTP and RPT? How did the ratio of RPT have the effect on the model performance?\n\n2. The window size w is set to 3 during training while w=2 during sampling, why did not keep them the same?\n\n3. The RPT shares the similar format to FIM (Fill-in-the-Middle). The discussion of relation to FIM could be helpful.\n\n4. In Table 1, what is the meaning and choices of $n$? A further discussion and analysis is required.\n\n5. The paper continue-trained LLM on RPT for 100 tokens after AR training for 900B tokens. How is the scaling trend on pure RPT?\n\n6. How about mixing the NTP and RPT? What is the performance of it?\n\n7. How did the author conduct RPT during inference on the evaluated experiments? A comparison results among varying sampling hyperparameters such as $w$ is required.\n\n7. What is the weakness of RPT brought during pretraining?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Pros**:\n\n1. The paper addresses a significant issue in AR: error accumulation due to the fixed left-to-right generation process. RPT provides a novel and practical solution to enhance the model performance.\n\n2. RPT is straightforward and easy to implement in AR training. It adds minimal computational overhead.\n\n3. The paper provides empirical validation of RPT via reasoning and coding tasks, e.g., HumanEval+, GSM8K, MBPP. \n\n4. The authors give a solid theoretical foundation for the method, which provides an asymptotic error analysis comparing RPT with NTP sampling.\n\n**Cons**:\n\n1. The paper does not provide sufficient detail on how RPT training influences the model during inference.\n\n2. In Table 1, the meaning of $n=1.5, 2.5$ is not clearly explained. Further discussion is needed to clarify this term and provide more context. Does it refer to different values of the window size or some other hyperparameter? Also, the results would benefit from a comprehensive analysis of how varying $n$ impacts the model’s performance.\n\n3. It lacks comprehensive analysis on the introduced extra hyperparameters, such as $s$, $q$, $n$, $w$. The choices and selection strategy are required to be carefully compared.\n\n4. RPT introduces extra relative position embedding layers, which requires further discussion about its potential drawbacks. It introduces extra parameters to optimize. How are these embeddings injected into the model’s layer inputs (added on each layer or on the token embedding layer), and what are the implications of this additional layer in terms of optimization and model efficiency?", "was_revised": true, "final_justification": "The author states many misuse of hyperparameters in the paper, which requires modification in the next version. Also, it adds additional ablation study to demonstrate the hyperparameter choice. Overall, this work could provide insight to the LLM society, I am increasing my score.", "rebuttal": "> In Table 1, the meaning of $n=1.5,2.5$ is not clearly explained. Further discussion is needed to clarify this term and provide more context. Does it refer to different values of the window size or some other hyperparameter? Also, the results would benefit from a comprehensive analysis of how varying n impacts the model’s performance. \n\n> It lacks comprehensive analysis on the introduced extra hyperparameters, such as $s,q,n,w$. The choices and selection strategy are required to be carefully compared\n\n> How did the author conduct RPT during inference on the evaluated experiments? A comparison results among varying sampling hyperparameters such as $w$ is required\n\n* Fractional values of $n$ (we had this mistakenly called also $k$, will fix) refer to iterations that stop at the first token. For example, for $n=1.5$, we have: \n    * $x_i \\sim \\hat{p}(x_i | x_{< i})$\n    * $x_{i+1} \\sim \\hat{p}(x_{i+1} | x_{< i+1})$\n    * $x_{i} \\sim \\hat{p}(x_i | x_{< i}, x_{i+1})$\n* We will add this clarification to the manuscript.\n* Table 1 and the tables in Appendix B.1 and B.2 show results for $n\\in \\\\{0,1,2\\\\}$ under various window sizes $w$ and greedy decoding and confidence. \n* Following the reviewer's suggestion, for the rebuttal we added table A and B below, showing the effect of the permutation probability $s$, and $q$ - the probability of moving a token forward within a permuted sequence (see Section 2.1). As can be seen in these tables our method is robust to these hyperparameters. We will add these results to the revised manuscript.\n\n**Table A**\n\nFixing $q=0.02$ and \n\n| s           | iterations | HumanEval+ | MBPP | GSM8K |\n|-------------|------------|------------|------|-------|\n| 0.1         | 0          | 28.0       | 38.0 | 35.4  |\n|             | 1          | 30.5       | 38.8 | 36.2  |\n|             | 2          | 29.3       | 38.6 | 35.9  |\n| 0.5 (paper) | 0          | 27.4       | 39.6 | 35.5  |\n|             | 1          | 27.4       | 40.6 | 37.5  |\n|             | 2          | 28.0       | 39.0 | 36.3  |\n| 1.0         | 0          | 26.8       | 36.8 | 34.2  |\n|             | 1          | 28.0       | 39.0 | 34.8  |\n|             | 2          | 28.7       | 39.4 | 35.6  |\n\n\n**Table B**\n\nFixing $s=0.5$ and\n\n\n| q            | iterations | HumanEval+ | MBPP | GSM8K |\n|--------------|------------|------------|------|-------|\n| 0.01         | 0          | 26.8       | 38.2 | 34.0  |\n|              | 1          | 27.4       | 38.8 | 36.2  |\n|              | 2          | 28.6       | 40.0 | 35.9  |\n| 0.02 (paper) | 0          | 27.4       | 39.6 | 35.5  |\n|              | 1          | 27.4       | 40.6 | 37.5  |\n|              | 2          | 28.0       | 39.0 | 36.3  |\n| 0.05         | 0          | 25.6       | 38.0 | 34.2  |\n|              | 1          | 26.8       | 39.3 | 36.2  |\n|              | 2          | 26.2       | 38.3 | 34.3  |\n| 0.1          | 0          | 25.6       | 38.2 | 34.9  |\n|              | 1          | 26.2       | 38.6 | 35.6  |\n|              | 2          | 25.6       | 38.4 | 35.9  |\n\n> What is the ratio of NTP and RPT? How did the ratio of RPT have the effect on the model performance?\n\nIn practice, the ratio of NTP and RPT is fixed by $s$ and $q$ corresponding to the sequence permutation probability and token forward push probability. We add in Tables A and B their effect on the model performance.    \n\n> RPT introduces extra relative position embedding layers, which requires further discussion about its potential drawbacks. It introduces extra parameters to optimize. How are these embeddings injected into the model’s layer inputs (added on each layer or on the token embedding layer), and what are the implications of this additional layer in terms of optimization and model efficiency?\n\nThe relative position is added only to the token embedding layer, and introduces a very small additional parameter count: 12k parameters for relative positions of $\\\\{-1,0,1\\\\}$ where $w=2$ and an 8B model, versus 500m parameters used by the token embedding layer. In optimization, only $O(w)$ new vectors are added to the embedding layers, which is negligible compared to vocabulary size . During inference, this additional lookup table computation is parallelized with the token embedding layer and does not introduce any overhead in practice. We will add this discussion to the revised manuscript.\n\n\n> The window size w is set to 3 during training while $w=2$ during sampling, why did not keep them the same?\n\nSampling with $w=3$ (Appendix B.1) achieved comparable results to $w=2$ (Table 1, main paper) however $w=2$ is a faster sampling method, so generally preferable. We do believe that it is quite possible that in other benchmarks $w\\geq 3$ could improve over $w=2$, but for the benchmarks we tested in this paper that was not the case. \n\n\n> The RPT shares the similar format to FIM (Fill-in-the-Middle). The discussion of relation to FIM could be helpful.\n\n\nThe fill-in-the-middle method decomposes the sequence to prefix, middle, and suffix and then orders it in PSM (prefix, suffix, middle) or SPM (suffix, prefix, middle) fashion. During generation, the model is required to predict the ‘middle’ part, which can introduce completions of arbitrary length with known issues of when to terminate the middle part. Our corrector iterations do not suffer from this issue as they are trained on fixed token lengths. Another question is if arbitrary middle completion has theoretical error correcting properties as RPT, a good question for future work. We will add this discussion to the related work section. \n\n\n> In Table 1, what is the meaning and choices of $n$? A further discussion and analysis is required.\n\nThe choices of $n$ in Table 1, refer to the actual number of corrector iterations we employ. Please note that we have a typo where we used $n$ and $k$ with the same meaning, we will fix it in the revised paper.\n\n> The paper continue-trained LLM on RPT for 100 tokens after AR training for 900B tokens. How is the scaling trend on pure RPT?\n\nAs suggested by the reviewer (and reviewer 6ySe), we investigated the relationship between RPT training time and improvements in corrector iterations. Due to the limited time available for this rebuttal, we were unable to complete a full training cycle. Instead, we fine-tuned a model using 200B tokens and the results are presented in the table below. Our findings indicate that increasing training steps does not enhance the model’s correction capabilities. We also fine-tuned the model with fewer tokens, observing that while 30B tokens might be insufficient for corrector sampling, training with 60B tokens yields improvements similar to those reported in the paper.\n\n\n| # Tokens  | Iterations | HumanEval+ | MBPP | GSM8K |\n|---|---|---|---|---|\n| 30B | 0 | 26.8 | 36.8 | 34.2 |\n|  | 1 | 27.4 | 37.2 | 35.9 |\n|  | 2 | 27.4 | 37.0 | 35.6 |\n| 60B | 0 | 28.0 | 38.2 | 35.5 |\n|  | 1 | 29.2 | 39.4 | 35.9 |\n|  | 2 | 28.7 | 39.2 | 36.2 |\n| 100B (paper) | 0 | 27.4 | 39.6 | 35.5 |\n|  | 1 | 27.4 | 40.6 | 37.5 |\n|  | 2 | 28.0 | 39.0 | 36.3 |\n| 200B | 0 | 28.7 | 39.2 | 35.5 |\n|  | 1 | 30.5 | 39.0 | 36.3 |\n|  | 2 | 29.9 | 40.6 | 37.5 |\n\n> How about mixing the NTP and RPT? What is the performance of it?\n\nNot sure we fully understand the reviewer question here. Our training indeed mixes RPT and NTP (see Section 2.1, paragraph “Training”). The probability of permutation $s$ sets the ratio between native NTP training and RPT training; see also Table A in the general response for its effect on model performance. If the reviewer means something else, we would be happy to clarify further. \n\n> What is the weakness of RPT brought during pretraining?\n\n\nDuring pretraining, we add an additional relative position embedding layer (discussed above), however an efficient implementation of this layer does not introduce any overhead. On the flip side, please note that pretraining with RPT improves the results compared to only NTP training, see Table 1 $n=1$ (NTP sampling after RPT training) versus AR-F. We attribute this to the nature of our method, which introduces a non-trival data/task augmentation.\n\n> The paper does not provide sufficient detail on how RPT training influences the model during inference.\n\nWe are not sure what influence the reviewer is referring to. As Table 1 shows (see, $n=0$ row), sampling purely in AR fashion after RPT training seems to somewhat improve results over training with AR and sampling with AR. If the reviewer has something else in mind, we would be happy to clarify further."}, {"review_id": "Ws3SGfalPe", "reviewer": "Reviewer_wdtD", "summary": "Corrector sampling proposes Resample Previous Tokens (RPT): a novel language modelling task and sampling strategy that facilitates revision of previously generated tokens within a fixed sized lookback window conditioned both on the prior context and subsequently generated tokens. The authors formally introduce their method and provide a robust theoretical analysis of RPT and Next Token Prediction (NTP) asymptotic errors. The error analysis is empirically validated with a synthetic experiment which agrees with their analysis that RPT exhibits a lower error than NTP. Finally, the authors pre-train a Llama-3-8B model from scratch for 900B tokens with the classic NTP task and use this checkpoint to further fine-tune two variants: 1) AR-F which is fine-tuned on an additional 100B tokens in the pure NTP task; and, 2) a model trained for RPT by randomly swapping target indices and carefully manipulating positional encoding. The results demonstrate that even in the purely auto-regressive sampling setting, RPT improves model quality on a variety of reasoning and coding benchmarks. Further, when using iterative RPT sampling, the model quality is further improved.", "questions": "## Key questions\n* Could the authors speculate on the cause of surprising results in Table 1. I would have naively expected AR-F to outperform AR-C in all cases and similarly was surprised to see AR-C significantly outperform in the PHP category. \n* Was only pre-training data used? Any SFT or RL training? Please discuss the high-level content of the corpus used; was it curated for the benchmark tasks specifically or drawn from a larger, more general dataset?\n* Can RPT training be extended to previously instruct-tuned checkpoints? What are some potential challenges to extending RPT training to this paradigm? \n* What is the computational overhead of RPT samping? Am I correct in my understanding that for a single RPT iteration (k=1) and a window size w, each output token requires w forward passes? \n* Please clarify the imprecise or misunderstood notation as requested in the weaknesses section. \n\n## Minor/Typos:\n* L129: “next-token-**rediction** (**NPT**)”\n* L132: NPT used in lieu of NTP twice. \n* L193: Should be revised as “left” and “right” instead of “top” and “bottom”, respectively. \n* L304: “**fo**”", "limitations": "* **No**. The authors stated in the checklist that limitations were discussed; however, I could not find any related discussion. I encourage the authors to consider the limitations of their method in particular from the perspectives of computational efficiency, performance (latency/throughput), and extending RPT training to previously instruct-tuned models.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 2, "significance": 4, "originality": 3, "strengths_and_weaknesses": "## Strengths\n* The proposed methodology appears to be novel. \n* The submission includes both high-quality theoretical analysis and compelling empirical evidence that the proposed technique materially benefits the RPT trained models compared to pure NTP models. \n* The submission is relatively clear; however, some notation could be made more consistent and a few statements clarified as outlined below. \n* The results are very promising and suggest a potential improvement to current the pre-training paradigm. In particular, it is very significant that RPT trained models improve in quality *even when auto-regressively sampled* (i.e., with no RPT iterations to actually revise past tokens). This suggests that the representations learned by the model during RPT are beneficially even for NTP. This is a very important outcome since, if I understand the sampling procedure correctly, RPT requires significantly more computation than pure auto-regressive decoding. \n\n## Weaknesses\n* The main weaknesses are some unexpected results in the empirical results where AR-C out performs AR-F and/or the RPT variants; for instance, the C# and PHP tasks for AR-C outperforming AR-F and RPT, respectively. \n* Some additional information on the empirical experiments would benefit the reader and improve reproducibility. \n* While the authors demonstrate the benefit of RPT training using only 10% of the pre-training budget; this still represents a very significant amount of training (100B tokens). It could be particularly impactful if RPT training could be extended to a more lightweight fine-tuning setting of instruct tuned model checkpoints. \n* Further discussion on the computational requirements of RPT sampling would improve the work and help clearly establish some of the limitations of RPT sampling. \n* Additionally to the above, since RPT requires additional computation, an interesting ablation would be to compare RPT with compute-matched test-time-scaling such as beam search or Best-of-N parallel sampling with the NPT trained checkpoint. Does RPT sampling provide a benefit in a strict compute-matched setting?\n* Some potentially imprecise notation:\n\t* L72: Shouldn’t $w=1$ represent the NTP sampling procedure as defined? The most recent token generated in equation 6 is $x_{i+w-1}$. This seems to be the case in algorithm 1 where w is explicitly restricted to $\\geq 2$. \n\t* L80: It appears that $k$ should be restricted to the domain $k \\in \\{1, …, n-w+1\\}$ since we can’t permute tokens within a window size of the last token in a sequence; there would be no future tokens to condition RPT sampling with. \n* In Table 1, $n$ is used to signify the number of RPT iterations. However, elsewhere in the paper such as section 5.3 $k$ is used. Unifying the notation between these sections would improve readability. Further, it’s unclear what n=0.5 or 1.5 refer to; does 0.5 represent a sliding window of 2 with a stride of 2 such that only every second token is resampled?\n* In Appendix B.2, the authors refer to “naive” sampling. It seems like this is simply nucleus samping. If so, please simply state nucleus sampling as it is less confusing to the reader than introducing a second, less common term. \n* L263 states that “greedy decoding is always used in the main paper”, but based on section B.2 it seems like what was actually reported in table 1 is greedy decoding **+ confidence**. Please clarify.", "was_revised": true, "final_justification": "The majority of my concerns were addressed during rebuttal. The AR-C variant outperforming AR-F remains a somewhat surprising result; however, the author's explanation regarding some variability in accuracy during AR training is sufficiently explanatory. The comparison between BoN sampling and corrector sampling suggests that corrector sampling offers a net benefit even in compute matched settings, further convincing me that the proposed method is a significant contribution worthy of inclusion in the proceedings.", "rebuttal": "> The main weaknesses are some unexpected results in the empirical results where AR-C out performs AR-F and/or the RPT variants; for instance, the C# and PHP tasks for AR-C outperforming AR-F and RPT, respectively.\n\n> Could the authors speculate on the cause of surprising results in Table 1. I would have naively expected AR-F to outperform AR-C in all cases and similarly was surprised to see AR-C significantly outperform in the PHP category.\n\nWe agree this is surprising, but we thought it’s useful to report this nonetheless. Our only explanation is that we notice in general that benchmarks results for AR training are not always monotonic during the pretraining. Consequently, also RPT is not always monotonic. However, note that we do compare AR-F and RPT after 1T token (at the final state of the cosine scheduler).\n\n> Some additional information on the empirical experiments would benefit the reader and improve reproducibility. \n\nIf the reviewer can specify where they feel information is missing we will happily elaborate in the text. \n\n\n> While the authors demonstrate the benefit of RPT training using only 10% of the pre-training budget; this still represents a very significant amount of training (100B tokens). It could be particularly impactful if RPT training could be extended to a more lightweight fine-tuning setting of instruct tuned model checkpoints. \n\nFollowing the reviewer's suggestion, we conducted additional experiments with 30B and 60B tokens. The results, presented in the table below, indicate that 30B tokens are insufficient to fully unlock RPT's capabilities, while 60B tokens appear to be adequate.\n\n| # Tokens  | Iterations | HumanEval+ | MBPP | GSM8K |\n|---|---|---|---|---|\n| 30B | 0 | 26.8 | 36.8 | 34.2 |\n|  | 1 | 27.4 | 37.2 | 35.9 |\n|  | 2 | 27.4 | 37.0 | 35.6 |\n| 60B | 0 | 28.0 | 38.2 | 35.5 |\n|  | 1 | 29.2 | 39.4 | 35.9 |\n|  | 2 | 28.7 | 39.2 | 36.2 |\n| 100B (paper) | 0 | 27.4 | 39.6 | 35.5 |\n|  | 1 | 27.4 | 40.6 | 37.5 |\n|  | 2 | 28.0 | 39.0 | 36.3 |\n| 200B | 0 | 28.7 | 39.2 | 35.5 |\n|  | 1 | 30.5 | 39.0 | 36.3 |\n|  | 2 | 29.9 | 40.6 | 37.5 |\n\n> Further discussion on the computational requirements of RPT sampling would improve the work and help clearly establish some of the limitations of RPT sampling.\n\n> What is the computational overhead of RPT sampling? Am I correct in my understanding that for a single RPT iteration ($k=1$) and a window size w, each output token requires w forward passes?\n\nCorrect. We will add the following to the revised paper (and mention it as a limitation). Each RPT iteration adds $w$ (window size) additional forwards. \n\n\n> Additionally to the above, since RPT requires additional computation, an interesting ablation would be to compare RPT with compute-matched test-time-scaling such as beam search or Best-of-N parallel sampling with the NPT trained checkpoint. Does RPT sampling provide a benefit in a strict compute-matched setting?\n\nAs suggested by the reviewer, we incorporated best-of-n results. In a non-verifiable setup (as we report in the paper), we generate $n$ samples and select the one with the highest model-assigned likelihood. While this baseline does improve benchmark performance, corrector iterations consistently yield better results with more significant improvements.\n\n| Best of n | HumanEval+ | MBPP | GSM8K |   |\n|-----------|------------|------|-------|---|\n| 1         | 25.6       | 39.0 | 35.2  |   |\n| 2         | 25.6       | 38.6 | 35.5  |   |\n| 3         | 26.2       | 39.5 | 33.3  |   |\n| 4         | 26.2       | 40.6 | 36.3  |   |\n\n> Some potentially imprecise notation:  1) L72: Shouldn’t $w=1$ represent the NTP sampling procedure as defined? The most recent token generated in equation 6 is . This seems to be the case in algorithm 1 where w is explicitly restricted to >= 2) L80: It appears that should be restricted to the domain $k\\in \\\\{1,..,n-w+1\\\\}$ since we can’t permute tokens within a window size of the last token in a sequence; there would be no future tokens to condition RPT sampling with.\n\nThanks for pointing it out, we will fix it! \n\n> In Table 1, $n$ is used to signify the number of RPT iterations. However, elsewhere in the paper such as section 5.3 $k$ is used. Unifying the notation between these sections would improve readability. Further, it’s unclear what $n=0.5$ or $1.5$ refer to; does $0.5$ represent a sliding window of $2$ with a stride of $2$ such that only every second token is resampled?\n\n**Typo of $n$ and $k$**: Thanks for pointing out, we will fix it in the revised manuscript. \n\n**Fractional iterations**: Fractional values of $n$ refer to iterations that stop at the first token. For example, for $n=1.5$, we have: \n* $x_i \\sim \\hat{p}(x_i | x_{< i})$\n* $x_{i+1} \\sim \\hat{p}(x_{i+1} | x_{< i+1})$\n* $x_{i} \\sim \\hat{p}(x_i | x_{< i}, x_{i+1})$\n\nWe will add this clarification to the manuscript.\n\n> In Appendix B.2, the authors refer to “naive” sampling. It seems like this is simply nucleus sampling. If so, please simply state nucleus sampling as it is less confusing to the reader than introducing a second, less common term.\n\nThanks, we will fix it.\n\n> L263 states that “greedy decoding is always used in the main paper”, but based on section B.2 it seems like what was actually reported in table 1 is greedy decoding + confidence. Please clarify.\n\nAll main Table results are with greedy sampling and confidence. We will clarify in the manuscript.\n\n> Was only pre-training data used? Any SFT or RL training? Please discuss the high-level content of the corpus used; was it curated for the benchmark tasks specifically or drawn from a larger, more general dataset?\n\nThroughout the paper, we used a fixed pretraining web-scale datamix (for reference, similar to DCLM [1]). We did not use any high-quality curated data such as being used in SFT or RL training. \n\n> Can RPT training be extended to previously instruct-tuned checkpoints? What are some potential challenges to extending RPT training to this paradigm?\n\nRPT adoption for such models should be straightforward. One thing to note is that, similarly to AR SFT training, we should take RPT loss only on certain parts of the data sequences. \n\n> Please clarify the imprecise or misunderstood notation as requested in the weaknesses section.\n\n> Typos: L129: “next-token-rediction (NPT)” L132: NPT used in lieu of NTP twice. L193: Should be revised as “left” and “right” instead of “top” and “bottom”, respectively. L304: “fo”\n\nThanks for pointing out, we will fix these in the revised manuscript.\n\n[1] DataComp-LM: In search of the next generation of training sets for language models, Li et al., 2024."}, {"review_id": "9jZTmrv75c", "reviewer": "Reviewer_6ySe", "summary": "This paper introduces a new sampling method, resample previous token, to allow the LLM to correct the previous tokens by themselves. The authors design a new training objective by modifying the positions/orders of tokens in a sequence to make it compatible with existing training infrastructure for next token prediction in autoregressive model. The authors provide theoretical analysis on why the correction prediction (conditioned on the future token) is easier than standard next token prediction. The experiments on coding benchmark show that the proposed method improves over the next-token-prediction baseline in pretraining.", "questions": "- What is the definition of $\\hat{p}$ and $p$ in Eq.(4)? Is the former an approximation and the latter the ground-truth distribution? It will be more clear if there are some explanations.\n- Most evaluation comparisons are on coding and math (only GSM8K). What's the performance of proposed method on other domains? E.g., general QA, science.\n- Why do you use 90% data for normal next-token prediction training and only 10% for RPT? Do you have any ablation of data proportion on final performances? For example, what if you use the full dataset for RPT training?\n- In general, how long will it take in RPT inference compared to NTP for window size=2?", "limitations": "yes.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**strengths**\n- Overall, the paper is well written and the illustrations such as fig.1 and 2 are clear.\n- The idea of RPT sampling to allow the LLM to self-correct the generated token is novel. Moreover, the authors also explain the effectiveness by an asymptotic analysis on approximation error.\n\n**weaknesses**\n- My major concern on this paper is on the training paradigm of RPT: In inference, RPT corrects token in the middle (e.g., when predicting $x_3$ by $x_1,x_2, x_4$, the $x_1,x_2,x_4$ are at the positions $1,2,4$). However, in training, RPT uses different positions of tokens after the corrected token (e.g., the $x_1,x_2,x_4$ are at the positions $1,2,3$). Meanwhile, the training objective of $x_5$ also changes in fig.2: the input becomes $x_1,x_2,x_4,x_3$ instead of $x_1,x_2,x_3,x_4$. Although the authors introduce new position embedding, I still feel there is large gap between training and inference.\n- In experiment, the authors train the model by perturbing data with window size=3. However, the table in Appendix B.1 shows that window size =3 achieves similar or even worse performances than window size =2, which also suggests the training inference gap.", "was_revised": true, "final_justification": "My major concern on training-inference gap is resolved so I will raise my score to 4.", "rebuttal": "> My major concern on this paper is on the training paradigm of RPT: In inference, RPT corrects token in the middle (e.g., when predicting $x_3$ by $x_1,x_2, x_4$, the $x_1,x_2,x_4$ are at the positions $1,2,4$). However, in training, RPT uses different positions of tokens after the corrected token (e.g., the $x_1,x_2,x_4$ are at the positions $1,2,3$). Meanwhile, the training objective of $x_5$ also changes in fig.2: the input becomes $x_1,x_2,x_4,x_3$ instead of $x_1,x_2,x_3,x_4$. Although the authors introduce new position embedding, I still feel there is a large gap between training and inference.\n\nWe believe there is a misunderstanding here. The training and inference of RPT are performed  **exactly the same way**. For the example in Figure 2: assuming we have $x_1,x_2,x_3,x_4$ and want to resample $x_3$ we push the **permuted sequence** $(x_1,x_2,x_4)$ (**in positions $1,2,3$, exactly as in training**, along with the positional encoding) through the network to get a prediction for $x_3$. \n\nTo clarify this we will add the following explicit inference pseudo-code to the revised paper. \n\n| |\n| :--- |\n| Require: Trained parameters $\\theta$, prompt $x_{<1}$, window size $w$, number of iterations $k$. |\n| Sample with NTP $x_1,\\ldots,x_{w-1}$ |\n| **For** $i = w, w+1, \\dots, n$ |\n| &nbsp;&nbsp;&nbsp;&nbsp;Sample with NTP $x_i$ |\n| &nbsp;&nbsp;&nbsp;&nbsp;**For** $j = 1, \\dots, k$ |\n| &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;**For** $s = w-1, w-2, \\dots, 0$ &nbsp;&nbsp;&nbsp;&nbsp;% we want to replace $x_{i-s}$ |\n| &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;Create permutation &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;$\\sigma = (1, 2, \\dots, i) \\setminus \\{i-s\\}$ |\n| &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;Create target index &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;$\\tau = (2, 3, \\dots, i-s)$ |\n| &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;Create permuted sequence $x_\\sigma = (x_{\\sigma_1}, \\dots, x_{\\sigma_i})$|\n| &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;Feed into model $(x_\\sigma, \\tau-\\sigma)$ and sample $x_{i-s}$ |\n\nRegarding the $x_5$ target in the example of Figure 2: Indeed in each block there is one token that predicts $w$ locations into the future. One could mask out this token in the cross-entropy loss, which we indeed tested while working on the submission and found it has no influence on results. This can be explained by noting that in this case the model gets a unique relative positional encoding of $w$ for these tokens, see Section 5.1.\n\n\n> In experiment, the authors train the model by perturbing data with window size=3. However, the table in Appendix B.1 shows that window size =3 achieves similar or even worse performances than window size =2, which also suggests the training inference gap.\n\nFirst, as mentioned above, we respectfully claim there is **no training-inference gap**. Training with $w=3$ inherently covers all cases $w≤3$. Indeed, we found that $w=2$ and $w=3$ yield comparable performance on benchmarks. We highlight $w=2$ as it is more computationally efficient, requiring fewer forward passes during inference and therefore was overall favorable and made it into the main paper. \n\n> What is the definition of $\\hat{p}$ and $p$ in Eq.(4)? Is the former an approximation and the latter the ground-truth distribution? It will be more clear if there are some explanations.\n\nYes. $p$ is the ground truth join distribution and $\\hat{p}$ is the distribution learned by the model. We define the hat operator in L40 (see second parentheses where we define $\\hat{*}$) and $p$ in L41. If the reviewer feels more clarification is required, we are happy to elaborate in the revised paper. \n\n> Most evaluation comparisons are on coding and math (only GSM8K). What's the performance of the proposed method on other domains? E.g., general QA, science.\n\nTo address the reviewer's comment, we evaluated our model on additional question-answering datasets (TQA and NQ). The table below presents results consistent with the coding and reasoning findings from the original submission. We will incorporate these new results into the revised paper.\n\n\n| Iterations | TQA | NQ |\n|---|---|---|\n| AR-F | 61.48 | 31.65 |\n| AR-C | 61.30 (-0.18%) | 31.53 (-0.12%) |\n| 0 | 61.64 (+0.16%) | 32.01 (+0.48%) |\n| 1 | 62.85 (+1.55%) | 33.30 (+1.77%) |\n| 2 | 63.13 (+1.83%) | 33.41 (+2.11%) |\n\n> Why do you use 90% data for normal next-token prediction training and only 10% for RPT? Do you have any ablation of data proportion on final performances? For example, what if you use the full dataset for RPT training?\n\n\nAs suggested by the reviewer (and reviewer o26f), we investigated the relationship between RPT training time and improvements in corrector iterations. Due to the limited time available for this rebuttal, we were unable to complete a full training cycle. Instead, we fine-tuned a model using 200B tokens and the results are presented in the table below. Our findings indicate that increasing training steps does not enhance the model’s correction capabilities. We also fine-tuned the model with fewer tokens, observing that while 30B tokens might be insufficient for corrector sampling, training with 60B tokens yields improvements similar to those reported in the paper.\n\n\n| # Tokens  | Iterations | HumanEval+ | MBPP | GSM8K |\n|---|---|---|---|---|\n| 30B | 0 | 26.8 | 36.8 | 34.2 |\n|  | 1 | 27.4 | 37.2 | 35.9 |\n|  | 2 | 27.4 | 37.0 | 35.6 |\n| 60B | 0 | 28.0 | 38.2 | 35.5 |\n|  | 1 | 29.2 | 39.4 | 35.9 |\n|  | 2 | 28.7 | 39.2 | 36.2 |\n| 100B (paper) | 0 | 27.4 | 39.6 | 35.5 |\n|  | 1 | 27.4 | 40.6 | 37.5 |\n|  | 2 | 28.0 | 39.0 | 36.3 |\n| 200B | 0 | 28.7 | 39.2 | 35.5 |\n|  | 1 | 30.5 | 39.0 | 36.3 |\n|  | 2 | 29.9 | 40.6 | 37.5 |\n\n> In general, how long will it take in RPT inference compared to NTP for window size=2?\n\nEach RPT iteration adds $w$ (window size) additional forwards. We will make this clearer in the revised paper."}, {"review_id": "bbmazOfKjW", "reviewer": "Reviewer_eT7X", "summary": "This paper introduces Resample-Previous-Tokens (RPT), a new sampling method for autoregressive language models. RPT iteratively revisits and replaces tokens in a window to reduce error accumulation. Fine-tuning an 8B model with RPT for 100B tokens improves reasoning and coding benchmarks by 10%. RPT preserves next-token-prediction quality and speed, outperforming standard sampling.", "questions": "1. What does the parameter n in Table 1 represent?\n2. Llama-3 uses RoPE, which already has relative position information. Why still need to add a new positional encoding? Can you directly set the position to $\\sigma$?\n3. Why do the authors claim \"preserving their next-token-prediction speed\"? During inference, each token needs to be predicted at least twice. Does this mean that the decoding speed is twice as fast?\n4. In the Appendix. B, $\\omega=3$ is used for training, but $\\omega=2$ works better for inference. Is this in line with expectations? Does this mean that the scaling limit of this method is very low? After all, in code and mathematical reasoning scenarios, a large number of tokens often need to be modified.", "limitations": "see weaknesses", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "#### Strengths\n1. The paper proposes a good solution to save the token consumption of large models\n2. RPT maintains next-token prediction quality.\n3. Both theory and experiment have proved that the upper limit of RPT is higher than that of NTP.\n\n#### Weaknesses\n1. The processing of position coding is relatively rough, the sequence order is modified, but no ablation experiment is performed on position coding.\n2. The paper does not provide a larger scaling. Can the mode modify the previous 512 tokens?", "was_revised": true, "final_justification": "Considering that the effect of this method when there are many tokens to be modified has not been verified, I will keep my score unchanged.", "rebuttal": "> The processing of position coding is relatively rough, the sequence order is modified, but no ablation experiment is performed on position coding.\n\nBy rough as assume the reviewer means one of two things: 1) the positional encoding is costly; or 2) training-inference discrepancy (similar to reviewer 6ySe).\n\nFor 1) note that the relative positions is added only to the token embedding layer, and introduces a very small additional parameter count: 12k parameters for relative positions for $w=2$ and an 8B model, versus 500m parameters used by the token embedding layer. In terms of training complexity, only $O(w)$ new vectors are added to the embedding layers, which is negligible compared to vocabulary size. During inference, this additional lookup table computation is parallelized with the token embedding layer and does not introduce any overhead. We will add this discussion to the revised manuscript.\n\nFor 2) note that the training and inference of RPT are performed  **exactly the same way** and therefore there is no training-inference discrepancy. \n\n> The paper does not provide a larger scaling. Can the model modify the previous $512$ tokens?\n\nWhile training with large window sizes (e.g., $w=512$) is feasible, it leads to very intensive inference iterations, which is why it was not tested.\n\n> What does the parameter n in Table 1 represent?\n\nThe parameter $n$ in Table 1 is a typo and should be $k$. We will correct this in the revised manuscript.\n\n> Llama-3 uses RoPE, which already has relative position information. Why still need to add a new positional encoding? Can you directly set the position to $\\sigma$?\n\nThe RPT sampling requires the model to receive the extra information of which token it needs to predict next, which is not necessarily the next token. For example, for a window of size $w=2$ it can be the next token (with positional encoding “+1”), or the previous token (with positional encoding “-1”). RoPE works with input token locations rather than target locations and hence does not cover this use case. \n\n\n> Why do the authors claim \"preserving their next-token-prediction speed\"? During inference, each token needs to be predicted at least twice. Does this mean that the decoding speed is twice as fast?\n\nReading back this sentence we agree it is confusing. Our claim of “preserving next-token-prediction speed” (and quality) refers specifically to the standard autoregressive generation process using our model, where no corrector-iterations are applied. We will remove this sentence from the revised version.\n\n> In the Appendix. B, $w=3$ is used for training, but $w=2$ works better for inference. Is this in line with expectations? Does this mean that the scaling limit of this method is very low? After all, in code and mathematical reasoning scenarios, a large number of tokens often need to be modified.\n\nIndeed, sampling with $w=3$ (Appendix B.1) achieved comparable results to $w=2$ (Table 1, main paper) however $w=2$ is a faster sampling method, so generally preferable. We do believe that it is quite possible that in other benchmarks $w\\geq 3$ could improve over $w=2$, but for the benchmarks we tested in this paper that was not the case. An interesting future research venue could be to replace single tokens with blocks of tokens (or more generally, patterns of tokens) which can also have useful applications to code and math reasoning."}]

Author rebuttal:

Meta-review: (a) Summary: The paper introduces Resample-Previous-Tokens (RPT), a novel training and sampling method for autoregressive language models that mitigates error accumulation by allowing the model to revise previously generated tokens within a fixed-size window. Unlike standard next-token prediction, RPT conditions revisions on both past and future context, making prediction easier and reducing asymptotic errors. The authors provide theoretical analysis, synthetic validation, and large-scale experiments, showing that fine-tuning an 8B Llama model with RPT yields 5–10% gains on reasoning and coding benchmarks, while maintaining efficiency and compatibility with existing AR models.

(b) Strengths:
1. The paper proposes a promising solution RPT to tackle the major challenge of error accumulation in AR models caused by their fixed left-to-right generation, while maintaining the next-token-prediction quality.

2. The submission includes both high-quality theoretical analysis and compelling empirical evidence.

3. The paper is well-written and reviewers all agree that the idea is novel.

(c) Weaknesses (after rebuttal):
1. The matter of hyper-parameter choice should be clearly stated in the paper and tested in the ablation study.

2. The effect of this method when there are many tokens, such as 512 tokens, to be modified has not been verified.

(d) Why this decision:
Reviewers all agree that this is a good paper which proposes a promising method to challenge traditional next-token prediction. All reviewers vote consistently positive. AC would follow reviewers' recommendation and recommend for accept.

(e) Summary of discussions:
The rebuttal has done a good job to mitigate some major concerns in the first round of review such as training of the method. However, some others are not addressed, e.g., the effect of this method when there are many tokens, such as 512 tokens, to be modified has not been verified. Given these, AC would recommend Accept (poster) but not higher ranking.
Decision: Accept (poster)