Transformers Pretrained on Procedural Data Contain Modular Structures for Algorithmic Reasoning

TL;DR

The paper investigates how pretraining transformers on procedural, semantically neutral data shapes algorithmic reasoning and where in the network these benefits reside. Using small GPT-2–style models trained on six procedural tasks and evaluated with diverse diagnostic tasks, the authors show that procedural data induce distinct, modular inductive biases that are often captured in attention layers but can also reside in MLPs. They demonstrate that these biases are largely composable: combining components from different pretraining tasks yields robust, multi-task improvements. The findings point toward data-efficient, modular pretraining approaches that can separate reasoning abilities from semantic knowledge, with potential for more robust and adaptable language models.

Abstract

Pretraining on large, semantically rich datasets is key for developing language models. Surprisingly, recent studies have shown that even synthetic data, generated procedurally through simple semantic-free algorithms, can yield some of the same benefits as natural language pretraining. It is unclear what specific capabilities such simple synthetic data instils in a model, where these capabilities reside in the architecture, and how they manifest within its weights. In this short paper, we identify several beneficial forms of procedural data, together with specific algorithmic reasoning skills that improve in small transformers. Our core finding is that different procedural rules instil distinct but complementary inductive structures in the model. With extensive ablations and partial-transfer experiments, we discover that these structures reside in different parts of the model. Attention layers often carry the most transferable information, but some pretraining rules impart useful structure to MLP blocks instead. Most interestingly, the structures induced by multiple rules can be composed to jointly reinforce multiple capabilities. These results suggest an exciting possibility of disentangling the acquisition of knowledge from reasoning in language models, with the goal of improving their robustness and data efficiency.

Figures (4)

• Figure 1: (a) We pretrain small transformers on various forms of procedural data, then fine-tune them on a series of diagnostic tasks. The data is generated from formal languages (b) or simple algorithms such as elementary cellular automata (c). In $k$-Dyck examples, matching brackets are color-coded. For Stack, 'P' is the pop operation. For Stack, Identity, and Set, '|' acts as a separator between the input and the expected output, on which the loss is computed (bold tokens).
• Figure 2: Models pretrained on carefully-chosen procedural data (blue) significantly improve specific downstream diagnostic tasks, compared to a random initialisation (light gray) or pretraining on natural language (dark gray, shown for Multiplication). See Appendix \ref{['app:additional-results']} for full results, including cases where procedural pretraining brings no benefit.
• Figure 3: Results of the selective transfer of procedurally-pretrained weights. Each color is a different type of transfer: $\mathbf{\mathcal{T}_{\text{full}}}$ , $\mathbf{\mathcal{T}_{\text{mlp}}}$ , $\mathbf{\mathcal{T}_{\text{attn}}}$. The horizontal dotted line is the baseline performance with a random initialisation. See Appendix \ref{['app:additional-results']} for full results including the variance across seeds.
• Figure 4: Relative downstream improvement using pretrained weights with perturbations (noise or shuffling). See Appendix \ref{['app:additional-results']} for full results and details.