Gating and Normalization Interact in the Attention Block
Abstract We study how an elementwise sigmoid attention-output gate interacts with normalization placement in small causal transformers (29–36M parameters). A 54-model main sweep crosses three schemes (pre-norm; pre-norm + QK-norm; sandwich + QK-norm), three depths, three seeds, and models with and without the gate on WikiText-2. At the best validation checkpoint selected for each run, gating lowers mean validation perplexity in all nine normalization-depth cells. Averaged across depths and seeds, the gated models improve over their matched no-gate baselines by 1.03 PPL with pre-norm, 2.64 with pre-norm + QK-norm, and 3.35 with sandwich + QK-norm. The improvement is therefore 2.32 PPL larger for sandwich + QK-norm than for pre-norm (p < 0.001). A separate fixed-budget 8-layer experiment on WikiText-103 independently finds a gate benefit under all three schemes (all p < 0.01). Normalization placement also produces qualitatively different learned gate distributions: C1 and C2 become low-mean suppressors with essentially no fully-open values, while C3 becomes bimodal, with roughly 20% of values fully closed and 10% fully open. A single-layer ablation identifies the attention-output RMSNorm as necessary and sufficient for this bimodality at similar perplexity, and the distinction persists across the stable learning-rate range. Gate sparsity consistently peaks in the second transformer block. These results show that normalization placement controls both the performance effect and learned geometry of attention-output gating at this scale. Version 2 This version substantially expands the original manuscript with: a fixed-budget 8-layer experiment on WikiText-103; a single-layer ablation isolating the role of the attention-output RMSNorm; learning-rate robustness analyses; a parameter-matched no-gate control; expanded analysis of gate distributions and layer-wise sparsity; additional discussion, limitations, tables, and supplementary results. The experiments use small causal transformers trained from scratch and characterize behaviour in a compute-limited regime. The results should not be assumed to generalise to fully converged or large-scale language-model training without further evidence.
Authors
- Shubham Kumar (ORCID: https://orcid.org/0009-0009-6532-4712)
Publication Details
- Journal
- Zenodo (CERN European Organization for Nuclear Research)
- Published
- 2026-07-25
- DOI
- https://doi.org/10.5281/zenodo.21577775
- Primary Topic
- Advanced Memory and Neural Computing
- Type
- preprint