514 lines
33 KiB
TeX
514 lines
33 KiB
TeX
%% =====================================================================
|
|
%% CS6513 -- Introduction + Mini Literature Review (IEEE conference)
|
|
%% WORKING DRAFT -- ROUGH PASS (page limit deliberately ignored for now)
|
|
%%
|
|
%% STATUS: This is a SCAFFOLD, not final text. The bullets below are
|
|
%% ideas / papers / arguments to be turned into flowing prose IN YOUR
|
|
%% OWN WORDS over several commits. Do NOT submit the bullets as-is.
|
|
%% Remember: declare this AI-assisted outlining in the notebook GenAI log.
|
|
%%
|
|
%% Target final structure (2 pages incl. refs), weaving together:
|
|
%% 1. Context & Hook 2. Literature Context (mini lit review)
|
|
%% 3. Problem & Gap 4. Proposed Contribution 5. Roadmap
|
|
%% =====================================================================
|
|
%%
|
|
%% ABOUT THESE COMMENTS:
|
|
%% Anything after a "%" is a LaTeX comment -- ignored by the compiler, so it
|
|
%% NEVER appears in the PDF. The annotations below are learning notes for the
|
|
%% author; they are safe to keep, and can be stripped later. Convention used:
|
|
%% % DEFN: <term> -- plain-language definition of jargon
|
|
%% % EXPL: ... -- expands on a concept / why it matters
|
|
%% % NOTE: ... -- drafting instruction to self
|
|
%% % [TO VERIFY] -- citation detail to confirm against Scopus / the PDF
|
|
%%
|
|
%% ---- MINI GLOSSARY (core terms, defined once here) --------------------
|
|
%% Transformer : the neural-network design behind modern language models
|
|
%% (Vaswani et al. 2017, "Attention Is All You Need"). Processes all tokens
|
|
%% in parallel using "self-attention" (below).
|
|
%% Token : the unit of text a model reads/predicts -- a word or word-piece
|
|
%% (e.g. "walking" -> "walk" + "ing"). Text is a sequence of tokens.
|
|
%% Parameters : the model's learned numbers (weights). "Model size" = how many
|
|
%% there are. 100M = 100 million. Frontier LLMs have 10^9-10^12 (billions+).
|
|
%% Self-attention : the mechanism that lets each token "look at" other tokens.
|
|
%% For every token it builds three vectors -- a QUERY, a KEY, and a VALUE.
|
|
%% Query / Key / Value (Q/K/V) : query = "what am I looking for", key = "what do
|
|
%% I offer", value = "what I pass on". A token attends to others by comparing
|
|
%% its query to their keys.
|
|
%% Attention logit : the raw match score between a query and a key (their dot
|
|
%% product). Large logits -> the next step (softmax) becomes extreme.
|
|
%% Softmax : turns a list of scores into probabilities (all positive, sum 1).
|
|
%% Used to decide how much each token attends to each other token.
|
|
%% Learning rate (LR) : the step size during training. Too low = slow; too high
|
|
%% = unstable / diverging. The central variable in THIS study.
|
|
%% Ablation : an experiment that removes or toggles ONE component to measure
|
|
%% its specific effect -- the core method here (QK-norm on vs off).
|
|
%% Perplexity : a standard quality score for language models (= exp of the
|
|
%% average prediction loss). LOWER is better.
|
|
%% -----------------------------------------------------------------------
|
|
|
|
\documentclass[conference]{IEEEtran}
|
|
% DEFN: \documentclass[conference]{IEEEtran} -- load the official IEEE conference
|
|
% style (the IEEEtran.cls file in this folder). Sets the two-column layout,
|
|
% fonts and margins the assignment requires; do not change.
|
|
% DEFN: \usepackage{...} -- pulls in an add-on ("package") of extra commands.
|
|
% Each line below is one such add-on; comments say what it gives us.
|
|
\IEEEoverridecommandlockouts
|
|
\usepackage{cite} % tidies up citation numbers, e.g. [2]-[4]
|
|
\usepackage{amsmath,amssymb,amsfonts} % maths symbols: the arrows, <, mu, etc.
|
|
\usepackage{graphicx} % lets us include figures/images (none yet)
|
|
\usepackage{textcomp} % extra text symbols
|
|
\usepackage{xcolor} % coloured text (if ever needed)
|
|
\usepackage{url} % formats web links in the bibliography
|
|
|
|
\begin{document}
|
|
% EXPL: everything the reader SEES lives between \begin{document} and
|
|
% \end{document}. Above that line is the "preamble" (setup only).
|
|
|
|
% DEFN: \title{...} sets the paper title. The "\\" forces a line break.
|
|
% DEFN: QK-Norm (Query-Key Normalization) -- the technique under study: before
|
|
% attention compares a query to a key, rescale both to a controlled length
|
|
% so their match score (logit) cannot blow up. Full explanation in Sec. below.
|
|
% DEFN: "Isolated" -- we change ONLY QK-norm, holding everything else fixed, so any
|
|
% effect is attributable to it alone (see "ablation" in the glossary).
|
|
% DEFN: "Sub-100M parameters" -- models with fewer than 100 million weights; small
|
|
% enough to train on a laptop (frontier models are 100-1000x+ larger).
|
|
\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 -- rough draft)}}
|
|
|
|
\author{\IEEEauthorblockN{Cormac Shannon}
|
|
\IEEEauthorblockA{\textit{MSc Immersive Software Engineering} \\
|
|
\textit{University of Limerick}\\
|
|
Limerick, Ireland \\
|
|
23366303}
|
|
}
|
|
|
|
\maketitle
|
|
|
|
% DEFN: abstract -- the ~150-word summary at the very top of a paper. A reader
|
|
% decides from this alone whether to read on. Convention: write it LAST.
|
|
% NOTE(abstract): write LAST, once the argument is settled. Rough placeholder only.
|
|
\begin{abstract}
|
|
% ROUGH -- rewrite in own words.
|
|
% EXPL: "components validated only at large scale" -- modern models are built from
|
|
% many add-on parts (see the architecture list in Sec. Introduction) that
|
|
% were tested mainly on huge models; nobody checked they still help small ones.
|
|
% DEFN: \emph{...} -- emphasis (renders as italics).
|
|
Modern language-model architectures bundle many components validated only at
|
|
large scale; whether they help \emph{small} models is under-studied and, for
|
|
some components, actively disputed. This paper motivates a controlled,
|
|
isolated study of one such component -- Query--Key Normalization (QK-norm) --
|
|
at sub-100M-parameter scale, characterizing when (at which learning rates and
|
|
depths) it becomes necessary. % TODO: 1-2 sentences on method + what we expect.
|
|
% EXPL: "depth" = number of stacked transformer layers; "necessary" = the point
|
|
% where turning QK-norm off makes training fail/diverge.
|
|
\end{abstract}
|
|
|
|
% DEFN: keywords -- a handful of index terms so search engines/databases can file
|
|
% the paper. Pick the phrases a reader would search to find work like yours.
|
|
\begin{IEEEkeywords}
|
|
small language models, transformer architecture, training stability,
|
|
query--key normalization, learning rate, ablation study
|
|
\end{IEEEkeywords}
|
|
|
|
% =====================================================================
|
|
\section{Introduction}
|
|
% For a 2-page intro-only paper this whole section IS the paper. Below it is
|
|
% broken into the 5 required components as comment banners; in the final
|
|
% version they should read as ONE flowing argument (few/no subheadings).
|
|
|
|
%% -------------------------------------------------------------------
|
|
%% COMPONENT 1 -- CONTEXT & HOOK (broad domain + why relevant NOW)
|
|
%% -------------------------------------------------------------------
|
|
% NOTE: 1 short paragraph. Funnel from broad -> the tension that motivates us.
|
|
% DEFN: NLP = Natural Language Processing (getting computers to work with human
|
|
% language); ML = Machine Learning. LLM = Large Language Model.
|
|
% DEFN: "on-device / edge" -- running a model locally on a phone/laptop/sensor
|
|
% instead of a data-centre; motivates SMALL models (limited memory + power).
|
|
% EXPL: \textbf{...} = bold. This whole block is bullet notes; final version is
|
|
% ONE flowing paragraph that funnels broad -> specific.
|
|
\begin{itemize}
|
|
\item Broad domain: transformer language models now dominate NLP/ML; but the
|
|
field is increasingly moving toward \textbf{small / efficient} models
|
|
(on-device, edge, cost, energy, and \emph{reproducibility}).
|
|
% DEFN: reproducibility -- whether other researchers can re-run your
|
|
% experiment and get the same result. A core scientific value, and
|
|
% easier to guarantee with small models anyone can afford to train.
|
|
\item Modern LLM architectures are an accretion of components -- rotary
|
|
position embeddings, RMSNorm, grouped-query attention, SwiGLU,
|
|
\textbf{QK-norm} -- almost all validated at the billion-parameter scale
|
|
\cite{olmo2,vit22b}.
|
|
% EXPL: "accretion of components" = built up by stacking many small
|
|
% add-ons over the years. The four named ones (context only; NOT
|
|
% what we study):
|
|
% DEFN: rotary position embeddings (RoPE) -- how a model encodes WORD
|
|
% ORDER, by rotating the query/key vectors by an angle set by each
|
|
% token's position.
|
|
% DEFN: RMSNorm -- a "normalization" layer that rescales a vector to a
|
|
% standard size to keep numbers well-behaved during training (a
|
|
% cheaper variant of LayerNorm).
|
|
% DEFN: grouped-query attention (GQA) -- a memory-saving trick where
|
|
% several attention "heads" share one set of keys/values.
|
|
% DEFN: SwiGLU -- a high-performing choice of activation function (the
|
|
% little non-linear step inside each layer's feed-forward block).
|
|
% DEFN: \cite{key} -- inserts a numbered reference; "key" points at a
|
|
% \bibitem in the bibliography at the bottom of the file.
|
|
\item Hook / tension: it is \emph{not} established that these components carry
|
|
down to small models. Design rules for the sub-100M regime are thin,
|
|
and for several components the evidence is contradictory (see below).
|
|
% EXPL: THE HOOK. "regime" = a size range with its own behaviour. The
|
|
% whole paper hangs on this tension: proven-big does not mean
|
|
% proven-small.
|
|
\item Why now: small-model pretraining is a live research area
|
|
(sample-efficiency work \cite{imu1}; compute-constrained ablations
|
|
\cite{slmtradeoffs}); getting the component choices right matters most
|
|
exactly where the compute budget is smallest.
|
|
% DEFN: pretraining -- the big first training phase where a model learns
|
|
% language from lots of raw text (before any task-specific tuning).
|
|
% DEFN: sample-efficiency -- getting good results from LESS training data.
|
|
\end{itemize}
|
|
|
|
%% -------------------------------------------------------------------
|
|
%% COMPONENT 2 -- LITERATURE CONTEXT (mini lit review: what's been done)
|
|
%% -------------------------------------------------------------------
|
|
% NOTE: the meat. Synthesize (don't list). Group by theme, not by paper.
|
|
% EXPL: COMPONENT 2 is the "mini literature review": what is already known.
|
|
% Group by THEME (what QK-norm is / who uses it / small-scale evidence /
|
|
% the contradiction / why LR), not one-paragraph-per-paper.
|
|
\begin{itemize}
|
|
\item \textbf{What QK-norm is + why it exists.}
|
|
\begin{itemize}
|
|
\item Origin: Henry et al. \cite{henry2020} -- normalize the query and key
|
|
vectors before the attention dot product (for low-resource NMT).
|
|
% DEFN: "normalize ... before the dot product" -- rescale the query and
|
|
% key vectors to a fixed length first, so their dot product (the
|
|
% match score) stays in a controlled range. THAT is QK-norm.
|
|
% DEFN: dot product -- a single number measuring how aligned two vectors
|
|
% are; here it is the query-key match score = the attention logit.
|
|
% DEFN: NMT -- Neural Machine Translation (automatic translation). QK-norm
|
|
% was first proposed there, on small models -- useful for our story.
|
|
\item Mechanism/problem it solves: without it, query/key norms grow during
|
|
training $\rightarrow$ large attention logits $\rightarrow$ softmax
|
|
saturates (near one-hot) $\rightarrow$ unstable/diverging loss,
|
|
forcing smaller learning rates.
|
|
% EXPL: the core causal chain, in words --
|
|
% 1. "norms grow" = the Q/K vectors get longer (bigger magnitude);
|
|
% 2. so their dot products (attention logits) get very large;
|
|
% 3. SOFTMAX SATURATES: given one huge score, softmax puts nearly all
|
|
% weight on a single token -> output is "one-hot" (a spike);
|
|
% 4. at a spike the gradient (learning signal) ~vanishes, so training
|
|
% becomes unstable or DIVERGES (loss shoots up), which forces you
|
|
% to use a smaller, slower learning rate.
|
|
% DEFN: one-hot -- a vector that is 1 in a single spot and 0 everywhere
|
|
% else; here it means attention collapses onto ONE token.
|
|
% DEFN: diverging loss -- training "blowing up": the error metric
|
|
% increases uncontrollably instead of decreasing.
|
|
% EXPL: QK-norm breaks step 1->2, so the rest of the chain never fires.
|
|
% This is exactly why it should let you train at HIGHER LR (our H2).
|
|
\item Deeper mechanism (why logits grow): ``massive activations'' -- a few
|
|
tokens acquire extreme-magnitude activations that become dominant key
|
|
vectors and form ``attention sinks'' \cite{sun_massive,attnsink};
|
|
QK-norm bounds Q/K magnitude, limiting this. % [TO VERIFY]
|
|
% DEFN: activation -- an internal value a neuron/layer outputs as data
|
|
% flows through the network.
|
|
% DEFN: massive activations -- a handful of these internal values become
|
|
% hugely larger than the rest (extreme outliers).
|
|
% DEFN: attention sink -- a token that soaks up a large share of every
|
|
% other token's attention regardless of relevance (a side-effect of
|
|
% those massive activations). This is the mechanism QK-norm curbs.
|
|
\item Popularized at scale for stability: ViT-22B \cite{vit22b}.
|
|
% DEFN: ViT-22B -- a 22-billion-parameter Vision Transformer (Google).
|
|
% Widely cited as the work that made QK-norm mainstream for
|
|
% stabilizing very large models.
|
|
\item Complementary mechanistic account: ``attention entropy collapse''
|
|
-- low attention entropy co-occurs with instability
|
|
\cite{zhai2023}. (QK-norm bounds logits $\rightarrow$ keeps entropy up.)
|
|
% DEFN: attention entropy -- a spread-out measure of the attention
|
|
% weights. HIGH entropy = attention shared across many tokens;
|
|
% LOW entropy = concentrated on a few (spiky). "Collapse" = it
|
|
% drops too low, which coincides with unstable training. Same story
|
|
% as softmax saturation, viewed through an information-theory lens.
|
|
\end{itemize}
|
|
\item \textbf{QK-norm is now standard in frontier models.}
|
|
\begin{itemize}
|
|
\item OLMo 2 applies RMSNorm to Q and K and reports improved training loss
|
|
+ stability \cite{olmo2}. Also adopted by Gemma~2/3, Qwen3, Chameleon.
|
|
% TODO: pick 1-2 to cite precisely; verify each.
|
|
% DEFN: "frontier models" -- the current top-tier, largest LLMs. Naming
|
|
% several (OLMo 2, Gemma, Qwen3, Chameleon) shows QK-norm is not
|
|
% niche -- it is standard practice -> the study is relevant.
|
|
\end{itemize}
|
|
\item \textbf{QK-norm is one of several logit-control methods} (matters for the gap).
|
|
\begin{itemize}
|
|
\item Alternatives/complements: attention logit soft-capping (Gemma~2,
|
|
$z\!\leftarrow\!c\tanh(z/c)$) \cite{softcap}; z-loss, which keeps the
|
|
softmax normalizer near zero \cite{zloss}; NormFormer's extra
|
|
normalization \cite{normformer}. % [TO VERIFY]
|
|
% EXPL: these are RIVAL ways to tackle the same "logits blow up" problem:
|
|
% DEFN: logit soft-capping -- squash every attention logit through
|
|
% z <- c*tanh(z/c), which smoothly caps it at +/-c so it can never
|
|
% explode ($ ... $ is inline maths; \tanh is the tanh function).
|
|
% DEFN: z-loss -- a small extra penalty added to the training objective
|
|
% that keeps the softmax's normalizing sum near 1 (stops it drifting
|
|
% and destabilizing training).
|
|
% DEFN: NormFormer -- a recipe that adds a few extra normalization layers
|
|
% inside each block to steady training.
|
|
\item These are almost always studied \emph{bundled together}, so QK-norm's
|
|
isolated effect is rarely attributed. % [TO VERIFY framing]
|
|
% EXPL: KEY SETUP FOR THE GAP -- if papers turn on QK-norm AND z-loss AND
|
|
% soft-capping at once, you cannot tell how much QK-norm ALONE did.
|
|
% "attributed" = credited to a specific cause. Our isolation fixes this.
|
|
\end{itemize}
|
|
\item \textbf{Evidence at / near small scale (the closest work).}
|
|
\begin{itemize}
|
|
\item Wortsman et al. \cite{wortsman2024} -- instabilities reproduce in
|
|
\emph{small} models at high LR; qk-layernorm (+ z-loss) mitigate them
|
|
and ``transfer down''; introduces the \textbf{LR-sensitivity} metric.
|
|
KEY methodological template for us.
|
|
% EXPL: the closest prior work + the method we borrow. "transfer down"
|
|
% = the same fix that works on big models also works on small ones.
|
|
% DEFN: qk-layernorm -- Wortsman's name for QK-norm (LayerNorm on Q,K).
|
|
% DEFN: LR-sensitivity -- their metric for HOW MUCH the final loss worsens
|
|
% as you vary the learning rate. Low sensitivity = robust to LR
|
|
% choice = a flatter, more forgiving curve. We reuse this metric.
|
|
\item IMU-1 (Grigorev et al.) \cite{imu1} -- at 70M, QK-norm is the
|
|
\emph{single largest} architectural win (measured via reduced
|
|
activation kurtosis) -- but bundled with $\mu$P, NorMuon, value
|
|
residuals, etc.
|
|
% EXPL: strongest recent "QK-norm helps small models" datapoint (70M =
|
|
% 70-million-parameter model) -- BUT confounded (see the Gap).
|
|
% DEFN: kurtosis -- a statistic for how "outlier-heavy" a set of numbers
|
|
% is. High kurtosis = a few extreme values (the massive activations
|
|
% above). QK-norm LOWERING kurtosis = it is taming those outliers.
|
|
% DEFN: NorMuon / value residuals -- other training tweaks IMU-1 turned on
|
|
% at the same time; because they are bundled in, IMU-1 cannot isolate
|
|
% QK-norm's own contribution.
|
|
\end{itemize}
|
|
\item \textbf{The countervailing line (the contradiction).}
|
|
\begin{itemize}
|
|
\item Compute-constrained small-model ablations report that modern
|
|
components (RoPE, RMSNorm, GQA) give ``no measurable benefit at
|
|
$\sim$70M'' \cite{slmtradeoffs}. Directly disputes the above.
|
|
% EXPL: THE CONFLICT that motivates the paper: one camp says QK-norm is
|
|
% the biggest small-model win (IMU-1), another says modern parts do
|
|
% nothing at ~70M. Both look at the same size -> something is untested.
|
|
% DEFN: "compute-constrained" -- trained under a fixed, limited compute
|
|
% budget (so, e.g., only one learning-rate setting was tried).
|
|
\end{itemize}
|
|
\item \textbf{Why learning rate is the right axis.}
|
|
\begin{itemize}
|
|
\item $\mu$P / $\mu$Transfer \cite{mup}: optimal LR transfers across width;
|
|
LR is the hyperparameter these stability effects hinge on. Motivates
|
|
studying QK-norm \emph{as a function of LR}, and switching $\mu$P OFF
|
|
so LR effects are not masked.
|
|
% DEFN: $\mu$P (muP, "maximal update parametrization") -- a special way of
|
|
% scaling a network so the BEST learning rate stays the same as you
|
|
% grow model width; $\mu$Transfer = tuning LR on a small model then
|
|
% reusing it on a big one. ($\mu$ = the Greek letter mu.)
|
|
% EXPL: why we turn $\mu$P OFF -- it deliberately makes results LESS
|
|
% sensitive to learning rate, which would hide the very LR effect we
|
|
% want to measure. Off = a clean, unmasked view. (Also a difference
|
|
% from IMU-1, which left it ON.)
|
|
% DEFN: hyperparameter -- a training setting you choose by hand (like the
|
|
% learning rate), as opposed to a weight the model learns itself.
|
|
\item Recent work keeps reaffirming the LR link -- e.g. QK-Normed MLA reports
|
|
QK-norm ``permits significant increases in learning rate'' \cite{qknmla}.
|
|
% [TO VERIFY] -- supports H2
|
|
% EXPL: a 2026 paper independently restating that QK-norm lets you train at
|
|
% higher LR -- direct external support for our hypothesis H2 below.
|
|
\end{itemize}
|
|
\end{itemize}
|
|
|
|
%% -------------------------------------------------------------------
|
|
%% COMPONENT 3 -- PROBLEM STATEMENT & GAP ANALYSIS
|
|
%% -------------------------------------------------------------------
|
|
% NOTE: name the gap crisply; it must fall straight out of Component 2.
|
|
% EXPL: COMPONENT 3 = the hole in the literature the study fills. Every claim here
|
|
% should already have been set up in Component 2 (no new surprises).
|
|
\begin{itemize}
|
|
\item The positive small-scale evidence is \textbf{confounded}: Wortsman
|
|
bundles qk-norm with z-loss and frames small models as a \emph{proxy}
|
|
for large-scale behaviour; IMU-1 bundles it with $\mu$P + NorMuon +
|
|
value residuals and measures it via activation statistics, not an LR
|
|
sweep, at a single (fixed) depth.
|
|
% DEFN: confounded -- when two changes are made at once so you cannot tell
|
|
% which caused the effect; the classic flaw an ablation is meant to avoid.
|
|
% DEFN: proxy -- a stand-in. Wortsman used small models only as a cheap
|
|
% stand-in for big-model behaviour, not as the object of interest.
|
|
% DEFN: LR sweep -- training the same model repeatedly across a range of
|
|
% learning rates to map out how the outcome depends on LR. The thing
|
|
% prior work at this scale did NOT do.
|
|
\item The negative evidence \cite{slmtradeoffs} tests a \emph{single}
|
|
compute-optimal budget, does not isolate QK-norm, and does not sweep LR.
|
|
% DEFN: compute-optimal budget -- the amount of training data/steps that is
|
|
% "best value" for a given model size (Chinchilla-style rule of thumb).
|
|
% Testing only this one point can miss effects that appear at other LRs.
|
|
\item So the two literatures \emph{disagree at the same scale} while neither
|
|
does the clean experiment.
|
|
% EXPL: one-line crux -- the disagreement persists because the decisive
|
|
% experiment (isolate + sweep LR) has not been run.
|
|
\item \textbf{GAP:} no \emph{dense}, \emph{isolated} ($\mu$P-off)
|
|
learning-rate-sensitivity characterization of QK-norm at 10--100M, and
|
|
no map of how the LR at which QK-norm becomes \emph{necessary} moves
|
|
with model depth.
|
|
% EXPL: the headline gap, four qualifiers doing the work:
|
|
% dense = many closely-spaced LR values (a smooth curve, not 3 points);
|
|
% isolated = only QK-norm toggled, nothing else;
|
|
% muP-off = the masking factor removed (see above);
|
|
% at 10-100M = the small-scale band nobody has charted this way.
|
|
\item Secondary gap: reproducibility -- most evidence is expensive / CUDA-only;
|
|
a laptop-reproducible study is itself a contribution.
|
|
% DEFN: CUDA -- NVIDIA's GPU software stack; "CUDA-only" work cannot run on
|
|
% an Apple laptop, so it is hard for others to reproduce cheaply.
|
|
\end{itemize}
|
|
|
|
%% -------------------------------------------------------------------
|
|
%% COMPONENT 4 -- PROPOSED CONTRIBUTION (objective, RQs, method)
|
|
%% -------------------------------------------------------------------
|
|
% NOTE: be explicit and honest -- this is VERIFICATION + boundary-charting,
|
|
% not a novel-technique claim. Do not say "first".
|
|
% EXPL: COMPONENT 4 = what WE do about the gap: objective, questions, hypotheses,
|
|
% method. "boundary-charting" = mapping exactly WHERE QK-norm starts to
|
|
% matter (which LR, which depth), rather than claiming a new technique.
|
|
\begin{itemize}
|
|
\item Objective: cleanly isolate QK-norm and chart \emph{when a small model
|
|
needs it}.
|
|
\item Research question: at $<$100M params, does isolated QK-norm improve
|
|
convergence and raise the maximum stable LR -- and at what LR does it
|
|
become necessary? Sub-question: how does that threshold shift with depth?
|
|
% DEFN: convergence -- training successfully settling to a low, stable loss.
|
|
% "Faster convergence" = reaching a target loss in fewer steps.
|
|
% DEFN: max stable LR -- the highest learning rate at which training still
|
|
% succeeds instead of diverging. A key thing QK-norm should raise.
|
|
\item Hypotheses (state as such):
|
|
% DEFN: hypothesis -- a specific, testable prediction stated BEFORE running
|
|
% the experiment. H1-H3 are what we expect the data to show.
|
|
\begin{itemize}
|
|
\item H1 -- effect is LR-conditional: negligible/slightly negative at low
|
|
LR, decisive at high LR.
|
|
% EXPL: at gentle LRs nothing diverges, so QK-norm barely helps (or
|
|
% slightly hurts); at aggressive LRs it is what keeps training alive.
|
|
\item H2 -- QK-norm raises max stable LR and lowers LR-sensitivity.
|
|
% EXPL: i.e. it widens the band of usable LRs and flattens the loss-vs-LR
|
|
% curve (more forgiving to the LR you pick).
|
|
\item H3 -- the ``necessity'' LR falls as depth increases.
|
|
% EXPL: prediction that DEEPER models (more layers) destabilize sooner, so
|
|
% they need QK-norm at lower LRs than shallow ones.
|
|
\end{itemize}
|
|
\item High-level method: controlled ablation on \textbf{plain nanoGPT}
|
|
\cite{nanogpt}, $\mu$P off, AdamW; single independent variable
|
|
QK-norm\{on,off\} $\times$ dense LR sweep $\times$ depth; $\sim$10--15M
|
|
params; measures = LR-sensitivity \cite{wortsman2024}, max-stable-LR,
|
|
steps-to-threshold, perplexity, plus attention entropy / activation
|
|
kurtosis as mechanistic checks.
|
|
% DEFN: nanoGPT -- Andrej Karpathy's minimal, readable codebase for training
|
|
% small GPT models; "plain" = without the extra speed-run tricks that
|
|
% would confound the experiment. Our experimental substrate.
|
|
% DEFN: AdamW -- the standard optimizer (the algorithm that updates weights
|
|
% using gradients); "W" = with decoupled weight decay.
|
|
% DEFN: independent variable -- the ONE thing we deliberately change
|
|
% (QK-norm on vs off). "$\times$" here means we cross it with every LR
|
|
% and every depth -> a grid of runs.
|
|
% DEFN: steps-to-threshold -- how many training steps to first reach a chosen
|
|
% loss value; a concrete convergence-speed number.
|
|
% DEFN: mechanistic checks -- extra measurements (attention entropy, kurtosis)
|
|
% that reveal WHY an effect happens, not just THAT it happens.
|
|
\item Data + precedent: small models learn coherent language from tiny corpora
|
|
(TinyStories) \cite{tinystories}; single-GPU near-BERT training is an
|
|
established accessible-research precedent (Cramming) \cite{cramming}. % [TO VERIFY]
|
|
% DEFN: corpus (pl. corpora) -- the body of text a model trains on.
|
|
% DEFN: TinyStories -- a dataset of very simple stories on which even
|
|
% sub-10M-param models produce fluent English; justifies our tiny scale.
|
|
% DEFN: Cramming -- a paper training a BERT-quality model on ONE GPU in ONE
|
|
% day; the precedent for "serious results on modest hardware".
|
|
% DEFN: BERT -- an influential 2018 language model; "near-BERT" = almost as
|
|
% good as BERT.
|
|
\item Feasibility/repro: runs on a consumer laptop (Apple M4 Pro); code, seeds,
|
|
logs released.
|
|
% DEFN: seed -- the number that fixes the random choices in a training run;
|
|
% sharing seeds lets anyone reproduce the exact run. Running multiple
|
|
% seeds also shows an effect is real, not a fluke of one random draw.
|
|
\end{itemize}
|
|
|
|
%% -------------------------------------------------------------------
|
|
%% COMPONENT 5 -- PAPER ROADMAP (standard closing paragraph)
|
|
%% -------------------------------------------------------------------
|
|
% NOTE: rewrite to match the final section list once it exists.
|
|
% EXPL: COMPONENT 5 = the "roadmap" paragraph. A standard academic convention:
|
|
% one or two sentences telling the reader how the rest of the paper is laid
|
|
% out. "Section II ... Section III ..." are the numbered parts that follow.
|
|
% (Roman numerals are how IEEEtran auto-numbers \section headings.)
|
|
\begin{itemize}
|
|
\item ``The remainder of this paper is structured as follows: Section II
|
|
reviews related work on QK-norm and small-model training stability;
|
|
Section III details the experimental design; Section IV reports results;
|
|
and Section V concludes.'' % adjust to actual sections.
|
|
\end{itemize}
|
|
|
|
% =====================================================================
|
|
% NOTE(refs): inline thebibliography for now (matches template). Switch to a
|
|
% .bib file later if we want. EVERY entry below still needs its exact
|
|
% authors/venue/pages checked against the primary PDF -- during this project
|
|
% AI-surfaced citation details were wrong more than once. arXiv IDs are the
|
|
% most reliable field; treat author lists / venues as [TO VERIFY].
|
|
% =====================================================================
|
|
% DEFN: thebibliography -- the reference list. Each \bibitem{key} defines one
|
|
% reference; \cite{key} in the text above prints its number. The {00} just
|
|
% tells LaTeX the labels are up to two digits wide.
|
|
% DEFN: arXiv -- a free repository of (often not-yet-peer-reviewed) research
|
|
% papers. "arXiv:2010.04245" is a permanent ID: go to arxiv.org/abs/2010.04245.
|
|
% Reliable for finding the paper; the AUTHOR/VENUE text still needs checking
|
|
% (hence the [TO VERIFY] tags) -- run each through Scopus before submission.
|
|
\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
|
|
|
|
% --- added 18/08/2026: surrounding literature; ALL need Scopus/PDF checks ---
|
|
\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] find arXiv ID + exact citation for logit soft-capping
|
|
\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{qknmla} ``QK-Normed MLA: QK normalization without full key caching,''
|
|
2026. arXiv:2606.16310. % [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}
|