Organizations: Department of Mathematics, University of Bologna, Italy · Gatsby Computational Neuroscience Unit, UCL, UK · Department of Information Engineering, University of Florence, Italy
We study whether a transformer network can learn the deterministic sequence of trees generated by the iterated prime factorization of the natural numbers. Each integer is mapped into a rooted planar tree and the resulting sequence NT defines an arithmetic text with measurable statistical structure. A transformer network (the GPT-2 architecture) is trained from scratch on the first 1011 elements and evaluated on next-token prediction and masked-token prediction tasks, with a Hidden Markov Model as baseline and a scaling analysis over context window, dataset size, vocabulary size and model size. The model reaches a word accuracy of about 0.4, well above the baseline, and its performance remains stable on test blocks located at 1013-1015, far beyond the training interval. Moreover, the likelihood assigned by the trained model separates the arithmetic text from two controls: synthetic token sequences that preserve the marginal token frequencies but not their order, with a separation that widens with the context, and sequences containing more than three consecutive square-free integers, a configuration that arithmetic forbids. These results indicate that the transformer captures regularities of the arithmetic text that go beyond its frequency profile.
Figures & tables
Figure 1: Number of elements in the natural dictionary \mathdsNDn as a function of n . It is numerically found that ∣\mathdsNDn∣ scales sublinearly in n up to n=1012 Contucci et al., 2026 .
Figure 2: ( Left ) Jensen–Shannon divergence between chunks of 108 consecutive Dyck words in our dataset \mathdsNTn with n=1011 . The largest value of the divergence, reached between the very first chunk and the very last (upper right corner), is approximately equal to 0.0238 , i.e., the distribution difference remains quite small despite the 11 orders of magnitude span. ( Right ) Detail of the Jensen–Shannon divergence of the last chunk of 108 trees with respect to all the previous ones of equal size.
Figure 3: An example of a numerical sequence (integers between 5 and 14 ) converted into a sequence of Dyck words. The tokenizer splits the sequence of 10 Dyck words into 8 tokens, corresponding to the black boxes. Note that a token can consist of a single Dyck word, a concatenation of Dyck words or, generically, parts of Dyck words (in the example, the Dyck word 110010 representing 12 is split into two different tokens). After training, the attention matrix A can be used to compute, within a sequence, the relevance of each element of the context for a given token. In the picture, we focus on the last token, containing the Dyck words representing 13 and 14 . Arrows point towards the other tokens, with width proportional to the associated attention weight: thicker lines point towards more relevant tokens in determining its value. In the example, tokens t3 and t7 are the most relevant for t8 , as A83 and A87 are the largest entries of the attention matrix.
Figure 4: Scheme of the dataset management. The sequence of 1011−1 consecutive Dyck words has been split in 10 blocks of 1010 Dyck words (except for the first block, that is long 1010−1 ), each further separated into two portions: the first \nicefrac34 of each block is used for training and the rest for validation, except for the last sub-block used for testing. Three additional blocks of 106 consecutive elements were produced for testing in the range [1013,1013+106] , [1014,1014+106] and [1015,1015+106] respectively.
Figure 5: Loss curves for a model trained via NTP, obtained for different tokenizer’s vocabulary size D , namely D=26 , D=28 , D=210 . The loss value, computed as in Eq. ( 4 ), is rescaled by lnD , loss value corresponding to a uniform distribution.
Figure 6: Accuracy obtained from prompts of length L=210 , at β=+∞ , for different generated sequence length L^ . Each box corresponds to a value of L^ and presents three test sets (marked with different colors) obtained from the first 106 consecutive Dyck words after 1013 , 1014 and 1015 (red, blue, green) respectively. Each point represents the result from one out of 128 batches for each of the three test sets. Remarkably, accuracy seems not to be affected by the “scale” at which the testing sentence is sampled (i.e., if the test is obtained from sequences around 1013 , 1014 or 1015 ). Results obtained on an architecture with 12 layers, 12 heads per layer and D=28 , trained on our dataset of 1011−1 Dyck words.
Figure 7: Accuracy for words ( Aw , above ) and Kullback–Leibler divergence for word distribution ( KLw , below ) are reported at different values of β , over a set of 128 different inputs of 210 tokens. The generated outputs are L^=25 tokens long. The dashed horizontal line represents the HMM performance, whereas boxes correspond to the LLM performance. For each box (corresponding to a given value of β ) test results are given for token sequences sampled from the tokenization of the first 106 words following 1013 (red), 1014 (blue) and 1015 (green), respectively (see also the caption of Fig. 6 ). Despite the fact that the HMM is actually good at reproducing the statistics of the dataset has been trained on, as shown by the divergence KLw , it fails in predicting the order of Dyck words exhibiting an accuracy comparable to a random guess. Results obtained on an architecture with 12 layers, 12 heads per layer and D=28 , trained on our dataset of 1011−1 Dyck words.
Figure 8: Average precision, recall and F1 scores for various Dyck words over 105 test sentences of length L=210 at β=+∞ and dictionary size D=28 . Only words appearing at least 50 times are represented. The plot shows results for the LLM (circles) and the HMM (squares). The LLM clearly outperforms the HMM, whose precision and recall are particularly low. Markers with red frames corresponds to the word 10, i.e., to primes, whereas the points in the blue frame correspond to the most frequent Dyck word in the dataset, namely 101010.
Figure 9: Distributions of the words appearing in the true sequence when the LLM predicts a prime ( left ) and of the predicted words at true prime positions ( right ), obtained via NTP on a set of 105 different input prompts of L=210 tokens. The model was trained with a context window of L=210 , 12 layers with 12 heads per layer and trained on a dataset covering integers from 2 till 1011 with vocabulary size D=28 . For comparison, circles represent the performance of the HMM architecture.
Figure 10: Confusion matrix in the NTP task for the LLM at β=+∞ . Each row corresponds to a Dyck word and exhibits the logarithm of the frequency at which it is substituted by each one of the Dyck words listed along the x -axis. The LLM is trained with a context window of L=210 , 12 layers with 12 heads per layer and trained on a dataset covering integers from 2 till 1011 with vocabulary size D=28 . Frequencies are obtained by feeding the trained model 105 sequences of 210 consecutive tokens, and querying the model to predict the next tokens until a word is completed. For comparison, the confusion matrix for the HMM is given in B .
Figure 11: Mean ( left ) and variance ( right ) of the log-likelihoods in Eq. ( 8 ) obtained by sets of 105 synthetic prompts consisting of 2κ BPE tokens sampled independently from the empirical marginal token distribution of the training corpus ( circles ) and 105 prompts of equal length from the true sequence of tokens provided by the dataset ( squares ).
Figure 12: Cumulative distributions of the standardized variables for various values of κ , with colors matching the ones in Fig. 11 , obtained from sampling 105 sequences for each κ . For each κ , the continuous line refers to the cumulative distribution computed on true sequences, whereas the dashed line of equal color refers to the cumulative distribution computed on plausible, artificially generated sequences, showing that in some cases distributions also markedly differ. The LLM is trained with a context window of L=210 , 12 layers with 12 heads per layer and trained on a dataset covering integers from 2 to 1011 with vocabulary size D=28 .
Figure 13: Distribution of the log-likelihood ℓ obtained via kernel density estimation from 105 true sequences from the test set consisting of 20 Dyck words . Dots denote the average values of the likelihood obtained by sampling from the dataset 28 sentences ending in 1010 101010: such ending is substituted with sequences of square-free Dyck words, the log-likelihood of the unaltered sequences being marked in pure blue. It is shown that the modified sequences representing more than three consecutive square-free integers (reddish dots) are placed at very low values of ℓ , whereas true sequences and fake sequences with no more than three consecutive square-free (bluish dots) are assigned a higher log-likelihood.
Figure 14: Accuracy at β=+∞ for various combinations of the hyperparameters considered in the experiments. The reference model is described in the text and corresponds to scale ρ=1 , it is trained on the dataset of the first n=1011 Dyck words with context window size L=210 and vocabulary size D=28 . In the panels, one of such parameters is altered by keeping the remaining parameters fixed: from left to right, L , n , D and ρ . Each dot represents the accuracy of a single architecture trained with the indicated hyperparameters on a sentence of length L at a test scale ranging from 1010 to 1015 , each scale being highlighted by a different color. Each of the four plots shows the effect of each hyperparameter, obtained grouping all the set of accuracy values only by this latter.
Figure 15: Loss curves for an architecture trained via masked language modelling. Here, tokenizer vocabulary sizes equal to D=26 and D=28 . The loss value is rescaled by lnD , loss value associated to a uniform distribution.
Figure 16: Token accuracy as a function of masking probability and β within the Masked Language Modelling strategy on the test set.
Figure 17: Accuracy at β=+∞ for various combinations of the hyperparameters considered in the MLM experiments. The model size correspond to ρ=1 scale in the experiments above, it is trained on the dataset of the first n=1011 Dyck words. In the panels, one parameter is altered in turn, by keeping the remaining ones fixed: from left to right, we tuned the tokenizer dictionary size D , the context window L and the masking probability at training time pM . Each dot represents the accuracy of a batch of 27 sequences, with the architecture trained with the indicated hyperparameters on a sentence of length L at a test scale ranging from 1010 to 1015 , each scale being highlighted by a different color. Each of the three plots shows the effect of each hyperparameter, obtained grouping all the set of accuracy values only by this latter. The results highlight the detrimental effect of a large D ; as D increases, tokens often encompass multiple words, creating a harder masking task during training. Sequence length L does not significantly impact the training procedure, aside from requiring more refined learning rate tuning at L=210 . Optimal performance is achieved at a masking probability of pm=0.15
Appendix figures & tables1 asset
Supplementary material from the paper’s appendix.
Appendix
Figure 18: Confusion matrix for the trained HMM. It is reported the logarithm of the frequency at which a true word appears as replaced by a different word in the dictionary in the HMM generative process.
We study how we can leverage only a handful of characteristics of a transformer's architecture to closely predict the number of different sequences it can output, both qualitatively and quantitatively. We provide an upper bound depending on the length of the prompt, which we show empirically to be tight up to a factor less than 10, across architectures and model sizes. Our analysis also provides a theoretical explanation for previously observed empirical failures of transformers on simple sequence tasks, such as copying and cramming. Formally, we prove that (i) the maximal length of accessible sequences (those that the transformer can output for some prompt) grows linearly with the prompt length, (ii) beyond a critical threshold, the proportion of accessible sequences decays exponentially with sequence length, and (iii) the linear coefficient relating prompt length to accessible sequence length admits a theoretical upper bound. Notably, these results hold even with unbounded context and computation time.
Maxime Meyer, Mario Michelessa, Caroline Chaux +1
Department of Mathematics, National University of Singapore, Singapore, 117543 · IPAL, IRL2955, Singapore · School of Computing, National University of Singapore, Singapore, 117543 +2
Grokking in transformers trained on algorithmic tasks is characterized by a long delay between training-set fit and abrupt generalization, but the source of that delay remains poorly understood. In encoder-decoder arithmetic models, we argue that this delay reflects limited access to already learned structure rather than failure to acquire that structure in the first place. We study one-step Collatz prediction and find that the encoder organizes parity and residue structure within the first few thousand training steps, while output accuracy remains near chance for tens of thousands more. Causal interventions support the decoder bottleneck hypothesis. Transplanting a trained encoder into a fresh model accelerates grokking by 2.75 times, while transplanting a trained decoder actively hurts. Freezing a converged encoder and retraining only the decoder eliminates the plateau entirely and yields 97.6% accuracy, compared to 86.1% for joint training. What makes the decoder's job harder or easier depends on numeral representation. Across 15 bases, those whose factorization aligns with the Collatz map's arithmetic (e.g., base 24) reach 99.8% accuracy, while binary fails completely because its representations collapse and never recover. The choice of base acts as an inductive bias that controls how much local digit structure the decoder can exploit, producing large differences in learnability from the same underlying task.
Transformer Programs (TP) (Friedman et al., 2023) constrain a Transformer's architecture so that its learned computation can be extracted as readable Python code. TP expose learned computations, but token accuracy can hide incorrect complete answers. In this work, we examine how program inspection provides insights beyond accuracy scores and helps guide training. We first evaluated TP on nine algorithmic tasks. We compared token accuracy with accuracy on complete output sequences and checked for failures missed by sampled tests. Further, we assessed training data improvements informed by program inspection and failure analysis, finding gains in sequence accuracy on several tasks. We examined a counting task where a program achieved 100% accuracy on sampled test sequences but failed on another input using the symbols and lengths allowed during training. The program counted correctly, but the required answer was missing from its output labels. After adding training examples with larger counts, the retrained models passed exhaustive evaluation over the bounded domain. In the frequency ranking task, inspection showed how correct counts contributed to incorrect answers in programs trained before and after changes to the training data. Finally, we studied constraints imposed by TP training and operations. The TP training process can lead to accuracy losses late in training. Our reference constructions also show how composing the restricted TP operations can use many layers or large lookup tables.