Counting with Attention and LayerNorm: Norm Scaling and Representation Geometry
The study investigates how placing Layer Normalization (LN) affects parameter efficiency when counting tokens within a single attention layer using softmax attention. It finds that learned post-normalization offers the smallest parameter budget growth rate for counting, while bare pre-normalization is limited by a finite cap on the count range. This geometry dictates how parameters scale with input length and dimension.