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

Feb 21, 2026

Exact Attention Sensitivity and the Geometry of Transformer Stability

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)

Seyed Morteza Emadi

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

Transformer Attention Softmax Attention Sensitivity Analysis

Ask PDF

Jul 20, 2026

A Controlled Study of Attention-Only Transformers

Feed-forward networks hold two thirds of a transformer's non-embedding parameters, yet the architecture has not received a necessity test that controls parameters, compute, and depth at once. We pretrain attention-only decoder transformers (Simple Attention Networks, SANs) against standard transformers matched separately for parameter count, training FLOPs, and depth (2 to 48 layers), for up to 105B tokens at 6M to 87M parameters. Deleting feed-forward layers in place is costly: the standard transformer leads by 0.47 nats at matched depth and 0.26 nats at matched FLOPs. Reallocating the freed budget into attention depth closes the gap: at matched parameters the difference is 0.006 nats (0.27 percent of loss), reproducible to one part in ten thousand across seed pairs, shrinking across 5B, 30B, and 105B budgets, and holding near 0.02 nats across a 29x size range. Three measurements localize the remaining gap to parametric recall: attention-only models are better on context-grounded answers and worse where knowledge must come from weights. Weight spectra show why: routing matrices (Q/K) crystallize early, content matrices accumulate rank slowly, and removing feed-forward layers relocates this accumulation to the attention output projection. QK-normalization, not feed-forward layers or residual gating, keeps 48-layer attention-only stacks trainable. The deficit concentrates on low-context query prediction and localizes there entirely by the largest budget. A pre-registered test confirms the account: it predicts a 0.02 to 0.05 nat gap on knowledge-dense web text; a matched pair trained on fineweb-edu measures 0.040. Within the tested regime, attention does the rest.

Henry Ndubuaku, Karen Mosoyan, Jakub Mroz +5

Cactus Compute, Inc.

Transformer FFNs Transformer Decoder-Only Language Models Transformer Attention

Ask PDF

May 5, 2026

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

Large language models often fail at simple counting tasks, even when items to count are in the prompt. We investigate whether this failure occurs because transformers do not represent counts internally, or because they cannot convert representations to the correct output tokens. Across three model families: Pythia, Qwen3, and Mistral, ranging from 0.4B to 14B parameters, we find evidence for the second explanation. Linear probes recover the correct count from intermediate layers with

R^2>0.99

, showing that the information is present. However, the internal directions that encode counts are nearly orthogonal to digit-token output-head rows (

|\cos| \leq 0.032

). In other words, the model stores the count in a form that the digit logits do not naturally read out. We localize this failure with two interventions. Updating only the digit rows of the output head (36,864 parameters) substantially improves constrained digit prediction (60.7--100.0% on four tasks), but it does not fix unconstrained generation (0%); we do not claim that digit-row repair fixes open-ended text. By contrast, small LoRA on attention Q/V (7.67M parameters) improves upstream routing and achieves 83.1%

\pm

7.2% in true greedy autoregressive generation (deployable fix). Logit-lens at layer 35 (entity counting; correct-digit rank): (i) median over 3 seeds drops from order-

10^4

to 1; (ii) seed 42 shows

54{,}332 \to 838

(median top-1 while one seed stays far below). Norm, logit-lens, and cross-task analyses generalize the bottleneck to counting, addition, and list length; nulls on MMLU and GSM8K and limited DROP transfer. These results identify counting failure as a geometric readout bottleneck, not an internal-representation failure: the model knows the count but the output pathway is misaligned with tokens needed to express it.

Gabriel Garcia

Independent Researcher

Numerical Reasoning in Language Models Transformer Interpretability LLM Fine-Tuning Object Counting Language Modeling Representation Geometry in Language Models

Ask PDF