(How) Can Transformers Predict Pseudo-Random Numbers?

TL;DR

This paper investigates the capacity of Transformer models to learn and predict pseudo-random sequences generated by linear congruential generators, revealing their ability to develop algorithmic strategies for sequence prediction and generalization.

Contribution

It demonstrates that Transformers can learn to predict LCG sequences with unseen parameters and moduli, uncovering their internal mechanisms and generalization capabilities.

Findings

Transformers can predict LCG sequences with unseen parameters up to 2^{32}.

Models generalize to unseen moduli up to 2^{16} by estimating the modulus.

Prediction accuracy sharply improves at a model depth of 3.

Abstract

Transformers excel at discovering patterns in sequential data, yet their fundamental limitations and learning mechanisms remain crucial topics of investigation. In this paper, we study the ability of Transformers to learn pseudo-random number sequences from linear congruential generators (LCGs), defined by the recurrence relation $x_{t+1} = a x_t + c \;\mathrm{mod}\; m$. We find that with sufficient architectural capacity and training data variety, Transformers can perform in-context prediction of LCG sequences with unseen moduli ($m$) and parameters ($a,c$). By analyzing the embedding layers and attention patterns, we uncover how Transformers develop algorithmic structures to learn these sequences in two scenarios of increasing complexity. First, we investigate how Transformers learn LCG sequences with unseen ($a, c$) but fixed modulus; and demonstrate successful learning up to $m = 2^{32}$. We find that models learn to factorize $m$ and utilize digit-wise number representations to make sequential predictions. In the second, more challenging scenario of unseen moduli, we show that Transformers can generalize to unseen moduli up to $m_{\text{test}} = 2^{16}$. In this case, the model employs a two-step strategy: first estimating the unknown modulus from the context, then utilizing prime factorizations to generate predictions. For this task, we observe a sharp transition in the accuracy at a critical depth $d= 3$. We also find that the number of in-context sequence elements needed to reach high accuracy scales sublinearly with the modulus.