**Assumptions.** In Lemma 4.1, we assume that both the matrix with columns \(g_{\theta}\left(p_{i}\right)\) and the matrix \(\mathbf{W}\) have rank \(h\). These matrices have either \(h\) rows or \(h\) columns, so both have rank at most \(h\). Moreover, it is very unlikely that they have rank \(<h\): this would require the distribution of \(g_{\theta}\left(p\right)\) to be fully supported on a subspace of dimension \(<h\) across all \(p_{i}\) we query, or all \(h\ll l\) columns of \(\mathbf{W}\) to lie in the same \((h-1)\) dimensional subspace of \(\mathbf{R}^{l}\) (the output space of logits). In practice we find this assumption holds for all larger models (Table 2) and when different normalization layers are used (Appendix B.1).