Log In Sign Up

Unveiling Transformers with LEGO: a synthetic reasoning task

by   Yi Zhang, et al.

We propose a synthetic task, LEGO (Learning Equality and Group Operations), that encapsulates the problem of following a chain of reasoning, and we study how the transformer architecture learns 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 are able to understand (to some extent) some of the attention heads as well as how the information flows in the network. Based on these observations we propose a hypothesis that here pretraining helps merely due to being a smart initialization rather than some deep knowledge stored in the network. 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 ability to generalize to simple variants of the main task, and moreover we find that one can prevent such shortcut with appropriate architecture modification or careful data preparation. Motivated by our findings, we begin to explore the task of learning to execute C programs, where a convolutional modification to transformers, namely adding convolutional structures in the key/query/value maps, shows an encouraging edge.


page 7

page 9

page 27

page 28


The NLP Task Effectiveness of Long-Range Transformers

Transformer models cannot easily scale to long sequences due to their O(...

Pay Attention to MLPs

Transformers have become one of the most important architectural innovat...

Modifying Memories in Transformer Models

Large Transformer models have achieved impressive performance in many na...

Graph-to-Graph Transformer for Transition-based Dependency Parsing

Transition-based dependency parsing is a challenging task for conditioni...

Linguistic Knowledge and Transferability of Contextual Representations

Contextual word representations derived from large-scale neural language...

Thinking Like Transformers

What is the computational model behind a Transformer? Where recurrent ne...

LOFT: Finding Lottery Tickets through Filter-wise Training

Recent work on the Lottery Ticket Hypothesis (LTH) shows that there exis...