Counting in Small Transformers: The Delicate Interplay between Attention and Feed-Forward Layers

2024-07-16

AI Summary

This work investigates the learnable solution space of small Transformers when solving histogram tasks—i.e., counting token frequencies in sequences—a deceptively simple yet revealing probe of model internals. Methodologically, the authors integrate theoretical modeling (via linear-algebraic characterization), empirical training, and mechanistic reverse-engineering (including attention visualization and gradient probing). They formally distinguish, for the first time, two distinct counting strategies: relation-based and inventory-based. Results show fine-grained coordination between attention and feed-forward networks; minor architectural changes (e.g., replacing softmax) induce abrupt strategy shifts. Quantitative analysis uncovers nonlinear performance boundaries governed by vocabulary size, embedding dimension, FFN capacity, and attention design. Crucially, both strategies emerge empirically during training, with their prevalence determined jointly by hyperparameters and implicit couplings among model components.

Technology Category

• Machine Learning: Deep Neural Architectures and Foundation Models
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

How do different architectural design choices influence the space of solutions that a transformer can implement and learn? How do different components interact with each other to shape the model's hypothesis space? We investigate these questions by characterizing the solutions simple transformer blocks can implement when challenged to solve the histogram task -- counting the occurrences of each item in an input sequence from a fixed vocabulary. Despite its apparent simplicity, this task exhibits a rich phenomenology: our analysis reveals a strong inter-dependence between the model's predictive performance and the vocabulary and embedding sizes, the token-mixing mechanism and the capacity of the feed-forward block. In this work, we characterize two different counting strategies that small transformers can implement theoretically: relation-based and inventory-based counting, the latter being less efficient in computation and memory. The emergence of either strategy is heavily influenced by subtle synergies among hyperparameters and components, and depends on seemingly minor architectural tweaks like the inclusion of softmax in the attention mechanism. By introspecting models trained on the histogram task, we verify the formation of both mechanisms in practice. Our findings highlight that even in simple settings, slight variations in model design can cause significant changes to the solutions a transformer learns.

Problem

• Analyzing transformer solutions for counting items in sequences
• Exploring interplay between attention and feed-forward layers in counting
• Investigating impact of design changes on basic aggregation tasks

Innovation

• Analyzes attention and feed-forward layers interplay
• Identifies relation-based and inventory-based counting strategies
• Uses softmax and BOS tokens for robustness