Learning Elementary Cellular Automata with Transformers

TL;DR

This paper investigates whether Transformer models can learn and generalize the rules of Elementary Cellular Automata, demonstrating that with appropriate training strategies, they can abstract rules and improve multi-step reasoning.

Contribution

The study shows that Transformers can effectively learn and generalize cellular automata rules, especially when trained with future states or rule prediction, and highlights the importance of model depth for complex reasoning.

Findings

Transformers achieve high accuracy in next-state prediction.

Including future states in training improves long-term planning performance.

Deeper models enhance sequential reasoning capabilities.

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.