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

Aug 4, 2026

Architectural choices in transformer blocks strongly affect the solution space beyond scaling. This work analyzes how simple transformer blocks solve a histogram task: counting occurrences of each token in input sequences. Despite its apparent simplicity, the task exposes interactions among predictive performance, vocabulary and embedding sizes, token-mixing mechanisms, and feed-forward layer capacity. Two strategies emerge—relation-based and inventory-based counting—allocating functionality differently across attention and feed-forward layers, with robustness improvements from softmax and beginning-of-sequence tokens, confirmed by mechanistic inspection during training.

LanguageEnglish

CategoryResearch & Report

FormatPDF

Pages33

Reading Time1 h 23 min

AccessFree to download

Abstract

Introduction

Transformer architecture components and roles

Histogram task and motivation

Method overview and contributions

Counting in Small Transformers: The Delicate Interplay between Attention and Feed-Forward Layers Freya Behrens 1 Luca Biggio 2 Lenka Zdeborov 1 Abstract Next to scaling considerations, architectural design choices profoundly shape the solution space of transformers. In this work, we analyze the solutions simple transformer blocks implement when tackling the histogram task: counting items in sequences. Despite its simplicity, this task reveals a complex interplay between predictive performance, vocabulary and embedding sizes, tokenmixing mechanisms, and feed-forward layer capacity. We identify two theoretical counting strategies transformers adopt, relation-based and inventory-based counting, each defining distinct learning regimes for the task. These strategies dictate how functionality is distributed between attention and feed-forward layers. We further show that adding softmax and beginning-of-sequence tokens allow for more robustness when embedding dimensions are comparatively small. Empirical introspection of trained models closely confirms both the learning regimes of the various architectures and the formation of these strategies during training. We demonstrate how a basic task that requires only aggregation and selection is significantly impacted by minor design changes. 1. Introduction Transformers are the key neural network behind many recent deep learning advances, most notably large language models (LLMs) . Their success is partly due to their versatility in processing diverse data types, including text, images, and video, represented as sequences of tokens (Liu et al., 2021 ; Girdhar et al., 2019 ; Brown et al., 2020) . While scale has 1 Statistical Physics of Computation Laboratory, `Ecole polytechnique fdrale de Lausanne (EPFL), Lausanne, Switzerland 2Department of Computing Sciences, Universit Bocconi, Milan, Italy. Correspondence to: Freya Behrens <[freya.behrens@epfl.ch](freya.behrens@epfl.ch) >. Proceedings of the 42 nd International Conference on Machine Learning, Vancouver, Canada. PMLR 267, 2025 . Copyright 2025 by the author(s) . been a key factor in unleashing the potential of these models, it is remarkable that their architecture still largely follows the same simple template of the original transformer model proposed by Vaswani et al. (2017) . At its core, a single transformer block primarily alternates two basic components: the token-mixing attention mechanism and a standard fully connected multi-layer perceptron. At a high level, the attention mechanism mixes the tokens, while the multi-layer perceptron applies a nonlinear feature transformation identically to each token. Despite the widespread use of transformers, there is no clear consensus on the distinct roles of their components, how they interact, or if they can be substituted with alternative modules (Tolstikhin et al., 2021 ; Dordevicet al., 2024 ; Gu & Dao, 2024) . In particular, the specific contribution of each architectural element to the model’s hypothesis space –the range of algorithms it can learn and implement in practice– remains opaque (Weiss et al., 2021 ; Deltang et al., 2023 ; Abbe et al., 2023 ; Ouellette et al., 2023) . In this work, we investigate this question for algorithms that compare and aggregate information from a mechanistic perspective (Cammarata et al., 2020 ; Olah et al., 2020 ; Elhageet al., 2021 ; Michaud et al., 2023 ; Ouellette et al., 2023) and focus on the histogram task (Weiss et al., 2021) as a prototypical problem. It consists of predicting the number of appearances of each token in the input sequence processed by the model – counting-and solving it requires both comparison and aggregation. Even modern language models with up to 8B parameters currently fail to solve this task robustly and efficiently in-weights, rather than via chain-of-thought,(Appendix A) . This failure, despite the task’s apparent simplicity, motivates a study of the relative role of different architectural component

FAQs

Why can adding softmax and beginning-of-sequence tokens help when embeddings are small?

The study shows these additions improve robustness for comparatively small embedding dimensions, supporting more reliable formation of the learned counting strategies during training.

Aug 4, 2026

Upload by Valentina