Reproduce “Counting in Small Transformers: The Delicate Interplay between Attention and Feed-Forward Layers” on OpenResearch

About this paper

Freya Behrens, Luca Biggio, Lenka Zdeborová

This study analyzes how small transformer blocks solve the histogram task, identifying two distinct algorithmic strategies: relation-based counting (RC) and inventory-based counting (IC). The work demonstrates how architectural elements like attention mechanisms, feed-forward network capacity, and special tokens dictate the emergence and efficiency of these strategies, even when operating with non-orthogonal token embeddings.

Key insights

• Two primary algorithmic strategies for counting emerge: Relation-based Counting (RC), which uses attention for pairwise comparisons efficiently (requiring minimal FFN capacity), and Inventory-based Counting (IC), which uses the FFN as a memory/lookup table (requiring FFN capacity proportional to alphabet size).
• Architectural choices, such as the presence of a BOS token, the use of softmax, or the capacity of the feed-forward layer, profoundly determine which counting strategy a transformer adopts and its efficiency.
• The softmax function enhances robustness to "entanglement noise" from non-orthogonal token embeddings (when embedding dimension 'd' is less than alphabet size 'T'), allowing models to achieve high accuracy even with significantly reduced embedding dimensions.

Summary by alphaXiv.