Counting with Attention and LayerNorm: Norm Scaling and Representation Geometry

Ayyub Shaffy; Manoj Dhakal; Safal Shrestha; Keith W. Ross

cs.LGOct 12, 2026

Counting with Attention and LayerNorm: Norm Scaling and Representation Geometry

Authors: Ayyub Shaffy, Manoj Dhakal, Safal Shrestha, Keith W. Ross

Organizations: New York University Abu Dhabi

Abstract

Counting is a simple task on which transformers can nevertheless fail at long sequence lengths. We study the cost of counting in a single attention layer, up to a maximum length $L$ , focusing on how softmax attention and layer normalization shape the representation from which the count is recovered. We mainly focus on regression read-out but also consider classification read-out. For regression, we show unnormalized linear attention counts exactly at parameter norms independent of $L$ . For softmax attention, regression cannot count exactly with finite parameters; it recovers the count after rounding, but only when the vector parameter norms grow with $L$ . We measure this cost by $B_{\min}$ , the smallest bound that can be placed on every trainable vector to ensure exact counting after rounding, and derive tight rates. At fixed model dimension, $B_{\min}$ grows as $L$ without normalization, as $L^{3/2}$ under bare post-normalization, and as $L^{3/4}$ with learned post-normalization. The last rate is unexpected: post-normalization discards the magnitude that carries the count and might be expected to make counting harder, yet with a learnable gain it counts at a smaller norm than softmax without normalization. Bare pre-normalization is different: it bounds the attention scores, so beyond a finite length no parameter norm suffices. A classification read-out keeps every norm bounded under every placement, but a dense head stores one output vector per count and costs more in storage. Experiments on the post-normalized model reproduce this split between the two read-outs. Together, our theoretical and empirical findings characterize how self-attention, layer-normalization placement, and read-out choice shape the representation geometry and parameter cost of counting in the minimal single-head architecture studied here.

Figures & tables

Panel of Figure 1 — Figure 1: Boundary representations $h(n,0)$ on space-only inputs. (a) Linear attention gives equally spaced points on a ray. (b) Unnormalized softmax compresses them along a segment. (c) Post-normalization maps the centered numerator onto an arc. (d) Bare pre-normalization constrains the segment endpoints to a sphere. Open circles mark limiting endpoints.

Model	Budget $B_{\min}(L)$	Supported length $L_{\max}(B)$
Linear, no LN / bare pre-LN	$\Theta(1)$	$\infty$ once $B$ suffices
Softmax, learned post-LN	$\Theta(L^{3/4})$	$\Theta(B^{4/3})$
Softmax, no LN	$\Theta(L)$	$\Theta(B)$
Softmax, learned pre-LN	$\Theta(L)$	$\Theta(B)$
Softmax, bare post-LN	$\Theta(L^{3/2})$	$\Theta(B^{2/3})$
Softmax, bare pre-LN	$\infty$ beyond $L_{\mathrm{cap}}(d)$	$\leq L_{\mathrm{cap}}(d)$

Table 1: Counting costs under a separate norm bound on every trainable vector. Rates are asymptotic in

L

B

at fixed

d\geq 3

; the dependence on

d

is given with each theorem and in Appendix H . Softmax rows are ordered by how far they count at a given budget.

Figure 2: What each read-out costs, and how far each variant counts within a budget. Softmax attention,

d=16

, trained from random initialization by L-BFGS (Appendix J.2 ). Panels (a–c) show runs that pass the full-domain check (all seeds pass at

L\leq 300

, one of three classifier seeds at

L=1000

). (a) Regression needs a read-out norm that grows with

L

(gray: the range Theorem 2 (a) allows for

B_{\min}

); the classifier’s class vectors and biases stay

O(1)

under our parameterization. (b) Where the two heads put the count representations before the read-out, on the letter-free slice: regression on a

2\times 10^{-4}

rad arc, classification over

1.7

rad. (c) The classifier’s cost is elsewhere: the smallest winning score gap over letter counts, per count, shrinks with

L

. (d) Largest

L

that rounds correctly on the full domain when every trainable vector is capped at

B

. Dotted lines are the constructions; shaded regions the ranges the theorems allow.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: Training and representation dynamics of the unnormalized softmax model. Columns show held-out exact accuracy, theory-normalized parameter and gamma-score scales, and the high-/low-count adjacent-spacing ratio. Training is run for 3 seeds.

Dimension $d$	Maximum count $L$	Steps	Test exact accuracy	Full-domain worst error
4	250	20M	$0.9878\ [0.9815,\,1.0000]$	$0.650\ [0.444,\,0.757]$
16	250	20M	$0.9921\ [0.9902,\,0.9951]$	$0.668\ [0.638,\,0.687]$
4	500	100M	$0.9713\ [0.9598,\,0.9883]$	$0.833\ [0.703,\,0.915]$
16	500	100M	$0.9778\ [0.9763,\,0.9800]$	$0.793\ [0.778,\,0.803]$

Appendix

Table 2: Final performance after 20M steps for

L=250

and 100M steps for

L=500

. Entries report the mean and

[\min,\max]

across three seeds. Exact accuracy is measured on the held-out test split. Worst error is evaluated over the complete admissible input domain; rounding-exact counting requires worst error

<1/2

Figure 4: Training run for mini-batch Adam with batch size 1,024. Comparison with full-batch Adam, d=16 and L=250. The mini-batch run takes more updates in this comparison.

	no LN		bare post-LN		post-LN, learned gain
$B$	attained	constr.	attained	constr.	attained	constr.
10	10	15	20	11	100	43
30	50	50	50	24	500	186
100	100	169	100	54	2000	928
300	500	510	200	112	5000	4016
1000	1000	1705	500	251	10000	20000

Appendix

Table 3: Largest grid length passing the full-domain rounding check under a per-vector cap

B

, next to the largest integer

L

each construction guarantees at

d=16

(Appendix C.3 and Theorem 2 ). The next grid point above each attained value failed. Attained values are lower bounds on capacity; the

B=3000

row is excluded from the slope fits.

Figure 5: Layerwise geometry of the zero-shot assistant-boundary representation at $L=256$ . Blue shows the mean adjacent-centroid gap for counts $91$ – $100$ divided by that for counts $1$ – $10$ ; red shows the corresponding effective-rank ratio. Lines are means over three seeds after averaging four fixed token pairs, and shaded regions show the seed range. The horizontal line denotes equality between the high- and low-count windows.

Figure 6: PCA of zero-shot assistant-boundary count centroids for $n=1,\ldots,100$ at four residual locations. Centroids are first averaged across prompt configurations and the three seeds, and color denotes count. Each panel has its own PCA basis and axis scale, so only within-panel geometry should be interpreted.

Figure 7: IID rounded exact accuracy of the validation-selected zero-shot ridge probe, binned into ten-count ranges. Points show the mean across three seeds and shading shows the seed range.

Figure 8: Regularization–norm tradeoff for zero-shot probes at the assistant boundary after the first transformer block. Each colored point is one fitted probe: horizontal position gives its native $\ell_{2}$ norm, vertical position gives held-out rounded accuracy, and color gives the relative ridge coefficient. Black stars denote the three unregularized probes.

Explore similar work

Exact Attention Sensitivity and the Geometry of Transformer Stability
Feb 21, 2026Seyed Morteza EmadiTransformer Attention Softmax Attention
UNC-Chapel Hill, Kenan-Flagler Business School Chapel Hill, NC, USA
A Controlled Study of Attention-Only Transformers
Jul 20, 2026Henry Ndubuaku, Karen Mosoyan, Jakub Mroz +5Transformer FFNs Transformer
Cactus Compute, Inc.
The Right Answer, the Wrong Direction: Why Transformers Fail at Counting and How to Fix It
May 5, 2026Gabriel GarciaNumerical Reasoning in Language Models Transformer Interpretability
Independent Researcher

Feb 21, 2026J/K move · Enter open · S save

Exact Attention Sensitivity and the Geometry of Transformer Stability

Seyed Morteza Emadi

UNC-Chapel Hill, Kenan-Flagler Business School Chapel Hill, NC, USA

Transformer Attention Softmax Attention Sensitivity Analysis

Open paper AskPDF arXiv

We develop a sensitivity analysis for transformer attention in a geometry aligned with tokenwise computation. Our main result is the exact identity

\|J_τ(u)\|_{\infty\to1}=θ(p)/τ

for the Jacobian

J_τ(u)

of the tempered softmax

u\mapsto\mathrm{softmax}(u/τ)

, where

θ(p)=4\max_{S\subseteq[L]}p(S)(1-p(S))

measures how evenly the attention distribution can be bisected rather than how concentrated it is. We combine this identity with a block-

\infty

/RMS norm under which row-stochastic attention mixing is nonexpansive. This yields a distribution-aware local Jacobian bound for multi-head attention and a sequence-length-independent Lipschitz bound on bounded input sets, with explicit dependence on width, input magnitude, temperature, and projection norms. We also identify a structural distinction between normalization placements: a pre-LN residual-sublayer Jacobian contains an additive identity term, whereas a post-LN residual-sublayer Jacobian does not. A LayerNorm projection lemma gives a sufficient condition under which the LayerNorm-only term in the post-LN expansion contracts geometrically; the condition is not tested by our experiments. Across three Pre-LN early-training runs of

774

M-parameter models, attention becomes substantially more concentrated while the median lower-bound certificate for

θ(p)

remains near one at every sampled layer and checkpoint. This certifies near-maximal exact sensitivity for at least half of the sampled rows within each layer. A minority of rows enters a dominant-atom regime with lower exact sensitivity, consistent with the deterministic relationship between

p_{\max}

and

θ(p)

Abstract

Explore similar work

Exact Attention Sensitivity and the Geometry of Transformer Stability

A Controlled Study of Attention-Only Transformers

The Right Answer, the Wrong Direction: Why Transformers Fail at Counting and How to Fix It