cs.AIDec 1, 2025
SaveTesting Transformer Learnability on the Arithmetic Sequence of Rooted Trees
Organizations: Department of Mathematics, University of Bologna, Italy · Gatsby Computational Neuroscience Unit, UCL, UK · Department of Information Engineering, University of Florence, Italy
Abstract
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 defines an arithmetic text with measurable statistical structure. A transformer network (the GPT-2 architecture) is trained from scratch on the first 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 , well above the baseline, and its performance remains stable on test blocks located at -, 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 as a function of . It is numerically found that scales sublinearly in up to Contucci et al., 2026 .
Figure 2: ( Left ) Jensen–Shannon divergence between chunks of consecutive Dyck words in our dataset with . The largest value of the divergence, reached between the very first chunk and the very last (upper right corner), is approximately equal to , i.e., the distribution difference remains quite small despite the orders of magnitude span. ( Right ) Detail of the Jensen–Shannon divergence of the last chunk of trees with respect to all the previous ones of equal size.
Figure 3: An example of a numerical sequence (integers between and ) converted into a sequence of Dyck words. The tokenizer splits the sequence of Dyck words into 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 is split into two different tokens). After training, the attention matrix 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 and . 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 and are the most relevant for , as and are the largest entries of the attention matrix.
Figure 4: Scheme of the dataset management. The sequence of consecutive Dyck words has been split in blocks of Dyck words (except for the first block, that is long ), each further separated into two portions: the first 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 consecutive elements were produced for testing in the range , and respectively.
Figure 5: Loss curves for a model trained via NTP, obtained for different tokenizer’s vocabulary size , namely , , . The loss value, computed as in Eq. ( 4 ), is rescaled by , loss value corresponding to a uniform distribution.
Figure 6: Accuracy obtained from prompts of length , at , for different generated sequence length . Each box corresponds to a value of and presents three test sets (marked with different colors) obtained from the first consecutive Dyck words after , and (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 , or ). Results obtained on an architecture with layers, heads per layer and , trained on our dataset of Dyck words.
Figure 7: Accuracy for words ( , above ) and Kullback–Leibler divergence for word distribution ( , below ) are reported at different values of , over a set of 128 different inputs of tokens. The generated outputs are 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 words following (red), (blue) and (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 , it fails in predicting the order of Dyck words exhibiting an accuracy comparable to a random guess. Results obtained on an architecture with layers, heads per layer and , trained on our dataset of Dyck words.
Figure 8: Average precision, recall and scores for various Dyck words over test sentences of length at and dictionary size . Only words appearing at least 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 different input prompts of tokens. The model was trained with a context window of , 12 layers with 12 heads per layer and trained on a dataset covering integers from till with vocabulary size . 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 -axis. The LLM is trained with a context window of , 12 layers with 12 heads per layer and trained on a dataset covering integers from till with vocabulary size . Frequencies are obtained by feeding the trained model sequences of 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 synthetic prompts consisting of BPE tokens sampled independently from the empirical marginal token distribution of the training corpus ( circles ) and 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 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 , 12 layers with 12 heads per layer and trained on a dataset covering integers from to with vocabulary size .
Figure 13: Distribution of the log-likelihood obtained via kernel density estimation from true sequences from the test set consisting of Dyck words . Dots denote the average values of the likelihood obtained by sampling from the dataset 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 , it is trained on the dataset of the first Dyck words with context window size and vocabulary size . In the panels, one of such parameters is altered by keeping the remaining parameters fixed: from left to right, , , and . Each dot represents the accuracy of a single architecture trained with the indicated hyperparameters on a sentence of length at a test scale ranging from to , 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 and . The loss value is rescaled by , 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 scale in the experiments above, it is trained on the dataset of the first 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 , the context window and the masking probability at training time . Each dot represents the accuracy of a batch of sequences, with the architecture trained with the indicated hyperparameters on a sentence of length at a test scale ranging from to , 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 ; as increases, tokens often encompass multiple words, creating a harder masking task during training. Sequence length does not significantly impact the training procedure, aside from requiring more refined learning rate tuning at . Optimal performance is achieved at a masking probability of
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.
More in this topic
Transformer
Transformers are neural architectures that use attention to model relationships among sequence elements in language, vision, and scientific data.
See all 540 papers