Counting with Attention and LayerNorm: Norm Scaling and Representation Geometry

arXiv:2610.16653 · cs.LG · Submitted 2026-10-12 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Counting with Attention and LayerNorm".

Jane: The gist: The cost of counting in a single attention layer depends fundamentally on where layer normalization is placed,

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So we're talking about "Counting with Attention and LayerNorm: Norm Scaling and Representation Geometry." Let's talk about who wrote this stuff. The authors are Ayyub Shaffy, Manoj Dhakal, Safal Shrestha, and Keith W. Ross from New York University Abu Dhabi.

Jane: Those are some solid names in the representation learning space. What do you think their focus on counting specifically tells us about where the current weaknesses in transformer architectures lie?

Lu: Their work points out that layer normalization and softmax averaging are empirical obstacles to counting generalization, but they also show that layer normalization can actually help improve it under certain conditions.

Meng: If they can show how normalization placement affects the parameter budget so tightly, does that mean we have a clearer blueprint for designing efficient attention layers?

Tom: Exactly. The paper proves that the geometry of the representation is what shapes these tight budget bounds, and post-normalization has some surprising efficiency gains.

Lalam: They also show how dimension itself can substitute for norm in some cases, where the minimum dimensions needed grow polynomially with length L.

The paper's summary: Jane: So to summarize what they did, they looked at counting up to a length L and tested no normalization, pre-LN, and post-LN placements for both regression and classification read-outs.

Tom: That's the setup. The main result is that for softmax attention regression, you can't count exactly with finite parameters unless the vector parameter norms grow with L.

Lu: They derive some specific minimum norm budgets Bmin(L) under regression followed by rounding: they get theta(max L, d1/four*sqrt(log L)) for unnormalized softmax, and then the post-normalization ones like theta(L(three/two)/sqrt(d)) for bare post-LN <ref:2610.16653#pg1>.

Meng: Wait, so the learned post-LN with a gain and bias gets the slowest budget growth of all those softmax variants at fixed dimension, even though normalization discards magnitude. That's a key finding for cost analysis.

Lalam: They also show that unnormalized linear attention counts exactly at a length-independent norm budget, which is really strong evidence that linear attention is inherently more efficient for this task.

The paper's improvements: Tom: The paper suggests a few ways to improve things, mostly by tuning the normalization or using learned components. They show how a learned gain helps in both post-LN and pre-LN placements, lowering the budget from theta(L(three/two)/sqrt(d)) down to theta(L(three/four)/d1/four) for post-LN with a learned gain and bias <ref:2610.16653#pg1>.

Jane: That reduction in the exponent is significant. It means that if we can use a learned gain, we can count to longer lengths L at a smaller norm budget than if we stick to bare post-normalization.

Lu: They also show that learnable pre-LN has a budget of theta(L/d1/four), which is the best one they found for those softmax variants, but they mention bare pre-LN has a finite cap of theta(e*sqrt(d)) under regression <ref:2610.16653#pg1>.

Meng: So it seems like the improvement is less about just adding more parameters and more about using learned components to manage that representation geometry better.

Lalam: It's also worth noting that the classification read-out is asymptotically superior in terms of storage complexity compared to a regression read-out when you look at bits, although it still costs theta(L) times the storage of the regression read-out at a common precision.

Conclusion: Tom: So wrapping this up, the main implication is that the placement of normalization fundamentally shapes how efficient counting is. Learned post-LN needs the smallest budget among all softmax configurations, theta(L(three/four)) at fixed dimension <ref:2610.16653#pg1>.

Jane: And they show that learned gain helps in both post-LN and pre-LN to lower those budgets significantly compared to their bare counterparts, even though normalization throws away some of the magnitude carrying the count.

Lu: The takeaway is that representation geometry dictates these rates; softmax attention puts counts on a bounded segment, and post-normalization puts them on a circular arc that stays nearly straight for longer.

Meng: And for deployment, we have to be careful about length generalization because if we use bare pre-LN, the parameter norms are capped at theta(e*sqrt(d)), which means we can't count beyond that fixed length Lcap(d).

Lalam: And the paper on "Counting with Attention and LayerNorm: Norm Scaling and Representation Geometry" proves that a single attention sublayer without feed-forward sublayers or positional encodings can count up to any fixed maximum length L under regression followed by rounding, except for that finite cap of bare pre-LN softmax regression.

Tom: Pretty neat stuff. It really shows how the architecture choices we make affect the practical feasibility of counting at scale. Thanks for tuning in.

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

New York University Abu Dhabi

cs.LG

Submitted: 2026-10-12

Updated: 2026-10-12

Comments: 29 pages, 8 figures, 3 tables

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

The gist: The gist: The cost of counting in a single attention layer depends fundamentally on where layer normalization is placed, with learned post-normalization offering the smallest parameter budget growth

Key concepts

Softmax Attention Counting: This involves using attention mechanisms to predict the exact number of tokens in an input sequence. The research analyzes how different normalization placements influence the mathematical properties of this counting process, specifically looking at parameter growth required to achieve accurate counts.
Normalization Placements: The paper compares three ways to place Layer Normalization: none (no LN), before the attention mechanism (pre-LN), and after (post-LN). The placement fundamentally changes the geometry of the representation, which in turn determines how many parameters are needed to perform accurate counting.
Parameter Budget Growth Rate: This refers to how quickly the number of trainable parameters must increase as the maximum input length (L) grows. The research derives specific mathematical bounds for this growth rate under different normalization schemes, showing that learned post-LN minimizes this growth compared to other methods.
Representation Geometry: This describes the shape and structure of the data points in the model's internal representation space. The study shows that softmax attention creates a bounded segment, while post-normalization places counts on a circular arc, and this geometric difference dictates parameter scaling.

Summary

The gist: The cost of counting in a single attention layer depends fundamentally on where layer normalization is placed, with learned post-normalization offering the smallest parameter budget growth rate for softmax attention at fixed model dimension.

How it works

The study isolates the mathematical properties of counting within a single attention sublayer by analyzing space counting up to a maximum length L, evaluating three normalization placements—no LN, pre-LN, and post-LN—with bare and learnable normalization The goal is to predict the number of occurrences of a designated token in an input that also contains distractor tokens.

The model architecture consists of a single layer with one single-head attention sublayer, no feed-forward sublayer, positional encoding, residual connection, or separate read-out bias. The input is represented by type (n, m) inputs where n is the number of spaces and m is the number of letters. The goal is to predict the space count n.

Key Findings on Softmax Attention

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. The minimum norm budget Bmin(L, d) under regression followed by rounding is derived as Θ(max[L, d1/4√log L]) for unnormalized softmax.

The analysis shows that the geometry of the representation dictates these rates: softmax attention places the counts on a bounded segment, and post-normalization places them on a circular arc. The arc stays nearly straight for longer, so a read-out does not have to stretch the small gaps between counts as much.

Normalization Placements and Budget Growth Rates

The paper derives tight bounds for the minimum norm budget Bmin(L, d) under regression followed by rounding across different placements. The derived rates are:

)&bare post-normalization, Θ(L3/2√d) (page 4).

)&post-LN with a learned gain and bias, Θ(L3/4d1/4) (page 2).

)&learnable pre-LN, Θ(L/d1/4) (page 2).

Surprisingly, at fixed dimension post-LN with a learned gain has the slowest budget growth of all the softmax variants (page 2). Bare pre-normalization is the one placement that larger parameters cannot rescue, as normalizing the embeddings bounds the attention scores and caps the counting range at Lcap(d) = Θ(e√d) under regression (page 7).

Classification Read-out

A classification read-out keeps every norm bounded under every placement. Its dense head uses Θ(L) times the storage of a regression read-out at a common precision (page 2).

Storage and Complexity

Classification is asymptotically inferior to regression when it comes to storage (page 8). Classification keeps parameter norms bounded, but replaces one regression read-out with one output vector per count (page 8).

Implications for Length Generalization

The results suggest that a model trained on inputs up to length L can fail on longer inputs because the parameters can grow to the norm that counting to L requires, and once training stops, the norms are fixed (page 7). This differs from usual explanations of out-of-distribution failure by focusing on what exists within a budget rather than what training finds (page 7).

Conclusion

The analysis uncovers how representation geometry and normalization placement fundamentally shape the parameter efficiency and feasibility of counting with attention (page 7). Learned post-LN needs the smallest budget of all the softmax configurations, Θ(L3/4) at fixed dimension, even though normalization discards the magnitude that carries the count (page 7). Bare pre-LN is the one configuration that no budget can rescue (page 7). Learned gain helps in both placements, lowering the budget from L3/2 to L3/4 for post-LN (page 7). The placement of bare normalization also matters, as it caps the counting range beyond a finite length (page 7). The dimension can also substitute for norm, with minimum dimensions dmin(L; B) growing polynomially with length L (page 22).

The paper proves that a single attention sublayer—without feed-forward sublayers or positional encodings—can count up to any fixed maximum length L under regression followed by rounding, except for the finite cap of bare pre-LN softmax regression (page 2). Unnormalized linear attention counts exactly at a length-independent norm budget, whereas the softmax regression models require growing norms (page 2).

The post-normalization construction maps the centered numerator onto a circular arc (page 4). A linear read-out measures a projection of this arc, and setting ω = L−3/2 makes the error strictly below 1/2 at norm L3/2√d (page 4).

The classification head keeps every class vector bounded under every placement (page 7). Storage cost for a dense classification head is Θ(L) times the storage of a regression read-out at a common precision (page 8).

The high/low gap ratio decreases through the network and remains well below one in later layers, suggesting representation crowding (page 27). The final residual location shows a ratio of 0.448 on average across seeds and token pairs (page 27).

The paper uses L-BFGS to train both read-outs under post-normalization from random initialization (page 4). In the unnormalized model trained with Adam, the saturation scale γ, not the parameter norm, is the quantity that lags during training (page 4).

The paper uses L-BFGS to train both read-outs under post-normalization from random initialization (page 4).

Improvements for AI systems

Bold header: Minimizing Parameter Norm Budget for Counting

The improved system can achieve counting with minimal parameter norms by leveraging learned post-normalization, which lowers the budget from L3/2 to L3/4 compared to bare post-LN. This allows architectures to count up to length L at a smaller norm budget, directly addressing the constraint that a single large entry can exceed the range of a low-precision format.

Bold header: Exploiting Representation Geometry for Efficient Readout

The system can be designed using read-outs that exploit representation geometry; specifically, post-normalization places them on a circular arc which stays nearly straight for longer, allowing the readout to achieve counting with less magnification than softmax without normalization.

Bold header: Implementing Length Generalization via Learned Pre-Normalization

The system can be configured with learned pre-normalization and a learned gain/bias to remove the finite cap, as the learnable gain removes the cap, allowing for counting beyond a fixed length Lcap(d) = Θ(e√d). This enables models to count correctly up to longer lengths based on training data.

Bold header: Optimizing Classification Storage and Accuracy Tradeoff

For classification tasks, the system can utilize Od(log L) bits for all parameters storage by using a bounded-norm classifier, which is asymptotically superior to regression storage cost of Od(Llog L) bits. This allows for storing more counts in a given memory budget while maintaining correct classification.

Bold header: Mitigating Representation Crowding in Large Models

The system can be trained to handle representation crowding by observing that adjacent high-count representations are more crowded than neighboring low-count representations and adjusting the readout or regularization accordingly, as evidenced by the high/low gap ratio decreases through the network.