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

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Behrens, Freya, Biggio, Luca, Zdeborová, Lenka
Format: Preprint
Veröffentlicht: 2024
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866914151682539520
author Behrens, Freya
Biggio, Luca
Zdeborová, Lenka
author_facet Behrens, Freya
Biggio, Luca
Zdeborová, Lenka
contents 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, token-mixing 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.
format Preprint
id arxiv_https___arxiv_org_abs_2407_11542
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Counting in Small Transformers: The Delicate Interplay between Attention and Feed-Forward Layers
Behrens, Freya
Biggio, Luca
Zdeborová, Lenka
Machine Learning
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, token-mixing 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.
title Counting in Small Transformers: The Delicate Interplay between Attention and Feed-Forward Layers
topic Machine Learning
url https://arxiv.org/abs/2407.11542