Unveiling Transformers with LEGO: a synthetic reasoning task

Abstract: We propose a synthetic reasoning task, LEGO (Learning Equality and Group Operations), that encapsulates the problem of following a chain of reasoning, and we study how the Transformer architectures learn this task. We pay special attention to data effects such as pretraining (on seemingly unrelated NLP tasks) and dataset composition (e.g., differing chain length at training and test time), as well as architectural variants such as weight-tied layers or adding convolutional components. We study how the trained models eventually succeed at the task, and in particular, we manage to understand some of the attention heads as well as how the information flows in the network. In particular, we have identified a novel \emph{association} pattern that globally attends only to identical tokens. Based on these observations we propose a hypothesis that here pretraining helps for LEGO tasks due to certain structured attention patterns, and we experimentally verify this hypothesis. We also observe that in some data regime the trained transformer finds ``shortcut" solutions to follow the chain of reasoning, which impedes the model's robustness, and moreover we propose ways to prevent it. Motivated by our findings on structured attention patterns, we propose the LEGO attention module, a drop-in replacement for vanilla attention heads. This architectural change significantly reduces Flops and maintains or even \emph{improves} the model's performance at large-scale pretraining.

• The paper introduces the LEGO task to evaluate Transformer reasoning and highlights ALBERT's advantage in handling length extrapolation.
• The paper demonstrates that pretraining boosts both classical generalization and structured iterative reasoning through refined attention mechanisms.
• The paper identifies crucial attention patterns, such as association and manipulation, which underpin robust performance on synthetic reasoning tasks.

Analyzing Transformers through the LEGO Task: Insights and Implications

The paper "Unveiling Transformers with LEGO: a synthetic reasoning task" presents an innovative approach to understanding Transformer architectures by introducing the LEGO (Learning Equality and Group Operations) synthetic task. This task is designed to encapsulate reasoning paradigms and dissect the learning dynamics of Transformers. By focusing on both architectural choices and data compositional effects, the research offers granular insights into how Transformers manage reasoning tasks.

At its core, the LEGO task serves as a controlled experimental setting to probe the reasoning capabilities of Transformers. The emphasis is on the sequence of reasoning, with key aspects such as variable assignments and operations defined by group actions. The ability to generalize, especially under distribution shifts that vary chain length, becomes a focal point of the study, enabling a nuanced understanding of extrapolation beyond classical generalization.

Key Findings

1. Performance and Generalization: Both BERT and ALBERT architectures demonstrate strong performance in classical generalization. However, ALBERT exhibits a marked advantage in length extrapolation, attributed to its iterative design which mirrors the iterative nature of reasoning tasks. The study suggests that ALBERT is inherently more suited for tasks that can be algorithmically reduced to iterative operations akin to a "for loop."
2. Effect of Pretraining: Pretraining significantly aids both classical generalization and length extrapolation in LEGO tasks. Interestingly, the pretraining advantage appears to stem more from structural learning (attention patterns) rather than direct knowledge transfer, challenging conventional views on pretraining's role.
3. Attention Patterns: Specific attention patterns, notably association (long-range identity matching) and manipulation (short-range operations), emerge as crucial in the successful performance of Transformers on LEGO. The paper introduces the LEGO attention module, which leverages these patterns to reduce computational costs without sacrificing performance.
4. Shortcut Solutions and Robustness: Transformers sometimes resort to shortcut solutions, achieving correct outcomes through unintended paths, which may hinder robustness. Preventative measures, such as pretraining that encodes structured attention patterns, help models avoid these pitfalls, underscoring the importance of architectural and training strategies in fostering robust reasoning capabilities.

Implications and Future Directions

The introduction of a synthetic reasoning task such as LEGO provides a valuable framework for dissecting the behaviors of Transformer models in a controlled setting, offering implications for both theoretical understanding and practical applications. The findings suggest that iterative architectures like ALBERT could be preferable in environments requiring structured reasoning. Additionally, the identification and exploitation of specific attention patterns provide pathways to designing more efficient models, potentially applicable beyond synthetic tasks to real-world scenarios where reasoning and generalization beyond training data are required.

Future research could expand by exploring larger task sizes or more complex group operations, potentially increasing the applicability of the insights gained. Moreover, as our understanding deepens, these insights could inform the development of novel Transformer variants or alternative architectures focused on reducing model size while maximizing generalization capabilities, a key consideration in resource-constrained applications.

In conclusion, the study provides a profound exploration of how Transformers learn reasoning tasks, with the LEGO task serving as an exemplary model for this inquiry. Through its systemic investigation of architecture, data effects, and pretraining, the research informs current practices and sets the stage for future advancements in AI models that more closely mimic structured reasoning processes.