🤖 AI Summary
This work investigates the information bottleneck inherent in modern deep learning architectures when performing indexing tasks, with a particular focus on performance disparities between indexing at the beginning versus the end of a sequence. By introducing the notion of causal complexity, the authors theoretically demonstrate that low-parameter RNNs, state space models (SSMs), and masked linear-attention Transformers cannot solve the tail-indexing task within a constant number of layers, whereas a Softmax-based Transformer can accomplish it in a single layer. Combining theoretical analysis under infinite precision with empirical validation for sequence lengths up to \( n \leq 64 \), the study shows that architectures with low causal complexity learn indexing tasks efficiently, while those with higher causal complexity degrade significantly as sequence length increases.
📝 Abstract
We study information bottlenecks in modern deep-learning architectures -- RNNs, softmax transformers, linear-attention transformers and state-space models -- through the lens of the indexing primitive. In this primitive, the input consists of $n$ bits and one integer $i$ from $1$ to $n$ called the index, and the output equals the value of the $i$-th bit.
We introduce causal complexity for masked architectures. We show that architectures with low causal complexity cannot solve the indexing primitive in any constant number of layers when the index appears at the end of the input. In particular, this limitation applies to low-parameter RNNs, SSMs and masked linear-attention transformers. In contrast, small softmax transformers can solve it in one layer, while non-masked linear-attention transformers can solve it in 2, which separates them from their masked counterparts. In turn, when the index appears at the beginning, we show that small RNNs are capable of solving this task in 1 layer, while all the other architectures require 2.
All our impossibility results are unconditional and apply even to models that employ infinite-precision real arithmetic. Moreover, experiments for up to $n=64$ qualitatively align with our theory: configurations with low-parameter theoretical solutions learn the indexing task easily, while configurations that do not admit such theoretical solutions struggle to learn as the sequence length grows.