Learning Elementary Cellular Automata with Transformers

Authors: mikhail  burtsev
Venue: NeurIPS_cc_2024_Workshop_MATH-AI
Published: 2024-09-26 11:26:25

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.

Main paper: 
 introduction: 
 Large Language Models (LLMs) have become valuable tools in mathematics, demonstrating impressive capabilities in problem-solving and reasoning tasks. Notably, OpenAI's o1 model achieved a ranking among the top 500 students in the US in a qualifier for the USA Math Olympiad (AIME) [1] . Despite these successes, LLMs still face challenges in reasoning [2] [3] [4] [5] [6] and planning [7] , particularly when required to infer and apply underlying rules from data.These observations raise a question: Do the limitations of LLMs in reasoning stem from the nature of their training data, the procedures employed during training, or inherent architectural constraints? Transformers [8] , which form the backbone of many modern LLMs, are universal approximators [9] [10] [11] [12] and theoretically capable of simulating Turing machines using intermediate computational steps, making them Turing complete and, by extension, capable of formal reasoning [13] [14] [15] [16] . Earlier studies found that transformers can be trained to perform symbolic integration and solving differential equations [17] , as well as symbolic regression [18, 19] .In our study, we continue the line of research exploring the trainability of Transformers for mathematical reasoning tasks by focusing on Elementary Cellular Automata (ECAs). The simplicity and clarity of this "toy" problem make it an ideal testbed for assessing the ability of Transformers to abstract and apply logical rules. By demonstrating that Transformers can learn and generalize Boolean functions of fixed arity inherent in ECAs, we aim to evaluate their ability to infer, generalize, and apply logical rules solely from observed data, without relying on memorization.  
methods: 
 An Elementary Cellular Automaton (ECA) is a one-dimensional, dynamical system in which space and time are discrete. Let r ∈ N : r ≥ 1 be the neighbourhood radius in space represented by a regular lattice of W ∈ N : W ≥ 2r + 1 identical, locally-interconnected cells with a binary state space, S = {0, 1}. The ECA's global state, x ∈ S W , is a lattice configuration specified by the values of all the states of all cells in the lattice at a given time. This state evolves deterministically in synchronous, discrete time steps according to a global map g ρ : S W → S W defined by a local rule ρ : S 2r+1 → S, so [g ρ (x)] w = ρ(x w-r , . . . , x w , . . . , x w+r ). The sequence of states an ECA passes through during its space-time evolution, O T (x) = [x, g ρ (x), g ρ (g ρ (x)), . . . , g oT -1 ρ (x)], defines its trajectory or orbit from an initial condition (configuration) x for T ∈ N : T ≥ 1. Examples of ECA orbits are visualized on the figure 1A.Let consider four modifications of learning tasks designed to evaluate different aspects of predictive modeling and rule inference in ECAs (see Fig. 1B ).Orbit-State (O-S), given an orbit O T (x) = [x (0) , x (1) , . . . , x (T -1) ] where x (0) ∈ S W , the objective is to predict the state x (T ) at time T . Orbit-Orbit (O-O), given an orbit O k (x) for some k < T predict the subsequent states up to time T , generating We generated a dataset with the CellPyLib [20] 1 for fixed lattice size W = 20 and neighborhood radius r = 2. This configuration results in a total of 2 2 2r+1 ≈ 4.3 × 10 9 possible boolean functions defining local rules. For each sample in the dataset, both the initial state and the local rule ρ were generated randomly. We then computed the orbit for T = 20 time steps using these parameters. The training dataset consists of 9.5 × 10 5 and the test of 10 5 samples. Importantly, the local rules included in the test set are exclusive and not present in the training set. This separation ensures that the model's performance reflects its ability to generalize to unseen rules, rather than simply memorizing the training data. To measure the quality of the model's predictions, we use per-bit accuracy averaged over free runs, which calculates the proportion of bits correctly predicted in the output sequences.EQUATION 
results and discussion: 
 We first assess whether the samples in our dataset provide sufficient information for the model to learn the underlying dynamics successfully. The minimal number of time steps T min needed to recover the local rule ρ from the observed orbit can be estimated analogous to the coupon collector's problem T min = 2 2r+1 (ln 2 2r+1 + γ)/W ≈ 6.47, where γ is the Euler-Mascheroni constant (γ ≈ 0.5772). Training of the model for the next state prediction (O-S task) confirms this theoretical estimation by showing that the accuracy plateaus after 8 time steps (Fig. 2A ). Therefore, we choose to use input orbits of 10 steps O 10 as our main setting.Successful learning of the O-S task demonstrate that the Transformer model is capable of generalizing not only over initial conditions for a particular function -commonly the focus in studies of transformer trainability in CA domain [21] [22] [23] [24] [25] [26] [27] [28] [29] [30] [31] [32] [33] [34] [35] -but also across different Boolean functions of fixed arity (5 in our case). This indicates that the model has learned to abstract a class of rules.Next, we investigated the Transformer's ability to plan ahead by predicting future states beyond the immediate next state. Specifically, we trained the model to predict the state at time x (T +k) for look-ahead steps k ∈ {1, 2, 3}. As presented in Figure 2B , this task proved to be significantly more challenging. While the average accuracy for next-state prediction (O-S task) was 0.96, it dropped to 0.80 for k = 1 and fell below 0.75 for k = 2 and k = 3.To determine whether this decline was due to the Transformer's architecture or the training objective, we explored whether accuracy could be improved by training the model to predict intermediate steps. This approach is analogous to the "chain-of-thought" method [36] used in large language models for in-context learning. We employed the Orbit-Orbit (O-O) task, training the model to predict the next four states in parallel. The results, also shown in Figure 2B , indicate that the model retains predictive abilities when skipping one step but struggles with skipping two or three steps. These results suggest that learning to store a hidden representation of intermediate states (as in the O-O training) is easier for the model but harder to generalize over multiple time steps. In contrast, developing a hidden representation of the underlying rule (as in the O-SR training) is more challenging initially but facilitates better generalization to longer planning horizons. This implies that explicitly encouraging the model to infer the generating rule can enhance its ability to make longer-term predictions by reinforcing the internalization of the system's dynamics.Finally, we explored the scenario where the local rule ρ is explicitly provided to the model, corresponding to the Rule and Orbit-State (RO-S) task. Intuitively, this should be the easiest task for the Transformer, as it eliminates the need to infer the rule from the observed data. As shown in Figure 2B , the Transformer indeed learns to apply the given rule for next-state prediction with near-perfect accuracy for k = 0 and k = 1. Surprisingly, however, the performance for look-ahead steps k = 2 and 3 drops to approximately 0.75, similar to the O-O training scenario where the rule is not provided. This unexpected decline hints that even with explicit access to the rule, the model struggles to predict future states beyond the two next steps. We hypothesize that this limitation arises from difficulties in learning effective hidden representations of the intermediate states required for multi-step predictions. Despite having the rule, the Transformer may not adequately capture and propagate the necessary state information over multiple time steps without explicit intermediate context. Average per bit accuracyStep 1Step 2Step 3Step 4 2C ). As expected, the success of these models in this task correlates with their next-token prediction accuracy (refer to the first group of bars in Figure 2B ). The RO-S model exhibits the best performance, followed closely by the O-O and O-SR models, while the O-S model shows significantly weaker performance.When comparing autoregressive generation (AR) for rollouts of 2, 3, and 4 steps to planning performance for the same look-ahead steps (LA) (Figure 2D ), we observe that all models perform better in the autoregressive inference mode. This suggests that the models are more adept at short-term state-bystate prediction than at planning multiple steps ahead without intermediate context. The superior performance in the autoregressive mode highlights the models' reliance on immediate past states to inform their predictions.Our planning results (Figure 2B ) demonstrate a sharp decline in performance across all training methods starting from a look-ahead of 2 steps. This decline may be due to the Transformer's inability to store more than one ECA state in its hidden representations or because the task requires sequential application of the local rule equal to the planning horizon. The O-O training method explicitly provides memory for intermediate states through tokens and associated hidden vectors, effectively ruling out the first explanation.The alternative explanation points to the limitations of sequential computation within the Transformer architecture, which is bounded by the number of layers. To test this, we examined how the ability to predict future states scales with network depth. Figure 3A shows that the O-O model begins to accurately predict the next state with 2 layers. Predicting two steps ahead requires at least 4 layers, three steps necessitate 7 layers, and four steps demand 10 layers. These results suggest that each additional planning step requires more computational layers, highlighting a direct relationship between the Transformer's depth and its capacity to model longer sequences of ECA dynamics. 
conclusions: 
 In this study, we have demonstrated that Transformer models possess the ability to learn and generalize the underlying dynamics of Elementary Cellular Automata. By designing specific tasks and training regimes, we showed that Transformers can abstract the governing rules and predict future states with notable accuracy. However, their performance diminishes when required to plan multiple steps ahead without intermediate context, highlighting limitations in storing and propagating state information over longer sequences. Our analysis reveals that including future states or rule prediction in the training loss improves the performance of next-state prediction and, as a result, enhances the overall performance in autoregressive generation. This finding might guide training strategies for LLMs to improve their perplexity and reasoning capabilities. We also confirmed that the model's depth plays an important role in extended sequential computations required for complex reasoning; thus, recurrence and adaptive computation time are promising directions for dynamic control of the model depth. These insights contribute to a deeper understanding of how neural networks can abstract rules and point toward future research directions in improving their planning and generalization skills. 
 
 Appendix: Dataset and a source code for experiments are available at https://github.com/burtsev/ TransformerECA. 


Decision: The reviewers are generally positive about the work and excited about the insights that it uncovers about the learning processes of transformers. We recommend that the authors include a more comprehensive comparison with related work in their revision.

Reviews: review: title: title: Good analysis on the learning dynamics and capabilities of Transformers
main_review: main_review: **Summary**

In this paper the authors have tried to study the ability of Transformers (and in turn LLMs) to infer and apply abstract rules for long-term planning tasks. They use Elementary Cellular Automata as a “toy example” to test the trainability of Transformers for mathematical and reasoning tasks. 

**Strengths**

The different experiments conducted and the related plots were well explained. The authors did a nice analysis of the internal learning dynamics and learning capabilities of Transformers. The work tries to hypothesize and conduct some experiments to prove that the Transformers find it difficult to learn effective hidden representations of the intermediate states required for multi-step predictions. The Transformers get relatively more capable at short-term state-by-state prediction than at planning multiple steps ahead without intermediate context. 
The authors conduct experiments to claim that Transformers can be made better at long-horizon multi-step predictions by making them learn the dynamics explicitly by making them predict the underlying rules. 
This is aligned to learning some world representations or world dynamics in order to be able to do longer horizon state predictions and planning. 
This also leads the authors to conduct experiments that try to explain the relationship between capability and capacity of the Transformers. The experiments try to explain that storing additional context for additional planning steps requires more computational layers and can be done by increasing the depth of the model.

These results might help in designing training pipelines for larger architectures and models for long horizon planning and reasoning. 

**Weaknesses**

The paper can mention some of the training details along with some hyper params (although the authors mention sharing the code upon acceptance). 

**Limitations**

It would have been nice to have an additional supporting test-bed/toy example to make the claims more effective and broadly applicable.
Was there enough data provided to test the generalization capability of the models?
title: title: Review
main_review: main_review: This paper studies the learning of elementary cellular automata (ECA) with Transformers. More precisely, the authors evaluate their models on two tasks: predicting the next state of ECAs and planning ahead i.e. predicting the steps ahead from the observed ECA. They consider different training tasks: Orbit-State (O-S), Orbit-Orbit (O-O), Orbit-State and Rule (O-SR) and Rule and Orbit-State (RO-S) and report the performance of the model on the evaluation tasks. They basically observe that when training with RO-S, they got the best performance on predicting the next state. However, the model then is very bad at planning. They find that the best model for planning is O-SR and hypothesize that forcing the model to learn to generate the rule is the reason for these benefits. Lastly, they do a study with respect to the depth of the model and show that deeper networks are better at planning. 

Overall, I really liked this paper and this is why I advocate for acceptance of this paper. It is well explained, studies a problem that has not been studied in Transformers and math, and the findings are interesting. 

A few suggestions that may be interesting to the authors: 

- First, maybe I am wrong but is this ECA learning analogous to learning automata? There are some works on learning automata with Transformers and I think it would be good to cite them / reflect your work with respect to these works. 

- I am not familiar at all with ECAs. Maybe it would be useful to add a small example / drawing in the Methods section to give concrete examples of an orbit and a local rule? This would greatly help the understanding.

- An interesting method to try is a more explicit comparison with Chain-of-Thought (CoT). The authors claim that O-O is the same as CoT but maybe we can make a more explicit CoT? For instance, putting the local rule output next to the intermediate steps to guess. The reason I am insisting on CoTs is because I really believe that if a model learns with CoT, it should be able to do planning. Not entirely related to this but [2] show that with CoT, they can achieve reasonable length generalization on any task. 

- I believe that there are other interesting tricks to maybe achieve better planning (inspired from the automata and Transformers literature):

* data mixing: here, it means that i would create a dataset made of O-S tasks but for some samples, the state to predict is x^T, others we need to predict x^T+1, or x^T+2 etc. This is also another implicit way to enforce the model to infer the rule. 

* recurrent architectures: how would LSTMs work on this task? I believe that they should do great even when trained on O-S? this would be interesting to see. Other variants would also be worth studying like Universal Transformers [3] or Looped Transformers [4]. 



[1] Liu, B., Ash, J.T., Goel, S., Krishnamurthy, A. and Zhang, C., 2022. Transformers learn shortcuts to automata. arXiv preprint arXiv:2210.10749.

[2] Hou, K., Brandfonbrener, D., Kakade, S., Jelassi, S. and Malach, E., 2024. Universal Length Generalization with Turing Programs.

[3] Dehghani, M., Gouws, S., Vinyals, O., Uszkoreit, J. and Kaiser, Ł., 2018. Universal transformers.

[4] Giannou, A., Rajput, S., Sohn, J.Y., Lee, K., Lee, J.D. and Papailiopoulos, D., 2023, July. Looped transformers as programmable computers.
title: title: This paper studies whether Transformer models can learn Elementary Cellular Automata. The authors find that the model's performance drops quickly in multi-step planning tasks. The paper also propose including future states and developing a hidden rule to facilitate better in long planning horizons.
main_review: main_review: This paper explores whether Transformer models can learn and generalize the rules of Elementary Cellular Automata (ECA). By training Transformers on sequences of states, the authors found that the models can generalize across Boolean functions and effectively abstracting the underlying rules.  Although the models exhibit high accuracy in predicting the next state, the paper finds that their performance drops sharply in multi-step planning tasks. By including future states or hidden rule predictions in the training loss, the models' ability to form internal representations of the rules is enhanced, leading to improvements in longer planning horizons.

Strengths:

1. The empirical analysis of the algorithm is interesting. The author compares different methods (O-S, O-O, O-SR, etc) to illustrate the 
 relationship between accuracy and planning steps. Moreover, The authors also study the role of depth in modeling long sequences of ECA dynamics.

2.  The research results provide guidance for improving training strategies of large language models when facing long term planning tasks, such as including longer horizon hidden rules in the loss function.

Questions and Suggestions: I'm not an expert in Transformer Model so please forgive me if my questions are naive.

1: May one ask in your setting, how will the lattice and radius sizes influence the results? Thanks!

2: Are there any theoretical guarantees in learning theory between transformer’s ability of planning and the number of layers?

3: Just out of curiousity, have you ever considered more complicated learning tasks rather than Elementary Cellular Automata? In other words, do you expect that your finding could be generalized to other kinds of multistep learning tasks?
title: title: A comprehensive work, but not novel enough.
main_review: main_review: This paper explores the potential of Transformers to learn and generalize the rules governing Elementary Cellular Automata (ECAs). By training on sequences generated with random initial conditions and local rules, the study examines whether Transformers can abstract from specific instances to a generalized understanding of Boolean functions. Despite achieving high accuracy in next-state prediction, the models face significant challenges in multi-step planning tasks, which highlights the limitations in their capacity for abstract reasoning and dynamic problem-solving. The paper suggests improvements in model architecture, such as increasing depth and incorporating dynamic elements, to enhance performance on complex reasoning tasks.

Strengths:

1.	The study innovatively uses ECAs as a testbed to examine the generalization capabilities of Transformers across different Boolean functions, highlighting a novel way of evaluating deep learning models' abstract reasoning capabilities.

2.	The paper provides a thorough examination of the models' performance, including different setups such as next-state prediction and multi-step planning. This comprehensive approach allows for a detailed understanding of the models' strengths and limitations.

Weaknesses:

1.	While ECAs provide a clear and controlled environment for testing, they might not capture the complexity of real-world scenarios that Transformers are expected to handle, potentially limiting the broader applicability of the results.

2.	The suggestion to increase model depth might lead to more resource-intensive models, which could be impractical for certain applications due to increased computational and energy costs.

3.	Although the paper aims to teach models to generalize, the specific focus on ECAs and Boolean functions might lead the models to overfit to these types of tasks, potentially reducing their effectiveness on more varied real-world problems.

4. Limited Novelty, learning a Turing Machine has been studied by previous work[1-3], this is a more simple case.



[1] Wei C, Chen Y, Ma T. Statistically meaningful approximation: a case study on approximating turing machines with transformers[J]. Advances in Neural Information Processing Systems, 2022, 35: 12071-12083.

[2] Ryoo M S, Gopalakrishnan K, Kahatapitiya K, et al. Token turing machines[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023: 19070-19081.

[3] Pérez J, Barceló P, Marinkovic J. Attention is turing-complete[J]. Journal of Machine Learning Research, 2021, 22(75): 1-35.