Learning Elementary Cellular Automata with Transformers

TL;DR

The paper investigates whether Transformers can learn abstract, generalizable rules of Elementary Cellular Automata ($ECAs$). Using a dataset with $W=20$, $r=2$, and $T=20$, plus four predictive tasks (O-S, O-O, O-SR, RO-S) and unseen test rules, the authors study rule inference and short- versus long-horizon forecasting. They find that next-state prediction generalizes across Boolean functions of arity $5$, but longer-horizon planning requires additional context or explicit rule representations; depth growth is necessary for extended sequential reasoning. The results suggest training strategies for longer-horizon reasoning in LLMs, including recurrence and adaptive computation time to improve planning and rule-based generalization.

Abstract

Large Language Models demonstrate remarkable mathematical capabilities but at the same time struggle with abstract reasoning and planning. In this study, we explore whether Transformers can learn to abstract and generalize the rules governing Elementary Cellular Automata. By training Transformers on state sequences generated with random initial conditions and local rules, we show that they can generalize across different Boolean functions of fixed arity, effectively abstracting the underlying rules. While the models achieve high accuracy in next-state prediction, their performance declines sharply in multi-step planning tasks without intermediate context. Our analysis reveals that including future states or rule prediction in the training loss enhances the models' ability to form internal representations of the rules, leading to improved performance in longer planning horizons and autoregressive generation. Furthermore, we confirm that increasing the model's depth plays a crucial role in extended sequential computations required for complex reasoning tasks. This highlights the potential to improve LLM with inclusion of longer horizons in loss function, as well as incorporating recurrence and adaptive computation time for dynamic control of model depth.

Figures (3)

• Figure 1: Learning Elementary Cellular Automata (ECA) with Transformers.A. Examples of training samples. Orbit of ECA is a sequence of binary strings of size $W = 20$. First $k = 10$ states marked by red rectangle encode Transformer input. B. Given a part of the orbit Transformer with full attention learns to predict the next state (O-S), the next few steps (O-O), the next state and a rule (O-SR), or predict the next state given a rule and an orbit (RO-S).
• Figure 2: Transformer learns to predict the next state of ECA but struggles to plan ahead. A. Accuracy of the next state prediction for ECA orbit with size 20 generated by Boolean function of 5 arguments for different input state lengths. B. Planning accuracy for different training settings (see the main text for details). C. Accuracy for autoregressive generation of ECA orbit. D. Comparison of autoregressive (AR) and look ahead (LA) predictions.
• Figure 3: Adding layers improves prediction of ECA orbit. Accuracy of O-O training for different number of layers.