Is logical analysis performed by transformers taking place in self-attention or in the fully connected part?

TL;DR

This paper investigates whether logical analysis in transformers occurs in self-attention or in the fully connected layers, revealing that self-attention can perform logical operations, challenging traditional views.

Contribution

The study introduces a handcrafted encoder layer performing logical analysis within self-attention and analyzes how models learn to utilize self-attention versus fully connected layers.

Findings

Self-attention can perform logical analysis, not just information aggregation.

Gradient descent can get stuck at undesired zeros, affecting learning.

Explicit methods to avoid these zeros improve model training.

Abstract

Transformers architecture apply self-attention to tokens represented as vectors, before a fully connected (neuronal network) layer. These two parts can be layered many times. Traditionally, self-attention is seen as a mechanism for aggregating information before logical operations are performed by the fully connected layer. In this paper, we show, that quite counter-intuitively, the logical analysis can also be performed within the self-attention. For this we implement a handcrafted single-level encoder layer which performs the logical analysis within self-attention. We then study the scenario in which a one-level transformer model undergoes self-learning using gradient descent. We investigate whether the model utilizes fully connected layers or self-attention mechanisms for logical analysis when it has the choice. Given that gradient descent can become stuck at undesired zeros, we explicitly calculate these unwanted zeros and find ways to avoid them. We do all this in the context of predicting grammatical category pairs of adjacent tokens in a text. We believe that our findings have broader implications for understanding the potential logical operations performed by self-attention.