Proposition 5. Let $`l`$ and $`\overline{l^\star}`$ denote the mean output logits before and after $`\mu`$-centering, respectively.

1. *The mean output logit after $`\mu`$-centering is zero:

   MATH

   Collapse Copy

   \begin{align}
      \overline{l^\star}= 0
   \end{align}
   ```*

   Click to expand and view more

2. *The output logits standard deviation is not affected by $`\mu`$-centering:

   MATH

   Collapse Copy

   \begin{align}
   \sigma_{l^\star}= \sigma_l
   \end{align}
   ```*

   Click to expand and view more

3. The output probabilities and the loss are not affected by $`\mu`$-centering.

Proposition 7. *An infinite $`\mu`$-loss $`\mathcal{L}_\mu`$ corresponds to

\begin{align}
\max_j \big| l_j \big| \to \pm \infty
\label{theorem_muloss}
\end{align}
```*

</div>

<div class="proof">

*Proof.* Follows directly from
Eq. (<a href="#eq:mer" data-reference-type="ref"
data-reference="eq:mer">[eq:mer]</a>). ◻

</div>

Note the subtle difference compared to z-loss and
Proposition <a href="#theorem_zloss" data-reference-type="ref"
data-reference="theorem_zloss">2</a>: Absolute logits
$`\big| l_j \big|`$ appear in the limit instead of the logits $`l_j`$
themselves. Hence, $`\mu`$-loss suppresses the positive or negative
divergence of any single logit.
Tab. <a href="#tab:regularization_methods_effect" data-reference-type="ref"
data-reference="tab:regularization_methods_effect">1</a> summarizes the
methods discussed in this section, and the means by which they prevent
logit divergence.

<div id="tab:regularization_methods_effect">

<table>
<caption>Overview of methods and means by which logit divergences are
suppressed. Note that suppression of single divergences implies
suppression of collective divergences, but not vice versa.</caption>
<tbody>
<tr>
<td style="text-align: left;"></td>
<td style="text-align: center;"></td>
<td colspan="2" style="text-align: center;">suppressed divergence</td>
</tr>
<tr>
<td style="text-align: left;">name</td>
<td style="text-align: center;">type</td>
<td style="text-align: center;">positive</td>
<td style="text-align: center;">negative</td>
</tr>
<tr>
<td style="text-align: left;">z-loss</td>
<td style="text-align: center;">regularization</td>
<td style="text-align: center;">single</td>
<td style="text-align: center;">collective</td>
</tr>
<tr>
<td style="text-align: left;"><span
class="math inline"><em>μ</em></span>-loss</td>
<td style="text-align: center;">regularization</td>
<td style="text-align: center;">single</td>
<td style="text-align: center;">single</td>
</tr>
<tr>
<td style="text-align: left;"><span
class="math inline"><em>μ</em></span>-centering</td>
<td style="text-align: center;">centering</td>
<td style="text-align: center;">single</td>
<td style="text-align: center;">single</td>
</tr>
</tbody>
</table>

</div>

The theoretical advantages of $`\mu`$-loss and $`\mu`$-centering over
z-loss are the suppression of single negative logit divergences, their
simplicity, and the fact that they have a theoretical foundation that
addresses the root cause of the problem. Potential additional advantages
of $`\mu`$-centering over the regularization methods are that it is
hyperparameter-free and deterministic instead of stochastic. In
contrast, the regularization methods might offer more flexibility
compared to $`\mu`$-centering.

# Experiments

Our approach to studying training stability with regard to output logit
divergence primarily follows . In particular, we train dense decoder
models with a modern Transformer architecture on 13.1 billion tokens for
100000 steps, using 7 different learning rates:
``` math
\begin{align}
\eta \in \{ \text{3e-4, 1e-3, 3e-3, 1e-2, 3e-2, 1e-1, 3e-1} \}
\label{eq:eta}
\end{align}

However, there are also a number of key differences. We use FineWeb and the GPT-2 tokenizer with a vocabulary size of $`V = 50304`$. Our 5 model sizes,

\begin{align}
N \in \{ 16{\rm M}, 29{\rm M}, 57{\rm M}, 109{\rm M}, 221{\rm M} \}
\label{eq:model_sizes}
\end{align}

and the corresponding specifications (e.g. widths, number of layers and attention heads) are taken from . In addition, we use SwiGLU hidden activations and a non-truncated Xavier weight initialization . Further details on model architecture and hyperparameters are provided in App. 8. For each of the 7 x 5 = 35 combinations of learning rate and model size defined by Eq. ([eq:eta]) and Eq. ([eq:model_sizes]), we train four different models: A baseline model with the standard language modeling head (Sec. 1), and models using z-loss, $`\mu`$-loss as well as $`\mu`$-centering (Sec. 2). In order to compare the variants, we evaluate the dependency of the test loss on the learning rate and the dependency of learning rate sensitivity on the model size, with the latter defined as in :

\begin{align}
        {\rm LRS} &= \mathbb{E}_{\eta} \left[ \min (\mathcal{L}(\eta) , \mathcal{L}_0) - \min_\eta \mathcal{L} \right]
        \label{eq:lr_sensitivity}
\end{align}

Here, $`\eta`$ are the learning rates from Eq. ([eq:eta]) and $`\mathcal{L}_0`$ denotes the loss at initialization time. Additionally, we investigate the dependency of a few other metrics on the learning rate for the purpose of analyzing the functionality of the different methods. Firstly, we consider the norm $`\| \mu\|`$ of the mean embedding (see Eq. ([eq:mu])). Secondly, we compute sample estimates for the mean logit $`\overline{l}`$ (see Eq. ([eq:meanlogit])), the logits standard deviation

\begin{align}
\sigma_l&= \frac{1}{V} \sum_{j=1}^V \left( l_j - \overline{l}\right)^2 \ ,
\end{align}

as well as the maximum absolute logit

\begin{align}
\max_j |l_j| \ ,
\end{align}

using $`5 \cdot 10^5`$ logit vectors created from the test data. Finally, the time $`t`$ to train a model on 4 A100 GPUs using data parallelism is compared.

Results

Training Stability

The main results of our experiments are shown in Tab. 4 and Fig  2.

The top table (i) demonstrates that the optimal loss $`\min_\eta \mathcal{L}`$ for each model size is virtually the same for all methods. As expected, the top figure shows that the non-regularized baseline is the first to diverge with larger learning rates. Interestingly, z-loss leads to occasional divergences as well, given a large enough learning rate². Meanwhile, none of the models using our methods diverge to any significant extent. This is also reflected in subtable (ii) of Tab. 4, which shows that $`\mu`$-loss and $`\mu`$-centering exhibit a lower learning rate sensitivity than z-loss, for all models sizes. In addition, subtable (iii) reveals that our methods are computationally cheap, such that the training time is minimally affected.

Analysis

The additional metrics mentioned at the end of Sec. 3 are visualized in Fig. 3.

Firstly, regarding the logits mean (top left), we find that $`\mu`$-centering and $`\mu`$-loss center the logits at and around 0, respectively. Similarly, z-loss indirectly controls the logits mean, although at negative values. In contrast, the logits mean diverges at higher learning rates for the baseline, in accordance with the loss divergence observed in Fig. 2. Secondly, the standard deviation (top right) is the same for $`\mu`$-centering and the baseline barring slight statistical differences, at least for lower learning rates for which the baseline training converges. This is consistent with the theoretical prediction, see Proposition 5. In contrast, z-loss and $`\mu`$-loss—since they are regularization methods—change the logit standard deviation slightly. Thirdly, the mean embedding norm is shown on the bottom left. As expected, $`\mu`$-centering maintains a norm of zero while both baseline and z-loss grow at higher learning rates, indicating that z-loss fails to prevent anisotropic embeddings. Meanwhile, $`\mu`$-loss constrains the mean embedding norm to relatively small values. Finally, as predicted by Theorem 6, both $`\mu`$-centering and $`\mu`$-loss restrict the logit bound such that the maximum logit remains stable. Similarly, z-loss also implicitly restricts the maximum logit, albeit to a lesser degree than our methods, which explains the divergence observed for training using z-loss. In contrast, the maximum logit grows extremely large for the baseline models. In summary, these results are in accordance with the theoretical predictions from Sec. 2.

Hyperparameter Sensitivity

So far, the regularization hyperparameters have been set to their default value $`\lambda = 10^{-4}`$ for both regularization methods, z-loss (cf. Eq. ([eq:zloss])) and $`\mu`$-loss (cf. Eq. ([eq:lambda_default])). We now vary the regularization hyperparameter

\begin{align}
\lambda &\in \{ 10^{-7}, 10^{-4}, 10^{-1}, 10^{2} \}
\label{eq:lambda_ablations}
\end{align}

for those methods, and determine the optimal loss and learning rate sensitivity as in Sec. 4 for each choice of $`\lambda`$. The results are presented in Tab. [tab:overview_comparison] and Fig. 8.

For $`\mu`$-loss, hyperparameter tuning is notably straightforward: the regularization coefficient only needs to be sufficiently large to enforce the centering effect. In fact, for larger values ($`\lambda \ge 10^{-4}`$), the training is stable and does not exhibit a strong dependency on the exact value of $`\lambda`$. Only when $`\lambda`$ is too small ($`\lambda = 10^{-7}`$), we observe that the loss diverges for large learning rates across all model sizes.

This behavior stands in contrast to z-loss, which requires more careful tuning. Severe divergences appear for $`\lambda=10^2`$, but also for lower values of $`\lambda`$ in conjunction with large learning rates. Our results indicate that the optimal value for z-loss is $`\lambda=10^{-1}`$, which is significantly larger than the previously assumed optimal value of $`10^{-4}`$. Importantly, however, even for the optimal $`\lambda`$, z-loss is outperformed by both $`\mu`$-loss and $`\mu`$-centering. This performance gap is evident in the learning rate sensitivity values for the largest model size $`N=221`$ in Tab. [tab:overview_comparison], as well as in the comparison of the rightmost points—corresponding to the largest model size—across the learning rate sensitivity plots in Fig. 8.

Conclusions

This paper establishes a link between the problems of anisotropic embeddings and output logit divergence. We have identified the former as the cause of the latter, and introduced $`\mu`$-centering and $`\mu`$-loss as theoretically well-founded mitigation strategies. Our experiments show that our methods outperform z-loss in terms of training stability, learning rate sensitivity and hyperparameter sensitivity. The code to reproduce our results is available at github.com/flxst/output-embedding-centering  .

Limitations

We have only trained models up to a size of 221M parameters. In addition, our experiments use a fixed dataset, vocabulary size, token budget and set of hyperparameters. Hence, the same limitations as in apply. We have not investigated the dependency of the results on these factors, so we cannot make any reliable statements about their generalizability. Finally, while we have discussed the theoretical pros and cons of $`\mu`$-centering or $`\mu`$-loss in Sec. 2, we do not provide a clear recommendation on which method is to be preferred in practice.

Hyperparameters

All our experiments use the architecture and hyperparameters specified in Tab. 5.

Results for $`B_{\rm ratio}`$

As described in Sec. 3, we trained a total of 35 baseline models with a standard language modeling head (see Sec. 1), using 7 different learning rates (see Eq. ([eq:eta])) and 5 different model sizes (see Eq. ([eq:model_sizes])). Tab. 6 lists $`B_{\rm ratio}`$, as defined in Eq. ([eq:theorem_oec_condition]), individually for each of these models, while Fig. 9 shows a histogram of all its values.

/>

For each model, we find that the condition for Theorem 6 is fulfilled: $`B_{\rm ratio}\leq 1`$. Tab. 6 also shows that $`B_{\rm ratio}`$ tends to decrease with a larger learning rate. This indicates that the beneficial effect of $`\mu`$-centering (or $`\mu`$-loss) on the output logit bounds becomes larger, which is also in accordance with our results in Sec. 4.