1. We merge the final LayerNorm weights \(\gamma\) into \(\mathbf{W}\) by linearity.3 Footnote 3: For a full explanation of this method of rewriting the unembedding matrix, see Appendix A.1, ‘Folding LayerNorm’ in Gurnee et al. (2024). 2. We assume the numerical precision is high enough that after the final LayerNorm, the hidden states are on a sphere. 3. There is no degenerate lower-dimensional subspace containing all \(g_{\theta}(p)\) for all our queries \(p\). 4. We assume the \(\varepsilon\) in RMSNorm/LayerNorm is \(0\). This is not a critical assumption.