Normalization kernels: compute the right statistics and reuse the row

For the LayerNorm variant used here, each chosen row contains n feature values. Compute its mean and population variance using denominator n, add positive epsilon inside the square root, then apply per-feature scale and bias. PyTorch 2.14 documents population variance for LayerNorm and normalization over the configured trailing dimensions. Dividing by n−1 would instead use a sample-variance correction and change the operation. Epsilon is a numerical and model parameter, not a value that may be moved outside the square root while preserving the same function.

LayerNorm uses each instance's selected feature population, unlike BatchNorm's different statistical axes and training/evaluation behavior. RMSNorm omits the mean subtraction and uses a different normalization statistic; it is not merely a faster spelling of identical LayerNorm semantics. Affine parameter shapes also matter when several trailing dimensions are normalized. Confirm broadcasting, dtype conversion, and output layout. Constant rows and single-feature rows are useful semantic tests: with positive epsilon, their centered values are zero and the output follows the bias rather than requiring division by zero.

Computing variance as mean of squares minus square of mean can lose significant digits when the mean is large and variation is small. A centered two-pass method or a Welford-style count/mean/centered-square summary can improve stability. Neither method is exact in finite precision, and accumulation dtype remains important. A parallel implementation must combine summaries with their counts rather than average two means equally when their populations differ. Empty masked groups require an explicit identity case before divisions by count.

A useful summary stores count n, mean μ, and M2, the sum of squared deviations from that mean. When two groups merge, the difference between their means contributes an additional centered-square term. After the merge, divide the combined M2 by the intended population count for this LayerNorm variant. Padding a row with zeros and then treating padded lanes as genuine observations changes both mean and variance. Masked lanes must affect the reduction identity and count consistently, including the final partial tile.

A fused row kernel can reuse loaded feature values while computing statistics and normalized output, avoiding some intermediate writes and rereads. A wide row can require too many registers or too much shared memory, leading to spills or less parallelism. Multiple-pass designs may then be better despite extra traffic. Affine scale and bias also require access and can change the bottleneck. Benchmark complete forward semantics and any required backward work separately; a forward-only improvement does not imply the same training-step improvement.

The function lab represents a small exact variance family. For the symmetric two-value row [−x,x], the mean is zero and population variance is x². Reset, set a=1 and x=0.50 in y=ax², predict y=0.25, then Play or Step. Separately Reset, set a=1 and x=1, predict y=1, then Play or Step. The graph cannot normalize an arbitrary row, set epsilon, or model rounding and GPU execution. Use it to see how small feature spread changes the variance scale, then compute the epsilon and affine effects by hand.