How reinforcement learning after next-token prediction facilitates learning
Authors: Nikolaos Tsilivis
Eran Malach
Karen Ullrich
Julia Kempe
Venue: ICLR 2026 Poster (2026)
Keywords: large language models
reinforcement learning
length increase
theory
Primary area: learning theory
TL;DR: We study and explain how and why RL after next-token prediction enables learning in cases where mere next-token prediction fails to improve performance.

Abstract: Recent advances in reasoning domains with neural networks have primarily been enabled by a training recipe that optimizes Large Language Models, previously trained to predict the next-token in a sequence, with reinforcement learning algorithms. We introduce a framework to study the success of this paradigm, and we theoretically expose the optimization mechanisms by which reinforcement learning improves over next-token prediction in this setting. We study learning from mixture distributions of short and long “chain-of-thought” sequences encoding a single task. In particular, when the task consists of predicting the parity of $d$ bits and long sequences are rare, we show how reinforcement learning after next-token prediction enables autoregressive transformers to generalize, whereas mere next-token prediction requires extreme statistical or computational resources to do so. We further explain how reinforcement learning leverages increased test-time computation, manifested in longer responses, to facilitate this learning process. In a simplified setting, we theoretically prove that autoregressive linear models following this training recipe can efficiently learn to predict the parity of $d$ bits as long as the proportion of long demonstrations in the data mix is not exponentially small in the input dimension $d$.
Finally, we demonstrate these same phenomena in other settings, including the post-training of Llama-series models on mixture variations of common mathematical reasoning benchmarks.

Published as a conference paper at ICLR 2026
How REINFORCEMENT LEARNING AFTER NEXT- TOKEN PREDICTION FACILITATES LEARNING
Nikolaos Tsilivis
Eran Malach Harvard University
Karen Ullrichi FAIR, Meta
New York University
nt2231enyu.edi
Julia Kempe' New York University & FAIR, Meta
ABSTRACT
Recent advances in reasoning domains with neural networks have primarily been enabled by a training recipe that optimizes Large Language Models, previously trained to predict the next-token in sequence, with reinforcement learning algo- rithms. We introduce a framework to study the success of this paradigm, and we theoretically expose the optimization mechanisms by which reinforcement learn- ing improves over next-token prediction in this setting. We study learning from mixture distributions of short and long "chain-off ofoiinhhtthouhtt sequences encoding single task. In particular, when the task consists of predicting the parity of bits and long sequences rare, show how reinforcement learning after next-token prediction enables autoregressive transformers to generalize, whereas mere next- token prediction requires extreme statistical or computational resources to do so. We further explain how reinforcement learning leverages increased test- cest-time com- putation, manifested in longer responses, to facilitate this learning process. In a simplified setting, we theoretically prove that autoregressive linear models follow- ing this training recipe can efficiently learn to predict the parity of bits as long as the proportion of long demonstrations in the data mix is not exponentially small in the input dimension d. Finally, we demonstrate these same phenomena in other settings, including the post-training of Llama-series models on mixture variations of common mathematical reasoning benchmarks.
INTRODUCTION
The application of reinforcement learning techniques large neural networks has led many ma- chine learning advances and successes (Silver et al., 2016 2017 Berner et al. 2019 Bakhtin et al.) 2023 In the realm of Large Language Models (LLMs), recent paradigm that leverages reinforce- ment learning consists of two main ingredients: training a large model, such as an autoregressive transformer Vaswani al. (017)) on diverse types of data (sequences) from the web with a next- token prediction objective (Radford et al., 2019). followed by fine-tuning stage with a reinforce- ment learning algorithm that seeks to improve model generations with respect to reward functior The latter part is, simply put, guess-and-check procedure (Rechi) 2225);: the model generates one or more guesses per prompt, a reward function evaluates them (usually for correctness), and train- ing proceeds on (prompts, guesses) with next-token prediction loss that weights each sequence according to its reward. When this procedure is applied on a frontier model using a mathematical, logical reasoning reward, it leads to rapid increase in the model's capabilities in related domains and is often accompanied by significant length increase in model response. Prominent examples of such "reasoning" language models include OpenAl's ol model (Jaech et al.) 2024) and Deepseek's 's RI model (DeepSeck-Al et al. 2025). The prolonged LLM generation that precedes an answer has
Currently at Apple.
Joint senior authors.
'There is sometimes an intermediate stage of supervised fine-tuning that aims to hone specific model skills, but we treat it as part of the first stage of next-token prediction optimization in our discussion.

---

Published as a conference paper at ICLR 2026
Pre-training & Posterraminng curves Partry, opt2, dose
Duta
Shor
Iono
Nest-token prodiction (jvestealviing)
Reinforemecl learuing (pas--cuining)
AL Straiion
Figure 1: Left: An illustration of our main learning setting: a mixture of long and short sequences encoding the parity of d bits, along with representation of pre-training and post-training for d=5. Right: Advantage of next-token prediction followed by reinforcement learning over mere next- token prediction in predicting the parity of d-50 bits with a transformer trained from scratch. The red line corresponds to pre-training using next-token prediction, while blue lines correspond to the same pre-training runs, each followed by GRPO (with final token accuracy reward) from a different checkpoint. Lines correspond to median across 3 seeds. Top: Test accuracy under greedy decoding- Bottom: Median length of model response under greedy decoding. The inset figures zoom-in on the curves during post-training. Accuracy plateaus around 50% (random guess) during pre-training and then immediately grows during post-training Also, length increases during post-training.
been dubbed the "thinking process" of the model, as it often resembles the "chain-of offttouuhttth Nye et al. 2021 Wei et al., 2022) of a human expert solving said task.
In this paper, we study how autoregressive models succeed in solving challenging prediction tasks by following this training recipe (next-token prediction followed by reinforcement learning D) To provide a theoretical account, we make a central simplification: we assume that pre-training data already contain correct and elaborate but perhaps rare demonstrations for a task of interest. We argue this is a natural and productive way to view internet-scale data, even for the most challenging tasks. Based on this modeling assumption, we study and explain:
1. Why in certain cases it is difficult for a model to generalize during pre-training.
2. How reinforcement learning (RL) leads to a rapid improvement in terms of samples, and
3. What optimization pressures cause increase in the length of the response.
Specifically, we study learning from distribution that contains both short and long. chain-of- thought-like sequences encoding a single task. The first part of the paper (Sections 2 devoted to the task of predicting the parity of d bits, while later this learning setup. Predicting the parity, or XOR, of objects has been an influential learning prob-
R1 is
consider more variations of
lem in the study of artificial neural networks (Minsky & Papert 1969 Shalev-Shwartz al. 2017
Daniely & Malach| 2020) Abbe et al. 2023 Barak et al. 2022] Glasgow (202)), statistical learning more broadly (Kearns 1998) and biological learning
animals (Howard et al. 2022). Through
extensive empirical simulations on transformers trained from scratch (Section 3 we show that if
the task complexity is high (large input dimension d) and long demonstrations are scarce, then mere next-token prediction results in model that fails to predict the parity on unseen examples. On the other hand, when switching to an RL algorithm with a correctness reward, model performance improves rapidly and the length of the generated sequences increases. The results hold for large array of model choices, problem and optimization hyperparameters, and reinforcement learning al- gorithms. Figure | (Left: Illustration, Right: Simulations) summarizes the landscape.
We demonstrate that this learning behaviour is primarily driven by two factors: (a) during pre- training, the model learns much faster from long demonstrations than from short ones and remains
In the argot of I.I. M.M these two stages are often called pre-training and post-training, respectively.
2

---

Published as a conference paper at ICLR 2026
calibrated in how often it generates long versus short responses, and (b) during post-training when we reward successful generations and train with a seward-weighte loss, we tend to amplify the presence of long responses in the training batches, as the model has much higher probability of being correct when the response is long rather than short. This is the mechanism that leads to length increase which in turn enables generalization. We must emphasize that the transformer's inability to learn during pre-training in this setup is not due to insufficient model size or depth, but rather due to insufficient number of samples; an estimation, not an approximation, limitation.
In Section we theoretically capture most of the previously observed phenomena in a simplified setting. Let Peot be the proportion of long data in the parity mixture distribution ((Poot). We analyze pre-training (next-token prediction) with Stochastic Gradient Descent on D(Peot) for an autoregres- sive architecture that consists of a series of linear models (Malach, 2024), while for post-training, we consider the STaR algorithm (Zelikman et al. 2022) (which is REINFORCE, (Williams) 1992)- type algorithm) with a reward that verifies correctness the whole chain-of-thought.
this setting, we prove that the model: (a) fails to generalize under greedy decoding during the course of pre- training long demonstrations are rare (i.e., Peot < 1/3, which matches the empirical threshold observed with transformers) (Theorem D. (b) learns fast from long demonstrations and remains length-cute calgbbaade (Theorem []. and (c) succeeds in generalizing after the completion of both pre- and post-training in (ppoly[d) SGD iterations, provided that long demonstrations are not expo- nentially rare in the data distribution, while this is accompanied with length increase during RI. (Theorem 2). The analysis captures the immediate progress during RL seen in practice: as we - RL. rounds suffice to obtain a generalizing model. To the best of our knowledge, this suite of results provides the first theoretical separation between next-token predic- tion and next-token prediction combined with RL in the autoregressive setting, as well as the first optimization result demonstrating length increase during RI in L.MS..
show, only O 2
Poot
Finally, in Section [ we experimentally probe the same behavior in other settings, including GPT2 (Radford et al. 2019) models trained from scratch on the task of number multiplication, and pre-trained Llama 3-8B models (Dubey al.] (2024) that are fine-tuned (first to predict the next-token and then with RL) on a short/long mixture version of common reasoning benchmarks
(GSM8K (Cobbe al. 202T and MATH (Hendryeks et al. 202T datasets).
SETUP: MIXTURE OF LONG AND SHORT SEQUENCES ENCODING PARITY
In this section, we present our main learning problem and the models and algorithms we consider.
Data We study the task of predicting the parity of d bits given access to a source of sequences which either consist of: (i) input bits and their parity, or (ii) input bits. intermediate computations and the final parity- see Figure I for a demonstration. Let X' {+1]d be the input space of d bits and y 111, <BOS>}" be the output space of sequences, where <EOS> is special symbol denoting the end of a string. Let D(pcot) be a distribution over X' x V parameterized by Peot E [0, 1]. such
that T1..... Fd (12(1 and (.l....,9J+t))
7 22. II I2; <EOS>) + (1
7) [T: 32 <ROS>) where 7 Ber(Peot). We will primarily be focused on the case where the mixture weight Peot is small, i.e., where the distribution contains more short responses than long.
Predicting the parity of uniform random bits is difficult learning problem for neural networks, as is believed to require exponentially many samples/iterations (Shalev-Shwartz et al. 2017 Daniely & Malach, 2020, Shoshani & Shamir 2025). On the other hand, when intermediate computations are available in the form of an elaborate chain of thought, the parity function, as any other binary computable function, can be learned efficiently (Joshi et al. 2025 Malach 2024 Wies et al. 2023).
Models In our experiments, we train decoder-only transformers (Vaswani et al. 2017). 7) standard approach for language and sequence modeling. We experiment with the GPT-2 (Radford et al.) 2019) and Mistral (Jiang et al. (202)) variants of the transformer architecture. See Appendix C. for further details on the specifics of the architectures.
We denote Rademacher random variables (uniform +1) as Rad(1/2) and Bernouli random variables with bias as Ber(p). Our notation is described in Appendix
3

---

Published as a conference paper at ICLR 2026
Parity, medel=gp12, depth-4, dena Ranas
Test Accuracy (di 50)
Probabiltty long generation (d-50)
Rinal Accuracy vs (dm25)
100
102
Acouracy
10000
200.00
30000
4000
50600
33800
20000
30000
45000
30303
0.2
1.4
0.2
(pre (rsir hersttons
lare (pre-tnannt trenaioss
greedy (heepod
Pat
Figure 2: Pre-training of transformers on mixture long and short sequences encoding the parity of d bits. Left: Test accuracy during the course of pre-training. Center: The probability that a model's generation length is equal to the maximum length present in the training distribution (which equals d). Solid lines correspond to greedy decoding, while dashed lines correspond to sampling with temperature 1. Each color denotes a training distribution with a separate mixture coefficient Peot- Figure shows average and standard deviation across 3 seeds. Right: Test accuracy with greedy decoding at the end of pre-training (50k iterations) versus mixture coefficient Pent- The red dashed line corresponds the critical threshold. Each bullet is the median of 3 runs.
Pre-training In the first stage of training (pre-training), we train a transformer to estimate the proba- bility of the next token sequences drawn from (Ppcot (Radford et al. 2019). We consider online learning, sampling a fresh batch of data at each iteration.
Post-training After several steps of pre- pre-training, we switch training algorithm and further optimize the model with reinforcement learning methods. We consider three reinforcement learning algo- rithms: STaR (Zelikman et al. 20022), REINFORCE (Williams) 1992) and GRPO (Shao al., (024). At bigh level, all three algorithms seek to maximize the reward of model generations. We defer their description Appendix due to space constraints. We focus on two types of rewards:
End-to-end correctness: The reward function assesses whether the last token is equal to the correct
answer: (x2: y) = { - - -I= } where v[-1] E {=1} denotes the token appearing
right before the <EOS> token in sequence y.
Chain-of-though correctness: The reward function assesses whether the whole sequence is valid: Tem(X; (x, y) y = Ii, 22,2 I Vi, <ROS> V y = [ Hy <FOS>) For all post- train methods, we can optionally set the sampling temperature TRI. to value different than 1.
3 EXPERIMENTS ON PARITY
We present our experiments with autoregressive transformers trained the mixture distribution (ppcot) of long and short sequences encoding the parity of d bits. Throughout training, we evaluate (test) accuracy in the prediction task, by measuring whether the token generated before <EOS> equals the parity of the input. We also measure the length of the generated sequence.
3.1 PRE-TRAINING VS (PRE- & POST-TRAINING)
First, we fix the input dimension d a large value (e.g. d=50) and consider pre-training (next-token prediction) on data sampled from D(pcot) ) for various values of mixture weight Poot. We present results in Figure 2 for GPT2 architecture. As we can see, performance under greedy decoding is determined by the value of Peot: if the percentage of long sequences in the data is large enough (e.g. Poo)=..75), then the model quickly learns to predict the parity the input correctly by generating long responses and continues to generalize perfectly until the end of training. On the other hand, when the long sequences are under-represented in the source distribution, the model - under greedy decoding generates short responses and performs par with random guessing even after training on several millions of sequences (256 sequences per iteration for 50,000 iterations). In fact, the critical threshold is around poot=1 1/3 see Figure (right).
4

---

Published as a conference paper at ICLR 2026
Pariby, 0025 Posttrainingg2kk d-50, model-qpt2. depthiw4, 420,s=188, Rhauce=32
Accuraty
Longthe (e2e), temp-1.0)
Sequencss
The miw0ds
Accurscy
Lengthh
Accuraty
Longtha temp-1.0)
Accurscy (M
Sequeincs
Tar=1.0
Lengzhh
Accuraty
Longthe (e2e), temp-1.0)
Accurscy a 8 8
Sequeicss
Length
The
RL Iterations
Beeration
Figure 3: Post-training of transformers on mixture of long and short sequences encoding the parity of d bits with various RL methods and generation temperatures (TRI.). Left: Test aceu- racy during the course of post-training with greedy decoding (solid lines) and sampling with tem- perature 1 (dashed lines) for various RL methods. Figure shows average and 1 standard deviation across 3 seeds. Right: Length of generated response (sampled with temperature 1) during the course of a post-training run (GRPO with end-to-end reward) for 640 test inputs, after 20k pre-training iterations. Note:
The sample size n of each RL iteration differs amongst the RL algorithms: n=64 for GRPO, REINFORCE and n=3,200 sequences for STaR.
For the distributions with small Peot where the models did not manage to generalize during pre- training. we consider switching to post-training: at some iteration, we stop training with the next- token objective of equation 5 and resume model training with reinforcement learning algorithm. We show results in Figure for Peat=0.255, where we plot post-training accuracy as a function of RL iterations for all post-train algorithms described in Sections 2 A and for a few different values of sampling temperature TRL. Progress during RL is immediate: after training on only handful of sequences, model performance under greedy decoding grows from random guessing to 100% accu- racy. This improvement is accompanied with concurrent increase in the length of the model's gen- erations, which grows from (min train length) 50 (max train length). The previous observations are robust with respect changes on the reinforcement learning algorithm and the hyperparameters of post-training (e.g. sampling temperature), with the notable exception of STaR with high sampling temperature TRL and end-to-end reward, where the lack of negative rewards and the sparsity of the reward function Te2e cause unstable post-training.
To further emphasize the difference between the two paradigms, we show aggregated results (pre- & post-training combined) in Figure [ for d=50, =0.5 and post-training with GRPO starting from various checkpoints. Observe that if post-training starts early, then it might not lead to gen- eralizing model despite the length increase. However, GRPO, when started from later checkpoint, improves performance rapidly, and the number of post-training samples required for generaliza- tion a miniscule fraction of the overall number of samples. On the other hand, mere next-token prediction does not result in generalizing model (and the median greedy response remains short) even with access to ~0' training samples. This demonstrates the enormous sample-complexity gap between the two approaches for this task. Note that when the input dimension d is smaller, pre-training eventually does lead to geleraenera model (if we optimize for sufficiently long) and post-training does not seem strictly necessary (see Figure I1 for d (11, 12}). However, for larger values of d, the amount of SGD iterations required is unreasonably large. In other words, the transformens' inability to succeed during pre-training does not stem from insufficient model ca- pacity (small number of parameters or lack of ). but it is rather because next-token prediction on this distribution requires many samples to lead to a good predictor. In fact, as the complexity of the task increases, does the advantage of reinforcement learning after next-token prediction
A shallow transformer can approximate the parity function for any input dimension d; uniform attention for aggregating the bit values and a 2-layer MLP for computing their XOR see Lemma 6 (Liu et al. 2023)
5

---

Published as a conference paper at ICLR 2026
over mere next-token prediction (Figure 12
Appendix C.1
contains more experimental results for
various model and task hyperparameters.
3.2 DIFFERENT SPEEDS OF LEARNING & GROWTHI OF TEST-TIME COMPUTE
We now describe what makes reinforcement learning effective and why post-training leads to length increase in our simple setting. First, notice the curves in Figure FigureQithtt correspond to sampling with temperature during pre-training. Even when Poot < 1/3, the accuracy of the models at the end of pre-training is greater than 50%- Figure 2 ( In fact, it is equal Peat+1 100% for all cases of Poot. Furthermore, notice in the length plot of the same figure (Figure models are calibrated with respect length, i... they generate long responses with probability Peot- These two simple observations directly suggest that the models learn the two parts of the mixture "in parallel" during pre-trainn vre-rraining: after few training iterations, when prompted with d bits 21,.....4 Idd the models learn to generate long and correct responses containing the parity II Hy with proba- bility Poot and short responses with probability 1 Poot- As mentioned before, learning from the short sequences requires learning the parity of d bits in a single prediction step, which belongs to class of computationally difficult problems (Kearns| 1998) and is believed to be hard for standard neural networks to learn in practice (Shalev-Shwartz et al. 2017: Shoshani & Shamir (025).. ) As result, we do not expect the model learn using any reasonable amount of samples, and hence Learning from long demonstrations,
S (Center) that the
it resorts to short, random guessing with probability 1 Pcot-
on the other hand, can be performed efficiently (Malach) (024). This asymmetric learning difficulty
leads to accuracy equal to Poot 100% + (1 PPoo)) x 50% Peet x 100%.
The learning process during pre-training sets the stage for what follows in post-training. It is perhaps simpler to understand the RL dynamics in the case of the STaR algorithm with the chain-of-thought correctness reward coot (cyan lines in Figure3). The objective of equation O implements next-token prediction solely on moodel-generated sequences which contain a correct chain of thought and a correct final answer. If, at the end of pre-training, the prohability of long, correct generation is Peot and the probability of a short, correct one is (1 - Post) /2 (as suggested by Figure 2). then the model fits a data mixture which contains both long and short, correct sequences with proportions
equal pl = +11-++(1Ppoo and Pcl - )22 respectively. If, further, the model succeeds in fitting them, then the conditional distribution model generations at the end of the first round of STaR has the same weights, provided the model remains length-calibrated and did not manage to learn from the short sequences. Continuing inductively, at the start of the n 'hth round, we effectively sample from distribution (pp,), where Po = Peot and Pn = 2p.,
Prot
()-Peor)/2
14pp-1 converges to 1 exponentially
fast. Once the effective coefficient Pr. becomes larger than 3/ (which happens when accuracy with
temp. exceeds 1/3+1 2 red dashed line in Figure 3) the model starts generating long, correct responses with greedy decoding, thus generalizing perfectly at the task. This is the mechanism that causes length increase in model response, which in turns allows the model to learn to generalize.
In Appendix C.2 we consider a variation of our setting, where the data include a third type of sequence consisting of a partial cot. We demonstrate how this affects length distributions at the end of RL. and the effectiveness of length penalties during RL as a function of prior pre-train iterations.
THEORY
Next, we theoretically capture the main empirical findings of the previous section and prove the identified mechanisms. We seek theoretical model that, when Poot is small, fails to generalize during pre-training, while being capable of learning the long component of the distribution. As we show, the analysis of a linear autoregressive model (Malach 2024) satisfies the above desiderata.
Architecture We study an autoregressive model that consists of series of d+1 linear predictors. We reduce next-token prediction on (x, y) ~ D (Pcot) to a series of binary classification problems. At high-level, the feature embedding used in each one of them involves (at most) second-degree monomials of the input, which is canonical embedding choice. In particular, for the first output position, we consider a linear hypothesis class H, = {x (W],x) X E {+1]d, w E Rd). For the second output position, we need to learn a mapping from [#1]d+] to {+1, <EOS>}, }. as all three characters valid second output tokens under distribution D(peot). We build such a mapping in

---

Published as a conference paper at ICLR 2026
two stages: first, linear model decides whether the halting symbol needs to be outputted and, if not, we then make a decision between +1 with a different linear model. This makes the analysis of the output of SGD simpler than directly considering single multiclass model. That is, define H2 =(x++ (W2a: (b(x)) + bza x E [+1]a=1 W2a E Id22+1 b E RR, where @2 is a feature map that augments the input with all second-degree monomials involving the last input bit. For the second piece of the second position and for the rest of the positions, we proceed in a similar fashion we Wt E 2 J. The embedding function is de- define H4- = {x + (Wi: (x())) x E [=1])t/-1,
fined as di: function h that belongs to the product of these classes H M/ N2a x H2... ... x Ha- For all d+1 models, the output space is {=1} and we map the label to a character in {#1, <EOS>,<) accordingly. Note that the token position d+1 is a constant function of the input and hence omit it from our analysis. Furthermore, we define GREEDY to be the operation of greedy decoding from the model logits (takes the sign of each model's output). By composing each individual com- ponent of / with the sigmoid function, we also define the probability measure induced by h, namely TA(: x) E A()). This corresponds to sampling from the model with temperature Despite their apparent simplicity, linear autoregressive models can be efficient universal learners (Malach [024] Joshi et al.) (025). Our architecture with our embedding choices, in particular, is capable of identi- fying and learning any sparse parity function, hence we argue the model is not particularly tailored to the full parity function. Appendix E contains background and formal definitions.
Id F: - - for 2 < 1 S d. Our final model is
Our first result characterizes the model during pre-training. We consider minimization of the (regu- larized) next-token prediction objective (logistic loss evaluated at all output positions) with SGD.
Theorem 1. (Pre-training, Informal) Let d 2 2. Poot E (0, 3/4). Consider distribution D(peot), as defined in Section and linear hypothesis classes H1, Hzas Hz...., H. as defined earlier. Consider running Stochastic Gradient Descent (SGD) minimizing the next-token prediction objective over Hhin -H1 x H2a X H2... ... x Ha with an additional & regularization term. Then, for any O < E < (d + min { (2 In 1+poot In 1- PPoo
2 } and after (pooyydd, 1/e, 1/Peot)) iterations, with high probability over sampling from D D(poot) SGD returns model hpre E HhR for which the following hold:
1-Peoi
1. Length calibration:
When prompted with input x {+1]d, model hire generates a long, correct sequence with probability approximately equal to Poot and a short, correct one with probability approximately equal to 1-Peon That is, for any x E (+134.
There
31,5112,
Il
Tiy <BOS>
x
Poot S e,
(1)
d
Thpre
II Di. <BOS>
x
1 Poot
Se.
2
2. Failure of greedy decoding: For any E {+1]d, it holds:
GREEDY((ppc) (x)) ={ (()) <EOS>), Peot < 1/3,
(2), ....,
11 II I Th <EOS>). otherwise.
(2)
That is, test accuracy under greedy decoding is on par with random chance if Peot 1/3, and perfect otherwise.
Theorem D proves the main phenomena identified in the pre-training of transformers in Section 3 i..., length calibration and failure of greedy decoding, in our linear autoregressive setting. Its formal statement (Theorem 4 and proof can be found in Appendix E The proof proceeds by analyzing the output of SGD each binary classification problem independently. Precisely, standard convex learning result (Shalev-Shwartz & Ben-David 2014) implies that SGD's 's output has generalization error comparable to the best-in-class. Then, strong convexity (afforded by the regularization term) and case-to-case analysis establishes calibration and performance under greedy for each model.
For notational brevity in the proofs, subscript to denote the second piece of the second position.
7

---

Published as a conference paper at ICLR 2026
Remark The model's length calibration at the end of pre- pre-training arises from the use of the logistic loss, which is a proper loss, in the next-token prediction objective.
Remark 2. The critical threshold Poot=1/3, which is the same as the empirical critical thresh- old observed in transformers
Section + arises from the model's greedy decision the first
and second token. Indeed, in absence a parity feature II TT, it holds: IP [8/1 = x1] =
Poot + 1-Peat =
L+Beot >
I-Pees = IP [8 = --i]. Thus, a calibrated model's "hard" decision
for the first token will be Φ1 for any Poot. For the second step, we calculate the conditional proba- bilities: [y/2 = 21 = x1] = Plwa =2) =x1] 2poot_ while P [y/ <BOS> 8/1 = x1] =
Fy
1+Peat"
2prtt PPPt That is, as long as Poot < 1/3, the prediction for the second token after 21 P+Pe P+
will be <EOS>. That is, greedy decoding will generate a short sequence (I,, <EOS>). See also Appendix
E.2 Theorem | asserts that SGD after sufficient many iterations returns model that approximately matches these conditional probabilities.
Next, we study reinforcement learning with (regularized) STaR objective and with chain of thought reward Tcot (defined in Section ( Theorem 2 below the main result of the paper: it shows that pre-training followed by post-training can be an efficient recipe for learning the parity from D(pent). and an indispensable element this success
increase of model response during RL.
Theorem 2. (Post- (Post-training training, Informal) Let hyre be the output of SGD after pre-training under the conditions of Theorem for Peot < 1/3 and S CO Peut for sufficiently small constant co > O. Consider post-training of hove with the STaR algorithm using reward Teot. Suppose each STaR round uses (poly((d, 1/e, 1/cot)) SGD iterations, and let h(m) denote the model after rounds. Then, there exists an integer n = O (lo 1-Peot) such that with high probability over sampling from (poot), and the model-sampled outputs used by STaR in rounds ,,..., n - 1, the following hold:
Poot
1. Length increases: The probability of a long generation increases:
Th(n) ((01 211,
II i=1
Fi, <EOS>
Se. (+1)
(3)
where (4n pal S g and pn+1
2pn I+pm
for all n S n² with po Poot-
2. Perfect generalization: After n RL rounds, the model under greedy decoding is only gen- erating long responses and generalizes perfectly:
21, 111 21 -(1))
d
GREEDY (n17)(x))= =
II 04s <EOS>>
(4)
The proof, which appears in Appendix E.44 proceeds in two steps. We first observe that the STaR objective for the n th round is equivalent to next-token prediction with respect to a re-weighted version D(gn) of the original distribution. Then, we bound the deviation of sequence 4n from 2p--1 that describes the evolution of the proportion long data in "noise-less" sequence p.. = 1+Pa-1
the data mix.
Remark 3. Consider the case where long demonstrations are very rare, that is Peot + 0 as d + which is perhaps the most interesting case. Theorem 2 shows that as long as Poot = (d for some constant K E N, then we leam the parity using (ppoly(d) SGD steps. This should be contrasted with the hardness of learning parities (Kearns] 1998; Shoshani & Shamir 2025).
5 EXPERIMENTS WITH MATHEMATICAL REASONING BENCHMARKS
Next, we consider the same training recipe (next-token prediction followed by RI.) in less idealized settings (including large pre-trained models) and for mixture distributions encoding tasks computa- tionally deeper than parity, such as 5-digit multiplication and mathematical reasoning benchmarks.
Numbers Multiplication First, we study the task of multiplying n-digit numbers. Following our previous setup, we construct datasets that contain the digits of the multiplier and multiplicand,
8

---

Published as a conference paper at ICLR 2026
GSMBK, model 3.2 3B (base), Pctt 0.05
Numbers Mutipiication 5x5, model=apt2, Par
Sacidre 03
Acturrrry (80
Jeeepg Leaoth (nsoone
X % 19
Duan sequerces Fos Inlee
bruin 9056-97669 aOc Ry langhh Au
57 SowarES
mere next-token prediction in (Leff) multiplying two 5-digit numbers with a GPT2 model trained from scratch, and in (Right) solving grade-school math problems with Llama 3.2 3B (base). Left: Red line: pre-training on a 5x5 dataset with Peor=0.25 using next-token prediction. Blue lines: the same pre-training runs, each followed by GRPO from different checkpoint. Top: median test accuracy under greedy decoding (3 seeds). Bottom: median average response length (median over seeds, averaged over sequences). Right: markers: supervised fine-tuning checkpoints on GSM8k with Pror=0.05 (epochs 1-30). Colored curves: GRPO post-training from these checkpoints. Training sequences are reused across epochs; during post-training, multiple generations per input are counted.
Figure 4: Advantage of next-token prediction followed by reinforcement learning
together with the execution trace of the grade-school multiplication algorithm and the final answer, with probability Peot. and just question, answer with probability Peot (Figure 22). We train GPT2 models from scratch on these datasets. Aggregated results for n=5, Peot Pcot=.. =0.25 are shown in Figure A (left). Full description and results (n=4, 7, Peot E {0.1, 0.25, 0.5)). along with detailed figures, are deferred Appendix C.4 In summary, we observe similar learning dynamics in the previous sections, except that here accuracy can even increase from ~0% to 100% when switching to RL.
Grade School & High School Mathematics (GSM8K & MATH datasets) Finally, we experiment with pre- tra-trained LLMs and mathematical reasoning benchmarks. We train Llama pre-trained models (Dubey et al.) 2024) variations of two standard mathematical benchmarks (GSM8k (Cobbe al. 2021) and MATH (Hendrycks et al. (02T)).
I We format the data by surrounding the chain of thought of each sample with special "thinking" tokens. As previously, the cot is replaced by the empty string with probability Poot- We first perform supervised fine-tuning (SFT) on these mixture datasets, and then switch to RL with GRPO and reward that assesses both correctness of the response and consistency with the specified data format. For experimental details, see Appendices B.3 B.4
We plot results for Llama 3.2 3B (base) on GSMSK, Peo:=0.05, in Figure (right). We show test accuracy and average response length for the SFT checkpoints (stars) and the post-training curves starting from these checkpoints. As baseline, we also plot the best SFT model trained with full chain of thought data; the Pcot=1.0 runs can be found in Figure 27 Early in SFT, accuracy is low and the average length of the model response is much smaller than the baseline. As training with the next-token objective continues, accuracy improves and length increases, but accuracy does not reach the baseline and appears to saturate at a lower value. Once we switch to RL with GRPO, the model rapidly reaches baseline accuracy while increasing its average response length. Notice how few samples RL requires to reach the baseline, even when starting from the first epoch check- point. This sample count should be contrasted with the number of SFT updates needed to reach comparable accuracy. Furthermore, in Figures 30 to 32 we provide examples of model completions during RL for this first checkpoint. We observe that early on, the model respects the SFT format but primarily generates short responses (Figure 30). As RL progresses, the model learns to produce longer responses with greater probability (Figure 3T), eventually reaching point where consis- tently generates correct elaborate responses that mimic the chain of thoughts in the long-form part of the training distribution (Figure F2). These training patterns are consistent with the mechanisms observed in the parity setting of the previous sections. In Figure 27 we repeat this experiment for various Poot values. Beyond the point where accuracy exceeds the SFT baseline, the model contin- ues to improve and its response length grows even further. This phase of post-training differs from

---

Published as a conference paper at ICLR 2026
the situations observed in the previous sections and is likely due to the model leveraging pre-training data. Indeed, in Figure 33) we show that model generations at the end of RL are qualitatively dif- ferent from the train sequences. Based these observations, phase of RL reflects in-distribution learning (where the model "mines" SFT data), while the later phase reflects out-of-distribution gains. MATH experiments with Llama 3.2 3B, 3.1 8B (instruct) are deferred to Appendix C.5
hypothesize that the initial steep
6 DISCUSSION
In this work, we introduced a theoretical framework to study the success of reinforcement learning applied after next-token prediction in Large Language Models. We demonstrated and proved that when the data contain rare elaborate sequences encoding a challenging target function, RL can help the model to learn by effectively up-sampling the presence of the long demonstrations in the data mix. Future work can address the limitations of our setting, by understanding better, for example, how noisy chains of thought affect the conclusions, as well as considering separate pre-train and post-train target functions.
Learning with chain of thought data Recently, a few theoretical works have attempted to capture the success of autoregressive modeling with LLMs (Wies et al., 2023 Malach] 2024 Joshi et al. 2025). In particular, Joshi et al. ): Malach (2024) showed that next-token prediction can lead computationally efficient learning of any efficiently computable binary function, which stands in sharp contrast to standard supervised learning where only limited class of functions can be learned efficiently. However, these results rely on the assumption that a learner has access with perfect chain of thought data. While this has turned out to be very productive assumption for the theoretical study of autoregressive learning. it nonetheless constitutes strong assumption. Our work takes step forward in relaxing this requirement; it instead assumes that the dataset contains a small but nonzero (possibly polynomially small in the context length) fraction of chain of thought data. This is arguably more natural assumption for modeling the presence of elaborate "good" demonstrations in the vast ocean of internet text. As we showed in Section this is enough to guarantee efficient learning of the parity function and, interestingly, the guarantee is achieved through a variation of the popular post-training recipe used in state-of-the-ar LLMs. We believe our proof strategy can be modified to show that pre-training followed by post-training can lead to efficient learning for other "hard" functions. We leave such a study for future work.
dataset
Length increase as learning phenomenon The common wisdom in the literature has been that the chain of thought during RL grows to enable the approximation of complex algorithmic tasks. Indeed, some computational problems require large computational depth to be executed with reason- able resources (Hastad) (187) and the standard decoder-only transformer architecture with context length n appears limited in what can represent exactly with constant depth. On the other hand, a transformer augmented with (p)oy((n) chain of thought can simulate any function described by a circuit of polynomial size (Feng et al. 2023] Li et al., [2024] Merrill & Sabharwal 2023). as the additional chain of thought provides computational depth to the model. As such, it seems theoreti- cally satisfying that when chain of thought grows during RI., the model's performance improves on complex tasks. One shorteoming of this perspective is that it does not aim to explain why how optimization pressures lead to length increase, which was the main focus of our work. On a more abstract level, the approximation advantage of long responses suggested by prior work indeed fundamental, as one cannot hope to learn task without the capacity to represent i Nevertheless, the learning advantage we captured in this paper can be more prevalent as it even appears in cases where representation is not an issue (such as the parity task we considered). Based on the above, we suggest interpreting the growth of the length during reinforcement learning of autoregressive models as bona fide learning phenomenon.
'The actual situation is bit more nuanced, as exact representation a strict requirement for learn- ing through approximation (which is the goal in e.g. PAC learning). See Kamath et al. (2020); Malach & Shalev-Shwart (2022) for some alternative definitions of approximation which can be fruitful for understand- ing learning.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENTS
JK thanks the Simons Foundation for support through the Collaborative Grant "The Physics of Learning and Neural Computation". NT and JK acknowledge support by the NSF through NRT Award 1922658.
REFERENCES
Emmanuel Abbe, Enric Boix Adserå, and Theodor Misiakiewicz. SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics. In Gergely Neu and Lorenzo Rosasco (eds.). The Thirty Sixth Annual Conference on Learning Theory, COLT 2023, 12-15 July 2023, Bangalore, India, volume 195 of Proceedings of Machine Learning Research, pp. 2552-2623. PMLR, 2023.
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303. 087744, 2023.
Anton Bakhtin, David Wu, Adam Lerer, Jonathan Gray, Athul Paul Jacob, Gabriele Farina, Alexander H. Miller, and Noam Brown. Mastering the game of no-press diplomacy via human- regularized reinforcement learning and planning. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023, 2023.
Boaz. Barak, Benjamin I.. Edelman, Surbhi Goel, Sham M. Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: SGD learns parities near the computational limit, In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Process- ing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28- December 2022, 2022.
Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemyslaw Debiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Christopher Hesse, Rafal Józefowicz, Scott Gray, Catherine Olsson, Jakub Pachocki, Michael Petrov, Henrique Pondé de Oliveira Pinto, Jonathan Raiman, Tim Salimans, Jeremy Schlatter, Jonas Schneider, Szymon Sidor, Ilya Sutskever, Jie Tang, Filip Wolski, and Susan Zhang. Dota 2 with large scale deep reinforcement learning. CoRR, abs/1912.06680, 2019. URL URL http: / arxiv. org/ abs 1912.1121. 06680
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. CoRR, abs/2110.144 14168, 2021. URL
https://arxx org. abw////s//111 14168
Amit Daniely and Eran Malach. Learning parities with neural networks. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS '20, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 9781713829546.
Dee,Seck-AI, Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, Xiaokang Zhang Xingkai Yu, Yu Wu, F. Wu, Zhibin Gou, Zhihong Shao, Zhuoshu Li, Ziyi Gao, Aixin Liu, Bing Xue, Bingxuan Wang, Bochao Wu, Bei Feng, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, Damai Dai, Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, Fucong Dai, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li, H. Zhang, Han Bao, Hanwei Xu, Haocheng Wang, Honghui Ding, Huajian Xin, Huazuo Gao. Hui Qu, Hui Li, Jianzhong Guo, Jiashi Li, Jiawei Wang, Jingchang Chen. Jingyang Yuan, Junjie Qiu, Junlong Li, J. L. Cai, Jiaqi Ni, Jian Liang, Jin Chen, Kai Dong, Kai Hu, Kaige Gao, Kang Guan, Kexin Huang, Kuai Yu, Lean Wang. Lecong Zhang, Liang Zhao, Litong Wang. Liyue Zhang, Lei Xu, Leyi Xia, Mingchuan Zhang, Minghua Zhang, Minghui Tang, Meng Li, Miaojun Wang, Mingming Li, Ning Tian, Panpan Huang, Peng Zhang, Qiancheng Wang, Qinyu Chen, Qiushi Du, Ruiqi Ge, Ruisong Zhang, Ruizhe Pan, Runji Wang, R. J. Chen, R. Jin, Ruyi Chen, Shanghao Lu, Shangyan Zhou, Shanhuang Chen, Shengfeng Ye, Shiyu Wang, Shuiping Yu, Shunfeng Zhou, Shuting Pan, and S. S. Li. reepseek-rl: Incentivizing reasoning capability in llms via reinforcement learning. CoRR, abs/2501. 12948, 2025. doi: 10. ASS50//RXIV.2555 12948. URI //do:i http:///oooooo//////// 8550 x v. 2501 12991
11

---

Published as a conference paper at ICLR 2026
Yuntian Deng, Yejin Choi, and Stuart M. Shieber. From explicit cot to implicit cot: Learning to internalize cot step by step. CoRR, abs/2405, 2024. doi: 10.45550ARKXV244454443
URL httpts::s/so..oooo//ooooootttttt http:///00..0g111..455555/0///// wtw..0..55500/.x..22.5 14838
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurélien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Rozière, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer. Cyrus Nikolaidis, Damien Allonsius, Daniel Song. Danielle Pintz, Danny Livshits. David Esiobu, Dhruv Choudhary, Dhruy Mahajan, Diego Garcia-Olano, Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Frank Zhang, Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Ander- son, Graeme Nail, Grégoire Mialon, Guan Pang, Guillem Cucurell, Hailey Nguyen, Hannah Ko- revaar, Hu Xu, Hugo Touvron, Iliann Zarov, Imanol Arrieta Ibarra, Isabel M. Kloumann, Ishan Misra, Ivan Evtimov, Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Ma- hadeokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong, Jenya Lee. Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasuden AI- wala, Kartikeya Upasani, Kate Plawiak, Ke Li, Kenneth Heafield, Kevin Stone, and et al. The Ilama 3 herd of models. CoRR, abs/2407 217833 2024. doi: 10.45550ARX...240.22277 URL https: /tg/10. 48550/5550B88 arxiv 2407 21783
Guhao Feng. Bohang Zhang, Yuntian Gu, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: A theoretical perspective. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural In- formation Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
Margalit Glasgow. SGD finds then tunes features in two-layer neural networks with near-optimal sample complexity: A case study in the XOR problem. In The Tweifth International Conference Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024, 2024.
Johan Hastad. Computational Limitations of Small-Depth Circuits. MIT Press (MA), 1987.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song. and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. NeurIPS, 2021.
Scarlett R. Howard, Julian Greentree, Aurore Avarguès- Avareu Jair E. Garcia, Andrew D. Greentree, and Adrian G. Dyer. Numerosity categorization by parity in an insect and simple neural network. Frontiers Ecology and Evolution, Volume 10- 2022, 2022. ISSN 2296-701X.
Aaron Jaech, Adam Kalai, Adam Lerer, Adam Richardson, Ahmed KI-Kishky, Aiden Low, Alec Helyar, Aleksander Madry, Alex Beutel, Alex Carney, Alex Iftimie, Alex Karpenko, Alex Tachard Passos, Alexander Neitz, Alexander Prokofiev, Alexander Wei, Allison Tam, Ally Bennett, Ananya Kumar, Andre Saraiva, Andrea Vallone, Andrew Duberstein, Andrew Kondrich, Andrey Mishchenko, Andy Applebaum, Angela Jiang, Ashvin Nair, Barret Zoph, Behrooz Ghorbani, Ben Rossen, Benjamin Sokolowsky, Boaz Barak, Bob McGrew, Borys Minaiev, Botao Hao, Bowen Baker, Brandon Houghton, Brandon McKinzie, Brydon Eastman, Camillo Lugaresi, Cary Bassin, Cary Hudson, Chak Ming Li, Charles de Bourcy, Chelsea Voss, Chen Shen, Chong Zhang, Chris Koch, Chris Orsinger, Christopher Hesse, Claudia Fischer, Clive Chan, Dan Roberts, Daniel Kap- pler, Daniel Levy, Daniel Selsam. David Dohan, David Farhi, David Mely. David Robinson, Dimitris Tsipras, Doug Li, Dragos Oprica, Eben Freeman, Eddie Zhang, Edmund Wong. Eliz- abeth Proehl, Enoch Cheung, Eric Mitchell, Eric Wallace, Erik Ritter, Evan Mays, Fan Wang, Felipe Petroski Such, Filippo Raso, Florencia Leoni, Foivos Tsimpourlas. Francis Song, Fred von Lohmann, Freddie Sulit, Geoff Salmon, Giambattista Parascandolo, Gildas Chabot, Grace Zhao, Greg Brockman, Guillaume Leclere, Hadi Salman, Haiming Bao, Hao Sheng, Hart An- drin, Hessam Bagherinezhad, Hongyu Ren, Hunter Lightman. Hyung Won Chung, Ian Kivlichan, Ian O' O'Connell, Ian Osband, Ignasi Clavera Gilaberte, and Ilge Akkaya. Openai ol system card. CoRR, abs/2412.16722 2024. doi: 2..
12

---

Published as a conference paper at ICLR 2026
Samy Jelassi, David Brandfonbrener, Sham M. Kakade, and Eran Malach. Repeat after me: Trans- formers are better than state space models at copying. In Forty-first International Conference on Machine Learning. ICML 2024. Vienna, Austria, July 21-27, 2024, 2024.
Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de Las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mistral 7b, CoRR, abs/2310.06823, 2023.
Nirmit Joshi, Gal Vardi, Adam Block, Surbhi Goel, Zhiyuan Li, Theodor Misiakiewicz, and Nathan Srebro. A theory of learning with autoregressive chain of thought. Nika Haghtalab and Ankur Moitra (eds.), The Thirty Eighth Annual Conference on Learning Theory, 30-4 July 2025, Lyon, France, volume 291 of Proceedings of Machine Learning Research, pp. 3161-3212. PMLR, 2025.
Pritish Kamath, Omar Montasser, and Nathan Srebro. Approximate is good enough: Probabilistic variants of dimensional and margin complexity. In Conference on Learning Theory, pp. 2236- 2262. PMLR, 2020.
Michael J. Kearns Efficient noise-tolerant learning from statistical queries. J. ACM, 45(6):983- 1006, 1998.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
Zhiyuan Li, Hong Liu, Denny Zhou, and Tengyu Ma. Chain of thought empowers transformers to solve inherently serial problems. In The Twelfth International Conference on Learning Represen- tations, ICLR 2024, Vienna, Austria, May 7-11, 2024, 2024.
Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. In The Eleventh International Conference on Learning Representa- tions, ICLR 2023, Kigali, Rwanda, May 1-5, 2023, 2023.
Iyy Loshchilov and Frank Hutter. Decoupled weight decay regularization. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019, 2019.
Eran Malach. Auto-regressive next-token predictors are universal learners. In Forty-first Interna- lional Conference on Machine Learning, ICMI. 2024, Vienna, Austria, July 21-27, 2024, 2024.
Eran Malach and Shai Shalev-Shwartz. When hardness of approximation meets hardness of learn- ing. J. Mach. Learn. Res., 23:91:1-91:24, 2022. URL
20 940. html
William Merrill and Ashish Sabharwal. The parallelism tradeoff: Limitations of log-precision trans- formers. Trans. Assoc. Compul. Linguistics, 11:531-545, 2023.
Marvin Minsky and Seymour A. Papert. Perceptrons: An Introduction to Computational Geometry. The MIT Press, 1969.
Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar. Foundations of Machine Learning.
Adaptive computation and machine learning, MIT Press, 2012.
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, Charles Sutton, and Au- gustus Odena. Show your work: Scratchpads for intermediate computation with language models. CoRR, abs/21 12.00114, 2021. ht /ttp:///aivvviviiii org/abs/ abs 2212. 00114
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Ed- ward Z. Yang, Zachary Detito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjic Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer,
13

---

Published as a conference paper at ICLR 2026
Florence d' d'Alcéé-Buc, Emily B. Fox, and Roman Gareet (eds.), Advances in Neural Informa- tion Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019. December 8-14, 2019, Vancouver, BC, Canada. pp.
2019.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. OpenAI, 2019. URL https: wtt.///.////..e. Jea.tter...
s/ hall//aaaggaaaaaaaeaaaeaaa are_
Ben Recht. Patterns, predictions, and actions (revisited). https::/www.argg:: net/pt patterns prtnass-rediictione aedd-actinns -ebb 2025. Substack post, published August 26, 2025.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. CoRR, abs/1 7674776477 2017. URL http://arxiv, org/ /
1707. 06347
Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. In Katrin Erk and Noah A. Smith (eds.), Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Longg Papers), Berlin, Germany, August 2016. Association for Computational Linguistics.
Shai Shalev-Shwartz and Shai Ben-David. Understanding Machine Learning From Theory to Algorithms.. Cambridge University Press, 2014.
Shai Shalev-Shwartz, Ohad Shamir, and Shaked Shammah. Failures of gradient-based deep learning. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning. ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pp. 3067-3075. PMLR, 2017.
Zhihong Shao, Peiyi Wang. Qihao Zhu, Runxin Xu, Junxiao Song, Mingchuan Zhang, Y. K. Li, Y. Wu, and Daya Guo. Decpseekmath: Pushing the limits of mathematical reasoning in open language models. CoRR, abs/2402.033000, 2024.
Ruoqi Shen, Sébastien Bubeck, Ronen Eldan, Yin Tat Lee, Yuanzhi Li, and Yi Zhang. Positional description matters for transformers arithmetic. CoRR, abs/2311 2023. doi: 10.48550/ ARXIV.2311.1477
2311.
Itamar Shoshani and Ohad Shamir. Hardness of learning fixed parities with neural networks. CoRR, abs/2501 0t//22511.00111 2025. doi: 10.48550/ARIV..550.11 10. URL 18550/arXiv> 2501 C0817
David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre. George van den Driess- che, Julian Schrittwieser, loannis Antonoglou, Vedavyas Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy P. Lillicrap, Madelcine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of go with deep neural networks and tree search. Nature, 529(7587):484-48) 2016.
David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, Yutian Chen, Timothy P. Lillicrap, Fan Hui, Laurent Sifre, George van den Driessche, Thore Graepel, and Demis Hassabis. Master- ing the game of go without human knowledge. Nature, 550(7676):354-33) 2017.
Jianlin Su, Murtadha H. M. Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
torchtune maintainers and contributors. torchtune: Pytorch's finetuning library, April 2024. URL https//githun com/ http://ooooct///occcctt.cco
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp- 9998-6008 2017.
14

---

Published as a conference paper at ICLR 2026
Leandro von Werra, Younes Belkada, Lewis Tunstall, Edward Beeching, Tristan Thrush, Nathan Lambert, Shengyi Huang, Kashif Rasul, and Quentin Gallouédec. Trl: Transformer reinforcement learning. https:// hwwehhh...oooouuuheeeeeee com wwt.o//hhggiggaaceetteee 2020.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed H. Chi, Quoc V. Le, and Denny Zhou. Chain-of-thought prompting elicits reasoning in large language models. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 December 9, 2022, 2022.
Noam Wies, Yoav Levine, and Amnon Shashua. Sub-task decomposition enables learning in quence to sequence tasks. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023, 2023.
Ronald J. Williams. Simple statistical gradient-following algorithms for connectionisstreiiforeeee learning, Machine Learning. 8:229-255 1992.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Remi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick Von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander Rush. Transformers: State-of of-thhe-rtt natural language processing. In Qun Liu and David Schlangen (eds.). Proceedings of the 2020 Confer- ence on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 38-45, Online, October 2020. Association for Computational Linguistics.
Jingfeng Wu, Peter L. Bartlett, Matus Telgarsky, and Bin Yu. Benefits of early stopping in gradient descent for overparameterized logistic regression. CoRR, abs/2502.132883 2025.
Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah D. Goodman. Star: Bootstrapping reasoning with reasoning. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neurul Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 December 9, 2022, 2022.
15

---

Published as a conference paper at ICLR 2026
Reproducibility Statement To allow independent reproduction of experimental results, provide extensive experimental details in the main paper and in the appendix: (i) description of syn- thetic datasets, model and algorithms hyperparameters (Section 2 and Appendix B.2 B.44 (ii) number runs per result, with mean/median + std at each figure; (iii) compute details and es- timates of GPU hours needed for reproduction, (iv) reference theoretical assumptions are mentioned in several places. Appendix contains theoretical details, which were not possible cover in the main text due space constraints.
any open-source code used. Our
A MODELS AND ALGORITHMS
Herein, describe the main models and training algorithms consider in the parity experiments.
A.1 ARCHITECTURE
A transformer is te-sueence neural network, which in its simplest form consists of lay- ers of self-attention followed by a position- wise feed-forward network. The term "autoregressive" indicates that the prediction of the model for each new token is conditioned only on tokens that appeared earlier in the sequence. To obtain distribution over the next token, common com- pose the output of the transformer with soft-max layer across the vocabulary index, which can be potentially parameterized by a temperature value.
A.2 PRE-TRAINING
The next-token prediction objective consists of the log loss of the model token distribution evalu- ated at output positions 1 to d. We omit the first d input positions of the sequence from the loss objective, input bits 21,-.., Ft distributed uniformly at random. If we denote the predicted sequence-level distribution by no (y x) I و | x, Meey) E (0, 1), which is induced by the prediction of the transformer at each position composed with the soft-max, then training corresponds to:
M
min E ((x.))>DDeed
Tay (U)
x, yes)
(5)
A.3 POST-TRAINING
We consider three main reinforcement learning algorithms: Self-Taught Reasoner (STaR) Ze-
likman et al. (022), REward nocrement=NNomegative Factor times Offset Reinforcement times Characteristic Eligibility (REINFORCE) (Williams) 1992) and Group Relative Policy Optimiza- tion (GRPO) (Shao et al. 2024). r X' x {0, 1] be binary reward function. We consider outcome-based rewards, which assign one reward value per sequence.
STaR In STaR [A first query the pre-trained model to generate responses for inputs. We then filter out those with 0 reward and fine-tune model on the rest using the next-token prediction objective. We repeat this process for several rounds, each time sampling from the model returned by the previous round. The algorithm over one round corresponds to:
min E
Radd(1222) pprev (1x)
In May (y) (cy)) r(x, v)
where i Uprev E RR corresponds to parameters returned in the previous round of the algorithm.
The pseudocode of STaR can be found in Algorithm
REINFORCE REINFORCE is a standard policy gradient algorithm. It seeks to maximize the expected reward of sequence generated by the model:
max E
Rx 1/2 [r(x, v)]
(7)
9
g (-1x)
16

---

Published as a conference paper at ICLR 2026
Algorithm Self-Taught Reasoner (STaR) Algorithm
Require: Pre-trained model parameters db. RL rounds n, Fine-tuning epochs per round E, Input distribution xx Rad(1 12) 228d, Reward function r(x, v). Number of samples to generate per round N.
1: Set 0 1/
2: for r = to do
3:
Set S 0
4:
for = to N do
5:
Sample x xx
6:
Sample ~ Tt(- x).
7:
if r(, y) then
8:
Set S - S u v).
9:
end if
10:
end for
for epoch = E do
12: for each (x, y) E S do
Update by taking gradient step on the next-token prediction loss for (x, y).
14: end for
15:
end for
16:
end for
17:
return Final model parameters d.
The gradient of the objective function can be written as:
Rad R2d(1/2)04 /2y0d,
[V ( v)1
(8)
y-<o (1)x)
which implies that the algorithm can be recast as:
min E
Rad(1/2)e |- In mo (3) x)) xxx,
(9)
go (-(x)
where the random variable is treated as constant and not as a function of 9. This is the surrogate objective that is optimized practice. Observe the similarities with the STaR objective of equa- lo Furthermore, a widely adopted option is to center the rewards in order to reduce the variance of the updates during the execution of stochastic optimization algorithm. This leads the final version of the REINFORCE method:
min E Rtd 84 g~
Σ
In may (y5 x, A(x,y)
(10)
where A(x,v) v) r(x, y) -E Bx [r(x, y) is often called the advantage function. grrto (-1x)
Group Relative Policy Optimization, GRPO for short, is policy gradient algorithm, motivated by the class of Proximal Policy Optimization (PPO) algorithms (Schulman et al. 2017). that is widely used for post-training large language models with reinforcement learning. Its objective amounts to:
GRPO
E Rod Rad 1/2) ()) ..... (N) mosa
min (rij clip(ri,, 1 - Celip- 1 + Celip) A.}
N
1g()))
7=1
where A, rxx, mean (( x (1) (x...((x.,(("))
is normalized reward, éclip is hyperparameter and oold a predicted distribution that corresponds an earlier model checkpoint during the exe-
stdl ()) ) ..... (x."(())
(M) (x
cution the algorithm. We defined rı of group of responses per input.
The term "Group" indicates sampling
Wold. ( xxy (
17

---

Published as a conference paper at ICLR 2026
B EXPERIMENTAL DETAILS
The experiments are implemented using PyTorch (Paszke et al. 2019)
B.1 PARITY
We use the Transformers library (Wolf et 2020) for our parity experiments.
Architecture The transformer is initialized with context length d'=2d, to be able to process the whole sequence, where recall d is the number of input bits. The two architectures we consider primarily differ in the type of positional encodings: GPT2 uses absolute positional encodings of the input sequence, while Mistral uses relative positional encodings (Su et al.) 2024). We use the GPT2Con f i g, Mi Con config classes define the models. We consider several model hyperparameters: depth (number of transformer blocks) L E {2, 4, 8), embedding dimension demdb {128, 256}, number of heads Theada {4, 32). By the default convention, each head has dimension demdb/ Pheads-
Pre-training
We consider distributions with
Poot 0.05, 0.1, 0.25, 0.3, 0.32, 0.33, 0.34, 0.35, 0,,,,,,,, 1.03.
We use the Adam (Kingma& & Ba 2015) optimizer with learning rate 7pr==10 and batch size Bpre= Test statistics are estumated using 2,560 new samples.
In Figure 2 (right), each bullet corresponds to the final test accuracy with greedy decoding over random seeds. The final accuracy of each run is defined as the average accuracy from the last checkpoints of the pre-trained run (that is, average accuracy iterations 47,500 to 49,500).
Post-training We continue post-training from the last checkpoint of pre-training. For STaR, we generate 3, 200 samples per RL iteration and train on the correct ones for 3 epochs with batch size 64. For REINFORCE and GRPO, we use batch size 64. We found it necessary to use larger sample size per RL iteration for STaR. The GRPO experiments use N 4 rollouts per input and clip parameter Éelip
0.2. No KL penalty is applied (as outlined in Section (A). We use the Adam (Kingma & Ba Test statistics are estimated using 640 new samples.
2015
optimizer with learning rate 710
B.2 NUMBER MULTIPLICATION
Architecture We use the GPT2 (Rafford et al., 2019) architecture which contains 124,439,800 parameters and the default GPT2 tokenizer that includes vocabulary of size 50,257 and utilizes Byte Pair Encoding (Sennrich et al. 2016).
Pre-training We consider datasets with Peor (0.11 0.5, 1}. The dataset and the pre-training code is based on the public repository of (Deng et al. 2024). We train with the next-token prediction objective applied to the chain of thought (if exists) and the answer tokens of the sequence. We use the AdamW W (Loshchilov & Hutter 2019) optimizer with learning rate 75 10 batch size Bpre=32, and maximum gradient norm equal to 1.0. We train for 15 epochs for the 4 x 4 task and 50 epochs for 5 x 5 and 7 7. All 3 training datasets consist of 808,00 008,000 unique samples. We train each model (3 random seeds per task) in a single GPU (NVIDIA A 1100-XXM4-00GB For x 4, each run takes approximately 3-5 hours to complete (depending on the value of Poot). For 5 x 5 and 7 x 7, the completion time ranges in 21-38 hours.
Post-training We train with GRPO with end-to-end reward function (that is, 1 if the kens generated after # A ## correspond to the correct answer] and 0 otherwise). We use the AdamW (Loshchilov & Hutter 2019) optimizer with learning rate Tppos=3 1 (and Hugging- face's s Transtormers' (Wolf et al. 2020) default rest of hyperparameters), batch size =1. group size (for GRPO) equal 4, steps of gradient accumulation (which makes the effective batch size equal to 32), no KL penalty, éelip 0.2. We use generation temperature TRI. = 1.0. We use the
That is, "complete number match".
18

---

Published as a conference paper at ICLR 2026
TRL library for efficient vost-training (von Werra et al. 2020 We found it necessary implement a wrapper around the model forward's 's function to force the model to be in evaluation mode dur- ing the generation of the rollouts for RL (that is, to turn off any sources of non-determinism in the model, such as dropout layers). Otherwise, under TRL's default implementation, the model was not able to generate correct responses, even during the very first RI. iteration. We perform post-training for each random seed in a single GPU (NVIDIA A A100-SXM4-80GB) which concludes in about 3.5 hours.
Test accuracy is estimated using set of 1000 unseen pairs of numbers.
Estimate of total GPU hours needed to reproduce the pre-training and post-training results: 900
B.3 GSM8K
Our code is based on the torchtune library (torchtune maintainers & contributors 2024).
Architecture We use the Llama 3.2 3B base model for our GSM8K experiments.
We split each dataset into equal parts, perform supervised fine-tuning (SFT) on the first half with the next-token prediction objective applied to the answer portion of the sequence, and then switch to RL on the second half.
Supervised fine-tuning We train with the next-token prediction objective ap- plied to the chain thought and answer tokens of the sequence. We struc- ture the data with preamble prompt (not calculating loss): "A conversation between User and Assistant. The user asks question, and the Asaistant solves it. The assistant first thinks about the
reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within <think></think> and <answer></answer> tags, respectively, i.e., <think>reasoning process here</think> <answer>answer"
here</answer>, User: (question] Assistant:
and structure the fol- where cot is
lowing sequence as: <think>co<<tthing everything that exists before characters #### in the original GSM8k dataset. As before, we drop the cot from each sample with probability Peot-
fanswer>answer//nnww
We use the Adam W (Loshchilov & Hutter 2019) optimizer with learning rate 1===10-5, hatch size Bat==56 and bfl6 precision. We train for 30 epochs. We consider datasets with Poot 0.25, 0.5, 1}. We use PyTorch for distributed training. Each run is performed in 8 GPUs (NVIDIA A100-SXM4-8000
and takes 1.5 hour to complete.
Post-training We train with GRPO with reward function that has stepwise structure: points for the presence of the answer tag, 5 points for the presence of the thinking tag, 20 points if it contains the correct answer in part of the response, and 100 points for correct answer the appropriate place. We found it important using the math-verifs module Huggingface for simplifying mathematical expressions and accurately rewarding generations. We use the AdamW (Loshchilov) & Hutter 2019) optimizer with learning rate 1 cosine learning rate scheduler, batch size Bpost= =, group size (for GRPO) equal to 32, no KL penalty, <elip 0.2. We generation temper- ature TRL = 1.0. We perform post-training trattrraiiig in 8 GPUs (NVIDIA (X000SXM4-80G) and each run concludes in 31-35hours and 20hours for Peot 0.
Test statistics are estimated on the test set of GSM8k using 4 generations per question and taking the average.
Estimate of total GPU hours needed to reproduce the pre-training and post-training results: 11k
B.4 MATH
Our code is based on the torchtune library (torchtune maintainers contributors 2024).
Architecture
We use the Llama 3.2 3B, and 3.1 8B instruct models.
19

---

Published as a conference paper at ICLR 2026
We split each dataset into equal parts perform supervised fine-tuning (SFT) the first half with the next-token prediction objective applied to the answer portion of the sequence, and then switch to RL on the second half.
Supervised fine-tuning We train with the next-token prediction objective applied to the chain of thought and answer tokens of the sequence. We structure the data with preamble prompt (not calculating loss):
< <begin of text starstart header id id/>system<| end header id/> Cutting Knowledge Date:
"December 2023
Today Date: 26 Jul 2024
A conversation between User and Assistant. The user "asks a question, and the Assistant solves it. "The assistant first thinks about the reasoning "process in the mind and then provides the user with the answer. " "The reasoning process and "answer are enclosed within <think></think> "and <answer></answer "<thinnk>reaonning process heres/think> <answer'answer "here</answen Inside the answer tag, put only
tags, respectively, i.e.,
"the answer and no additional commentary. <|eot id|><sstas
"header id |>user<|end header id/> (question)
"< id/><|start header id salsssssantt//em header id/>
and structure the sequence as:
<answer>answrr<//nswwee where cot is everything that exists in the original MATH sequence before the boxed answer. We drop the cot with probability Poot-
We use the AdamW (Loshchilov & Hutter 2019) optimizer with learning rate 7/a==10 batch size B2n=24 for the 3.2 3B model and =1n==6 for the 3.1 8B model, and bf16 precision. We train for 20 epochs. We consider datasets with Peot (0.1, 0.25, 0.5, 1}. We use PyTorch for distributed training. Each run is performed in 8 GPUs (NVIDIA A100- (00-SXXM-88GGB and takes 1.5 hour to complete.
Post-training We train with GRPO with a reward function that has stepwise structure: 5 points for the presence of the answer tag, 5 points for the presence of the thinking tag, 20 points if it contains the correct answer in part of the response, and 100 points for correct answer at the appropriate place. We found it important using the math math-ver veriiyy module Huggingface for simplifying mathematical expressions and accurately rewarding generations. We use the AdamW (Loshchilov) & Hutter 2019) optimizer with learning rate 1ppos=10033 (and Huggingface' 's default rest of hyperpa- rameters), batch size =l, group size (for GRPO) equal 20, 2 steps of gradient accumulation (which makes the effective batch size equal to 2), no KL penalty, Ediip = 0.2. We use generation temperature TuL. 1.0. We perform post-training in 8 GPUs (NVIDIA A A100-SXM4-80GBB for 300 steps which concludes in about 14-29 hours.
C ADDITIONAL EXPERIMENTAL RESULTS
C.1 MORE EXPERIMENTS ON MAIN PARITY SETTING
Model hyperparameters ablation We present aggregated results for pre-training and post- training on D(pcot), ((Peot),Peot=0... =0.25 for GPT2 and Mistral architectures of varying embedding dimension and number of heads in Figure (deeth L=2), Figure (igure L=4) and Figure (depth L=8). Post- training uses GRPO with end-to-end reward function Fe2e and sampling temperature 7RI.=1.0 for all configurations. Notice that the Mistral architecture exhibits more unstable learning (pre-training) than GPT2. We observe that pre-training alone does not induce a model capable of consistent gen- cralization in any of the cases. We plot number of training sequences on the x- x-axis (as opposed to,
We split the datascts, since they are of small size and we perform multi-epoch training during SFT.
20

---

Published as a conference paper at ICLR 2026
say, SGD iterations allow simpler comparison with post- post-training). For the number of post-traie train samples, we take the multiplicity of GRPO samples into consideration.
Comparison of RL algorithms for more values of Poot In Figure we presented test accu- racy and response length during post-training for all RL algorithms (STaR, REINFORCE, GRPO) and for few different values of sampling temperature used for generating training sequences for Pcor=0.25. We show now additional results for distributions where data are even more rare (Pcot=0.1 and P.0)= In Figure 8
we observe that all RL algorithms induce generalization for most val- ues of sampling temperature (post-training starts after 20k pre-training iterations) when Peot=0.11 1. On the other hand, when Peot is even smaller (Poo= post-training after 20k iterations does not reliably lead to well-generalizing models for any of the RL algorithms (Figure ( If, instead, we consider more pre-training iterations (50k) before the start of post-training. then we observe greater probability of success with all RL algorithms for reasonable (==1) sampling temperatures (Figure (0). This illustrates, amongst others, that the learning behaviors STaR, REINFORCE and GRPO are similar in our simple setting.
Smaller input dimension In Figure 11 we demonstrate that pre-training sometimes leads to generalizing model if we wait long enough (example seeds shown for d=11 and d=12). However, as d increases, this waiting time becomes prohibitive.
Sample complexity gap We design an experiment to try estimating the sample complexity gap between the two paradigms (next-token only vs next-token followed by RL). We consider GPT2 ar- chitecture of depth 4, number of heads 32 and embedding dimension equal to 128. We fix P0.2= and we vary the input dimension d {6, 8, 10, 11, 12). We train for at most 100,000 iterations (of batch size 256) and we consider switching to post-training (GRPO with e2e reward) at several checkpoints. We consider checkpoint be successful if its test accuracy is greater than 95%. In Figure [2] we show the gap between the number of samples required for any pre-trained checkpoint to generalize vs the number of samples required by any pre- and post-trained model (combined) for several input dimensions d. For the number of post-train samples, we take the multiplicity of GRPO samples into consideration. This gap seems be growing exponentially fast. Some lim- itations of this simple experiment: the pre-trained checkpoints from which we consider starting post-training are at iterations {1000, 1500, 2000, 2500, 3000, 5000, 10000, 15000, 20000, 30000, 50000, 100000}. Also, for d=12, only one of three seeds ended up generalizing during pre-training (so by definition we could only calculate the gap for that seed and we discarded the rest).
C.2 PARTIAL CHAIN OF THOUGHT
The purpose of this subsection is relax our main setting, and consider a more complex variation which might better capture "real-world" data. In particular, we consider a setting, where the data include third type of sequence consisting of medium-sisee sized chain of thought, in addition to short and long sequences.
We consider the distribution Da(Poot Poda) over {+1)d d x <FI,< parameterized by Peot, Podd E (0, 1), such that: Ti,... .,, Fd ~ Rad(1/2) and
21, IL Z4, <EOS>
w.p. Peot
(y1...., Wd+1)
21, x 2122233115
I 33, <BOS>), w.p. Podds
(11)
[[' Tr <EOS>) w.p. 1 Peot Podd
The medium-sized sequence consists of the same random variables as before input positions to d, while the chain of thought skips the intermediate computations that involve partial parities ending at an even position (with the exception of the final parity, if d is even). This results in sequences that contain partial chain of thought with omitted reasoning steps. The purpose of this sequence is to provide the model with some samples of intermediate complexity between the two extremes of D(Pcot). Indeed, learning predict each term of the chain of thought after the first position requires a leap (Abbe et al. 2023) of order 2, as opposed to either leap or d for short and long sequences, respectively. For instance, predicting 213223 from the history of the sequence (I1,..., Td, x1) is
21

---

Published as a conference paper at ICLR 2026
Pre-rrainng & Post-trainng curves
Pre-trainnng Posteraining curves
Partry, 5-50 model depii darerr122
Party, d-50, model mittlll depth dear1120, Roarr
Fraht
training Post training ourves
training& training curves
1arty.d-50, 5-50. model 1p12. 2epth-2, 128, Thom 225
Partly, 5-50, motel maainl, depth?2, 2, deau-128 Puse 2225
Sidde Unde pran
Meddon
Under
RCS
47gt 35
Fre tratrainng&& training
Pre troining Post training curves
Partiy. d-50, model-gol2 depth-2, 456 R A25
Partyy d-50 model-mitinal degttm 250, Resa Pur
Pre training & Post training Ou DUI
Pre traininng S Poct training curves
Partyy
Panty
W AcurTe 2.R.
Lncll presey droourig
Figure 5: Pre-training and post-training curves combined on the parity task for GPT2 and Mistral architectures. We vary the embedding dimension and the number of heads. Depth L=2.
22

---

Published as a conference paper at ICLR 2026
Pre-rapingg & Post-trainng curves
Pre-training Posteraining curves
Pertry, 5-50 model
depth darerr122
Party, d-50, model mittll depth Sesr1120
Fraht
training Post training Ourves
training Post curves
1arty.d-50 5-50 model 1p12. depthnd, 128, Thoe -22 225
Partly, 5-50, motel manni depthrn, deau Pusee 2225
Meddon seculing under
Sidde unde pran
RCS
47gt 35
Pre tratraiing&& training
Pre troining Post
Partiy. 4-50, model-gpl2 depsi-d, 256 R A25
Partyy d-30 model-mitiaa, degttme, 250, Resa Pur
Pre training & Post training Du NUs
Pre training S Poct training curves
Partyy modalngat2 de5th-4,
Panty
desth-d,
Figure 6: Pre-training and post-training curves combined on the parity task for GPT2 and Mistral architectures. We vary the embedding dimension and the number of heads. Depth L=4.
23

---

Published as a conference paper at ICLR 2026
Pre-trainng & Post-trainng curves
Pre-training Posteraining curves
Pertry, 5-50 model depii darer
Party, d-52, model mittll depth desr-120, Roaar
Fraht
training Post training Durves
training& craining curves
1arty.d-50, 5-50. model 1p12. Septh it 128, Thoe -22 D25
Partly, 5-50, motel maini depthnf & 1eau Pusee -32, 2225
Sidde Unde pran
Meddon
Under
eygt 35
Pre tratrannng&& training
Pre troining Post training curves
Partly d-50, model-gol2 depsi-d, 456 R A23
Partyy d-50 model-mitinal degttmi, d20, R Pur
Pre training & Post training Ou cui
Pre training S Post training curves
Partyy modalngat2
Panty
Figure 7: Pre-training and post-training curves combined on the parity task for GPT2 and Mistral architectures. We vary the embedding dimension and the number of heads. Depth L=8.
24

---

Published as a conference paper at ICLR 2026
Pariyy P.. 1. Post-traininga2220 dm50,
cmm5, modelngpt2, depthmn, dessm122 m32
That Accuracy
Lenithh (Stait (=2ei, temp-1.0
Sequeences
Accurscy
Lengths tomp-1.01
To-00.5
Text Accuracy
Lengths (Stait (o2e), temp-1.01
Accuracy
399
Lengths lemp=1.01
Accuraoy
That Accursy
Lengths (StaR le2e), temp-1.01
ta.e1.5
Accurasy (N)
Sequerccs
RL Reeratons
RL lttrations
creedy 20mpp-0:)
Immp-11
536 Urbel
STA (v4))
REA (000
GARO UNAr
Figure 8: Post-training of transformers on mixture of long and short sequences encoding the parity of d bits with various RL methods and generation temperatures (TRL). Mixture coeffi- cient: P..1=0.1. Pre-traiming iterations: 20k.
Parkyy -0.05 05, Post-rr Nest-traininggp20tt d-50, model-gp12, depthm4, devas-122 Pheas
Text Accunacy
Lengths (Stait (=2e), lemp=1.01
Length
te.00.55
390
1emp-1.01
Accraccy 188)
Length
The 75
Tist Accuracy
Langth (STaR (42e), temp=1.01
Accuracy
299
Lengths
1.25
Acuracy
Test Accunacy
Lenghs (Sta (=2e), temp=1.0 1.01
tare1.5
Accurasy
330
greed Swapp 0.0)
2TA (04l
REIATOACE (000
GRPD (AR
Figure 9: Post-training of transformers on mixture of long and short sequences encoding the parity of d bits with various RL methods and generation temperatures (TRL). Mixture coeffi- cient: =0055, Pre-training iterations: 20k.
25

---

Published as a conference paper at ICLR 2026
Partly Do.m0.05 Pest-traiinggg90k. die56, modelngp12,, depthm4, 6l Rewam
Test Accurasy
Lengths (StaR (o2e), temp=1.01
Secuennoss
0..00.5
Lengths
699
Sequeenes
Acuracy
Lenath
The Tee
Text Accuracy
Lengths (Staif (+2e), temp=1.0
Accuraa
Accunacy
lemp=1.01
Ter-1.25
Accuracy
Test Accuracy
Langts (Stait (620), lump-1.0
ta.11.5
Accuracy
299
Marations
gready Sumpeo B)
Sump--2
STA (624)
STA Scol
REWTORCE 120
SAPO 624
Figure 10: Post-training of transformers on mixture of long and short sequences encoding the parity of d bits with various RL methods and generation temperatures (TRI). Mixture coefficient: =0.55 Pre-training iterations: 50k.
Pu- 2el2, vedioga2 dep Sou P seddnd
Dasttude
Iaparton (lampnnd
Figure 11: Pre-training the parity task leads to a generalizing model after many iterations. The average response is short (length equal to 1). Top row: d=11. Bottom row: d=12.
26

---

Published as a conference paper at ICLR 2026
Post-training Benefit (Parity, Pac (0.25,, model=gpt2. depth
10
125
104
10
11
12
Input dimension
Figure 12: Sample complexity advantage of pre-training post-training over mere pre- training vs input dimension. Note the logarithmic scale in the y-axis. The figure shows mean and one standard deviation across random seeds.
Parity, thought. (Pox (00.. 0.11, d-15,, depthw4,
REINFORCE/<2IIBS<
RENFORCE/<22/8500
Tex: Accunccy
Proportion responses by length
Text: Accuraty
Probertion of reigarte wnghh
13 Fall Atunee
n. tevvlsrs
mezzranng
201, lom pe0p
Dow 132
Hamppnl.
Vow Rar 110
Bration
RBNFFORCE/22 w/ peralty)@@k
REINFORCH wi pena Fy1850k
Accuracy
Proportion
Accuracy
Proportion
100
101
w5gh
MCn 0.2
Figure 13: Parity with partial chain of thought data and effectiveness of length penalties. Post- training experiments using REINFORCE with (bottom row) and without (top row) length penalty on pre-trained models that correspond the training shown in the middle. Notice that leaps in accuracy during pre- pre-training control length distributions at the end of post- post-training and success of length-penalized RL.
27

---

Published as a conference paper at ICLR 2026
leap of order 2. The leap between terms characterizes the additional sample complexity for learning each
We repeat pre- pre-training and post-training on this distribution and present results in Figure 13 for d=15 and (Peot Poda (0.1), 0.1).
Pre-training stairease As before, improvement during pre-training happens in stages: early training. model accuracy under sampling with temperature 1 is around 50%, yet it suddenly jumps to Pous+11 x 100% as the model learns from the long demonstrations. Later in training, there is an- other phase transition indicating the model has learned the medium-sized sequences, the accuracy leaps to PoottPoaat X 100%. As before, if Poot. Podd not large enough, then the models under greedy decoding do not generalize in reasonable time. In fact, the critical values of Ponts Podd can be calculated in our main setting. It holds:
2
1i I
21 +
ri
TP [v = + 1 = (Peat + Podd)
+
Poot
Podd)
(12)
2
2
In absence of a parity feature, this conditional probability simplifies to (see Section
E.2
for details):
P Py1 = +1|x] (Peot + Poda) 21 + 1
(13)
2
or in other words:
Poot + Podd +
P [y x1 [x]
(14)
2
This implies that first token output of a greedy decoder will always be æ1, as long as Poot + Podd > 0. For the second position we have:
1/2 21
21
P 12/2 I
22
x1
P [vV x1 x1|x
Poot + #1 2
Podd
(15)
Peot +odd
2pcot + (x3 + 1) Podd
Poot + Podd
and likewise:
2podd (I's + 1) Peot
P [y2 x11223191 -1,
(16)
Poot + Podd +
1- (x3 + P) + Podd)
P 1y22 <BOS>|un = 11 x]
(17)
Peot + Podd +1
Thus, model under greedy decoding will be stopping generation at the second token long as:
(#3 + (Peot + Poaa) 2 max (2peot + (Ig + Podd, 2podd + (xr + 1) Poot)
(18)
This certainly shows that model trained on distribution endowed with additional "partial" chain of thought data generate long responses and generalize during pre-training already, even if Peot 1 1/3 which was the critical threshold in our main setting. Such an example is presented in Figure 14 for (Peot, Poda (00.),,,,,
Pre-training checkpoints & length penalties In this setting, the point when we switch to post- training, can not only affect the eventual success of RL, but also the length distribution of the model responses at the end of it:
Post-training at 5k iterations yields a model that generalizes, yet it learns to almost exclu- sively generate long responses (Figure 13 top left).
Post-training at 50k iterations, the other hand, results in model that not only manages to reach 100% test accuracy but also has more diversity in its length distribution, generating both long and medium-sized responses (Figure [3] top right).
28

---

Published as a conference paper at ICLR 2026
Farty Vurnd poand deik reddlettinn digiteed Gearrillyy ceadeds
Marro
Prpprrden lenght Compendde
Figure 14: Pre-training on the parity task with "partial" chain of thought data leads to
generalizing model even for Poot 1/3. The average response length "medium".
The previous observation is both good and bad news: good news, as the 50k checkpoint ends up being cheaper at inference than the post-trained 5k counterpart; but bad news, as the model could have been even faster at inference without sacrificing accuracy. Indeed, after 50k pre- gre-training itera- tions the model has learned to generate correct medium-sized responses, hence, principle, need not learn from the long self-generate ones. To test this, we consider reinforcement learning with pegelleee reward. In particular, we consider the following
reward:
d
yyl Alen d
Aeen (x,9) -1 11-11|
36
Aten E [0.
(19)
We use A0.4 As we confirm in Figure 13 (bottom row), a length penalty applied during post- training of a late checkpoint is effective, leads to model that is both accurate and fast. On the contrary, we apply the same length penalty at an earlier checkpoint, post-training fails. This is because the pre-trained model has managed to learn from the medium-sized sequences early in pre-training.
C.3
NOISY DATA
In this subsection, we introduce and study variations of our main setting in which noise may be present in the data. Specifically, we consider cases where the noise affects the pre-training dis- tribution. We observe that many of the previously observed learning phenomena continue to hold in the presence of moderate amounts of noise, even in cases that noise seems be "dominating" pre- pre-training performance.
We consider different noise models.
Final token noise In this case, the final token is flipped with probability v) E (0, That is, Ti...., Zd Rad(1/2) and
/ 81, 01,4112, I Iy, Ed, <BOS>), w.p. Poots
(y1...., 12811)
(20)
(Ea, <BOS>), w.p. Peot,
d i= Zi, w.p. 1 - 17, For this type of noise, we do not anticipate a change in II 24, w.p. 17.
where Ea
the critical threshold pre-training, since there no noise the first token of the long reasoning chain (and in absence of parity feature, the noise in the short path does not affect things). On the other hand, the pre-training accuracy under sampling with temperature is expected change to (1 P)PP + 1-Post (1-2r)peat+1. We present pre-training simulations that confirm the above predictions for r) E (0.1, in Figures 15 and
2
This small change in the pre- rre-training accuracy does not alter significantly the post-training be- haviour. Indeed, pre-training noise this type only affects the starting point of the model during reinforcement learning, and as result it does not affect significantly the sample complexity of post- training. Experimental results for Pcmt=0.25 and {0.1, 0.2) can be found in Figures) and 18 respectively.
29

---

Published as a conference paper at ICLR 2026
Parity, d-25, medel-gpt2, depth-4, Marte=128, 22, noise type-final, g=0.11
Test Accuracy
Probsbilty generation
type-final
0U8
Accurste
Freebae
80
kay prrteery
42-
50
50303
20003
30000
50000
Part
Semp=0.0)
Povie.5
temp-).
Rumm 625
Figure 15: Pre-training of transformers on mixture of long and short sequences encoding the parity of d bits under "final" type noise for 700.1. Left: Test accuracy during the course of pre- training. Center: The probability that a model's generation length is equal to the maximum length present in the training distribution (which equals d). Solid lines correspond greedy decoding. while dashed lines correspond to sampling with temperature 1. Each color denotes a training dis- tribution with a separate mixture coefficient Poot. Figure shows average and standard deviation across 3 seeds. Note the change in the legend from the noise-less figures. Right: Test accuracy with greedy decoding at the end of pre-training (50k iterations) versus mixture coefficient Peot The red dashed line corresponds to the critical threshold. Each bullet is the median of 3 runs.
Parity, d-25, model-gpt2, depth-4, 10ee-112800 noice type-firal, p-02
Test Accuracy
Probability generation
Final Arcurscy (noike bypesfinl,, na022
0.8
W/ groedy serorro
Accurco
0.6
80
40
50
20000 iterations
terraioos
Puu
(lamswon
Poym
temp-7.0
Aur 0:35
C:77
Figure 16: Pre-training of transformers on mixture of long and short sequences encoding the parity of d bits under "final" type noise for 700.2. Left: Test accuracy during the course of pre- training. Center: The probability that a model's generation length is equal to the maximum length present in the training distribution (which equals d). Solid lines correspond greedy decoding. while dashed lines correspond to sampling with temperature Each color denotes a training dis- tribution with a separate mixture coefficient Poot Figure shows average and standard deviation across 3 seeds. Note the change in the legend from the noise-less figures. Right: Test accuracy with greedy decoding at the end of pre-training (50k iterations) versus mixture coefficient Peot The red dashed line corresponds to the critical threshold. Each bullet is the median of runs.
30

---

Published as a conference paper at ICLR 2026
Parity, P5x Pest trakn d--5,, model-gpk2, depth-4, 1288, Rwess-32, nolse-(final, 0.1)
tempa1
8 8
Lemgth
Accuress
Lengths temp-11
Accuracy 1o
Te=11.
KCe
Lengoh
8
temp=1
rem11.5
Accuracy
K00
838 8
RL Iterations
AL Iterabors
Figure 17: Post-training of transformers on mixture of long and short sequences encoding the parity bits under "final" type noise (7=00.1) with various RL methods and generation temperatures (TR.) Left: Test accuracy during the course of post-training with greedy decoding (solid lines) and sampling with temperature (dashed lines) for various RL methods. Figure shows average and I standard deviation across 3 seeds. Right: Length of generated response (sampled with temperature during course post-training run (GRPO with end-to-end reward) for 640 test inputs, after 20k pre- pre-training iterations. Note: The sample size of each RL iteration differs amongst the RL algorithms: n=64 for GRPO, REINFORCE and n=3,200 sequences for STaR.
Parity,Pxx:0.2 0,2, Pest-triining330k. d=25, model=gptz, deptho:4, desss=120, Rbeas noise=-final 0.2)
temp=]
Seng Ceeeeee
3.8
temp-1
Accuray No
Th 1.0
8 a 88888 8 88
kre
Length
Seng
temp-1
1.511.5
Accuracy
RL Ieertions
AL Ittrrborss
Figure 18: Post-training of transformers on mixture of long and short sequences encoding the parity bits under "final" type noise (770.2) with various RL methods and generation temperatures (TRL). Left: Test accuracy during the course of post-training with greedy decoding (solid lines) and sampling with temperature (dashed lines) for various RL methods. Figure shows average and I standard deviation across 3 seeds. Right: Length of generated response (sampled with temperature during the course of post-training run (GRPO with end-to-end reward) for 640 test inputs, after 20k pre- pre-training iterations. Note: The sample size of each RL iteration differs amongst the RL algorithms: n=64 for GRPO, REINFORCE and 3,200 sequences for STaR.
31

---

Published as a conference paper at ICLR 2026
Patly, 0x23, makii gez, destsed, 120 Seaa rebs bopeet, p=0.00
Perts, 5-25, -352 oithted Soc nawe lopeet ge03
brg prreriion
Lyne
Figure 19: Pre-training of transformers on mixture of long and short sequences encoding the parity of d bits under "all" type noise for 7=0.01 and 7=0.1. Left: Test accuracy during the course of pre-training. Center: The probability that a model's generation length is equal to the maximum length present the training distribution (which equals d). Solid lines correspond to greedy decoding, while dashed lines correspond to sampling with temperature 1. Each color denotes a training distribution with a separate mixture coefficient Peot- Figure shows average and 1 standard deviation across 3 seeds. Note the change in the legend from the noise-less figures. Right: Test accuracy with greedy decoding at the end of pre-training (50k iterations) versus mixture coefficient Poot. The red dashed line corresponds to the critical threshold. Each bullet is the median of 3 runs.
All tokens noise In this case, there is noise in each token of the output that gets amplified ("snow- balls") in the case of long reasoning chains. That is, ,1,-..,, Jdd ~ (1d(1/2) and
($1, E2...,....... .....,,,, Ea, <EOS>), w.p. Peots
(21)
(Ea, <ROS>), w-p. 1 Pent:
I1 w.p. 1 17, and & = [1: w.p. 1 17. for i It is clear that this 1- w.p. 17.
where E
T1, w.p. 17.
noise disproportionately affects the long reasoning chains. The critical threshold of pre-training here will change, as the noise affects the distribution of the first and second output tokens. Additionally,
the pre-training accuracy under sampling with temperature will now equal PP
Ea = II Di +
1-Peot
1+(1-2n)* + 1-Peat 2 Poot (even sucs in d trials] + 1-Peet Poot 2
(1-2n) PPot- We
present pre-training simulations in Figure 19 for 1) E {0.01, 0.13. We confirm the analytical calcu- lation for the pre-training accuracy under sampling, while we observe that when the noise level is moderate (1=0.1) pre- pre-training seems be failing to induce any correlation with the target. How- ever, and perhaps surprisingly. post-training with reinforeement learning can still lead to generaliz- ing models. Post-training experimental results for Poor=0.22. and rj E (0.01, 0.1} can be found in Figures 20 and 21 respectively.
2
In both cases, we observe that there exists combination of post-training algorithm and sampling temperature that yields a generalizing model. We observe moderate increase in the sample com- plexity in comparison to the noise-less setting, and a lower tolerance to large sampling temperatures, even with "perfect" (i.e., Feot) verification. We defer a theoretical understanding of these interesting phenomena to future study.
C.4 NUMBERS MULTIPLICATION
Here, we describe the setting and our results on the task of multiplication in more detail.
It has been reported that even LLMs as capable as GPT4 (Achiam et al., 2023) have struggled with accurate multiplication of 4-digit numbers (Shen al., 2023) without the use of external calculator program We encode the task in sequences of characters, including the digits of the multiplier and multplicand (both in reversed order). the execution trace of the grade-school multi- plication algorithm, and the final answer. Prior work has reported that reversing the input order helps transformers with absolute positional embeddings generalize better on arithmetic tasks (Shen et al.) 2023). As stated in the main text, we construct datasets that contain the full sequence, including the
Otherwise, the problem reduces to copying input (from the calculator program) to output and transformers well-suited for it (Jelassi et al. 2024p.
32

---

Published as a conference paper at ICLR 2026
Parily, Poe- 0.25, Rost traininggSs d-25, model-gpts depth-i, 12n, fread=322, noise-(all, 0.01)
Lengthe (ete), temp=11.0
Length
Lengthe temp==1.p
7e=11.0
Length
Lengths temp=]
Accuracy
623
Lemgth
RL Ioe ratinns
AL heraations
Figure 20: Post-training of transformers on mixture of long and short sequences encoding the parity d bits under "final" type noise (v/0.011) with various RI, methods and generation temperatures (TR.) Left: Test accuracy during the course of post-training with greedy decoding (solid lines) and sampling with temperature (dashed lines) for various RL methods. Figure shows average and I standard deviation across 3 seeds. Right: Length of generated response (sampled with temperature during course post-training run (GRPO with end-to-end reward) for 640 test inputs, after 20k pre-training iterations. Note: The sample size of each RL iteration differs amongst the RL algorithms: n=64 for GRPO, REINFORCE and n=3,200 sequences for STaR.
Parity, Pss: Pest-rrainngggokk d=25,, model-gat2, depth=4, desss=120, Phaas=32, 32, noise=(all, 0.1) Accurary
(GRPP tempm]
190
m0.5
Length
Accurary
(GIPP tempm]
Accureey 19
1.0
Length
Accurary
tempm]
1.5
Accuracy
Length
RL Ittrrttions
RL Itrrations
Figure 21: Post-training of transformers on mixture of long and short sequences encoding the parity of d bits under "all" type noise (1=00.1) with various RL methods and generation temperatures (TRL)- Left: Test accuracy during the course of post-training with greedy decoding (solid lines) and sampling with temperature (dashed lines) for various RL methods. Figure shows average and I standard deviation across 3 seeds. Right: Length of generated response (sampled with temperature during the course of post-training run (GRPO with end-to-end reward) for 640 test inputs, after 20k pre- pre-training iterations. Note: The sample size of each RL iteration differs amongst the RL algorithms: n=64 for GRPO, REINFORCE and 3,200 sequences for STaR.
33

---

Published as a conference paper at ICLR 2026
chain of thought, with probability Peot and just the problem and answer with probability Peot for various values of Pcot see Figure 22
We pre-train randomly initialized GPT2 transformers either 4 x 4, 5 5 x 7 datasets for multiple epochs over a training set of 808k examples using the next-token prediction objective. At various checkpoints, we switch to reinforcement learning (with GRPO) and compare perfor- mance during pre-training vs post-training. Full experimental details are provided Appendix B.2 Figures 23 to 25 show pre-training results for all values of E 4, 5, } and Poot E {0, 0.1, 0.25, 0.5, 1.0]. When Pcot=1.0 (i.e., the algorithm's trace is included in all sam- ples), models during pre vreetrainnng quickly learn to answer correctly with long responses, leveraging the intermedi- ate calculations present in the training data. By contrast, when Poot is small, models struggle to generalize: per- formance under greedy decoding remains close to 0% throughout training. some runs, we observe sudden in- crease in accuracy under sampling with temperature 1.0, jumping to approximately Peot x 100%. This outcome is consistent with the parity experiments: the objective en- forces length calibration and, once the model learns from long sequences, its output becomes mixture of short ran- dom guesses (with probability of success 10 10-2xn) and long correct responses (with probability of success 1). Similar to the parity experiments, we observed unstable behavior under greedy decoding when training with Peot close the critical threshold (which is close to 0.5 here).
Numbers Multiplication
6 ****
5
Short: 44d4
5
3
2
4 3 1.
II
Figure 22: Example of 4x4 sequence, encoding 2365#4374 (digits appear in reverse order). Top row: long format with full sequence. Bottom row: short format without chain of thought.
After switching to reinforcement learning (Figure (left) for (n=5, Pr.t and Figure 26 for the rest of the values), the model's accuracy improves rapidly and the average response length in- creases, provided the pre-training checkpoint has developed some correlation with the target. These observations hold for most values of and Peot For the most challenging setting
considered, (n, Pcot) ) = (7, 0.1). the model failed to learn from long sequences under any of the random seeds, after 38 GPÚ hours, preventing post- trst-training from being successful in this case. Note that this section provides an example of a task where the capability obtained during RL (accuracy under greedy decoding on task of multiplication) appears to be almost absent during pre-training (accuracy goes from 2 to 100%)- for example, when Peot=0.11
Miscellaneous observations Notice that the length of the response is very large for all runs at initialization, as the model has not learned to generate the end-of-sequence token. For the 4 x 4 task, we observe success after RL for at least one checkpoint for values Peot- If post-training starts early, the model is not able to benefit from reinforcement learning. As anecdote, some of the early checkpoints in Poot 0.5 (top left corner) seemingly failed to yield a generalizing model, evaluation protocol and
but found that this was actually due to slight mismatch between the reward function; the model did learn to generate correct answers to the multiplication questions but did not learn end the generation afterward, generating tokens until the maximum limit was exceeded (indeed, see in the plot that the average length is much larger than in the rest of the checkpoints). The evaluation code, however, deemed such responses incorrect, and thus RL was recorded as a failure. This illustrates the various ways that RL might not go as planned if applied to an "immature" checkpoint.
C.5 MATH BENCHMARK
MATH results with the instruct version of Llama 3.1 8B and Llama 3.2 3B models are presented in Figures 28 and 29 As in the GSM8k experiments, we observe that SFT on a dataset with small Poot does not enable the model to generalize well SFT with full chain-of chain-of-thought data. During RI., checkpoints with smaller Poot benefit the most in terms of relative performance gains, which are accompanied by an increase in response length. However, we also find quantitative evidence that this improvement is sometimes (Figures 34 and 33) (but not always Figures 36 to 33) due to out-of- distribution gains, as model completions can differ significantly from the SFT dataset. We believe
34

---

Published as a conference paper at ICLR 2026
Bumbers cattl dall saad-42
Lengle Carpee)d))
Keselty of eepinn Sampe1
Anruncy
Fresshity orgenen
Avenda
Dreien Femricn
En-aal cenaler
Thes Accursyy
182
Treesbi by of orginont hamp-1
AcAMPec
by of orgonent bempl1s
Avensa
angh 13
Bambers Mutppication 4ol, 30nd+43
Lenotle Canp-).01
Sempel
Annumcy
Fressh
ThE ACAMSsy
122
Tresabi of orpanott bamped
Pysslaky oegioen Sempe1
angh 13
Narben Walbo tcaiin 4e4. 3000e44
"ccuraay
1200
Probash boy bhe tamp-1
Robul by bngolee thrree
lant beicas
femoad)
6ngth33
Figure 23: Pre-training GPT2 transformers mixture of and short sequences encoding the multiplication of 4-digits numbers for 3 different seeds and for various values of Peot- Left: Test accuracy during the course of pre-training. under greedy decoding and sampling with temperature 1. Center, Left: The length of the greedy response for 1000 test samples during the course of pre-training. Center, Right: The length of the sampled response (temperature of for 1000 test samples during the course of pre-training. Right: The probability that the length of model's s autoregressive generation equals the maximum or minimum length present in the training distribution.
35

---

Published as a conference paper at ICLR 2026
Bumbers cuttl Sc. sadd-42
Lengle Campeedd)
Kesdlty of eepinn Sampel
2
5
Jrrunyy
Fresshity orgensn
Avenda
LAND
En-aal benties Treesbi by of orginont hamp-1
Thee Accurbcy
182
Eagi
AccMrecy
byssbalyy of orgónett bampll
Avens
angh 12
Bambers Murpliction 565 3unde43
Canp-).01
Sempel
Annumyy
Fressh hame
Legps
hentier Treseb by of orpanort bamped
The Acuussy
122
orginen Sempedt
Avendz
angg 13
Narben Mallo tcation 545, accde64
Auturuursy
1200
Probash bopsher tampel
T- 1
boLacurE boy
Robul by bngolee tarred
Lane's teicas
bnghh35
Figure 24: Pre-training GPT2 transformers mixture of and short sequences encoding the multiplication of 5-digits numbers for 3 different seeds and for various values of Peot- Left: Test accuracy during the course of pre-training. under greedy decoding and sampling with temperature 1. Center, Left: The length of the greedy response for 1000 test samples during the course of pre-training. Center, Right: The length of the sampled response (temperature of for 1000 test samples during the course of pre-training. Right: The probability that the length of model's s autoregressive generation equals the maximum or minimum length present in the training distribution.
36

---

Published as a conference paper at ICLR 2026
Bumbers catto TV7, oadd-42
Lengle Herpp110
bargshat Sempme
Arouncy
144
Smoshd iirgitu pamp-ll
Aventa
LANSS
The Accurscy
Tronshd by langinen hampet
larponort bampee
Avensa
lngh
Bumbers Mutplication TVT 3onde43
dorosttt Sempw).
1200
Drocahi inreiiee
The Accurscy
Tropsbt of largginnnt hamper
PM
borgshan Sempei
Avendz
Narben Mdlbb lcaiion 767.
Trye ingther: (tmp-1
andies Trobel by of Vngtiee: Hempei
Figure 25; Pre-training GPT2 transformers mixture of and short sequences encoding the multiplication of 7-digits numbers for 3 different seeds and for various values of Peot- Left: Test accuracy during the course of pre-training, under greedy decoding and sampling with temperature Center. Left: The length of the greedy response for 1000 test samples during the course of pre- preitrrin Center, Right: The length of the sampled response (temperature of 1) for 1000 test samples during the course of pre-training. Right: The probability that the length of model's s autoregressive generation equals the maximum or minimum length present in the training distribution.
37

---

Published as a conference paper at ICLR 2026
Numbers Multiplication, 4x4, model-gpt2, Pot 0.5
Numbers Mutiipicaaioo, 4x4, model=gpt2, Powe 0.25
Kocurecy (NJ ai gieedyy peeel--derddree
Rocuracy d0000
62-
0.0
3.3
04 04 0.5
30
6.3
32
64
3.6
CA
10
Arsrape Secddeg
Ausage Langh
Dun Requences
Erain sequances
nac--oaan..o o1ly
reot/ictan 6
nodit only
cota frr
Numbers Multiplication, 4x4, model=
Numbers Multiplication, 5x5, model= gpt2,
100
Accurso Eccoddre
3422
Accurery 081
25-
Avstage deccddmng
Avstiga Langen
42-
Numbers Multiplication, 5x5, model-gpt2, Puot 0.1
Numbers Multiplication, 7x7, model: gpt2, Pox 0.5
M greedyy decodng
Accurery (4)
ACCUICC (N)
0.0
6.5
13
15
25
3.5
3.0
3.5
3.0
Aaarage Lenge
Resrape Lonath
22-
2,
4 20. Scgucnccs
caan ScRUCnoos
neniookarert only
resticher
mod-- 0nly
wAw.AAaaa Our
Numbers Multiplication, 7x7. model=gpt2, 0.25
Numbers Multiplication, 7x7, model=gp12, 0.1
Steedy deroudio 10
Accurer (86)
Accuracy (6))
00
0.0
10
25
0.5
15
2.5
3.5
120
-250
Nuerap lengle
Under geeeey derod
Figure 26: Pre-training and post-training combined of a randomly initialized GPT2 on num- bers multiplication for various values of numbers of digits n and fraction of chain of thought data Peot in the mix.
38

---

Published as a conference paper at ICLR 2026
GSMBK, Lame 32 38 (base), 0.0
GSMBK, model=Llama 3.2 38 (base).
20006
$2000
250900 aEQuEnces
10000
38.6
GSMBK Lame 3.2 38 (base), Post = 0.1
GSMBK, model=Liama 3.2 30 (base), 0.9
50.30
up.
49.38
9
1230
1223
12003
GSM8K, Lldea 3.2 3B (base), 1.0
100
5.
10
20
23
Test Accurasy (5)
tempp=1 1.0
40
Aoeerage Length [lokens)
350
502-
200
50309
100000
1500.0
200000 250000
train sequences
checksoints
Figure 27: SFT and GRPO of Llama 3.2 3B (base) on the GSM8k dataset for various values of Poot- Some observations: For Pco:=0, only two checkpoints succeed to generalize during RL. Interestingly, the length does not grow in these cases. The model obviously leverages some kind of pre-training data. For =0011, we observe length increase for almost all checkpoints and accom- panying performance gains. When the length does not grow (SFT epochs and 10), test accuracy platcaus to small value (in comparison to the other checkpoints). RL consistently induces gener- alizing models for larger values of Poot. We observe that test accuracy and average token length are very similar across these checkpoints.
that these models were heavily finetuned on the MATH dataset prior to release, introducing additional confounding factors.
39

---

Published as a conference paper at ICLR 2026
MATH, model-llsma.. 36 (instruct), Pox
MATH, model=Llama3.1 38 (instruct), Re
8
R
The Prc.Arcuncyy
N
2000
400E
0006
E3X
13000
220000
34.00
300093
200%
400ce
$2.00
30500
LOCK
MATH, modelmtJmme3.1 finstruct), 8avm0.05
MATH. modelnLlama3.1 (instruct), Aaa
Ted Aocuu (88
359
600c
46890
C5000
12300
LICK
146300
28.
MATH, model= 3B (instruct), Pur= 1.0
50. 10.
15
50. 20
op.
Test Accuracy
Median (talen))
400
100
250
200
1:0
40630
reñoo
80000
120030
train sequences
sft &
sft checkpoints
sft acc
1.0
kinth
1.0
Figure 28: SFT and GRPO of Llama 3.2 3B (instruct) the MATH dataset for various values of Pent- The situation is similar to the GSM8k experiments, yet RI. gains might be due to out of distribution gains see qualitative analysis. We early stopped the curves before the RL collapse, which we attribute to large learning rate.
40

---

Published as a conference paper at ICLR 2026
MATH, model-Llsma.. 86 (instruct), Pox
MATH, model=Llama3.1 88 (instruct), Re
2000
52.3
EXP
LOCKCE
100000
34000
20000
46900
60300
53000
LOCK
1200
14000
MATH, modelmtJmme3.1 finstruct), Aa.m0.05
MATH. modelnLlama3.. (instruct), Aar
28
Thd Aocuna (8)
350-
359
63000
33090
2203
1000K
MATH, model= 8B (instruct), 1.0
10
60.
Test Accuracy
(ermu]]
Median (talen))
400
100
250
200
1z0
20000
48503
mino
#8630
1000000
120000
140000
train sequences
5ft &
sft checipoints
sft acc
1.0
kinth
1.0
Figure 29: SFT and GRPO of Llama 3.1 8B (instruct) the MATH dataset for various values of Pent- The situation is similar to the GSM8k experiments, yet RI. gains might be due to out of distribution gains qualitative analysis. We early stopped the curves before the RL. collapse, which we attribute to large learning rate.
41

---

Published as a conference paper at ICLR 2026
GSM8K Completions (pr... 0.05, epoch=1, RL step=1)
Question:
In four years, Peter will be twice as old as Harriet. If Peter's age is currently half of his mother's age, who is 60, how old is Harriet now?
Ground truth cot: If Peter's age is currently half of his mother's age, who is 60, Peter is 1/2*60 <<30=30>>3> years old. In four years, Peter will be 30+4 << <<<0+4=34>>> years old. Since Peter's age in four years will be twice Harriet's age, Harriet will be 34/2 <<34/2=17>>> 17 years in four years. Currently, Harriet
<11==<<17-4=13> 13 years old. Answer: 13
Generation I; reasoning process here</think> <answer>32</answer> Reward: 10
Generation 2; reasoning process here</think> <answer>35<<ansswer Reward: 10
Generation 3; ...1</think>> <answer>24<<answe>> Rewand: 10
Generation 4; reasoning process here</think> <answer>34< Reward: 10
Generation 5; ?</think> <answer>4D</answer Rewand:
Generation 6; procassing online </think> <answer355<aanswen Reward: 10
Generation 7; reasoning process heres/thinks (naxer>18</ansses Reward: 10
Generation 2020-11-23111:49: htt.::21111......999000000000 <answer240</aneeee Reward:
10
Generation 9; thinking... </think> <answer>2l</answer Reward: 10
Generation 10; Mazriet's age right now <</think> <answer>2</answer Reward: 10
Generation reasoning process here</think> <answer>32</ansseer Reward: 10
Generation 12; 00t2000000-11100000000////hh---1
sanswer>2<</answer: Rewand: 10
13; 34/think> canswer336<<answer>> Reward:
Generation 14; reasoning process here</think> <at.wer>40</acsee //nswerr//// Reward: 10
Generation Harriet 4</think> <answerr84<<nnsweer Reward: 10
Generation 16; zeasoning process here</think> ranswer>4<</anskee Reward: 10
Generation 17; </think> <answer>25</enswer Reward: 10
Generation 18; Por Peter to be twice old Harriet, laariiet's age has to be half the current age of Peter. Start with the current age of Peter as years. Then half of half of mother's age=60, 1 60
60 x
120(0] Harriet's age-z - = 120 - 118 [1] Harriet's a age is 118 yeara, and 120 years.
No]eevre 2021 Rewand: 0
Generation 19; (please wait a few seconds) </think nk> (onswer>18</anse Reward: 10
Generation my reasanning heres/thinks eosxer>>8</anssee Reward: 10
Generation 21; </think> Canswer320</answer, Reward: 10
Generation 22; could not solve the problems/think> 1nnswer>13</anser> Reward:
Generation 23; 2016-03-11 20:54:37</think> <answer>36</ answer> Reward: 10
Generation 24; reasoning process here</think> Canswer>4D</answee Reward: 10
Generation
25;
IN FOUR YEARS, PETER WILL BE TWICE AS OLD</chink>
<answer's r>8 < answer> Rewar:: 10
Generation 26; He is fourty years now</thinks <answer>35</ 3 Reward: 10
Generation 27; 10</think> <anawer>25</anewe> Reward: 10
Generation Hazriet's current age?</think> <atswerr24</anssee
Reward: 10 Why is Harriet's s age h? Who is w? Who is d? <think>
Generation 29;
<answer>> 374 answer> Reward:
Generation 30; </think> <answer>36</anseee Reward: 10
Generution 31; include forgive-me. fre/think <answer>32</answee Rewand:
Generation 32;
Ernesto thought about reasoning process here </think>
<answer> / answer> Reward:
Figure 30: Llama 3.2 3B completion on a GSM8k prompt for Poot 0.05, SFT cpochss RL Most completions are short (following the SFT format of and inaccurate.
42

---

Published as a conference paper at ICLR 2026
GSM8K Completions (Pcos 0.05, epoch=1, RI. step=58)
Question: Mrs. Watson is grading 120 final exams from her American History class, On Monday, she grades 60% of the exams. On Tuesday, she grades 75% of the remaining exams. On Wednesday, how many more exams does she need to grade?
Ground truth cot: On Monday, Mrs. Watson grades 120 x 60%
11120"60*001722>>3 final exams. At the end of Monday, she has 120 72 <<<120-72=48> >48 exams left to grade. On Tuesday, she grades 48 x 75%
<=<48+75*.01=36>>> final exams. On Wednesday, she has 48 36 <<48- 36=12>>> > 12 exams left to grade.
Answer: 12
Generation 1; </think> <answer>30</anske Reward: 10
Generation 2;
She graded 608 of 120 72 exams on Monday conday.</thinbl </think>
fanswer>244 24< answer> Reward: 10
Generation 3; </think> <answer>15</answez Reward:
Generation 4;
Let's get variable for Mrs. Watson's total number of exams. </think nk> <answer>30</ansses Rewand: 10
Generation 120 final
left grade, 30. She grades 608 of them on Monday, which is 18.</think> <answer>12<<answer> Reward: 100 Generation
Mrs. Watson needs to grade of the exams</zaink>
<anawer> answer> Reward: 10
Gemeration 7; ... </think> <anawer>30</anaee
Reward:
Generation 8; She's finished 60% of the total. How many are left? We have 100+60=60 exams left. Next, on Tuesday, she grades another of those. She now has finished 60+75=135 out of 120 exams. So hox many exams are left? There are 120-135-15 exams left.</think> <anawer24<</anawe Reward: [...]
608 completion rate of all exams. On Monday, 608 of all grades were given. This means that 608 of 120, or 12, were graded so far. Thus, 48 have yet to be graded. 758 of the remaining exams that haven't been graded will graded Tuesday, There then 48tstimes; 75% 36 more exams to be </aded.</think> <answer>36<<ansser Reward: Gemeration ...</think>> <answer>10</answee> Reward: 10 Generation 24; 120</think> <anawer>30</anawet Reward: 10
Generation 22; There
good thing worked through so many examples last night. </thik</think>> <answer>3D</anseee Reward: 10 Generation 26;
Generation That's really easy problem.
She has graded them at rate of and then of them remaining, so she will need to grade... </adbn..//think> <anawe:>30</anawer: Reward: 10 ...once she grades 608 of the exams, she grades the remaining </tthink> kanswer>30</answer: Reward: 10
Generation
27;
Generation 28; Well, we know that we have 120 exams that need to be graded. Mrs. Watson graded 608 Monday that 60 exams. That means that she graded 60 120 = 60 exams on Monday.</thinis. <arswer>30</anssee Reward:
Generation Let's see </tt../k <answer>15</answe>>> Reward: 10
Generation 30; ....<//tinkk <anawer>15</answee Reward: 10
Generation ... </think> <ansver>20<<answer: Reward:
Generation 32; Nee asked how many more exams she needs to grade...wish she had thought this one.. </thinkk Smnwer>><//nswer Reward: 10
Figure 31: Llama 3.2 3B completion on GSM8k prompt for Peot = 0.05, SFT epochs=1, RL. step=58. The model starts generating longer responses that mimic the long SFT demonstrations, while continuing to respect the SFT answering format.
43

---

Published as a conference paper at ICLR 2026
GSM8K Completions (Pcos 0.05, epoch=1, RI. step=118)
Question: Ethan is reading sci-fi book that has 360 pages. He read 40 pages Saturday morning and another 10 pages at night. The next day he read twice the total pages on Saturday. How many pages does he have left to read?
Ground truth cot: Ethan read total of 40 pages 10 pages = < <<<++000=00>00 on Saturday. On the next day, he read total of 50 pages x <<50=>=100>> 100 pages. The total number of pages that Ethan had read is 50 pages + 100 pages <<50+100=150>>> 150 pages. Therefore, Ethan has 360 pages 150 pages <<360-- 150=210>>200 210 pages left to read. Answer: 210
pages
Generation I; Ethan was reading sci-fi book that had 360 pages. Ethan read 40 pages or: Saturday morning and then another pages on Saturday night. far he had read 50 pages Saturday. the next day be read 2+50 100 more pages So far he had read
100 150 pages So the next day he has read
total of 150 pages.Ethar has 360 150 210 pages left to read.</think
<answer>210</ ///a//wr>>>>> Reward: 100
Generation 2; Our book had 360 pages in total Rthan read 40 peges Saturday 30 there are 320 pages left. Ethan read 10 nore pages at night so there are 330 pages left.
if read one day on Saturday then be must have read 25 on Sunday. So now there are 320-25 295 pages left </tink/ kanswer>295</ansker Rewand: Generation 3; Ethan is reading sei fi book that has 360 pages. He read 40 pages or Saturday morning and pages at night.
far he has read 50 pages in days.
the total numbar of pages he read on Saturday was 40 10 = 50 pages. The next day he read twice the number of pages he read on Saturday. So the number of pages he read on Sunday was 50 100. So the :otal number of pages he read Saturday and Sunday was 50 + 100 150 pages. So the total number of pages he read was 50 100 150. so the remaining rumber of pages were 360 150 Rewand: 100
210 pages. </think> <answer>21D</anske [...]
Ethans book has 360 sages.Ethan read 40 pages on Saturday morning. Ethan read additional 10 pages night. Ethan read Saturday 50 pages - 40 + 10 50.8thans book has on the next day twice as many pages as on Saturday. on next day Eshan would
Generation 30;
read
100 210 pages he had left to Btges. book had a Total of 360 50 read.</think> <answer>210<</ansees Reward: Generation 31; Ethan has 360 page science fiction book. He read 40 pages Saturday morning and then 10 pages on Saturday night. So far he has read 50 pages. The next day he zead 2 times 50 which equals 100 pages. so he has read a total of 150 pages. Se he has 360 150 210 pages left to read.</think> canswer>210</ansven Generation 32; Rthan is reading sci-fi book that has 360 pages. He read 40 pages Saturday morning and another 10 pages Saturday night. So far today he has read 50 pages. The next day he read times 50 which equals 100 pages has read se today and the next day he read total of 150 pages. Bo far he has read 360 minus 150 which equals 210 pages. So he has 210 pages left to read.</thinko <answer>210</ answer>> Reward: 100
Reward: 100
Figure 32: Llama 3.2 3B completion on GSM8k prompt for Peot = 0.05, SFT epochs=1, RL. step=117. After only 100 RL steps, the overwhelming majority of model generations consist of long, elaborate responses that give the correct answer.
44

---

Published as a conference paper at ICLR 2026
GSM8K Completions (pcos 0.05, epoch=1, RL step=4676)
Gabe has three songs on his playlist. "The Best Day" is 3 minutes, "Raise the Roof" minutes, and "Rap Battle" minutes. How many times can Gabe listen to his entire playlist on the 40-minute ride to his wrestling match?
Question:
Ground truth cot: If you add all the songs together, Gabe's playlist + 3 <<3+2+3=8>>> minutes long, During the 40-minute car ride, Gabe can listen to his 8-minute playlist 40
40/8=5>> > times.
Answer: 5
Generation 1; Be have the following information: Gabe has three songs his playlist.
"The Best Day" is 3 minutes
Raise the Roof" is minutes
"Rap Battle* is 3 minutes
total length of three sengs minutes So the rumber of minutes in the ride to his wrestling match is 40 mi minutes So the rumber 40-minute ride 40 = plays of the playlist
minutes
that the
playlist can played during the
</think> <answe>>5</enawer> Reward:
Gemeration 2; We have the following information:
Gabe has three songs on his playlist.
"The Best Day" is 3 minutes long
"Raise the Roof" is 2 minutes long
"Rap Battle" is 3 minutes long
So the total length of the three songs on Gabe's s playlist - minutes for "The Best Day"
2 minutes for "Raise the Roof"
minutes "Rap Battle"
So the total length of the playlist - 2 3 8 minutes
So the length the ride to wrestling match ninutes
So the number of times that the entire playlist can fit into the 40-minute ride 40 - total plays the playlist
</think> <anawer>5</answer Reward: 100 [...]
Generation 32; We have the following information: Gabe has three songs on his playlist.
\The Best Day" is 3 minutes long
\Raise the Roof" is 2 minutes long \Rap Battle* is 3 minutea long So the total length of the three songs on Gabe's s playlist - minutes for Day"
2 minutes for /Raise the Roo
minutes for \Rap Pattle"
So the total length of the playlist = 2 minutes
Gabe's s ride his wrestling match lasts 40 ninutes.
So the number of times that the entire playlist can fit into the 40-minute
ride 40 /
complete plays of the playlist
So the the 40-minute ride to his wrestling match 5 plays
how many times Gabe lister. to his entire playlist
</think> <anawe>>S</answer Reward: 100
Figure 33: Llama 3.2 3B completion on GSM8k prompt for Poot 0.05, SFT epochs=1, RL step=4676. After 1000s of RL steps, model responses cease to resemble SFT data and are much more verbose.
45

---

Published as a conference paper at ICLR 2026
MATH Completions (pest 0.01, epoch==0,, RL step=1)
Question: In triangle ABC, AB = 13, BC = 15 and CA = 17. Point D) is on AB, E is on BC, and F is CA. Let A/ = p- AB, BF= = BC, and CF =r =r-CA CA, where p4 ς, and positive and satisfy ++ 2/3 and p² + + the area of triangle DEF to the area of triangle ABC can be written the form m/n, where and n are relatively prime positive integers. Find +
2/5. The ratio of
Ground truth cot: [asy] arbitrary values, I couldn't find nice values for pqr please replace if possible */ real 0.5, q 0.1, 0.05; /* arbitrary values, couldn't find nice values for pqr please replace if possible - */ pointpen = black; pathpen = D=A+p*(B-
linewidth(O.7)) black; pair "DD"::""""""""""A"""DM"""""""""" -DDA DDMMMI""""" C(,)) A), E-B+--B---BB,,,,,- F=C+r DDDMP""" cycle); D(D(MP("D DDDMP(""" DDD
[...] denote area; then the desired value is (FNN)""""FFEW)ccyycye" [/asy] We let
DEF
[ABC]- [BDE] (CFF)
[UBC]
[ABC]
Using the formula for the area of a triangle Labsiin C, we find that ABC
IADF]
(1-- /CAB
p(1 - r) and similarly that BDE
q(1 p) and
F.AB.AC-sM sin /CAB
ABC
[CEP] [ARC] = 9). Thus, wish to find
DEF [ABC]
= 1 [ADF] [BDE]
CEF
[ACC] [ABC]
[ABC]
=1 =1 p(1 - 7(1 p) r(1 - 4)
(pq + qr + (p (p+ +)
We know that p+t-+- = and also that (p+q+r)² = p² +"
-3
pq + Qr rp
2
45
Substituting, the answer 45
+
45
and
n
61
Answer: 61
Generation 1; </think> canswer>103</arssear Reward:
Generation 2; </think> <answer>103</ /answrr>> Reward: 10
Generation 3; </think>
Reward:
Generation 4; </think> kanswer>103</ansue Reward: 10
Generation 5; DSE2NIT</think: <answer>123</anssee Reward: 10
Generation 6; </think> <answer>103<<annwer>> Reward: 10
Generation 7; ></think> <enswe>>103</annsee Reward: 10
Generation 8; forcedettemll.ant///hint <anawer>1214/ answer> Reward: 10
Generation 9; thinks/thinks <onswer>103</answee Reward: 10
Generation 10; </think> <answer339</answert Reward: 10
Generation 11; <think> <chinb></thinkk</thhii
Canswer>137<c anseers Reward: 10
Generation 12; <answer>>13<</nsserr>>> Reward: $
Generation 13; </think> <answer>103<<anssee Reward: 10
Generation 14; </think> <answer>937</anse Reward:
Generation 15; ></think> <answer>101</anssee: Reward: 10
Generation 16; </think> <answar>103</anawk Reward: 10
Generation 17; ></think> canswer>1<</answe Reward: 10
Generation </think> <answer>244</anssee Reward: Reward: 10
Generation 18; comnands- i</think> <answer>9</answer
Generation 20; </think> amsser>103</ansxer> answer> Reward: 10
Figure 34: Llama 3.1 8B (instruct) completion on MATH prompt for Peot 0.01, SFT epochs=20, RL step=1. The model generates short inaccurate responses.
46

---

Published as a conference paper at ICLR 2026
MATH Completions (Poot 0.01, epoch=20, RL step=181)
Question: Let b, c be distinct complex numbers such that
b
k.
Find the sum of all possible values k. Ground truth cot: From the given equation,
a = k(1 (1 - b):
b k(1 1 c):
k(1 a)
Then
a k( b)
k(1 - k(1 c))
k(1 k(1 k(1 a))).
Expanding, we get 0h + a 3. + 122 = 0, which factors as
(h: k + 1)(kk
If ak + a - = 0, then a = k+1 in which case c k+1 This not allowed, as a,
and are distinct, so k22 k 0. The sum the roots 1. Note: The roots of k.2 = 0 are
1iv3
2
For either value of k, we can take a 0, b 1, and k. Answer: 1
Generation I: ></think> <answer>Pron the given equation, we have: a
k(1 b).b k(1 c).c h(1 a). Substituting b k(1
into
k(1 b)
gives: = k(1 - k(1 - c)) = - lPc. Substituting
= k(1 b)
into k(1
gives:
k(l k(1 -(1-))) b)) k 12 + k2b. Equating
= k- 12 + Rec and
k =++ 12b
get: kfe = kb so
b.
Substituting
b into
k x2 Ae gives:
22 k.b.
A1so, we know that = - k2 + Ka. Equating = k - ka + lAb and k 12t Aaa we get:
= a. Since a, are distinct, we have
contradiction. Therefore, the initial assumption that there are distinct complex numbers a, is incorrect. There can be only one value of k. Adding the equations
k 12 b26,b - k- +Pa,c
k
k k² + 12a,
we get: a 3k(1 - 0. Solving for h, get:
= 3k 312 + la + + +). Equating 3k 312 = 0,
k 1. We are only interested in non-trivial solutions, so the sum of all possible
values k is: 0+1=1. </answer> Reward: 100
[...]
Generation ></think> <answer>
Therefore, t
=
be any value that the equation I
k holds. It can be shown that the sum of all possible values of ④ is </answer> Reward: [...]
Figure 35: Llama 3.1 8B (instruct) completion on MATH prompt for Poot = 0.01, SFT epochs=20, RL s1ep=181, The model starts generating longer responses, which however do not respect the SIT format. This behavior is likely due to out of distribution generalization.
47

---

Published as a conference paper at ICLR 2026
MATH Completions (Peot 0.01, epoch=1, RL step=1)
Question: In triangle ABC, /B 600 and LC 45 ratio 3. Find
The point D divides BC in the
sin LBAD
sin LCAD
Answer: yi
Generation 1; Categories problem. Tnclines RC4/thinks <answer3y22//nnwerr Reward:
10
Generation 2; </think>
<answer22</answer> Reward:
Generution 3; </think> <answer>2<<answer> Reward: 10
Generation 4; </think> <answer>2</amsseer Reward:
Generation 5; </think> <answer> </ answer>>> Reward: 10
Generation 6; ></think> <answer>3</answerr Reward: 10
Generation 7; ></think>
<anawer2<</answr>> Reward: 10
Generation 8; </think>
<answera1</answerr Reward:
Generation 9; </think> <answer>2</answer Reward:
Generation 10; </think>
<answer> 24v3
/ </answer> Reward: 10
Generation 1I; </think> <answer>y3</anskee Reward: 10
Generation 12; </thinks <answer>v3, </answer> Reward: 10
Generation 13; No </think> /x
Reward: 10
Generation 14; </thinks <answer> v3 </answer> Reward:
Generation 15; </think> <answer>2 / Reward: 10
Generation 16; </think>
<answer> 2-13 end answer> Reward: 5
14y3
Generation 17; </think>
<answer>/7<<anavee Reward: 10
Generation 18; </thinks
<answer>7</answer Reward:
Generation 19; </think> amsser>3/3<//ansee Reward: 10
Generation 20; </think>
<answer>
</answer> Reward:
Figure 36: Llama 3.2 3B (instruct) completion on a MATH prompt for Poot = 0.01, SFT epochs=1, RL step=1. The model generates short inaccurate responses, which nevertheless adhere to the SFT format.
48

---

Published as a conference paper at ICLR 2026
MATH Completions (Peot 0.01, epoch=1, RL step=87)
Triangle ABC has sides of 6 units, 8 units, and 10 units. The width of a rectangle, whose area is equal to the area of the triangle, is 4 units. What is the perimeter of this rectangle, in units?
Question:
Answer: 20
Generation 1; </think> <answer>2
28</answer> Reward: 10
Generation 2; </think> (an5 28</answer> Reward: 10
Generation 3; </think> (4 44</answer> Reward: 10
Generation 4; </think>
<answer>44<<ansswer> Reward: 10
Generation 5; </think>
(answer22-6--8-- 28<answeer> Reward:
Generation 6; </think>
<ansswrr2((662)22222/aaaxexe Reward: 10
Generation 7; </think>
Reward:
Generation 8; </think>
22<anssers> Reward: 10
Generation 9; </think>
<answer>20 8</answer> Reward: 20
Generation 10; </think>
<answer>2 x 8) = 40 /answer Reward: 10
Generation 11; </think>
<answer>2(6+8) 28</answer> Reward: 10
Generation 12; </think>
<answer>2
324/answer> Reward: 10
Generation 13; </think>
<answer>2
10 4) = 76</answer> Reward: 10
Generation 14; </think>
<answer>2 [10 6) 32</answess Reward: 10
Generation 15; </think>
<answer>2:10 + 6) 32</answer> Reward:
Generation 16; </think>
kansmer>21<+12+))
52</answer> Reward:
Generation 17; </think>
<answer>2 - (a(r4)))336//aaeee /
Reward: 10
Generation 18; </think>
<answer>2
8</answer> Reward:
Generation 19; </think>
<answer>2(445)=1 B</answer>> Reward: 10
Generation 20; </think>
<a 8) 28</answer> Reward: 10
Figure 37: Llama 3.2 3B (instruct) completion on MATH prompt for Poot 0.01, SFT epochs=1, RL 87 The model learns generate some reasoning in the answer part of the response, which does not seem to be very successful strategy.
49

---

Published as a conference paper at ICLR 2026
MATH Completions (Pest 0.01, epoch=1, RL step=92)
Question: Simplify
10 2 + 1
log10 +
logs +
Answer: 2
Generation 1; //></hhin>> canswer>1</answee> Reward: 10
Gemeration 2; </think> <answer>1<<answer> Reward: 10
Generation 3; </think> <anawer>i</anawee Reward: 10
Generation </think> <answer21<<ansser> Reward: 10
Generation 5; </think> Sanswer>l<<answer> Reward: 10
Generation 6; </think> <answer>1</anssee Reward:
Generation 7; </think> <answer>1<<ansswr> Reward: 10
Generation 8, Recall the change-of-base formula for logarithma:
logg a
log. e log.
for any positive numbers a.b, and Also that x+1 zewrite each tern of the given expression using theae properties. After the rewriting, we obtair
log(n+
1)
log 15
log 10
log log log
log 15 + log log 10 log
where xe have renoved the denominators to make the argunents in the denominator each equal logıs 2, logio 3,and logg since the 2nd and 3rd terms are already in lewest terms, we rocus on simplifying the 1st cerm:
log 15
log 15
log 15 + log
log(15
logao
Thus, the given expression be written
logio 3
loga logg 5
log30 15
lugio 10 + log10 loga
By using change-of-baae formula, find that logao reault is the following: Reward: 0
log1s 30. The final
Generation 9; </think> <answera1</answer: Reward:
Generation 10; ></think> <enswer>lc/answer> Reward: 10
Generation II; </think> <answer>1<<answer> Reward:
Generation 12; </think> <answer>l</answer> Reward: 10
Generation 13; </think> <answer>1</anwwe Reward:
Generation 14; Let's use the change-of-base formula to simplify each of the
logarithmic expressions in the denominators of the fractions. </think> <
<answer>1 </aswer>>>>> Reward: 10
Generation 15; </think> <answer>1</answe Reward:
Generation 16; </think> <answer>1</answer Reward:
Generation 17; </think> <answer>l</answer Reward: 10
Generation 18; Using identity log, a
log, ab, expression be
simplified as logis 30
logi0
loge 30
then applying the change of base
formula, logs a log,
get
log30 log
then log + log
log
log 30
log 30
log 10
log
and after cancelling out common terms, we Eind
lox 15+ 10+lag and finally log
log(15-10:8) which is log 30
log 900 log 30
equal to log(3u2)
equal to 2, and the final value
log 30
2 </thikk <answer>2</answerr Reward: 100
Generation 19; </think> <answer>1</ansves Reward: 10
Generation 20; </think> canswer>1</answer Reward:
Figure 38: Illama 3.2 3B (instruct) completion on MATH prompt for Poot 0.01, SFT epochs=1, RI. step=92. The model starts exploring longer responses in the thinking part of the response.
50

---

Published as a conference paper at ICLR 2026
MATH Completions (Pest 0.01, epoch=1, RL step=102)
Question: Let G be the centroid of triangle ABC. If GA + GB + GC = 58, then find AB + AC BC
Answer: 174
Since G is the centroid of triangle ABC, by the properties of centroid, G divides the medians of AABC inte six equal segmenta, with GA = {AD, GB BE, and GC JCF, where D, E, and F the vartices of trisections of the medians AD, BE, and CF. Then by the Midaoint Theorer applied to triangles AGD, CFC, we have CD = 2(AC), GE = 2(BC), and FC = 2(CF), and since AD = AE, BE BG, CF CE, we get GD = 4AG2, GE? = 4BG², and FC
Generation
4CP where D,E and F are the midpoints of sides AB, BC, (HAD)* +(BBE) ([BE))+((CFT) = = + BE + CF
and CA. Now, GBB+GB++GGR Thus, 9 (GA4 GB CC*) AD BE 1FF =3 (GA" CB GC4) 3 58, then AB I AC BC = AD | BE + CF2, </thinks fact that the centroid divices each median in:o two segments with 2:1 ratio,
1onswer>174</nnswer Reward: 100
Gemeration 2; Using
have AG JAD, BG BE, and CG CFF where D, E, and F are the midboints of BC AC, and AB saspetiiely,</thhae <thinkeTherefore, GA2 GB2 + GC2 (AD2 + BE CF²). Since AD AC,
BE 1A and CF = BC, we have GA + GB + CC = AAAC + AB + BC
58.</think> <thinksIt follows that AC + AB + BC = = 58 = 132 </ </think>
sanswer>13<</asswe
Reward: 5
Gemeration 3; By the zelationahip between the centroid and the circumcenter of triangle,
know that G is the circumcenter of triangle ABC. Thus, AG = BG = CG, since the circumcenter is equidistant from the vertices of the triangle. Letting AG =G=CG- CG, we that
2 58
3r 58
2² 58/3.
</think> <anawer>58 58/3</answer> Reward:
Generation 4; By the configparser propesition, we know that GA? + GB + GC2 (AB* AAC2 BC*), so AB AC BC = 3-58 1.58-174</hhnnk <answer>174<aanswer
Reward: 100 [...]
Generation 17; By the formula for the centroid, GA GB2 GC AB*+BC2+AC (AB² + BC + AC*) = 58 * AB + BC + AC = 174 < /think > <anawer>174< < /
3
Reward: 100
Generation 18; By Apollonius's Theoren, AB AC#
2GA | CB²) and
2GA + GC3). Adding these together, we have AB AC BC = 2 + GB + GC2) = 2(58) = 116.</chnnk>
AB2 BC
2(GB2 + GC and AC? + BC
<answer> 2011//////werr Reward: 10
Generation 19; By the definitian of the centroid, we know that G is the point 1) in the 3-2-1 coordinate system. Also, by Stewart's Theorem, GA2 GB2 + = (AB² + 3C2+ CAP). Ne want to solve for AB + BC + CA3, so 3(GA2 + GB2 GC2) AB2 + BC CA2. 58, it follows that BC +
:
multiply both sides by and Since are given that GB"+ Rewand: 100
3(58) = =14.c/think> <answer>i7d (aae=>114///aseeeee Generution 20; Racall that the centroid divides the median each trianqle into segments whose lengths hold the weights of 2:1, with the longer length being towarda the vertex. Then, by the properties
medians, we know
GA + GB2 + GC = (AB² + AC2 BC
If know this, then the answer is
AB + AC + 2C = S(GA + GB2 + GC = 8(58) = 464
</think> <answe=>464</ answerr Reward: 10
Figure 39: Llama 3.2 3B (instruct) completion on MATH prompt for Peot = 0.01, SET epochs=1, RL 10ep=102, After 100 RL steps, the model generates long correct responses similar SFT samples.
51

---

Published as a conference paper at ICLR 2026
D USE OF LARGE LANGUAGE MODELS
LLMs were used at several stages during the development of this work:
1. for assistance with coding and figure creation (especially in the Llama experiments);
2. for checking grammar and syntax in parts of the paper;
3. for brainstorming different directions for the proofs of Section
4. for writing this Section.
E
PROOFS
Here, we present our theoretical analysis. For pre-training, we consider learning the mixture of long and short demonstrations of the parity task with linear (time-dependent) autoregressive hy- pothesis class. That is, we consider a series of d+1 linear models. For post-training, we consider reinforcement learning with the STaR objective and a chain-of-thougee correctness reward. The goal is to show two main learning results: a negative one during pre-training and a positive one after post-training.
Remark 4. Linear autoregressive architectures were introduced in (Malach] 2024) to demonstrate the power of autoregressive learning, independently of the specifics of selfaatentiion networks. Leverag- ing the equivalence between binary computable functions and binary circuits, Malach| (2024) proved that such simple autoregressive models can approximate and learn any function efficiently computed by Turing machine given dataset that contains appropriate "chain-of- oftn--ugtthhhht data.
We first cover some preliminary results. We then define our architecture series of linear models, and then define pre- & post-training algorithms.
F. PRELIMINARIES SETUP
We make use of standard definition and theorem for convex learning presented in (Shalev- Shwartz & Ben-David 2014). First, we recall the definition of a convex (and Lipschitz) learning problem, which based on Definition 12.12 in (Shalev-Shwartz & Ben-David, 2014).
Definition 1. (Adaptation Definition 12.12 in "Shalev-Shwarde & Ben Ben-David] 2014 Let H be
hypothesis set, Z be measurable instance set and H. x Z + E be measurable loss function. learning problem (H, 2, 1) is called Convex and p-1.ipschitz for p >> 0 if:
The hypothesis set H is convex.
For all E Z, the (partially applied) loss function ((:2) is convex and p-Lipschitz.
The original definition in (Shalev-Swwarzz & Ben-David 2014) about bounded hypothesis sets, but in our analysis we will use unbounded ones (as regularization term in the loss function will effectively bound the solution space).
The learning algorithm will be Stochastic Gradient Descent (SGD) for & regularized learning prob- lems (Algorithm 2 below). Note that the algorithm returns the average of the weights across T iterations, as it is typically the case in online learning.
Algorithm Stochastic Gradient Descent (SGD) for minimizing EzzD (!(w, 2)] + 3 /|/||2
Require: Integer T > 0
1: Initialize (1)
2: for 1, 1, ........ T do
3: Sample z ~ D
4: Set (t) = V w (wwe(((())
5: Set T)- Xt
6: Set (++1)
= (()) - ni (vit)_ + (w(())
7: end for
8: Output w
1 ΣT= w(t).
T
52

---

Published as a conference paper at ICLR 2026
We study strongly convex learning objectives, they permit tight bounds on the calibration of the output of SGD later on. Analyzing their convex (only) counterparts would have yielded weaker calibration guarantees see, for instance, Chapter 4.7 in Mohri et al. (2012) and the analysis of early stopped SGD in Wu et al. 2025).
We now state guarantee the output of SGD after T iterations for convex and Lipschitz objectives plus an additional l regularization term. The proof follows from Theorem 14.11 in (Shale-- Shalev-Shwartz & Ben-David (2014) by applying it learning objective and combining it with Markov's inequal- ity. We present it here for completeness.
Theorem 3. (Adaptation of Theorem 14.11 in Shalev-Shwariz Ben-David (014)) Consider a Convex, p-Lipschitz learning problem (H R 2,1) and a distribution DD over Z. For every 5 E (0, 1), if we run the SGD method (Algorithm 2 for minimizing LD.x(W) = Ez-D [l(w, =)] + Ž llwll? À > O, for iterations and with learning rate = 1 then the output w of SGD satisfies with probability at least 1 ô over the sampling ......... , ~ D:
Xt-
2p²
LD. (w) S min (w) +
(1+1
(22)
WEH
8XT
Proof. The objective (.D,(W): ED + >> 0 X-strongly convex. Let DD then v(t) Vwrol(((()) 2) and, from the Lipschitz condition, it holds: ||v(0)|| S p. Let the update direction be: g: = v(t) + (t), Then, we have for the weight vector (+++));
((+1) = w(( - 17 (v((0) + (wm(1))
= (!)
1 11_ 1
At (()
1
(23)
(())
At i=1
Hence, have || X[- (t 1)p G. which implies that for the update vector it holds:
gg. < p A 2p. From Theorem 14.11 (SGD guarantee for strongly convex functions)
in Shalev-Shwand & Ben- 2014), we get:
E E LLD (w) S min (D.(WW) + (2p)² wEH
2p² 2XT (1 (1+ min (.px(w)) + XT (1 + In (I+ITT).
wEH
(24)
Finally, apply Markov's inequality on the non-negative random variable LD, (w) minwen LD.((w)) Let 0 < 8 1, then by leveraging the above guarantee, we get:
P Pp
LD. (w)
min LD,2 (D,)(w) 27| < WEH
E Exp D LD, (w) - minwEN (D)
20 (1 (1+mn7)
X7
< 6,
T
(25)
then with probability less than 8: LD,A (W) min LD.> (w) 2
(1+In) X78
In other words,
with probability at least - 5 over 21,..., 2T ~ D, it holds:
2p2
LD.A (w) < min (w) + 8XT (1 + In7).
(26)
WEH
Remark 5. We assumed that the loss function is differentiable, but the proof goes through for any continuous function using subgradients.
Remark 6. It is possible to tighten the 8 dependency on the previous bound from @(1/8) to (In (1/8)) by using Azuma's instead of Markov's inequality.
53

---

Published as a conference paper at ICLR 2026
E.1.1 SETUP: DATA DISTRIBUTION, ARCHITECTURE AND LEARNING ALGORITHMS
Data As reminder, we consider parity mixture distribution. Let X' {+1} be the input space of d 2 2 bits and y = {=1, <EOS>} J be the output space of sequences, where <EOS> is a special symbol denoting the end of string. Let D(pco)) be distribution over X' V, parameterized by Poot E [0, 1], such that:
F1. Fd Rad(1/2)
(27)
(yl...., ......... y8+1) 2(11,
II n. <EOS>) + -1-36 2) ==1
Fzj <EOS> 1=1
II
where Z ~ Ber(Peot).
Model We consider an architecture that consists of d+1 linear models. For the first out- put position, we consider the hypothesis class Hi {x+ (W1,x) : e {#1)d, wi e Rd). We break the second decision into two parts and consider two separate linear classes: Hza = {x+ (Wza: (2(x)) + bza x {=1]d+1, W2a E R24+1 bbe E R and H2 {x+ (W2, (x))) xx E {#1]d+1, W2 E IR2d+1). where the first model decides between <EOS> and {+1} while the second model decides between -1 and +1, in case the first model did not predict <EOS>, For the rest of the positions, consider hypothesis classes Hi = fx+ (Wi, (x(x)) x E (#1]2d++11 W/ 12d+1-1]. The feature embedding defined as follows Φ1 x x [2] ... Fdti-1 1d+1-1#1
(=1]22th11 for < d. For position d + 1, the output is a constant function of the input, as the output symbol is always <EOS>. As learning is trivial in this case and, for not complicating the analysis any further, 2o
Id41-1
we assume access to this deterministic function o {+1] H <EOS> for which it holds: o(x) = <FOS> for any x E [+1]22.
Our final model is a function h that belongs to the linear autoregressive hypothesis class H defined as:
HAiR Hi H2a x H2... Hd-
(28)
Note that we can learn any sparse parity inside Hlin hence we argue that our definitions are not particularly tailored for the problem at hand. Given an h E H and the conresponding + models, W1, Ozo, W2, W3...., Wa. define deterministic autoregressive process h : R2 x Rd x 122++2 x IR 2d41 x .... x 1d3d-1 + 1-1, +1,6, <BOS>}* as follows:
h (1: sgn( (W1,
h2a 02a) diet sggn W2e: 02 h) (x. b2a)
Ozo = (wiar b2a),
e, if hea (x; O2a) = <BOS>,
he(x;; =
(x, o.w..
sgm
W2: 02
e, if hzo (x; O2a)
<E0S>,
ha(xx; wa)
sgn (( (ww. d3
x hx (x), (x) (x) o.w.,
e, if hea (x; 020) = <EOS>>
ha(x;;a))
sgm We. Ød x (x: w1). R(d-1) (x:Wa-1)))) o.w.
e, if hea (x; O2a) <BOS>,
hxt1(x)=
o(x), o.W.,
(x; {W1, Oza, W2,... wal)
(x: w1), (8),).))....,
(29)
54

---

Published as a conference paper at ICLR 2026
where dict {+1} - fe, <EOS>} is deterministic function (dictionary) that maps label -1 to the empty string and label +1 to the <EOS> token. Likewise, we define randomized sequence to sequence model h as follows:
1t (Ie-(wwis) (wix)) (
b1(x; ~ Rad
(2;(x; 02a) ~ dict
Rad
1 1te ( (x bya
0za (w lr2a))
if (x)(x)) = <EOS>,
(x(x;; w2)
Rad
1tc (x O.W..
c, if h2a <EOS>,
ha(x;; (x)))))) ~
(30)
Rad
14e (owm) O.W.
c, if hxa((x) = <EOS>,
ha(x;; we) ~
Rad
1te (M (K (x O.W.,
e, if hza(x)) <FOS2,
Ma41(x)
(x(x), o.W.,
[xx; [W1, O2a, W2. W wa)) = (11(x; (x; wi)...., 1x+1(x))
We adopt the convention that ae,3 = aB, i.e., concatenation with the empty string character e does not change string. We define the probability measure induced by as:
P [x(x: = v/
h2a 02n) <EOS> h (x(x;;))
1/
if |yl 2,
Th (y xx)
IP [xxxx vi] .... [axxx wa) Va (x -11 11.....
if [y] > 2.
fa-i(x; = ya-
(31)
For any other y (t11,c,KESS: is Tr. (y) x) =0
The two autoregressive models, h and h. correspond to greedy decoding and sampling with temper- ature 1, respectively.
Pre-training loss function During pre-training, the loss function is the next token prediction ob- jective, together with an le regularization term. Let (x, y) D(pcot). We define the loss functions corresponding to each linear model:
1. Position 1: (((w,, (x, v1)) In (1 e xwwx))))) + 2 w 1 21 > 0. In (1 + - u((( (h +
2. Position 2a: ((20) ((w2a bea): ((x, 81) (2))
Aan 2 21 - (2a)) with Aza > 0 and Tz2 1-1, U2 E (+1).
++,,
y/ = <BOS>,
3. Positions to d:
In (1 -ww. + A
1t) ((x, V/ ......., 10-,,,8111 -1-1),
with A for all sldd.
55

---

Published as a conference paper at ICLR 2026
+
BOS
+,
W.
w.
55
4
-1+
1,
350 1925 1ARA
Figure 40: An illustration of next-token prediction training (for d=5) with the linear (time- inhomogeneous) autoregressive architecture in the case of a short (left) and long (right) training sequence.
We included the regularization term in the definition of the loss functions. Then, we seek to solve the following optimization problem:
1l {lyl 23 (1() (x, + ((2a) ((w2ar b), ((x, M)), .52))
min
((x.y))((Peon))
W1, Was W..... Wid
+ 1 (1vl =d d 11 (1)) (x, + (wa, ((x, .i.-.. -dd-1) ,
(LIN-NTP)
Note that the regularization term can have a different coefficient for each parameter vector. Also, observe that Problem LIN-NTP corresponds d+1 binary classification problems with respect the logistic loss.
Pre-training Algorithm We minimize the previous objective with Stochastic Gradient Descent see Algorithm 3
56

---

Published as a conference paper at ICLR 2026
Algorithm 3 Stochastic Gradient Descent (SGD) for solving Problem equation LIN-NTP
Require: Integers Ti, T2a ....... ., Ta 0, Real numbers X A2as . Aa 0. ..... w (1)
1: Initialize w W 2a W2
(1) (1)
= 0, b2a = 0
2: Set Hlong =0
3: for 1........ ,,, do
4: Sample (x. y)
D (peot)
5: if S T1 then
6:
Set ? Aut
7:
Set w
(++1)
wi
ni V w/th wos)(t)) (w)" (x, v1))
8:
end
9:
if t S
T2a then
+1,
if <BOS>,
10:
Set 9/2
-1,
if 1/2 E {+1}
Set ?)
Aau
12:
Set 2a
[++1)
W
(t)
1i V w (2))
(f W 2b2 (1x, vil (22)
13:
Set (++1)
2a
bia
It Tbza
á ((2)) ((
w 2a
(e) (bx, V11 512)
14:
end
15:
if
| > 2 then
16:
Set tong = flong + for 1 2, d do
17:
18:
hong S Ti then
19:
Set Mionn Aebiong
Set w (Ponn
w, (1oong))()) 100) ( (Moou)) ((x,, y1. ..... ... 300 1), 311
W
21:
end
22:
end for
23:
end
24:
end for
25: Output W1
Ti Eth W
W1
W2a
Tza >tan
Wia Dan
T2a Eten bla
W2
W2
21=1
Post-training loss function For post-training, we consider the STaR algorithm (Zelikman et al., 2022). Recall that the STaR algorithm involves n reinforcement learning rounds, where each round involves optimization of next-token prediction loss model sampled responses (from the model of the previous round) that are correct according to some reward function. We use a reward rex X' y that assesses whether the whole sequence is valid:
d
d II
x4 <BOS>> }}
(xo y) {
F1, 11.
II 2. <EOS> y
(32)
Namely, at the ks round, we minimize the following objective:
1l flvl 2} (100) (11)((1,, (x, + ((20) (fW2.. b), ([x, 311 02))
E xx-Rad(1/223) y006 (00) (-Ix)
+ I + = d d (w)) (x,, + (2a) (1 b}, ([x, yn) (2)) +...
+... + (1)) (wa, ([x, "..., yd-1] ma))
reor Fxox(x, v) =
(LIN-RL)
where (R) is the model returned the end of the k'th round.
57

---

Published as a conference paper at ICLR 2026
Post-training Algorithm The algorithm consists of n rounds, where at each one minimize the previous objective with Stochastic Gradient Descent. At each round, we start optimization from a freshly initialized model at the origin.
E.2 WIYY IS THE CRITICAL TIRESHOLLD 1/3?
Before proceeding with the theoretical analysis of the pre-training and post-training stages, we ex- plain why should expect the critical threshold of pre-training to equal 1/3. Our argument is in some sense generic - see Lemma |
below.
Recall that from the definition of D(Peot), have for the distribution of the first token: y1 Zæz + (1 - 2) I[= Ts; where ~ Ber(prot). We have:
d
1y = +1 | = Poot 1 fx = +1] + (1 (1- Peot) 1
24 +1
(33)
+ 11
[ Iz +1
Poot
(1 Poot)
2
Any model that matches this conditional probability should be able to implement the parity function + Il 24. In absence of such parity feature, the conditional probability simplifies to:
I1 +
Poot
Pcot +
P vv +1 x
Peot
(34)
2
2
In other words, P Pyv æ1 x] Poatt1 1/2 for all pest Hence, for any model that matches this conditional probability, the "greedy" decision on the first token will be Ij, since 21 = arg (x P [y1
|x].
For the second position, we have:
P 1y22 21
x, 3/1 21]
P Ply2 31 91
31 x]
Peot
2pcct
(35)
P [/1
x
Pest +1
Peot +
and also P [y2 <ROS>
X. 1/1 -11] PP [y2 322
x,
an
]
1-Peot
Hence, long
Peoto1
as Pcot > 2poot -
Peot < 1/3, the "greedy" generation of a model will output the sequence
(I), <EOS>)
We formally explain now why all logistic regression models without a parity feature will match the conditional probability expression of eq. equation 34 The proof uses simple ideas from the Fourier analysis of Boolean functions.
Lemma 1. Denote the parity functions with support S S [d] by xs(x) II Let F
span (xs: S _ [d) & [d] be the class of functions spanned by all but the complete parity X(d) =
II 24. Then, if f = arg minjer Rad-R 22004, 1184.--Zayy + (1--2)1111 [log (1+ + te eu)(())] is the
output of logistic regression in class F, then it holds:
Poottt +1
Pr [v = +1|x] := 1+e-f
(36)
2
Proof. We write any function E F in its Fourier expansion:
f(x) Escla (xxs(x), where
cs = R2d (ftpg)/o_[[((((((0 (x)), Then, the objective becomes:
L (co.....qu-11) = E Ex- xx-Rad(1/2jo) u- +1--))11121)) -1) log 1 + - (= (37)
This strictly convex objective with respect to the Fourier coefficients. Denote by r7(x) =
[y = +1 x] = Peot
zr- + (1 - )cot) 2
[T #:+1
We calculate the partial derivatives for all
2
S C [d]:
AL dcg
Xx(x)e V Esscw (x)
=
F. x~
(1-2)
1 -V co (x)x)r(x)
(38)
Rad(1/)) (1
I+e Es'cin Estciageseeee (x) - n(x)) xs(x))
F
58

---

Published as a conference paper at ICLR 2026
Therefore, at the optimum, have 2d orthogonality conditions,
(S) 22))2/
Estcia (xxxsr(x) xx(
Rad (x(x)xxxxxx))x
for all S C [d]. In particular, if we denote (u)) I-e- then: oo 2.0 (39)
Pess
and a (S) 0 for S & {1). [d]). Observe that the full parity coefficient is unconstrained.
Due to symmetry, we can observe the form of the optimal solution (*(x = az1 + 3, then the optimality conditions yield:
E Rad(1/(122300
(40)
1 1
2
or:
=1 3) + +a0-a-a-0 - 9) =)=1
(41)
This implies that B = 0. From the second equation:
Poot
E Rxd
(42)
azı
2
which implies a In Root I-Poot
This the only possible solution as the objective is strictly convex.
Therefore, showed that:
(1 + Poot #1
Poot + 1
P.. = [+|x]
(43)
In P-Peot y"T
(1 Poot 10 + (1 -Poot)
2
1
The intuition is that "reasonable" architecture, such auto-regressive transformer, requires exponentially many iterations to form a parity feature (and therefore to "escape" F). As result, a pre-trained transformer will exhibit the same threshold.
E.3 PRE-TRAINING
We proceed by providing guarantees for each one of the models Wi, O2a. W2, W..... Wa indepen- dently, and then state and prove our main pre-training theorem for the behavior of induced models h, h, leveraging the per-position results. All results are stated for a mixture coefficient p E (0, 3/43. The upper bound needs to be greater than 1/2 to obtain the post-training result later on, but the specific 3/4 value is arbitrary.
F.3.1 POSITION
For the first position of the output, consider binary classification problem where X, {+1]d, y1 = [=1} with distribution (p) over X1 x 3 such that: II,..., Zd ~ (1)) and y = Zzi + (1 Z) II 24 where Z ~ Ber(p), < p < 3/4. We consider the hypothesis Rd and the logistic loss plus additional & regularization term: 200) (w. (x, v)) In (1 e ((wx)) + |/] A >
class H1 (x (w,x)
We prove the following guarantees on the hypothesis returned by SGD.
Proposition 1. Consider running SGD (Algorithm 2 for minimizing LD, (w) E(x.))/PiG) [In (1 + e =<ww.))]] 21 [[w]] with \ι > 0. Then, after Ti iterations and for any 8 E 1), with probability at least over the sampled (xx,,)))) Di(p), it holds for all {+1]d:
( W1 -ln
1+p
2d
1 + In
4ln I-p
In
e1, x
<
A1
1
X1
8T
1
p2
(44)
1
d
1 In T
In
1+p
X
1+e-(w..) (W)))
1-p
2X
8T;
p²
1 1
I+p
59

---

Published as a conference paper at ICLR 2026
where W1 is the output of SGD and e1 [1,0,... 0"" R..
Proof. We first show that learning 1 with respect (unregularized loss) In (1 e (()))) cor- responds to a Convex and Lipschitz learning problem. The loss is convex with respect to its first argument. For the Lipschitz constant, have for all [+1]dd E [+1} and Rd:
V In (1+
)
yx
S Vd.
(45)
1+ ev(w,x)
Therefore, applying Theorem 3 we have that SGD after Tj iterations returns hypothesis w1 such that for any 8 E 1) with probability at least - 6 holds:
2d
IDIG)), S (.(.)) (w)
(1 + ln
(46)
where w arg min (Di( (w). From the strong convexity of (Di()). this implies:
2 A1
4d
W1 wll
(w1) - LDI())A -LoGeA (w)) <
(1
(47)
X707i
The previous bound the parameter space be translated to guarantee the estimated proba- bility of the output being equal to 1. Recall that for a hypothesis w. this probability is defined as the output of the hypothesis passed through the sigmoid function (u) Ite-
Pi (i(w]] x) o(w,x)) = 1
(48)
We calculate the Lipschitz constant of pi(w x). For any x E {#1)d, di
Rd, we have:
xe (w.x)
5 Vd
VV (w
x)]
(1 allte-lwsh (1+e-(wwx)) (mxor))
4
(49)
Therefore, combining eqs. 47 49 we have:
Vil | wl 2A1 V
d +IIT
Ipi (w) x) pi (w x)]
(50)
8Ti
remains estimate the value of PI(w x). In order to find the minimizer arg min, (p) (w), we set the gradient of T. (p).A1 (w) to zero:
yx
E (xx,))DD. (p)
+ w = 0.
(51)
The objective LD, (p).A, (p),,,, (w) strongly convex, so admits a unique solution. We observe that this solution is of the form w ae1; E IR, where e1 0, 01. Indeed, the optimality conditions become:
E(x.g) (p)
+ aia = 0,
Ite
(52)
E E(x,)))D (p)
veL
= 0, = 2....,d. d.
IteoysI
The last d equations are satisfied as (x.x)) [yx:] for all 2....,d. The first equation simplifies to:
p + 0.
g(a) I-e-a + Aia
(53)
2
This equation has, indeed, a unique root as g is continuous, g (u) [1te + A > for all A > 0, u I and lim g(u) = -o, lim g(u) too. Furthermore, g = - 0, hence a > 0. This proves that w ae,, where a > is such that g(a:) = 0. Furthermore, let dg In +)) be the weight of the unregularized solution wo: that is, og is the solution of
g(au) 0 for ' = 0. We have g(a) - g(0o) = - 0 which implies that ag. From the mean value theorem there exists { E (a,, aw) such that:
d'(g) (oo) - a(a) \
(54)
Qg a Qg a
60

---

Published as a conference paper at ICLR 2026
But, observe that for any [0, aol. the derivative of the sigmoid is bounded as follows: 172 z d'(u) S 1 Therefore, for any X1 O, it holds:
4a
400
4 In
P
(55)
a aol 0ol -
Given the above, we can bound the calibration of w. For any A1 > 0 and x E {=1]² have:
1
S
aol
1 + em
4
(56)
In
\1
D²
Combining the above with equation holds:
50
we finally obtain that for any A1 0 and x E (+134, it
= (w1
x) Pi(Wo x)]
+ e
(Wix)
1-p I+p
m1
1
I 11 (wi
x) (w xx] + [Pi (w
x) -Pi x)] (57)
d
+ In
In
p2 A1.
2
8T,
Furthermore, for the optimality of w in parameter space, have:
[[w1
a0e1. x) = I( x) (aDe1 ae x)
(58)
vdll w/ll + lol.
From eqs. 47 and 55
obtain:
2d
1+ T
4
(wi ooee, x)
(59)
A1
87;
1 p2
The guarantees of equation 44 quantify the calibration of the model returned by SGD and its prox- imity the best-in-class model in parameter space. We obtain the following explicit corollary on the calibration and "hard" prediction of the model. Corollary
1. Consider running SGD (Algorithm
2 for minimizing LD.(I)A,
B(x.v)~D,()) (p) [In (1((1) e w(wx)x))]] A1 l/wl Then, for any = > 0, if =
6//2
and after
81/41//7
T. =
iterations, with probability at least - ô over the sampled {xx; v:) 1i Pi(p),
holds:
<E,
(60)
1 e- (w
21
1
I+p
for all x E [+1]d. If. additionally, E < In 11p 1-p) then with probability -
sgn ((w,,)) = 11
(61)
for all x E {+1]d.
61

---

Published as a conference paper at ICLR 2026
Proof. Recall that from eq. equation 44 with probability - 8 we have:
1 + p 1-p/
2d
+ I
4 4lm
1+2 1-p
W1 In
c1 x
A1
A1
8T1
1 - p²
(62)
In
1
d
1 In T1
X1.
1+ eee--MM (w)x))
di
2
871
p2
I+p
In 11
To instantiate the bounds, first observe that for any 0 p < 3/4, it holds:
16
1 1-p"
7
d
1+ 7
1/2
Furthermore, the first bound is always larger than the second one. We set \
32 ln
71
to
bound the two terms and solve the following inequality for T1:
d/22
+
In
1/4
512 In 7 1/2
<E
8T1
(63)
Ti
5122
(In 7)2
d
+ T1
49
8e4
This proves the first part of the claim. For the decision term sgn ( ((w1, x)). we have:
sgn ((1,,))
sgn
W1- ln
p
e1,
+ In
1 +
C1, xx
1-p p
(64)
+ p
1 + P 1 1-
sgm
In
e1,
+ + In
If
I-p
1+2
Wi - In
I+p i-p
e
xx
then the following holds:
sgn ((W1, x))
sgn
In
1+p
Z1. For this to happen, we additionally require having e <
T-p
1+2
E.3.2 POSITION
Position 2a We consider a binary classification problem where X20 = {+1]d+I, V/2 = {+1} with distribution (p) over Xzo V/2a such that: I1,..., Fd Rad(1/2) and (Id+1, v)
(z1 11) + (1 2) ( where Z ~ Ber(p), 0 < 3/4. The -1 label should be in- +1,
terpreted as the empty string e, while the +1 output corresponds to the <EOS> symbol. We consider non-homogeneous linear hypothesis class H2a {xx + (x)) + bb w E 2d++1 , E R). xama+1]" E [+1]2++ augments the input with the second degree monomials that involve the inputs bits and the bit from the previous position. In this way, the linear model is capable of approximating the target via the 21 feature which can help predict whether the generation is on the "Iong" or "short" path. Note that the dimension of the feature map - O(d) - is still polynomial the input dimension. The bias term is crucial for representing the best-in-class conditional probabilities. We consider learning with the ég-regularized logistic loss:
where the feature map Φ2 : X + [x]
Td+1 #1
T2
[(a)) (0- = bw, (x) In (1 ( ++ (1+e=M((,ooxxx)) -((m)02(x))+))) Ago ([]w]] b²). A2a > 0. (65)
Proposition Consider running SGD (Algorithm
2 for minimizing ToDza(V). (0 (w, (w,
+ Aga (l|w b²), Aza > 0. Then, after Tza iterations for any 8 (0, 1), with probability least over the sampled (xx,, v2)) it holds for all x E {+1]d=) with 2 21
((x.y)~Dz.(() (p) [ln (1 + --te-(lle
2(d + 1)
+ In T2a
Aza
+ t)
Aza
8T2a
+
V
1-p'
(66)
62

---

Published as a conference paper at ICLR 2026
while x {+1]++) such that 21
+1, it holds:
1 -
2(d 1)
+ In T20
A2a (1 + +p)
In
1-p
2p
+
e (W2
1+p p
A2a
8T2a
8p(1 p)
(67)
(W2o- 0ged+2 (x(x)) + bza bul
8(d 1)
+ In T20
12a (1 +1)
In
I-p
2p
2ua
8T2a
2p(1 - p)
where W2a: b2a is the output of SGD and āu, bu E IR such that: âu + bu = In 1-2 2p
Proof. We first show that learning H2a with respect to the unregularized loss In (1 + e (w((w)2()))++b) corresponds to a Convex and Lipschitz. learning problem. We will abuse notation and write E H2o for 0 = (w,b), where (w, b) E H2a- The loss is convex with respect to 0. For the Lipschitz constant, we bound the gradient with respect to O. For all x E {=1]d=l E {+1}, we have:
you (x)
V (20) (0, (x,
(68)
1+ c0((w,,z(()))-))
and for the bias term:
i (20) (0, (x. y))
y
(69)
+ co((m.0>(x))+b)
As result, we have for the Lipschitz constant:
(ol(20) (0, (x. (x,,)) w))|
|-
you (x)
+
S (2d + 1)
+1
(x((w,
(70)
which implies that the function is 2 with respect to O. Therefore, applying Theo- rem3 3 we have that SGD after T2n iterations returns hypothesis Oza such that for any 8 E (0, 1)
with probability at least - ö it holds:
1(d+1) Aza 8Taa
Da) (0za) S (D) (0)
(1 T2a)
(71)
where 1= = arg mingeH20 LDD (p) (0). From the strong convexity of (pz this implies:
(Lean(p)den (0za) - LDza (p) (6) S
8(d + 1)
020
(1 In
(72)
A2a
X2 872
We characterize now the loss minimizer θ = { (w,b) by setting the gradient of LD to zero:
E ((xx)))Dea()) (p)
(x)
A2a
(73)
E E(x,))) Dan (p)
+ A2ab O.
The objective T.Daa Aza (0) is strongly convex, so admits a unique solution. We observe that and i e R.
this solution of the form åea-2; E RR, where 6,42
Indeed, the optimality conditions yield:
Ala ( - i) 0,
(a
I-p
(74)
2p p-1 A2a (a+ b) 0.
Ite (148)
Itea=h
The functions 9 (u) T+e- Au, X 0 and g2(u) = Itc p-1 Au, À > 0 are continu-
ous, monotonically increasing and lim 91 (u) lim,- oc G2 (u) 00 lim 91 (u) lim,- (2(")= = +oo, hence they have unique roots. This proves that, indeed, w de where à, b are such that @1 (a- 6) 0 and g2 a 6) O. We will now bound the emror in the predicted probabilities, depending on what region of the distribution the input lies. We consider the following cases:
63

---

Published as a conference paper at ICLR 2026
If r1dd+1 -1, then we have:
p-1 (0aa Di 11
1 + (2)))
1 1 (2))
+ c(m,)2(x))++6
1 c(w))
v2(d+11 + 1)
| -0
1 c(M,)2(x)))
2(d + 1)
1 + T2a
A2a
8T20
+ 6(0162(0))4
(75)
where we used equation 72 and the Lipschitzness of the logistic function. For the bias error term, have:
1 + (x))+)
1 + mai 1
1 (a-b)
Aza
(76)
(from equation 74
p
Aza
V
as the solution of the gi (u) = 0 equation lies in (-v 1 0).
If x1 = 1, then denote by åg + by the solution of the unregularized problem g2(āg + Bu) for A 0 or, equivalently from equation 74 åg ho In
(1-p)
2p We bound the distance between the model-induced probabilities and the best in-class probabilities:
P1 O2a 21 = +1) Pi (6) [1
O2a I1J4+1 = +1) Pi (0 =+1|
pi
(77)
P1 0 21
+1)- -p1 P1 (0o 21 = +1)
v (2d /2d++1] + 1)
02a
61
1 =
(a à (d+b) b) -(āo + ho)
4
( + b) (ān bo) ) is bounded. Let denote the
It remains to show that the distance
unregularized part of 92 as g. that is g(u) = I + Itee First, observe that g2(0)
3--1, which means that the sign of + b depends on the value of p. It will be easier treat these three subcases separately:
If 1/3, then (0) 0 and g + bo) (āo +oo+0 bu) O. it holds 0
- 6 āo + bo. From the mean value theorem, there exists & E (ā + b, âo + bo)
such that:
Aza àa +
g((E) (ap + b0) - (
(78)
64

---

Published as a conference paper at ICLR 2026
But, g'(u) =
1e+ (1+e+)² and, in particular, o' (āo + b0) = g In
1- ) = 2p(1-1)
2p
1+p
This implies that for any E (o, In 2p )) it holds:
2p(1-p) 1+p
S g 1 1+p Thus,
1
Aza à (a+b)
(āo + b0) (a +b]
('(E)
X2a (1 + p)
(āg bo
(79)
2p(1 p)
A2a + p) In
1_
2p
2p(1 p)
p 1/3, then à + 6 g2(0) 0 and, also, âg + by In
1-p 2p
0. Hence,
the bias error term is 0.
- p > 1 1/3, then (0) > 0, (āo + -bo) = >za (āo + bn) < 0, so āo + bo < a- < 0.
From the mean value theorem, there exists & E (ān + bo, + b) such that:
Aza (a+t)
Aza
(ān bo)
F'(E)
(80)
à + 6) @s + 80)
a + b)
àn + bu
which, further implies:
Aea (1 + p)
In
3p
( a +)+ (āo + ho)
1-p
(81)
2p(1 - p)
from the same bound on ÿ' before.
Therefore, we showed that for any 0 < p < 3/4, we have:
(a+ 6) - (āu+ + bu)l A2a + p) In T-p
(82)
2p(1 - -p)
Combining this guarantee with the bound of equation 77 and the parameter space guarantee of equation 72 we finally have:
Aza (1 +
In
2p
pi (Zza (1a-11=
1 -
2(d+1)
/
1+ T2.
1 + In
T-p
1 1+p
Aza
8T2a
8p(1
p)
(83)
Furthermore, for the optimality of O2. in parameter space, we have:
(Waa - (0ex+2;02(()) + Tra - bn]
(W2a ded+2, x2(x)) + 6 b- 66((0002+2 ded42, (x)) + 6o - 6)
2 (d + 1) 02a - 011 (āo + bo)
à (a+b)
(84)
1(d+1)
1 In T20
Xza + p)
In
I-p
2p
Ago
8T2a
2p(1 1 - p)
where the last inequality follows from eqs.
82
L
As before, we obtain the following explicit corollary on the calibration and "hard" prediction of the second linear model.
65

---

Published as a conference paper at ICLR 2026
Corollary 2. Consider running SGD (Algorithm
for minimizing Lp)a(p), (w)
B(x.))DDa(()) [In (1 + +==((w,b2(0))40) + Aza
(Iww)
2
+ b²). Then, for any > 0, after T2a
2
Ô (max 82 prbct) }) iterations with A2a = θ min
dāpt
}). with probability at 8-1312
gara
54/4124
least 1 - 8 over the sampled {(x,, / D2a (p), it holds:
1
<E, VxE (#1]d+) S.t. 21 +1,
e (Wza ben
1+1
(85)
<e,V (+1]0+1 S.I. 31 = -1.
1 e x) +b2
If additionally
1-p) In
then with probability least 8 we have:
2p
1 (1-1)
-1, if 2 1/3,
sgn W2ay 62 (x)) + b2a)
sgn
In
(86)
2p
+1, ifp 1/3,
for all x E {#1]d+l with F1
Proof. Recall that from eqs. 66
67
we have: for all x E {+1]d+1 with 21 -1:
2(d+1) 1)
1 + In T2a
Aza
+
1 + +b2
A2a
8T2
1-p
(87)
2(d + 1)
/
1 + In T2a
+ 22vza = Qa(1zar T2a);
A2a
8720
while for x E {=1]d+l such that 21
+1, it holds:
A2a (1 +p) In
2p
1 1
2(d + 1)
1 + In T20
1-p/
+ e- (W2o x22
1 1+ p
\za
8T2a
8p(1 p)
(W20- - 20ea+2, (2(x)) + bza
bol
<
8(d + 1)
In T20
A2a + p)
In
1-p
2p
(88)
A2a
812a
2p(1 p)
Qa(Aza T2a)
We want to upper bound all three quantities by 5 E (0, 1). In equation 88 the second inequality
upper bounds the first one. Thus, it suffices to only consider that one. Let Aa As be the optimal Aza's 's for the two previous expressions (i.e., quantities Qu (Ara ), Qo T2a)) and Ta. To the
22/3 corresponding number of iterations to get the expressions less than e. We have: Aa =Ō (
while As è
p2dz (by balancing the two terms of each of Qa. Qb). For these values of
Xa; As, the two expressions are equal to:
d1/3
-
/51/672/6
(89)
12
(b(),TTb) =
/2/01/47711 1/4
This implies that it sufficient to take To Ô and To = ò p to satisfy
Qa Ta) < E and Qo(As, Tb) E. respectively. In other words, is sufficient to take 2/3
Aza ê min at atpd 81/41/44 31/0T//3 and T2a Ô maxx { de de })- This proves the first
8e 73aa
part of the claim. For the decision term sgu ((w.., tbe (x)) + bra), we have:
sgn ((W2a) 62 (x)) + b2a) =gn
(W2. agea+2, x2(x)) + bza - bo + (ā0e d2(x)) + bo) sgn (W2a agedt2, x2(x)) + bea - bo agz + bu).
(90)
66

---

Published as a conference paper at ICLR 2026
In partieular, when 2124+1 1, we have:
sgn ((Waa: (2())) + b2a) = sgn
(W2a doCd+2, &2 (x)) + bza - by + âo + bo)
(91)
1 p (1-1)
sgn
(22a
0ood+2: 1b (xx) + bza bo +
2p
(W22 0gea+2, 00ex42, + b2a boo. then the following holds:
If In
sgn (W2a: 022 (x))
+ (2a)
sgn
In
1-2
sgn (1/3 p)- For this to happen, it suf- 20
fices additionally have <
In
1-2
2p
Position 2b For the second part of position 2, we consider binary classification problem where all the data come from the "long" path. As this is similar to positions 3 to d, we treat all of these positions together next. We will denote the second part of position 2 simply as position in the next subsubsection for ease of presentation.
E.3.3 POSITIONS 2 TO D
For the l'th position of the output, 2 < 1 I d we consider a binary classification prob- lem where Xi (+1]dt-11 .) {+1} with distribution Di(p) over Xi Vi such that: .I,...., ... I wi and y II x;. We sider linear hypothesis class H4 {x (w, 61(x)) w [2d++-1). where the feature map [=1]2 augments the input with the second degree monomials that involve the inputs bits and the bit from the previous (I - 1'th) position. We consider learning with the logistic loss plus additional & regularization term: 110 (w. (x,z))
31, od Rad(1/2), (.44l,...,,,41111
Φ [x']
Id+i-1 Id+1-121 I..I
v)) = In (1 + ((=((((((())) + b l| λ O.
Proposition 3. For any 1 E (2....,d] d), consider rumning SGD (Algorithm 2 for minimizing Ay
LDD( (w) ((x,)))D(()) [In (1 + e x(ww,e(((i)) }
Z ||wl2 with A > 0. Then, after Ti it-
erations and for 5 (0, 1), with probability at least 5 over the sampled {(x,, vo it holds:
2d
In Ti
+ VA
(92)
1 e- (w).
2Ar
8Ti
for all E [=1]d+i-) such that T4+1-1
[1/=1 23
1-1
[ 24 1, and
2d 1
In
Ti
+ Ar.
(93)
1 1 (x))
2
8Ti
for all x (+1]dti-1 such that Ta+t-1
II T i= 2; 1, where w, is the output of SGD.
Proof. We first show that learning Hi with respect (00 (w, (x, (x) In (1 e (l (x))) + I
d ||| corresponds to a Convex and Lipschitz learning problem. The loss is Au-strongly convex
with respect to its first argument. For the Lipschitz constant, we have for all and w 122++111,
(x)
V In (1+c- ((w.))1)))))) 1 (x))
I 2d +
(94)
Therefore, applying Theorem we have that SGD after Ti iterations returns hypothesis w; such that for any 8 10 with probability at least - 5 holds:
2 (2d+1-1) (I+InTi).
LD. (wi) I (..iO),, +
(95)
X1ōTi
where w: = arg minwer, LD.(E).A, (w). From the strong convexity of ADi()).A this implies:
4 4(2d + 1)
[w 2 N
(LD(G)) (wi) - LD,()), (w:)) <
(1 In),
(96)
X781
67

---

Published as a conference paper at ICLR 2026
The previous bound the parameter space can be translated to guarantee the calibration of the model. First, we find the optimal solution W, = arg minweH, (P),A),,, We set the gradient of (eDi()) (w) zero:
(
E (xx)))D((())
AwW, =
(97)
1 + (
The objective (p) (w) is strongly convex, so it admits unique solution. We observe that this solution of the form Wi = a04421-1, a E R, where cd+21-1 [O,.... ......,,, ..........0 0] Indeed, the optimality conditions become:
E Ra
+ A1a = 0,
1 Miim =; Il
(98)
- (x); IT 82
E Ra
0, j & d + 21 1.
All but the d 21 1'th equation are satisfied as Rodd [T'm 2101(x),} (there is always at least one zero-mean bit that survives in the product). The d + 21 1'th equation can be simplified to:
g(a) ==
Apa =0 0.
(99)
1 + et
This equation has, indeed, unique root as g is continuous and g'(u)
a (Ite u)² A
0 for
all A 0, u E R. Furthermore, g(0) > O, while g (V))
1
I+eV Ai < as for any
X > 0 holds: 1+ VV
+2 From the intermediate value theorem, we get that
(0. 3).
Hence, for the output probabilities of the optimal model Wi, we have:
1
1
(wr x) =a((w) (x)))) = 1 (W1)01(0)
1 e 1 11
(100)
1
p-1 x) ((Wi, (x)))
+ ((i))r))))
ea II
For any x E [=1]d+l-1 with Id41-1 = 11
2ú. the above together with eq. equation 99 imply:
1
02 W
1 I+e ca = Ara < Vda.
1=1
(101)
p-1
24 1
tiae Vd
1 ea
Finally, combining with eq. equation 96 we have:
P1
W
x =
-1
+ (x))))(
1 dx
I 1 (((())((()))
(102)
2d
|w
will
Av
4
2d + 1
I
<
Vxv,
2
8Ti
One way to "guess" the optimal solution is the following. Suppose the solution had additional non-zero Take, for concreteness, j (the argument invariant
coefficient: Wi a0 + Se,, j a d
to the choice of j). Then. the optimality conditions yield:
Ex
+ λi = 0,
11 #1 #0))
and Ex
-x 11 Zi 14 7 s.tars)
+ A
0. By summing up and subtracting the equations, we get:
1 14.5 Ar + 8) and 14ea Au(a - 3) = 0. However, since the equation 14am Auu = 0 has
unique for any > 0, it must be: a + B, cquivalently, 8 0.
68

---

Published as a conference paper at ICLR 2026
and, similarly:
2d +
I Ti
p-1
Wt
1
+
Ar.
(103)
2
87;
1=1
We obtain the following corollary on the calibration and "hard" prediction of the hypothesis returned by SGD for the l'th position of the output.
Corollary Consider running SGD (Algorithm 2 for minimizing LD(())
((xx))D(()) [In (1+ ( + 2 llll. Then, for any E > O. if i
da/a and after
(1/8T/1/8
Ti =
iterations, with probability at least - 8 over the sampled {(x; 1i) 1i) ~ Di()),
82 we have:
1-1
+ 14--wwxx (wix) <,, Vx [=1]dti-1 s.A. Zd4t-1
Tj -
(104)
1-1
1 c(w,xx) E [=1]d+1-1 s.t. ad4i-1
a 1
==1
If additionally, & < 1/2, then with probability at least 1 8, we have:
sgn (( (x)))
24
(105)
for all x E (=1]d+t-1 with #d+1-1 i=1
Proof. Recall that from equation 92,
we have:
1
2d + -
In Ti
+ VA, vx ex (+13" s.t. II 24 1,,
2λ
8Ti
(106)
2d 1
In Ti
+
Mr Vx {+1]d s.t. II x 1,
+ e(W(x)
2A1
8Ti
Let \v
(2dt)/1))
2/3
1/3
then the upper bound becomes:
21
2d 1
1
InT;
2d 2d+1-11
1/3
1 In
1/6
(107)
2
8T:
2
8Ti
The first claim follows by solving the following inequality for Ti:
2d +1+111
1/3
(1+ In
1/6
2
<E
(108)
817
For the second claim, observe that by a standard property of the logistic function, we have:
sgn ((w,, (x))) sgn
(109)
1 (w),
Assume that & < 1/2, then we consider the cases:
If [T i= = 1, then we have:
<8 1/2,
(110)
+ e (x))
which implies: sgn
-11
itt (10)))000)
69

---

Published as a conference paper at ICLR 2026
If [T 2, = L, then have:
+ e(w),0)())) <E<1/2,
(1I1)
which implies:
2) sgn
1 (w 1+1 2
sgn
1 e- (W),
2
1
+ 1
(112)
sgm
1 + ((
2
Therefore, for all x E (+1)d+1-1 with Id+1-1 = I Fú. holds:
sgn((wr, (x)))
II
T
(113)
1=1
E.3.4 END-TO-END PRE-TRAINING RESULT
Theorem 4. Let d 2 2, e and O distribution over sequences defined as in equation
p
3/4. Let D(p) be (parameterized) Consider SGD (Algorithm 3 with per-
27
position iterations Ti =
oct T2a
Ô
max
Beu p=det }) .T.....,
d²
a2t total
iterations T = max [Ti. T2a:
T2a: (2...
Ti
In
2
and with regularization coefficients A1 =
8
13/4
A2a = - min dapa
A = 81/ST!73 for all < < d.
81/4T/77
81/47/4 81/97
175
Then, with probability at least 1 - 8, SGD returns hypotheses W1 E Hi, O2a E H2as W2 E
H2...., Wa E Ha that induce h E Hlin for which for all x E {+1}d holds:
d
Th
11, 12,
II
T4 <EOS>
x
p
E,
(114)
p
Th
34, <BOS>
x
Se.
2
Furthermore, ife < (d+ + 1) min
[ 2r
In
1+p I-p
In
1-p 2p
}
then it holds:
(oi, <EOS>) ifp < 1/3,
h (x;; (W1, 02,, W2. w2...., wal)
(
II.....,I_ 24, <BOS>) otherwise.
(115)
Proof. The proof strategy is to bound the probability of either not sampling enough 'long' sequences or SGD returning 'problematic' hypothesis at some position. We invoke Corollaries [ to 3 for
a+1 2[d-1) First position First, we calculate the probability of outputting the parity II. J in the first position for the unregularized solution Wo:
P (x;
= II
1 + pci I
3; +
+
-
p01
1==1 F4
d
2
2
2
(116)
+ p I[:
23
2
From Corollary after T 2 Ti
5 iterations for \1 = θ
8//
SGD returns W1
d+1 2(d+1))
51/49/77
such that the following hold with probability less than
5 2(d+1)
[x; w1) = z4
1 1
2 E
(117)
2
+
70

---

Published as a conference paper at ICLR 2026
and
P
b (x; W1)
II Zi
1 +
1i=2
I Tj +1 1
1
1 + par
1
i=
1
1 + pad
(w1,)
2
1 + e(w1,))
C d d+1"
(118)
and, since e
In
1+2 T-p
it also holds:
2+1
(1 W1) = sgn ((W1. = ø1.
(119)
Second position (2a) We calculate the probability of stopping the generation given that the first 2; for the unregularized solution 0u:
token has been equal the parity II
- P
IP
h2a (x: 0b) <EOS> hn W1)
II
T4
21
31
1 1+p
1=1
-
1 +
1im2
1
1
1+p
2
2
1
24
1
(120)
From
Corollary
after
T
T2a 041 2(d41]
iterations for A2a
d 81/47/7 8 13 }). SGD returns
é min
Wea bzn such that, it holds with probability
less than 2(d+1)²
P (2a O2a) = -x: = =x1
2p
2
E +
(121)
1
and
p
hza (x; (x: 02a) = <EOS>> h1(x:w1)
II 2
1
II
1 24
++P P
IT 1 34 +
1
p
1 [Td i= T2
1 + e (Wanx)) -1)
1 1 (Wao
1+
2
d +
(122)
and, since e In (1-0) it holds:
In (1-2) ))
-1, if 2 1/3,
sgn ((waa Oz (x)) + Iraa) = sgn
2p
=
+1, if < 1/3
(123)
for all x E (+1]d+1 with 21 1.
Count of long sequences Let ž1, 2T D(p) and let Tiong count of "long" samples. By the definition of D(v), it is: E [Tongl = Tp. By the multiplicative
ΣT. {[z]| > d d+2] 23 be the
Chernoff bound, holds:
Toong S
Tp
Se
TP
(124)
2
71

---

Published as a conference paper at ICLR 2026
hence with probability at least 1 6 2. we have:
Tp
Toog
(125)
as long as T > In
5. If further, T 2 (2....d))) then with probability at least
we have:
2
Tiog Ti for all I E (2....,d) d).
Positions 2 to From Corollary
3 after Tiong 2
Ti
5 d-I 2(d41))
iterations for A =
SGD returns such that, with probability less than 8/2d:
11/9T//3
P hx(x:; wi) = II
x4 (1) = æ1.... h (x: WI-1) = 21-1 -
2 E
(126)
d+1"
2m1
and, since _ < ::
sgn ((W,, 0i (x)))
(127)
for all x E (=1]24i-1 with Fd41 = 11--1
Union bound Let ELONG denote the event that Tiong > Ti for all 2 < < and denote by P. the event that the above guarantees hold for position Then, by union bound, have:
P [LLLNG A Pi A Pao A P3 A. Pal
P [LoNN]] P [PP A P2u P A PalELONG!
1
P Ei {1,2a, 2......................
........... ... -PP))
(1- 1-
1
8)
8 2(d 2(d+1) + +) 1E(1,20,2,... d..
Σ
2 1 8.
2
(128)
Thus, with probability at least 5, SGD returns a hypothesis h whose probability of correct, long and short, sequences is:
d
Thh
[11 211,
<BOS>
x =
=P [1((x: W1) =x1 ..... IP
ha(x; wa)
II x h (xx:ww) wi) I1..... <x: wi-1) 1=1
21-1
1 (1 <((x))
(1 1+ (220))
2p
+ E2a(x)) (1 Ex(x)) ... Ex((x))
(129)
and:
Th I
d
E4, <EOS>
x
=P
(x; w1)
II
IP
higa O2o) = <EOS> h (x: W1)
II
3;
(130)
1=1
1 1+
i=2 1:=2 (x()
/I-p
E2a(x)
2
+p
where [x(x)) (x., [a(x)) O(e/d) for all x. In other words,
d II c4 <EOS>>
Th
xx) - 2S5 c.
(131)
72

---

Published as a conference paper at ICLR 2026
min (122)
Lod
vsp
3.0
25
men In
Fuhczen
1.0
0.5
Figure 41: The functional form of the error term vs p (0, 1) in Theorem 4
and
(ii
-
Zi. <BOS>
S
(132)
2
Furthermore, from eqs. equation 119 equation
123 equation
127 for all x E {+1]d, we have:
hi(x;; sgn ((W1, x)) ø1.
((w2er dz x, w:))) +b20) = sgn
1 (1-P)
2p ))
sgn
in
(133)
-1, if p 2 1/3,
+1, 1f3
hence
<BOS>, if > 22 1/3, if < 1/3,
(x; 0za)
and
I1'm Fi, if (x; 02a) 4 <BOS>, c, if (xx 02a) <EOS>,
(x(x; wi)
for <l2d
This implies:
(x1, (B0>))))>),))) if < 1/3,
h (x; (W1, O2a.
wal)
( ... I 24, <EOS>) otherwise.
(134)
Theorem in the main text follows as corollary by Theorem functional form the error term sp)) particular, the functional form also explains the difficulties faced in training transformers for Peot & 1/3.
4 for p Peot See Figure 41 for the
73

---

Published as a conference paper at ICLR 2026
E.4 POST-TRAINING
We complement the analysis of the previous subsection to show now a positive learning result on the combined pre- & post-training procedure. Precisely, show that, when Poor < 1/3 which is
RL rounds suffice for the pre- pre-trained model to reach perfect accuracy. Furthermore, we prove that the length of the answer grows predictable way.
the hard regime for pre-training, only O (log 1-Best
Peos
Theorem 5. Let d > 2, 6 e (0,1) and Poot (0,1/3).
Let D(pcot) be (pa- Poot Let E S Cu 1-
rameterized) distribution over sequences defined as in equation
27
for some constant co > 0. Consider pre-training (Algorithm
per-position iterations
T1
Ö
87
T2a
(max
de Be
.T..... .
d" Bell
total iterations
proct
max
Tí. T200
2 max (2..oa) Ti
In }} and with regularization coefficients 2
69/4
d Poot 81/422 4
. .v= .
A2a
min
5 175
81/2121/3 for all 2 S S d.
Furthermore, consider post-training with the STaR algorithm solving Problem each round and denote by 1())))) E Hiii the model returned at the end of the n-th STaR round. Then, there exists
n" =0 =0(10g - such that, if:
Poot
1. SGD steps per-position per-STaR round are as follows:
Ti
d' DE
T2a δ
max
d
Pocrde })
de
T2...., Ti Ö
d 5e
and total iterations T
max
Tí, T2a:
2 maxxeg
Ti In
logg
1-Poon
Poot
and
Poot
Poot
2. Regularization strengths per-position per-STaR round are as follows: x =
j0/4
113) }). N - 51/37775 for all 2l5dd. d.
A2a = min
min-
d Poot 51/47//7
51/31223
with probability at least 1-6 8 over sampling from D(peot) (during pre-training) and over sam- pling from models ho....,hw hn (during post-training), post-training returns w ...., w|"" (n") E Hi, O()) 0(n') E H2a, (h) .....
H2...., w ,..., w(n") E Ha that induce h(1) "U,,...,m'''' E Whin X' + y) for which for all x E X it holds:
w..,... W² w!"")
(mn)
Th(n) ((s1 [m1,
II 1=1
rs. <EOS>
x
In Se.
(135)
where 9n - Ppl S 2 and Pat1 =
20 Itp
for all n S n* with po = Peots and also:
d
h(aa)) (x: (xi(w/)' (hm (1n)) @(n") W2
(n))
w (0(")))
Ø1, 11
T <EOS>
(136)
Proof. The proof proceeds in two main steps. We first observe that the STaR objective for the 'th round can be re-written as next-token prediction with respect to shifted version D(g,) of the original distribution (P where < In < 1. We then bound the deviation sequence 4n from "noise-less" sequence Pn. that describes the evolution of the proportion of long data in the data mix, and whose closed form is available.
74

---

Published as a conference paper at ICLR 2026
From Theorem 4 for p Peot and å 8/2, with probability at least 5/2, obtain model (0)) H such that for all x {=1} it holds:
01,2112,
Bis <EOS>
Pcot
E,
i=1
(137)
1 Poot
T4. <EOS>
x
S E,
and
(0) (x; [W1, O2.. W2.
Wa)) (2), <BOS>)
(138)
The previous point-wise guarantee also implies:
E Rx-Rad(1
T)()))
F1 21221
3is <BOS>
xx
Poot
S E,
i=1
(139)
II
1 Pont
E Rad 11/2)ad
Ty:o:
F4 <EOS>
x
S E,
i=1
Recall the form of the STaR problem for the first round:
min Wi.W2u. b. W2...,
E
Rad((1 (12) ,00 (-1x)
[1 {lyl 2} Lehort
flyl = d + 1} (eotg|reo((x
(140)
where Labhott = (12) (w1, (x, + 1(2a) ((w2a: b). ((x, (1) 32)) and Liong = (w) (x, v1)) + ((a)) (1w2a) 6}, ((x, 31), (22) +... + ((d) (wa, ((x, 91: .... Md-1) Ba)). Note that, with probability at least 1 5/2, the objective can be re-written as:
Px,y
y =
Ti, T122
II 24 <EOS> 1=1
(co0((2 =
[Giong]
y= (1) IT <ROS>)
d
+ P
Fx,y
y
Fi. <EOS>
'oot (x. 11) (x. v)= =1
E x Rad(1
Lahortt]
y= 0= # IT (B0S)
Fx T1)))
F1,
11 Ti= T44 <EOS>
x
Ex
Th(0)
Fi, 2i 222
11
Ti, <EOS>
xx
+ Ex T/(0)) (II
[ Li, <EOS> x)
[Liong]
x- Rad()
y=
(o)
0102.....
I1 =3.0)
Ex (()) (II Fz; <BOS>>
+
Ex T
Ø1, 21,22122. .... [Lshort]
-i=1 Izi <BOS> [x)
d
Ex Th10) (IT <BOS> xx)
E
y= (TTL 21.5E00>)
Po + Co
Pa
+ Co
+ Eo
E
x- R2d
[Cion]]
2
y (1
11 (R00)
Po + So
E
[Lshort]
Pu+1
+ Co + Eo
x>Rad(1/2)04 -=(II:m (E00))
2
= ((xx)))D(a1)
[1 {lyl = 2} Cshort + (lvl = d + 1} Liong]
where Q1
with |Col, 1col $ (from eq. equation 139 and the first and second inequality holds from the law of total expectation and the definition of the conditional probability.
5o -
to
75

---

Published as a conference paper at ICLR 2026
In other words, showed that the reinforcement learning risk with reward Tcot during the first round of STaR equal to next-token prediction risk, where the distribution is the mixture of long and short strings but with shifted mixture weight 91- Thus, from Theorem 4 for p 91 and 6 = 81 E (0, 1) with probability (1 81) (1 6/2), SGD returns model (()) Alt such that:
21, 22:2
Th <BOS>
x
E
(141)
KA(D)
1I T4 <EOS>
)x
1 - 91
S
2
and
(1), <EOS>>) if < 1/3,
(,00) (x; (WI: Ozas W2.
Wa])
(142)
21,21122:1 II 24 <EOS>) otherwise.
Therefore, by induction, as long In < 1/2, invoking Theorem times for 8 = 5, E (0,1)
and p qs. with probability at least IT (1 6t) 2 Σ i= di post-training returns model
h(m) E HA such that:
(h()
21, 21,2122,1...
II
Ia <BOS>
x
9n
S E,
i=1
(143)
On
Th(n)
T4z <EOS> {
S 6,
and
((()) (x: { (11, Oza, W2.... wal) (1),, <EOS>), if On <1/3,
(144)
(11,01222 I1: 32, <BOS>) otherwise,
where In+1
2(m 1 G.) Nu+1+2(6, [.. with [6al. (knl S with qo = Po-
We now analyze the perturbed sequence du- Let f(ans Sni En) Taylor expand f around (gn, 0, 0):
du+1
+1+22((+t(m)
2(g.
We
0f OG,,
0f + En DEn
f(gn, SniEn) f(gn, 0,0) + Sn
+ Ry((u, Enr
(9m.0,0)
(gm.0,0)
2qn
+ Sn
2 2qn
En (pn + 1)2
4qn
+ Ry,uu, En.
(145)
On + 1
(4n + 1)2
2qn
2Cn 26n
-4En
On +
RRr(u, (ns
(qn + 1)2
where we used the Lagrange form of the remainder: Ru, En) ) = (0, u) (S + {) + where < |6n and |u2| 1&nl and H (q, Sns En) is the Hessian of f evaluated at (q, Sn, We now bound the deviation of the perturbed sequence qu+1 from the un-perturbed one Pat1:
u E IR? such that: |ur|
2pn
190t1 Pn+1l =
fgnn Sn. En)
Pn +
24n
2Cn
2Cn
4Endm + Ry(u, Sns
2pa
On 1
(gn + 1)²
Pn +
2qn
2pn
22n 2Cn - -En0dn
+
+ RR En))
9n +1
pn +
(qn + 1)2
2Cn 2(n
45m9n
< 2 2/4
pnl +
+ [Ry (u, Sue En)l. provided > 0.
(gn + 1)2
(146)
76

---

Published as a conference paper at ICLR 2026
Since In and 1enl O(e) and Ry(u, Sns En)l X E?,
if we denote the error sequence by
en = 19n px
then:
Cn-1 2en S E.
(147)
Since co 0, this implies that the emror at step n is bounded
en S e.
(148)
We want to find an integer n* such that model (nn)), generates long answers. For this, suffices to have 7n 1/3, equivalently:
Pn- 2n C'e
(149)
3'
2pn-1 We show that the unperturbed sequence Pn 1+1 then:
admits a closed form solution. Let uy
1
Un +
Up
Un
(150)
2
2
or, for the original sequence Pn:
1
2nPo Pc
Po
(151)
Pn. n
Po
2n + - po
+ (1 - po) 2-n"
Returning back to eq. equation
149 seek n² > 1 such that:
Po
2n C'E> 1
(152)
Po + (1 - Po) 2-
3
Assume that 2" C²e <
1 for some E to be specified later. Then, suffices to have:
15
Po
1
(153)
Po + (1 - (P) 2-n"
15
3'
2((1 - Po)
n" log
(154)
3po
In that case, the error parameter needs to be as small as:
1 Po
(155)
C1-Po Po
for some C O. Finally, we pick the smallest n such that On- > 1/3 μ4, where μ 2 w(e
Note that Pr. < 1/3 = Pat1 < 1/2, hence Theorem 4 can be invoked for 9n- < 3/4 if C is
sufficiently small. We set 60 = 6/2 and bj....,dn.. we have:
2n then with probability at least 1 - 6,
d
((n")) x: { ((m") )(n²)
(n')
W (a ((()))=
W2
= 01 22
F4s <EOS>
(156)
Remark 7. Note that the guarantees of Theorem for the number of SGD iterations required per STaR round are very pessimistic: they depend on the original mixture weight Peot rather than the per-round updated weight yn. and 2) they assume that the models are re-initialized at the origin prior to each STaR round, whereas of course in practice we continue fine-tuning the previously obtained model which speeds up convergence.
Remark 8. One can also analyze an (almost offline) version of REINFORCE algorithm with no reward normalization, instead of STaR, by noting that its loss also corresponds to the next-token prediction objective but scaled by different per-round coefficient.
In particular, if the proportion of long data in the data mixture is not exponentially small, i.c., there exists constant 6 E N such that: Peot E 52(d=*), then by following the previous pre- & post-training recipe on data coming from D (Poot). we can obtain model that emits long, correct sequence and predicts the parity of d bits after O (poly(d) SGD iterations.
We require qa* to not exponentially close 1/3, so that we do not pay 1/s sample complexity for 2g, E < (d + T-4.. x (d+ 1) (g.. 1/3)
Note that for
the regularization terms and SGD iterations grow together with factor
0 In 1-Peot ). but this gets suppressed by the Ō(.) notation.
Poot
77


Reviews:
[{"review_id": "xeref2dEv5", "reviewer": "Reviewer_StSb", "summary": "This paper studies why RL after next-token prediction helps large language models learn better reasoning. The authors build a simple setting where training data mixes short and long “chain-of-thought” examples. They theoretically and experimentally prove that next-token prediction alone often fails when long samples are less than 1/3 in pretraining dataset. And RL can quickly improve learning by focusing on longer samples, which leads to longer and more correct responses.", "questions": "1. How robust are the theoretical conclusions if long CoT samples contain noise?\n2. Can you still observe “length-driven generalization” when RL rewards contain a certain proportion of length penalty?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. Strong theoretical proof combined with solid experiments, showing when and why RL succeeds.\n2. The paper gives a simple and convincing explanation of why RL helps reasoning. Providing clear insights.\n3. The answer of  two core questions—\"why RL works\" and \"why length increases\" could offer a new design direction for LLM reasoning optimization.", "weaknesses": "1. The theory assumes fully correct CoT examples. However, the real pretrain data usually contains noisy or wrong data. And even positive RL trajectories could contain noise or false-positive ones. The paper lacks discussion of these factors.\n2. The theoretical proof could be written in a more organized manner, including formulas. This would make that part easier to understand.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "i9WLpwif1R", "reviewer": "Reviewer_pXVH", "summary": "This paper proposes a framework to theoretically understand the success the popular LLM training paradigm that RL post-training after SFT. This paper also uses the parity check experiments to show that RL enables the model to generalize to difficult tasks while SFT cannot. The paper also demonstrated this phenomenon on math problems.", "questions": "- Could similar phenomena be observed when in a worded version of a parity task? Such as giving the task description as language input of a LLM.\n- The theoretical analysis relies relies on a simple linear framework, how is it extended to non-linear structures like the Transformers?", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- the paper is well-motivated - trying to understand the reason behind the current successful LLM training paradigm.\n- the paper proposed a relatively comprehensive theoretical analysis framework.\n- the paper empirically validated its perspective through a cleverly designed and controlled parity check experiment.", "weaknesses": "1. The theoretical proof relies on a simplified linear autoregressive model.\n2. Some of the ideas of this paper are already pointed out in papers like DeepSeek-R1.\n2. This paer lacks of practical guidance for future LLM training and new algorithm.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nhIrK8OjOH", "reviewer": "Reviewer_Hty6", "summary": "This paper explains why large language models improve dramatically when reinforcement learning (RL) follows next-token prediction training. The authors prove that next-token prediction alone cannot generalize on hard tasks if long reasoning chains are rare, whereas RL amplifies these rare sequences and enables efficient learning. Experiments on both synthetic and real-world reasoning tasks confirm that RL boosts accuracy and increases response length. The results highlight RL’s role in unlocking generalization from limited but valuable reasoning data.", "questions": "* Theorem 1 gives $p_{cot} < 1/3$ as the exact point where greedy decoding stays short. Is this an artifact of the two-step linear decision model, or does it survive richer embeddings (e.g., transformers with non-linear MLPs)? An ablation that keeps the data distribution but increases model expressivity would clarify whether the threshold is distribution- or architecture-specific.\n\n* The post-training bound hides the per-round sample size inside Õ(·) and assumes fresh data every round. How many total unique prompts does the algorithm really need? \n\n* Parity has a single deterministic “long path”. Real reasoning data often contain many valid chains of varying length and quality. Does the two-component mixture still capture the dynamics, or does the presence of noisy/partial chains shift the critical $p_{cot}$ or require a different RL objective?\n\n* The theory 2 requires $p_{cot} \\in \\Omega (d^{-\\kappa})$. For $d\\approx 1,000$ (typical for LLM prompts) this seems prohibitive. Is the polynomial dependence tight, or do transformers empirically succeed with much smaller constants?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* The paper is the first to provide a formal separation between next-token prediction and next-token prediction plus RL in autoregressive models, giving a novel theoretical account of why the widely-used “pre-train then RL” recipe succeeds.  \n* By showing that RL can turn a sample-hard, exponentially data-hungry problem into a polynomial-time solvable one whenever long demonstrations are merely polynomially rare, the work directly informs how scarce reasoning data should be leveraged in large-scale model development.", "weaknesses": "* All conclusions of the paper are drawn around the task of predicting the parity of d bits given access to a source of sequences, and its generalization to more common reasoning tasks (science or open-ended) is questionable.\n* There is a lack of experiments on a broader range of and more recent LLMs, e.g., qwen3.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)