Transformers Provably Learn Sparse XOR with Polylogarithmic Parameters

TL;DR

This paper proves that single-layer, two-head Transformers can efficiently learn sparse XOR functions with polylogarithmic parameters, surpassing FFNNs, and demonstrates their rapid feature discovery and generalization capabilities.

Contribution

It provides the first theoretical analysis showing Transformers can learn sparse XOR with minimal parameters, breaking the FFNN parameter bottleneck.

Findings

Transformers can learn sparse XOR with O(polylog(d)) parameters.

Exact softmax attention is crucial for rapid feature discovery.

Transformers generalize well with finite data, as shown by sample complexity bounds.

Abstract

Learning sparse parity functions has become a theoretical testbed for studying feature learning in neural networks. However, existing analyses primarily focus on Feed-Forward Neural Networks (FFNNs). Meanwhile, theoretical understanding of Transformers in this setting remains limited, despite their empirical success and structural suitability for discovering sparse support over long sequences. To address this gap, we analyze how a single-layer, two-head Transformer learns the sparse XOR problem. Considering samples $(\mathbf{x}, y) \in \lbrace\pm 1\rbrace^d \times \lbrace\pm 1\rbrace$, where the label is defined by $y = -x_{i^*} x_{j^*}$ for some unknown $i^*, j^* \in [d]$, we prove that, with only $O(\mathrm{polylog}(d))$ trainable parameters, Transformers can successfully discover the relevant features and drive the loss for every input to nearly 0 with one gradient step. This result establishes that Transformers break the fundamental $\Omega(d)$ parameter bottleneck inherent to FFNNs for this problem. Furthermore, we empirically show that this rapid feature discovery is uniquely driven by the exact softmax attention, outperforming common substitutes such as linear or component-wise attention. Finally, we provide a theoretical sample complexity bound for learning from finite data, demonstrating the generalization ability of Transformers in this task.