%% ===================================================================== %% draft.tex -- LaTeX conversion of draft.md (the author's own prose) %% ===================================================================== %% This is a FORMAT conversion only: the Introduction/Abstract prose is %% reproduced VERBATIM from draft.md (Cormac's own words, typos included -- %% nothing here is copy-edited). Errors were REPORTED back in chat rather than %% silently fixed, so the repo keeps showing your own incremental writing (the %% point of the anti-AI audit). Only structure was added: markdown headings -> %% IEEE sections, and the numeric [1]..[17] citations -> \cite{} keys. %% %% CITATIONS FILLED (draft.md named the source but gave no [n]) -- both flagged %% inline with "% [FILLED CITATION ...]" so you can confirm/adjust: %% * "Wortsman ... methodological template" -> \cite{wortsman2024} %% * Gap para, "Wortsman pairs qk-norm with z-loss" -> \cite{wortsman2024} %% %% Number mapping (key -> printed number). qknmla (was [15] in main.tex) is %% never cited in this draft, so it was REMOVED; LaTeX renumbers the rest from %% \bibitem order, so tinystories/cramming are now [15]/[16] (were [16]/[17]). %% NB: this draft's numbering therefore no longer matches main.tex's 17-item list. %% [1]=henry2020 [2]=vit22b [3]=zhai2023 [4]=wortsman2024 [5]=olmo2 %% [6]=imu1 [7]=slmtradeoffs [8]=mup [9]=nanogpt [10]=sun_massive %% [11]=attnsink [12]=softcap [13]=zloss [14]=normformer [15]=tinystories %% [16]=cramming %% %% Title/author block carried over from main.tex (draft.md had none). %% Build: ./paper/build.sh draft %% ===================================================================== \documentclass[conference]{IEEEtran} \IEEEoverridecommandlockouts \usepackage{cite} \usepackage{amsmath,amssymb,amsfonts} \usepackage{graphicx} \usepackage{textcomp} \usepackage{xcolor} \usepackage{url} \begin{document} % Title/author carried over from main.tex (draft.md did not include them). \title{When Does a Small Language Model Need QK-Norm?\\ An Isolated Learning-Rate Characterization of Query--Key Normalization\\ at Sub-100M Parameters\\ % {\footnotesize \textsuperscript{}(working title -- draft)} } \author{\IEEEauthorblockN{Cormac Shannon} \IEEEauthorblockA{\textit{MSc Immersive Software Engineering} \\ \textit{University of Limerick}\\ Limerick, Ireland \\ 23366303} } \maketitle \begin{abstract} Modern language-model architectures ship many components which have been validated using large models, however, whether the same benefits are reproduced for smaller models is actively disputed. This paper aims to produce a controlled, isolated study of one such component: Query-Key Normalisation (QK-norm) against sub 100M parameter models, defining when it becomes necessary (at which learning rate and model depths). \end{abstract} \begin{IEEEkeywords} small language models, transformer architecture, training stability, query--key normalization, learning rate, ablation study \end{IEEEkeywords} \section{Introduction} % [Context & Hook] The transformer architecture has quickly become a dominant choice in the natural language processing (NLP) and machine learning (ML) field. While popularly used in large scale models boasting hundreds of billions to trillions of parameters, the transformer architecture has also seen growing success in smaller models, running inference on edge-compute. These models are cheaper to run and store on low-end hardware, efficient models can run with only a few tens of millions of parameters. Modern LLM Architecture are the culmination of many components, including but not limited to: rotary position embeddings, RMSNorm, grouped-query attention, SwiGLU, and QK-norm. Almost all of which have been validated against models billions of parameters in size \cite{vit22b, olmo2}. Although we know these techniques provided benefit in large models, it is not clearly established that these components carry their benefits to small models. Guidance for composing sub-100M parameter models is limited, and for several of the components mentioned above, the evidence from literature is contradictory regarding their benefits. Gaining clarity in this area is important, small-model pretraining is an active area of research (sample-efficient work \cite{imu1}, compute-constrained ablations \cite{slmtradeoffs}), choosing the right components matters when compute budget is tightest. % [Literature review] QK-Norm is a modification to a transformers self-attention step. First proposed by Henry et al. \cite{henry2020} for use in low-resource neural machine translation. The technique scales both the query and key vectors to a fixed magnitude before their dot product is calculated to create their attention logits. Without this step, the attention logits can grow without bound after training for many steps, causing the subsequent softmax step to create a near one-hot outcome. The learning gradients become vanishingly small, and the training loss can start spiking or diverge, requiring the learning rate to be reduced \cite{zhai2023, wortsman2024}. Recent mechanical studies associate the reason behind logit growth to ``massive-activations'', wherein a small number of tokens gain extreme-magnitude activations and become the domainant influence over the key vector, forming ``attention sinks'' \cite{sun_massive, attnsink}. QK-Norm resolves this issue by preventing unbound growth of these vectors, thus removing the failure mode. QK-Norm was popularised by the 22-billion-parameter Vision Transformer (ViT-22B) of Dehghani et al. \cite{vit22b}, the technique has since grown to become a near standard step in frontier language models. For example, OLMo 2 utilised RMSNorm to the Q and K vectors and found improvements to their models stability and training loss \cite{olmo2}. The technique has since been adopted by Gemma 2/3, Qwen3, and Chameleon. However, QK-norm stands among several logit-control techniques: Gemma 2 soft-capped their attention logits using a smooth saturating function \cite{softcap}, Z-loss maintains softmax normalisers near zero preventing drifting \cite{zloss}, and NormFormer adds an extra normalisation layer inside each block \cite{normformer}. Each of these techniques are frequently studied together, so the impact of just QK-Norm alone is difficult to observe. Literature regarding small models is sparse, and of what exists, often entangled. Wortsman et al. provides a solid methodological template \cite{wortsman2024}. % [FILLED CITATION -- draft.md named Wortsman but gave no [n] here] Wortsman found that instabilities found in large models can be reproduced in small models at high learning rates, QK-layernorm paired with z-loss can resolve these issues. They found this behaviour transfers across model sizes. They also introduce a learning-rate (LR) sensitivity metric that measures how far a model's final loss deviates from its optimum as the learning rate is adjusted. Recently, IMU-1 by Grigorev et al. found in 70M parameter models, the single most positively impactful architecture change was QK-Norm. This was measured via reducing the activation kurtosis \cite{imu1}. However, in their study, they paired QK-Norm alongside maximal update parametrisation ($\mu$P), NorMuon, and value residuals. It is important to note that these papers do not stand without objection. Other literature finds that some of the same ablations (RMSNorm, rotary position embeddings, grouped-query attention) do not transfer well to models sized at 70M parameters \cite{slmtradeoffs}. The two literatures disagree at effectively the same scale, while neither paper conducts the isolated learning-rate experiment that would settle the matter. From this, learning rate arises as the right axis to work along. Maximal-update parametrisation details that optimal learning rate can transfer across model widths \cite{mup}. Learning rate is the hyperparameter that these stability effects are dependant on. Therefore, a faithful examination of QK-Norm requires us to study it as it relates to learning rate while maximal-update parametrisation is disabled, ensuring that its effects do not mask the behaviour we are studying. % [Problem statement & gap] In regards to evidence inspecting the behaviour of small models, there are two limitations surrounding them. First, evidence studying the impact on small models is confounded, Wortsman pairs qk-norm with z-loss and uses small models as a proxy for larger models \cite{wortsman2024}, while IMU-1 pairs it with several other techniques % [FILLED CITATION -- draft.md named Wortsman but gave no [n] here] ($\mu$P + NorMuon + value residuals) measuring its effectiveness through activation kurtosis at a single fixed depth and never sweeping learning rate \cite{imu1}. Second, the opposing works test against a compute-optimal budget and does not isolate QK-norm or sweep learning as we desire \cite{slmtradeoffs}. The arising issue is that there is no study isolating the behaviour of QK-norm through different learning rates or at what model depth it becomes necessary. A second gap revolves around accessibility, most texts harness expensive GPUs, making reproducibility difficult. A laptop-reproducible study on its own is a worthwhile contribution. % [Proposed contribution: objective + research questions] This study therefore aims to isolate the effect of QK-norm and to determine at what scale a small language model benefits from it. It asks, at fewer than 100M parameters, does QK-norm --- on its own --- improve model convergence and raise the maximum stable learning rate? Additionally, at what learning rate does it become necessary? As a final subquestion, how does that threshold shift as the model depth is changed? % [Hypotheses] Three hypotheses can there be derived: First, the effect of QK-norm is learning-rate conditional, negligible impact or marginally harmful at low learning rates but highly impactful at higher ones (H1); that QK-Norm raises the max stable learning rate and lowers learning-rate-sensitivity (H2); and that the learning rate at which QK-norm becomes necessary lowers as depth increases (H3). % [Method] The study will investigate these hypothesis by conducting controlled ablations/changes on a nanoGPT \cite{nanogpt} project with maximal update parametrisation disabled, AdamW set as the optimiser, while toggling QK-norm as the sole independent variable through a dense sweep of learning rates, and additionally against a small range of model depths at a model scale of roughly 10-15M parameters. % [Metrics] Results will be measured through the learning-rate sensitivity \cite{wortsman2024}, the maximum stable learning rate, steps-to-threshold, perplexity, attention entropy and activation kurtosis. % [Data + feasibility] The project will be working with small models, which has a supported precedent to produce cohesive text by using the TinyStories \cite{tinystories} body of text, which can train on a single-GPU as established in the Cramming study \cite{cramming}. The project is designed to run on a consumer laptop (Apple M4 Pro chip), with code, seeds, and logs distributed. % ===================================================================== % Bibliography carried over from main.tex, minus qknmla (which this draft never % cites) -> 16 entries. LaTeX numbers them by \bibitem order, so removing the % uncited entry renumbers every \cite automatically. ALL entries still need % Scopus / primary-PDF checks -- see the % [TO VERIFY] tags. % ===================================================================== \begin{thebibliography}{00} \bibitem{henry2020} A. Henry, P. R. Dahl, et al., ``Query-key normalization for transformers,'' in \textit{Findings of EMNLP}, 2020. arXiv:2010.04245. % [TO VERIFY authors] \bibitem{vit22b} M. Dehghani et al., ``Scaling vision transformers to 22 billion parameters,'' 2023. arXiv:2302.05442. \bibitem{zhai2023} S. Zhai et al., ``Stabilizing transformer training by preventing attention entropy collapse,'' in \textit{ICML}, 2023. arXiv:2303.06296. \bibitem{wortsman2024} M. Wortsman et al., ``Small-scale proxies for large-scale transformer training instabilities,'' in \textit{ICLR (oral)}, 2024. arXiv:2309.14322. % [TO VERIFY full author list] \bibitem{olmo2} OLMo~2 Team, ``2 OLMo 2 furious,'' 2025. arXiv:2501.00656. % [TO VERIFY preferred citation format / authors] \bibitem{imu1} G. Grigorev et al., ``IMU-1: Sample-efficient pre-training of small language models,'' 2026. arXiv:2602.02522. % [TO VERIFY authors; preprint] \bibitem{slmtradeoffs} ``Architectural trade-offs in small language models under compute constraints,'' 2025. arXiv:2512.20877. % [TO VERIFY title + authors] \bibitem{mup} G. Yang et al., ``Tensor programs V: Tuning large neural networks via zero-shot hyperparameter transfer,'' 2022. arXiv:2203.03466. \bibitem{nanogpt} A. Karpathy, ``nanoGPT,'' GitHub repository, 2022--. \url{https://github.com/karpathy/nanoGPT}. % software; check IEEE format for repos \bibitem{sun_massive} M. Sun et al., ``Massive activations in large language models,'' in \textit{ICML}, 2024. arXiv:2402.17762. % [TO VERIFY] \bibitem{attnsink} ``When attention sink emerges in language models: an empirical view,'' 2024. arXiv:2410.10781. % [TO VERIFY authors] \bibitem{softcap} Gemma Team, ``Gemma 2,'' technical report, 2024. % [TO VERIFY] arXiv ID + exact soft-capping citation \bibitem{zloss} A. Chowdhery et al., ``PaLM: Scaling language modeling with pathways,'' 2022. % [TO VERIFY] z-loss source; confirm arXiv ID \bibitem{normformer} S. Shleifer et al., ``NormFormer: Improved transformer pretraining with extra normalization,'' 2021. arXiv:2110.09456. % [TO VERIFY authors] \bibitem{tinystories} R. Eldan and Y. Li, ``TinyStories: How small can language models be and still speak coherent English?,'' 2023. arXiv:2305.07759. % [TO VERIFY] \bibitem{cramming} J. Geiping and T. Goldstein, ``Cramming: Training a language model on a single GPU in one day,'' in \textit{ICML}, 2023. arXiv:2212.14034. % [TO VERIFY] \end{thebibliography} \end{document}