Research map / Architectures and efficiency / Attention and transformers
Counting in Small Transformers: The Delicate Interplay between Attention and Feed-Forward Layers
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Analyzes how one-layer transformer blocks (token mixing plus MLP) solve the histogram/counting task, identifying relation-based and inventory-based counting strategies and how embedding size, FFN width, softmax and BOS token shape the learned solutions.
Problem
Understanding which algorithms simple transformer blocks implement for counting (histogram task), and how architectural choices (token mixing type, softmax, BOS token, embedding dimension, feed-forward width) determine the solution space and feasibility.
Approach
Identifies two counting strategies, relation-based counting (pairwise token comparisons via dot-product attention, needs only one FFN neuron) and inventory-based counting (FFN memorizes the alphabet, needs p>=T), gives explicit constructions for each architecture, and bounds embedding dimension for non-orthogonal embeddings via mutual coherence and softmax error reduction.
Key results
- Relation-based counting with dot-product mixing and BOS token reaches ~100% accuracy with d>=T and p=1, and a single hidden neuron suffices (construction proven).
- Dot-product attention with softmax fails to reach near-100% accuracy for p=1, performing well only when both d and p exceed T.
- Softmax with BOS or temperature-based error reduction allows solutions with d=O(log T), e.g. d=7 for T=32, L=10 in theory, though no learned solutions were found in that regime.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase hidden-layer size from p=1 to p=T (ablation) | Learned dot model · Histogram (d≥T) | accuracy: 100% -> 99% | Hurt |
| Add softmax to dot-product mixing (ablation) | dot attention · Histogram (p=1) | accuracy | Hurt |
| Add a BOS token to softmax dot-product attention | dot+sftm · Histogram (d≥T, p=1) | accuracy | Helped |
| Use softmax in token mixing (ablation) | dot-product attention · Histogram (d<T) | accuracy | Helped |
| Use relation-based counting with dot-product mixing instead of inventory-based counting | Histogram transformer · Histogram (d≥T, p=1) | accuracy | Helped |
Surprising findings
- Dot-product attention without softmax reaches performance comparable to BOS-token models even without a BOS token.
- Llama-3.1-8B and Mistral-7B perform poorly on in-context counting of 10-item lists, mostly predicting 3.
Limitations
- Analysis limited to 1-layer, non-autoregressive models without causal masks or positional encodings.
- Numerical precision demands for the lowest-dimension softmax constructions may make them hard to learn; no learned solutions found there.
- Single task (histogram), so generalization to other algorithmic tasks is unclear.
Experimental setup
Datasets: synthetic histogram task sequences
Benchmarks: histogram task (T=32, L=10) accuracy counting prompts for Mistral-7B and Llama-3.1-8B
Metrics: accuracy
Baselines: mistral-7B-v0.1 Llama-3.1-8B
Builds on: RASP Welch bound
Closest papers by meaning
- Understanding the Expressive Power and Mechanisms of Transformer for Sequence Modeling (NeurIPS 2024 · similarity 0.89)
- Memorization Capacity of Multi-Head Attention in Transformers (ICLR 2024 · similarity 0.88)
- Language Models Need Inductive Biases to Count Inductively (ICLR 2025 · similarity 0.88)
- On the Optimal Memorization Capacity of Transformers (ICLR 2025 · similarity 0.88)
- What are you sinking? A geometric approach on attention sink (NeurIPS 2025 · similarity 0.87)
- Arithmetic Transformers Can Length-Generalize in Both Operand Length and Count (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.