Enhanced QKNorm normalization for neural transformers with the Lp norm

TL;DR

This paper tackles numerical instability in Transformer attention by generalizing Query–Key Normalization (QKNorm) from Euclidean to $L_p$ norms. It introduces $L_p$-norm attention, where queries and keys are normalized by $\|\cdot\|_p$ and attention logits are computed as $s^{(p)}_{ij} = \alpha\, \hat{\mathbf{q}}^{(p)}_i{}^T \hat{\mathbf{k}}^{(p)}_j$, with a learnable scaling $\alpha$; the final output uses softmax over these logits. Experiments on Tiny Shakespeare with a 6-layer, 6-head Transformer and 10-fold cross-validation across $p \in \{1.0,1.5,2.0,2.5,3.0,3.5,4.0\}$ show that $p>2$ yields faster convergence and lower average validation cross-entropy than the standard $p=2$ baseline, with negligible runtime overhead. The results suggest that non-Euclidean normalization can improve attention geometry, offering a tunable mechanism to shape the span of influential vector components and potentially enhance training stability in low-resource settings.

Abstract

The normalization of query and key vectors is an essential part of the Transformer architecture. It ensures that learning is stable regardless of the scale of these vectors. Some normalization approaches are available. In this preliminary work, a generalization of the QKNorm normalization scheme is proposed. The approach is based on the Lp norm, allowing non-Euclidean norms to be employed. Experimental results demonstrate the suitability of the method for a simple problem.

Figures (3)

• Figure 1: Per-$p$ validation loss curves, averaged by fold.
• Figure 2: Per-$p$ training time distribution across folds: the central line denotes the median, the box spans the interquartile range (IQR), whiskers extend to 1.5×IQR, and outliers are shown as individual points; the diamond marks the mean.
• Figure 3: Validation loss vs. iterations for each fold (KFold with $K=10$).