Better Embeddings with Coupled Adam
Authors: Felix Stollenwerk, Tobias Stollenwerk
Accepted at: acl25

Better Embeddings with Coupled Adam

Felix Stollenwerk*
AI Sweden

Tobias Stollenwerk
Forschungszentrum Jülich

Abstract

Despite their remarkable capabilities, LLMs learn word representations that exhibit the undesirable yet poorly understood feature of anisotropy. In this paper, we argue that the second moment in Adam is a cause of anisotropic embeddings, and suggest a modified optimizer called Coupled Adam to mitigate the problem. Our experiments demonstrate that Coupled Adam significantly improves the quality of embeddings, while also leading to better upstream and downstream performance on large enough datasets.

1 Introduction

Anisotropic Embeddings  Large Language Models (LLMs) take a sequence of tokens as input and predict the next token. An embedding matrix is used to map the input tokens to the hidden space of the model, while an unembedding matrix provides the inverse mapping to the output token space. Although the two matrices can in principle be different, it is common practice to apply weight tying (Press and Wolf, 2017) and use the transpose of the embedding matrix for unembedding. During training, the model learns an embedding vector in hidden space for each token in the vocabulary. However, it is observed that those embedding vectors are clustered in a small subspace away from the origin (Gao et al., 2019). This anisotropy limits the semantic usefulness of the embeddings and, in turn, the expressiveness and generalizability of the model. Multiple attempts have been made to both explain the root cause of the problem and alleviate it (more on this in Sec. 7). In particular, Biš et al. (2021) have shown that the problem can be traced back to a mere shift of the mean embedding vector away from the origin. With the mean embedding vector as reference point, the embeddings feature near-perfect isotropy. However, the role of the employed optimization algorithm has, to the best of our knowledge, not yet been investigated.

Optimization Algorithms  Optimization algorithms are an indispensable ingredient in the training of neural networks generally and LLMs in particular. While SGD is the foundational optimization technique, Adam (Kingma and Ba, 2014) is the most widely used optimization techniques for LLMs due to its superior performance and robustness. While it provides multiple conceptional advantages over SGD, see e.g. Ruder (2017) for a detailed discussion, the one that is particularly striking with regard to word embeddings is that Adam is well-suited for sparse data. More concretely, this means that using Adam, the embedding update vectors for rare words are scaled up in comparison to those of more frequent words. This is relevant in the context of LLMs as word frequencies in the training data are typically very skewed and may differ by several orders of magnitude. Formally, this is captured by the unigram probability distribution \( \tilde{p} \in [0, 1]^V \), which for a given dataset \( d \) and tokenizer \( t \) is defined by

\[
\tilde{p}_i \equiv \tilde{p}_i(d, t) = \frac{n_i}{\sum_j n_j},
\]

where \( i \in \mathcal{V} \equiv \{1, \ldots, V\} \) is the vocabulary index and \( n_i \) is the total number of occurrences of the \( i \)-th token in the tokenized dataset. A visualization of an example unigram probability distribution can be found in App. A.

Our Contributions  In this work, we combine the research areas of anisotropic embeddings and optimization algorithms and provide the following contributions:

• We show that the Adam optimizer plays a crucial role in causing anisotropic embeddings.
• We suggest Coupled Adam, an easy-to-implement yet efficient adjustment of the original Adam optimization algorithm, which is specifically designed for embedding parameters in order to alleviate the anisotropy problem.

• We demonstrate that our method not only significantly improves the quality of word embeddings, but also has a beneficial effect on upstream and downstream performance for sufficiently large datasets.

2 On the Root Cause of Anisotropic Embeddings

We study the collective shift of the embeddings (that underlies the anisotropy problem), by analyzing their vector updates based on the optimization algorithms SGD and Adam. Weight tying is assumed, but only contributions from the output layer are considered, following Biš et al. (2021). Our results apply to all model architectures with a standard language modeling head.

2.1 Language Modeling Head

The equations for the standard language modeling head read

\[
\mathcal{L} = - \log \left( p_t \right) \tag{2}\]
\[
p_t = \frac{\exp \left( l_t \right)}{\sum_{j=1}^V \exp \left( l_j \right)} \tag{3}\]
\[
l_i = e_i \cdot h \tag{4}\]

where \( \mathcal{L} \in \mathbb{R}_{\geq 0} \) is the loss for next token prediction, and \( p_t \in [0, 1] \) is the predicted probability of the true token \( t \in \mathcal{V} \). \( l_i \in \mathbb{R} \) and \( e_i \in \mathbb{R}^H \) denote the logits and embeddings for each token \( i \in \mathcal{V} \), respectively. \( h \in \mathbb{R}^H \) is the final hidden state provided by the model for a single token. Note that the operation in Eq. (4) is the dot product of two vectors in \( \mathbb{R}^H \). Backward propagation yields the following gradients with respect to the input vectors \( e_i \) and \( h \) of Eq. (4):

\[
g_i := \frac{\partial \mathcal{L}}{\partial e_i} = - (\delta_{it} - p_i) \cdot h \tag{5}\]

This result was first reported using a different notation in Biš et al. (2021), and is rederived in App. B for the reader’s convenience.

![Toy example of a hidden state vector h (shown in blue) and three embedding vectors ei (shown in red) in H = 2 dimensions. The gray arrows represent the embedding update vectors, for the SGD (dark) and the Adam (light) optimizer. The update vector of the true token is aligned with h, while the others point in the opposite direction, see Eq. (5). Note that the sum of embedding update vectors vanishes for SGD, while this is not necessarily the case for Adam, cf. Eqs. (11) and (16).](page_820_172_591_377.png)

Figure 1: Toy example of a hidden state vector \( h \) (shown in blue) and three embedding vectors \( e_i \) (shown in red) in \( H = 2 \) dimensions. The gray arrows represent the embedding update vectors, for the SGD (dark) and the Adam (light) optimizer. The update vector of the true token is aligned with \( h \), while the others point in the opposite direction, see Eq. (5). Note that the sum of embedding update vectors vanishes for SGD, while this is not necessarily the case for Adam, cf. Eqs. (11) and (16).

2.2 Vanishing Sum of Embedding Gradients

Optimization algorithms for neural networks usually update the model parameters iteratively, using an additive update vector that points in direction opposite to the gradient of the loss with respect to the parameters. In the case of embedding vectors, this can be expressed by

\[
e_i^{(\tau)} = e_i^{(\tau-1)} + u_i^{(\tau)} , \tag{6}\]

with

\[
u_i^{(\tau)} \propto - g_i^{(\tau)} , \tag{7}\]

where \( u_i^{(\tau)} \) is the update vector for \( e_i^{(\tau)} \) at time step \( \tau \). Eq. (5) implies that the embedding vector \( e_t \) of the true token is updated in direction \( +h \), while the update vectors \( u_i \) for all the other embedding vectors \( e_i \) with \( i \neq t \) are proportional to \( -h \), see Fig. 1. This circumstance is referred to in the literature as the "common enemy effect" (Biš et al., 2021), and regarded as the cause of the representation degeneration problem. However, as we will see in the following sections, this explanation is incomplete, as it does not take into account the scaling of the gradients with the predicted probabilities \( p_i \), see Eq. (5). The basis for our argumentation is the observation that the sum of embedding gradients vanishes, as the following simple calculation shows:

\[
\sum_{i=1}^V g_i^{(\tau)} \overset{(5)}{=} - \sum_{i=1}^V \left( \delta_{it}^{(\tau)} - p_i^{(\tau)} \right) \cdot h^{(\tau)} \\
= - \left( 1 - \sum_{i=1}^V p_i^{(\tau)} \right) \cdot h^{(\tau)} = 0 \tag{8}\]
Next, we will study how Eq. (8) translates to the sum \( \sum_{i=1}^V u_i^{(\tau)} \) of embedding update vectors, as well as the mean embedding vector

\[
\mu^{(\tau)} = \frac{1}{V} \sum_{i=1}^V e_i^{(\tau)}
\] (9)

Since the exact definition of the embedding update vector \( u_i \), i.e. the proportionality factor in Eq. (7), depends on the optimization algorithm, we discuss SGD and Adam separately.

2.3 Invariant Mean Embedding with SGD

We consider the application of the SGD optimization algorithm on the embedding vectors\footnote{Details are given in App. C.}. At each training step, an embedding vector is simply updated by adding the associated negative gradient \( -g_i \), multiplied by a global learning rate \( \eta \). Hence, Eq. (7) becomes

\[
u_i^{(\tau)} = -\eta \cdot g_i^{(\tau)}
\] (10)

Together with Eq. (8), this implies that the *sum of embedding update vectors vanishes* at any time step \( \tau \):

\[
\sum_{i=1}^V u_i^{(\tau)} \overset{(10)}{=} -\eta \sum_{i=1}^V g_i^{(\tau)} \overset{(8)}{=} 0
\] (11)

Consequently, the mean embedding vector will stay invariant during the training process:

\[
\mu^{(\tau)} - \mu^{(\tau-1)} \overset{(9.6)}{=} \frac{1}{V} \sum_{i=1}^V u_i^{(\tau)} \overset{(11)}{=} 0
\] (12)

This holds even though the different embeddings \( e_i \) will be individually updated in different directions with different magnitudes. Moreover, all of the above is true also in the case of SGD with momentum, which follows from linearity and mathematical induction. Eq. (12) has far-reaching implications with regard to the anisotropy problem. It entails that the embedding vectors do not collectively shift away from the origin if SGD (with or without momentum) is used.

2.4 Shifted Mean Embedding with Adam

In this section, we analyze the behavior of the mean embedding during optimization with Adam (Kingma and Ba, 2014), see Algorithm 1.

The update vector Eq. (7) for the Adam algorithm is given by

\[
u_i^{(\tau)} = -\eta_i^{(\tau)} \cdot \hat{m}_i^{(\tau)} ,
\] (13)

where we have introduced an \( i \)-dependent effective learning rate

\[
\eta_i^{(\tau)} := \frac{\eta}{\sqrt{\hat{v}_i^{(\tau)}} + \epsilon}
\] (14)

Note that \( \hat{m}_i^{(\tau)} \) and \( \hat{v}_i^{(\tau)} \) denote the exponentially averaged first and second moments, respectively, defined according to lines 4-7 in Algorithm 1. The \( i \)-dependent learning rate serves the purpose of individually normalizing the update vectors for different parameters in the Adam optimizer. However, it also has an unwanted effect specifically on the embedding vectors. While we know from Eq. (8) and Algorithm 1 (lines 4,6) that the *unweighted* sum over the first moments vanishes, \( \sum_{i=1}^V \hat{m}_i^{(\tau)} = 0 \), this is not true for the *weighted* sum,

\[
\sum_{i=1}^V \eta_i^{(\tau)} \hat{m}_i^{(\tau)} \neq 0 ,
\] (15)

unless \( \eta_i^{(\tau)} = \eta_j^{(\tau)} \) for all \( i,j \in \mathcal{V} \). Hence, the *sum of embedding update vectors does not vanish* in general,

\[
\sum_{i=1}^V u_i^{(\tau)} \overset{(13)}{=} -\sum_{i=1}^V \eta_i^{(\tau)} \cdot \hat{m}_i^{(\tau)} \overset{(15)}{=} \neq 0
\] (16)

<table>
  <tr>
    <th colspan="2">Input:</th>
    <th>η (lr), \( e_i^{(0)} \) (initial embeddings), \( \mathcal{L}(e_i) \) (objective), \( \beta_1, \beta_2 \) (betas), T (number of time steps)</th>
  </tr>
  <tr>
    <th colspan="2">Initialize:</th>
    <th>\( m_i^{(0)} \leftarrow 0 \) (1st moment), \( v_i^{(0)} \leftarrow 0 \) (2nd moment)</th>
  </tr>
  <tr>
    <th colspan="2">Output:</th>
    <th>\( e^{(T)} \) (final embeddings)</th>
  </tr>
  <tr>
    <td colspan="3">1: for \( \tau = 1 \ldots T \) do</td>
  </tr>
  <tr>
    <td colspan="3">2: for \( i = 1 \ldots V \) do</td>
  </tr>
  <tr>
    <td colspan="3">3: \( g_i^{(\tau)} \leftarrow \nabla_{e_i} \mathcal{L}^{(\tau)}(e_i^{(\tau-1)}) \)</td>
  </tr>
  <tr>
    <td colspan="3">4: \( \hat{m}_i^{(\tau)} \leftarrow \beta_1 \hat{m}_i^{(\tau-1)} + (1-\beta_1) g_i^{(\tau)} \)</td>
  </tr>
  <tr>
    <td colspan="3">5: \( \hat{v}_i^{(\tau)} \leftarrow \beta_2 \hat{v}_i^{(\tau-1)} + (1-\beta_2) \left( g_i^{(\tau)} \right)^2 \)</td>
  </tr>
  <tr>
    <td colspan="3">6: \( \hat{m}_i^{(\tau)} \leftarrow \hat{m}_i^{(\tau)}/(1-\beta_1^T) \)</td>
  </tr>
  <tr>
    <td colspan="3">7: \( \hat{v}_i^{(\tau)} \leftarrow \hat{v}_i^{(\tau)}/(1-\beta_2^T) \)</td>
  </tr>
  <tr>
    <td colspan="3">8: if coupled then</td>
  </tr>
  <tr>
    <td colspan="3">9: \( \hat{\nu}^{(\tau)} \leftarrow \frac{1}{V} \sum_{i=1}^V \hat{v}_i^{(\tau)} \)</td>
  </tr>
  <tr>
    <td colspan="3">10: for \( i = 1 \ldots V \) do</td>
  </tr>
  <tr>
    <td colspan="3">11: if coupled then</td>
  </tr>
  <tr>
    <td colspan="3">12: \( \hat{v}_i^{(\tau)} \leftarrow \hat{\nu}^{(\tau)} \)</td>
  </tr>
  <tr>
    <td colspan="3">13: \( e_i^{(\tau)} \leftarrow e_i^{(\tau-1)} - \eta \frac{\hat{m}_i^{(\tau)}}{\sqrt{\hat{v}_i^{(\tau)}} + \epsilon} \)</td>
  </tr>
  <tr>
    <td colspan="3">14: return \( e^{(T)} \)</td>
  </tr>
</table>

Algorithm 1: Pseudocode for the Adam algorithm and our extension, the Coupled Adam algorithm (highlighted), applied to the embedding vectors \( e_i \). Note that weight decay is not applied.
This, in turn, causes the mean embedding to change during training,

\[
\mu^{(\tau)} - \mu^{(\tau-1)} \overset{(9.6)}{=} \frac{1}{V} \sum_{i=1}^V u_i^{(\tau)} \overset{(16)}{\neq} 0 ,
\] (17)

which is in stark contrast to the case of SGD (cf. Eq. (12)). We have thus identified that an \( i \)-dependency of the second moment \( \hat{v}_i^{(\tau)} \) of the Adam optimizer leads to the observed collective shift of the embedding vectors away from the origin. Next, we will show that the second moment indeed depends on \( i \). More concretely, we will argue that its expectation value is proportional to the unigram probability\footnote{Note that from here until Eq. (23), the time index (\( \tau \)) is dropped for the sake of readability.} (see Eq. (1)),

\[
\mathbb{E}[\hat{v}_i] \propto \hat{p}_i
\] (18)

In App. D.1, Eq. (18) is derived using minimal assumptions and experimental input. Here, we restrict ourselves to confirming the relationship in a purely experimental manner. \( \mathbb{E}[\hat{v}_i] \) is estimated directly by measuring \( \hat{v}_i \) multiple times during training, using different models. We then perform linear fits of \( \mathbb{E}[\hat{v}_i] \) as a function of \( \hat{p}_i \). Indeed, the fits yield a high coefficient of determination, on average \( R^2 = 0.85(7) \), and a proportionality constant of

\[
A := \frac{\mathbb{E}[\hat{v}_i]}{\hat{p}_i} \approx 10^{-4}
\] (19)

Details about the exact procedure and plots showing the data and linear fits can be found in App. D.2.

3 Coupled Adam

In the previous section, we have identified the individual scales of the second moments \( v_i \) for different embedding vectors \( e_i \) as the root cause of the anisotropy problem. This implies that a solution to the problem is to enforce that the second moments are the same for every \( i \). The question arises whether and how this can be done in the best way, without harming the performance of the model. To answer this, we note that the normalization of the embedding update vector by the Adam second moment can be split into two parts:

\[
\mathbb{E}[\hat{v}_i] \overset{(19)}{=} A \cdot \hat{p}_i = \frac{A}{V} \cdot (\hat{p}_i V)
\] (20)

The first factor introduces a global scale to all update vectors simultaneously:

\[
\frac{A}{V} \overset{(19)}{\approx} \frac{10^{-4}}{5 \cdot 10^4} = 2 \cdot 10^{-9} ,
\] (21)

where the numbers correspond to our experiments from the previous section with \( V \approx 50000 \). The second factor scales the update vectors individually. It is one on average:

\[
\frac{1}{V} \sum_{i=1}^V (\hat{p}_i V) = 1
\] (22)

Our goal is to retain the first, global factor and get rid of the second, individual factor. The canonical way to do this is to simply take the average of the second moment over the vocabulary items \( i \):

\[
\frac{1}{V} \sum_{i=1}^V \mathbb{E}[\hat{v}_i] \overset{(20,22)}{=} \frac{A}{V}
\] (23)

In practice, the exponentially averaged second moments \( \hat{v}_i^{(\tau)} \) as they appear in Eq. (14) are replaced by their average:

\[
\hat{v}^{(\tau)} := \frac{1}{V} \sum_{i=1}^V \hat{v}_i^{(\tau)}
\] (24)

We call the resulting algorithm *Coupled Adam*, as it couples the second moments of the embedding vectors via Eq. (24). It is displayed in Algorithm 1. Evidently, with Coupled Adam, the effective learning rate in Eq. (14) that enters the update vector in Eq. (13) becomes independent of \( i \). Hence, like SGD but unlike standard Adam, the sum of embedding updates vanishes. However, like standard Adam but unlike SGD, Coupled Adam uses a second moment to normalize the embedding update vectors.

4 Experiments

Two types of experiments are conducted to study the impact of coupling the second moments of the embedding update vectors. First, a set of small-scale experiments (Sec. 4.1) with models and datasets of varying sizes up to 1B parameters and 20B tokens, respectively. Afterwards, we perform a few large-scale experiments (Sec. 4.2) to verify that the usefulness of our method extrapolates to the realm of large language models with more than 1B parameters trained on at least the corresponding compute-optimal (Hoffmann et al., 2022)
amount of data. In order to verify the generalizability of our method, the small- and large-scale experiments involve different datasets, training frameworks and dense transformer model architectures. An overview of the model and dataset sizes employed in our experiments is given in App. E.1. For each combination, two models are trained: one using standard Adam and one using Coupled Adam for the embeddings, see Eq. (24). Both variants use standard Adam for all non-embedding parameters. The various metrics we employ to assess both the general model performance and the quality of the model embeddings will be discussed in Sec. 4.3.

4.1 Small-scale Experiments

Our small-scale experiments use the OpenWebText Corpus (Gokaslan and Cohen, 2019) and the GPT-2 tokenizer (Radford et al., 2019). The model architecture also follows GPT-2, while the hyperparameter setup is taken from GPT-3 (Brown et al., 2020), see App. E.2 for further details. An implementation based on nanoGPT (Karpachy, 2022) is used. We define a grid \((D, N)\) with dataset sizes \(D \in \{5B, 10B, 20B\}\) and model sizes \(N \in \{125M, 355M, 760M\}\), and repeat each experiment \(S = 3\) times with different seeds in order to estimate uncertainties and assess statistical significance.

4.2 Large-scale Experiments

For our large-scale experiments, we use the SlimPajama dataset (Soboleva et al., 2023) and the GPT-2 tokenizer. A state-of-the-art dense transformer model architecture akin to (Touvron et al., 2023) is chosen, including e.g. RoPE embeddings (Su et al., 2023) and the SwiGLU activation function (Shazeer, 2020). Details can be found in App. E.2. The experiments are conducted using Modalities (Lübbering et al., 2024) as the training framework. We consider two model sizes, 1.3B and 2.6B. In order to cover the two common scenarios of compute-optimal training and overtraining, we conduct two sets of experiments: Firstly, we use near compute-optimal dataset sizes, 26B and 52B tokens, respectively. Secondly, we increase the number of tokens by a factor 4, resulting in 105B and 210B tokens, respectively. Each large-scale experiment is performed \(S = 1\) times.

4.3 Evaluation

Upstream performance is measured in terms of test loss, while downstream performance is evaluated using the Language Model Evaluation Harness (Gao et al., 2023) on the following tasks: ARC easy and challenge (Clark et al., 2018), HellaSwag (Zellers et al., 2019), LAMBADA (Paperno et al., 2016), RACE (Lai et al., 2017), TruthfulQA (Lin et al., 2022) and WinoGrande (Sakaguchi et al., 2020). More concretely, the considered metric is the average\(^3\) accuracy, which we will denote by Acc. To assess the quality of the embeddings, we first compute their isotropy, defined as (Arora et al., 2016; Mu et al., 2018)

\[
\text{Iso}(E) := \frac{\min_{c \in X} Z(c)}{\max_{c \in X} Z(c)},
\]

where \(E \in \mathbb{R}^{H \times V}\) is the embedding matrix, \(Z(c) = \sum_{i=1}^V \exp(c^T e_i)\) is the partition function and \(X = \{c\}\) is the set of eigenvectors \(c \in \mathbb{R}^H\) of \(EE^T \in \mathbb{R}^{H \times H}\). Secondly, the 2-norm \(\|\mu\|\) of the mean embedding, see Eq. (9), and the average 2-norm of the embeddings \(\|e_i\| = \frac{1}{V} \sum_{i=1}^V \|e_i\|\) as well as their ratio

\[
\|\mu\|^r := \|\mu\| / \|e_i\|
\]

are determined. In addition, we evaluate the models on embedding benchmarks for word similarity and relatedness, to assess how well they represent semantic meaning. Following Biš et al. (2021), we consider the benchmarks SimLex999 (Hill et al., 2015), MEN (Bruni et al., 2014), WordSim353 (Finkelstein et al., 2001) and Stanford Rare Words (Luong et al., 2013). Each dataset provides pairs of words labeled with a ground truth score that represents the words’ semantic similarity. We derive model scores from the cosine similarity of the corresponding embedding vectors, and report the Pearson correlation of the two scores averaged over the datasets, which we denote by \(\bar{\tau}\). Finally, some additional important properties of the embedding matrix are investigated. We study the correlation between the length of an embedding vector and the unigram probability,

\[
\rho := 100 \cdot \operatorname{corr}((\|e_i\|_{i=1}^V, \hat{p}) ,
\]

to measure how well the former represents the latter. Furthermore, the condition number \(\kappa\), defined as the ratio of the smallest and largest singular values of the embedding matrix, is determined in percent:

\[
\kappa := 100 \cdot \frac{\min_i \Sigma_{ii}}{\max_i \Sigma_{ii}}
\]

Here, \(E = U \Sigma V^T\) denotes the singular value decomposition of the embedding matrix.

\footnotetext{3Individual task performance is reported in App. G.3.}
<table>
  <tr>
    <th rowspan="2">D</th>
    <th rowspan="2">N</th>
    <th rowspan="2">Adam</th>
    <th colspan="1">L (↓)</th>
    <th colspan="1">Acc (↑)</th>
    <th colspan="1">Iso (↑)</th>
    <th colspan="1">\( \|\mu\|_1 \) (↓)</th>
    <th colspan="1">\( \|\mu\|^T \) (↓)</th>
    <th colspan="1">\( \overline{\tau} \) (↑)</th>
    <th colspan="1">\( \rho \) (↑)</th>
    <th colspan="1">\( \kappa \) (↑)</th>
  </tr>
  <tr></tr>
  <tr>
    <td rowspan="2">125M</td>
    <td rowspan="2">Standard</td>
    <td>3.14 (0)</td>
    <td>0.340 (2)</td>
    <td>0.31 (2)</td>
    <td>1.10 (6)</td>
    <td>0.67 (3)</td>
    <td>15 (3)</td>
    <td>-54 (3)</td>
    <td>0.6 (1)</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>3.12 (1)</td>
    <td>0.339 (2)</td>
    <td>0.94 (1)</td>
    <td>0.02 (0)</td>
    <td>0.01 (0)</td>
    <td>55 (0)</td>
    <td>87 (1)</td>
    <td>4.8 (2)</td>
  </tr>
  <tr>
    <td rowspan="2">355M</td>
    <td rowspan="2">Standard</td>
    <td>2.95 (0)</td>
    <td>0.352 (3)</td>
    <td>0.44 (2)</td>
    <td>0.81 (2)</td>
    <td>0.67 (1)</td>
    <td>16 (2)</td>
    <td>-47 (0)</td>
    <td>0.8 (0)</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>2.93 (0)</td>
    <td>0.350 (4)</td>
    <td>0.98 (0)</td>
    <td>0.01 (0)</td>
    <td>0.01 (0)</td>
    <td>56 (1)</td>
    <td>86 (0)</td>
    <td>6.8 (4)</td>
  </tr>
  <tr>
    <td rowspan="2">760M</td>
    <td rowspan="2">Standard</td>
    <td>2.85 (0)</td>
    <td>0.360 (3)</td>
    <td>0.43 (1)</td>
    <td>0.84 (1)</td>
    <td>0.63 (0)</td>
    <td>14 (3)</td>
    <td>-49 (2)</td>
    <td>0.7 (0)</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>2.86 (0)</td>
    <td>0.357 (3)</td>
    <td>0.97 (0)</td>
    <td>0.01 (0)</td>
    <td>0.01 (0)</td>
    <td>55 (1)</td>
    <td>85 (1)</td>
    <td>6.9 (3)</td>
  </tr>
  <tr>
    <td rowspan="2">125M</td>
    <td rowspan="2">Standard</td>
    <td>3.07 (0)</td>
    <td>0.343 (3)</td>
    <td>0.21 (3)</td>
    <td>1.58 (5)</td>
    <td>0.75 (0)</td>
    <td>9 (2)</td>
    <td>-64 (5)</td>
    <td>0.4 (0)</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>3.03 (0)</td>
    <td>0.343 (1)</td>
    <td>0.91 (1)</td>
    <td>0.05 (0)</td>
    <td>0.02 (0)</td>
    <td>57 (2)</td>
    <td>82 (0)</td>
    <td>3.6 (8)</td>
  </tr>
  <tr>
    <td rowspan="2">355M</td>
    <td rowspan="2">Standard</td>
    <td>2.86 (0)</td>
    <td>0.359 (2)</td>
    <td>0.35 (2)</td>
    <td>1.01 (4)</td>
    <td>0.74 (2)</td>
    <td>10 (3)</td>
    <td>-55 (3)</td>
    <td>0.5 (0)</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>2.83 (0)</td>
    <td>0.365 (2)</td>
    <td>0.96 (0)</td>
    <td>0.02 (0)</td>
    <td>0.01 (0)</td>
    <td>57 (1)</td>
    <td>83 (1)</td>
    <td>5.3 (2)</td>
  </tr>
  <tr>
    <td rowspan="2">760M</td>
    <td rowspan="2">Standard</td>
    <td>2.75 (0)</td>
    <td>0.375 (2)</td>
    <td>0.38 (3)</td>
    <td>0.97 (4)</td>
    <td>0.66 (2)</td>
    <td>11 (2)</td>
    <td>-56 (1)</td>
    <td>0.5 (0)</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>2.74 (1)</td>
    <td>0.372 (3)</td>
    <td>0.96 (0)</td>
    <td>0.02 (0)</td>
    <td>0.01 (0)</td>
    <td>57 (0)</td>
    <td>84 (0)</td>
    <td>6.2 (0)</td>
  </tr>
  <tr>
    <td rowspan="2">125M</td>
    <td rowspan="2">Standard</td>
    <td>3.03 (0)</td>
    <td>0.346 (1)</td>
    <td>0.10 (3)</td>
    <td>2.14 (7)</td>
    <td>0.82 (2)</td>
    <td>5 (1)</td>
    <td>-66 (2)</td>
    <td>0.3 (0)</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>2.97 (0)</td>
    <td>0.350 (1)</td>
    <td>0.83 (1)</td>
    <td>0.11 (0)</td>
    <td>0.03 (0)</td>
    <td>57 (2)</td>
    <td>77 (0)</td>
    <td>1.7 (5)</td>
  </tr>
  <tr>
    <td rowspan="2">355M</td>
    <td rowspan="2">Standard</td>
    <td>2.79 (0)</td>
    <td>0.366 (4)</td>
    <td>0.25 (2)</td>
    <td>1.32 (8)</td>
    <td>0.82 (2)</td>
    <td>5 (2)</td>
    <td>-65 (4)</td>
    <td>0.3 (0)</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>2.75 (0)</td>
    <td>0.372 (6)</td>
    <td>0.95 (2)</td>
    <td>0.04 (0)</td>
    <td>0.02 (0)</td>
    <td>57 (1)</td>
    <td>78 (0)</td>
    <td>4.1 (3)</td>
  </tr>
  <tr>
    <td rowspan="2">760M</td>
    <td rowspan="2">Standard</td>
    <td>2.68 (1)</td>
    <td>0.385 (3)</td>
    <td>0.28 (2)</td>
    <td>1.21 (8)</td>
    <td>0.73 (3)</td>
    <td>3 (4)</td>
    <td>-64 (2)</td>
    <td>0.3 (0)</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>2.65 (0)</td>
    <td>0.392 (2)</td>
    <td>0.94 (2)</td>
    <td>0.03 (0)</td>
    <td>0.01 (0)</td>
    <td>58 (0)</td>
    <td>81 (0)</td>
    <td>4.4 (2)</td>
  </tr>
</table>

Table 1: Results of our small-scale experiments. D and N denote the dataset and model size, respectively. L is the test loss, and the column Acc represents the accuracy averaged over the downstream tasks listed in Sec. 4.3. The other evaluation metrics are defined in the same section, see Eqs. (25)-(28). The arrow in parentheses indicates whether a higher or lower value is desirable. Every training was conducted S = 3 times with different seeds, and the numbers represent the (rounded) averages and standard deviations in the following shorthand notation format: 0.123 (4) ≡ 0.123 ± 0.004. For each combination (D, N) and each metric, the respective better value is highlighted in bold if the (unrounded) difference is significant according to Student’s t-test with a one-sided confidence level of \( \alpha = 95\% \) (see App. F for details). Plots for L and Acc are shown in Fig. 2.

![Line plots showing differences in loss and average downstream task accuracy between Coupled Adam and standard Adam, as a function of dataset size D.](page_1012_1042_420_312.png)

Figure 2: Difference in loss (left) and average downstream task accuracy (right) between Coupled Adam and standard Adam, for the different dataset sizes D (horizontal axis) and model sizes N (colors) of the small-scale experiments. The vertical bars indicate the one-sided 95% confidence interval for the difference to be significant. In order to avoid overlaps, the data points for \( N = 125M \) and \( N = 760M \) have been slightly shifted to the left and right, respectively.

5 Results

5.1 Small-scale Experiments

The results of the small-scale experiments (Sec. 4.1) are listed in Tab. 1 and illustrated in Fig. 2. We find that both upstream and downstream performance are better with Coupled Adam if the dataset size is sufficiently large. In fact, the improvement appears to increase monotonically with the dataset size D. In addition, the embedding-specific metrics benefit greatly from Coupled Adam. In particular, the isotropy reaches values above 0.90 (with a single exception), while \( \overline{\tau} \) and \( \kappa \) are hugely improved as well. The mean embedding is evidently close to the origin. Finally, Coupled Adam leads to a significantly stronger (positive) correlation \( \rho \) between the length of an embedding vector and its associated unigram probability.

5.2 Large-scale Experiments

The results of the large-scale experiments (Sec. 4.2) are shown in Tab. 2. We observe very similar patterns as for the small-scale experiments. Although upstream and downstream performance are worse with Coupled Adam for compute-optimal dataset sizes, they are better if 4 times larger datasets are used. Note that for the small-scale experiments, the
<table>
  <tr>
    <th>D</th>
    <th>N</th>
    <th>Adam</th>
    <th>L (J)</th>
    <th>Acc (%)</th>
    <th>Iso (%)</th>
    <th>\( \|\mu\|_1 \) (J)</th>
    <th>\( \|\mu\|^2 \) (J)</th>
    <th>\( \tau \) (%)</th>
    <th>\( \rho \) (%)</th>
    <th>\( \kappa \) (%)</th>
  </tr>
  <tr>
    <td rowspan="2">26B</td>
    <td rowspan="2">1.3B</td>
    <td>Standard</td>
    <td>2.433</td>
    <td>0.402</td>
    <td>0.50</td>
    <td>0.67</td>
    <td>0.45</td>
    <td>53</td>
    <td>-41</td>
    <td>1.2</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>2.448</td>
    <td>0.396</td>
    <td>0.96</td>
    <td>0.05</td>
    <td>0.03</td>
    <td>66</td>
    <td>74</td>
    <td>4.5</td>
  </tr>
  <tr>
    <td rowspan="2">52B</td>
    <td rowspan="2">2.6B</td>
    <td>Standard</td>
    <td>2.257</td>
    <td>0.451</td>
    <td>0.51</td>
    <td>0.68</td>
    <td>0.40</td>
    <td>55</td>
    <td>-44</td>
    <td>1.0</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>2.273</td>
    <td>0.441</td>
    <td>0.86</td>
    <td>0.10</td>
    <td>0.06</td>
    <td>66</td>
    <td>74</td>
    <td>3.0</td>
  </tr>
  <tr>
    <td rowspan="2">105B</td>
    <td rowspan="2">1.3B</td>
    <td>Standard</td>
    <td>2.278</td>
    <td>0.446</td>
    <td>0.40</td>
    <td>0.87</td>
    <td>0.43</td>
    <td>52</td>
    <td>-52</td>
    <td>0.6</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>2.277</td>
    <td>0.447</td>
    <td>0.82</td>
    <td>0.23</td>
    <td>0.09</td>
    <td>67</td>
    <td>71</td>
    <td>2.2</td>
  </tr>
  <tr>
    <td rowspan="2">210B</td>
    <td rowspan="2">2.6B</td>
    <td>Standard</td>
    <td>2.131</td>
    <td>0.490</td>
    <td>0.34</td>
    <td>0.96</td>
    <td>0.41</td>
    <td>54</td>
    <td>-52</td>
    <td>0.5</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>2.129</td>
    <td>0.492</td>
    <td>0.65</td>
    <td>0.49</td>
    <td>0.14</td>
    <td>67</td>
    <td>69</td>
    <td>1.5</td>
  </tr>
</table>

Table 2: Results of our large-scale experiments. See the caption of Tab. 1 for an explanation of the column names. For each combination (\( D, N \)) and each metric, the respective better value is highlighted in bold.

upstream and downstream performance were found to be better already for compute-optimal dataset sizes. We attribute this to the fact that the batch size for the large-scale experiments is five times larger (cf. App. E.2), which results in fewer optimization steps for the same dataset size. Regarding the embedding-specific metrics, we again find significant and consistent improvements throughout all experiments. However, we do observe a certain shift of the mean embedding vector away from the origin, even if Coupled Adam is used. The shift becomes more pronounced as the model and dataset sizes increase, and is also reflected in a reduced isotropy. As we shall see in the following section, it comes along with optimal model performance though. An obvious hypothesis in light of our analysis in Sec. 2 is that the residual shift of the mean embeddings is due to weight tying. This is supported by the results of Machina and Mercer (2024), who find improved isotropy for models without weight tying. We leave it for future work to verify the hypothesis.

6 Ablations

We perform some additional experiments to shed further light on how Coupled Adam works. A model size of \( N = 125 \mathrm{M} \) and the dataset sizes \( D \in \{5 \mathrm{B}, 10 \mathrm{B}, 20 \mathrm{B}\} \) from the small-scale experiments (Sec. 4.1) are used, and each experiment is repeated \( S = 3 \) times with different seeds.

6.1 Scaled Coupled Adam

While coupling the second moment of the embedding gradients using the average in Eq. (24) is the canonical choice, one could also use a multiple of the average. We conduct additional experiments where the coupled second moment is scaled by powers of 2:

\[
\hat{\nu}^{(\tau)} \rightarrow 2^{-n} \cdot \hat{\nu}^{(\tau)} ,
\]

with scaling exponents \( n \in \{z \in \mathbb{Z} \mid -5 \leq z \leq 5\} \). Note that using a scaling exponent \( n \neq 0 \) is equivalent to using a different effective learning rate for the embeddings than for all the other parameters, via Eqs. (24) and (14). In particular, a smaller scaling exponent \( n \) corresponds to a smaller effective learning rate and vice versa. The results for \( D = 20 \mathrm{B} \) are shown in Tab. 3, and the dependency of the loss on the scaling exponent \( n \) for that very dataset size is visualized in Fig. 3. Results for other dataset sizes and plots for the other evaluation metrics can be found in App. G.1. Our data shows that the loss reaches a minimum close to \( n = 0 \), with a rather weak dependence on the scaling exponent in its vicinity. Nevertheless, for the smallest and largest scaling exponents studied, we find that the loss gets significantly worse. Regarding downstream performance, we see indications of a similar pattern, although the statistical uncertainties are too large to draw definite conclusions. The semantic usefulness of the embedding vectors as measured by \( \tau \) seems to suffer from a scaling exponent \( n < 0 \). For the isotropy and the mean embedding, we observe the opposite behavior. They benefit from a smaller scaling exponent \( n \) and the associated smaller embedding updates, with the effect being more pronounced the larger the training dataset size \( D \). However, this also negatively affects the model performance. Hence, we conclude that, at least within the range of our experiments, the optimal setting is to have the same learning rate for the embedding parameters as for all the other model parameters, as implied by \( n = 0 \) and Eq. (24).

6.2 SGD

We train several models using SGD with momentum \( \gamma = 0.9 \) as the optimizer for the embeddings. Since Adam via the inverse square root of its second moment effectively scales the learning rate up by a factor comprising orders
<table>
  <tr>
    <th>n</th>
    <th>Loss (J)</th>
    <th>Acc (%)</th>
    <th>Iso (%)</th>
    <th>\( \mu \| \) (J)</th>
    <th>\( \mu \|^2 \) (J)</th>
    <th>\( \tau \) (%)</th>
    <th>\( \rho \) (%)</th>
    <th>\( \kappa \) (%)</th>
  </tr>
  <tr>
    <td>-5</td>
    <td>2.99 (0)</td>
    <td>0.349 (2)</td>
    <td>0.97 (0)</td>
    <td><b>0.01 (0)</b></td>
    <td>0.01 (0)</td>
    <td>55 (1)</td>
    <td>76 (3)</td>
    <td><b>6.0 (6)</b></td>
  </tr>
  <tr>
    <td>-4</td>
    <td>2.99 (0)</td>
    <td>0.348 (5)</td>
    <td>0.97 (1)</td>
    <td>0.02 (0)</td>
    <td>0.01 (0)</td>
    <td>55 (1)</td>
    <td>78 (3)</td>
    <td>4.4 (4)</td>
  </tr>
  <tr>
    <td>-3</td>
    <td>2.98 (0)</td>
    <td>0.352 (5)</td>
    <td>0.95 (1)</td>
    <td>0.03 (0)</td>
    <td>0.02 (0)</td>
    <td>57 (2)</td>
    <td>78 (2)</td>
    <td>2.7 (5)</td>
  </tr>
  <tr>
    <td>-2</td>
    <td>2.98 (0)</td>
    <td>0.352 (1)</td>
    <td>0.94 (2)</td>
    <td>0.04 (0)</td>
    <td>0.02 (0)</td>
    <td>57 (1)</td>
    <td>78 (2)</td>
    <td>2.3 (4)</td>
  </tr>
  <tr>
    <td>-1</td>
    <td>2.98 (0)</td>
    <td>0.348 (3)</td>
    <td>0.87 (2)</td>
    <td>0.07 (0)</td>
    <td>0.02 (0)</td>
    <td>57 (1)</td>
    <td>79 (2)</td>
    <td>2.1 (3)</td>
  </tr>
  <tr>
    <td>0</td>
    <td>2.97 (0)</td>
    <td>0.350 (1)</td>
    <td>0.83 (1)</td>
    <td>0.11 (0)</td>
    <td>0.03 (0)</td>
    <td>57 (2)</td>
    <td>77 (0)</td>
    <td>1.7 (5)</td>
  </tr>
  <tr>
    <td>1</td>
    <td>2.97 (0)</td>
    <td>0.351 (4)</td>
    <td>0.66 (6)</td>
    <td>0.20 (1)</td>
    <td>0.03 (0)</td>
    <td>56 (0)</td>
    <td>78 (2)</td>
    <td>1.6 (5)</td>
  </tr>
  <tr>
    <td>2</td>
    <td>2.98 (0)</td>
    <td>0.353 (2)</td>
    <td>0.47 (6)</td>
    <td>0.34 (1)</td>
    <td>0.04 (0)</td>
    <td>58 (1)</td>
    <td>78 (2)</td>
    <td>1.6 (9)</td>
  </tr>
  <tr>
    <td>3</td>
    <td>2.97 (0)</td>
    <td>0.352 (0)</td>
    <td>0.27 (4)</td>
    <td>0.54 (2)</td>
    <td>0.05 (0)</td>
    <td>58 (2)</td>
    <td>78 (1)</td>
    <td>2.0 (7)</td>
  </tr>
  <tr>
    <td>4</td>
    <td>2.97 (1)</td>
    <td>0.352 (1)</td>
    <td>0.09 (0)</td>
    <td>0.83 (2)</td>
    <td>0.05 (0)</td>
    <td>58 (1)</td>
    <td>77 (1)</td>
    <td>2.3 (2)</td>
  </tr>
  <tr>
    <td>5</td>
    <td>2.98 (0)</td>
    <td>0.349 (4)</td>
    <td>0.01 (1)</td>
    <td>1.32 (2)</td>
    <td>0.06 (0)</td>
    <td>57 (1)</td>
    <td>75 (1)</td>
    <td>2.1 (6)</td>
  </tr>
</table>

Table 3: Results of our experiments with Scaled Coupled Adam, for \( N = 125M \) and \( D = 20B \). Values are highlighted in bold if they are significantly better than all the other values in the same column, see the caption of Tab. 1 for more details.

Figure 3: Dependency of the loss on the scaling exponent \( n \), see Eq. (29), for \( N = 125M \) and \( D = 20B \). The plot shows the difference to the loss obtained for \( n = 0 \).

of magnitude (see Eq. (21)), we explicitly multiply the learning rate in SGD by a factor \( f \) of comparable size\footnote{Note that the difference between momentum in SGD and the first moment in Adam also plays a role here.}. A hyperparameter search using \( f \in \{100, 200, 300, 400, 500, 600\} \) is performed to search for the optimum with respect to upstream performance (loss), see App. G.2 for details. It is found at \( f = 300 \) for \( D \in \{5B, 10B\} \) and \( f = 400 \) for \( D = 20B \). The respective optimal model is compared to its counterpart trained with Coupled Adam in Tab. 4. The results show that, although SGD is advantageous with respect to isotropy, the mean embedding shift and the condition number, Coupled Adam consistently achieves better results on all upstream and downstream task metrics, while having one less hyperparameter to fine-tune.

7 Related Work

Gao et al. (2019) first described the anisotropy issue, which they referred to as *representation degeneration problem*, and suggested cosine regularization as a mitigation strategy. Alternative techniques to address the problem have been developed, including adversarial noise (Wang et al., 2019), spectrum control (Wang et al., 2020) and Laplacian regularization (Zhang et al., 2020). Bi et al. (2021) have shown that the anisotropy of embeddings can for the most part be traced back to a common shift of the embeddings in a dominant direction. They called this phenomenon *common enemy effect*, and provided a semi-quantitative explanation (Eq. (5)), which we developed further in the present work by including the optimizer in the analysis. In Yu et al. (2022), Adaptive Gradient Gating is proposed, based on the empirical observation that it is the gradients for embeddings of rare tokens that cause anisotropy. Our analysis conforms to this finding and attributes it to a massive up-scaling of the gradients for rare embeddings with Adam, cf. Fig. 1. Machina and Mercer (2024) have demonstrated that large Pytha models (Biderman et al., 2023) show improved isotropy compared to similar models, and attribute this to the absence of weight tying. This is in accordance with our analysis of the unembedding gradients in conjunction with Adam, Sec. 2. While all the previously mentioned papers use average cosine similarity (Ethayarajh, 2019) or Iso from Eq. (25) to quantify the geometry of embedding vectors, Rudman et al. (2022) deviate from this. Their notion of isotropy is based solely on the embeddings’ covariance matrix and embodied by the metric IsoScore. In particular, IsoScore is mean-agnostic, while Iso strongly correlates with the mean embedding (see e.g. Tab. 1). In a follow-up paper (Rudman and Eickhoff, 2024), IsoScore is used to regularize isotropy, which appears to benefit performance for fine-tuning tasks. Finally, concurrent to our work, Zhao et al. (2024) have investigated the importance of using the second moment in Adam with regard to performance and stability. They found that simplified variants of Adam that use the same effective learning rate either for the whole embedding matrix (Adalayer) or each embedding vector (Adalayer*) are slightly worse than Adam but better than SGD. Adalayer* is similar to Coupled Adam, but corresponds to the second moment averaged over hidden space instead of vocabulary space.

8 Conclusions

Our work addresses the well-known anisotropy problem for LLM embeddings. We have advanced the theoretical understanding of the phenomenon by showing that it is a combination of the common enemy effect and the individual second mo-
<table>
  <tr>
    <th>D</th>
    <th>N</th>
    <th>Optimizer</th>
    <th>L (J)</th>
    <th>Acc (%)</th>
    <th>Iso (%)</th>
    <th>\( \| \mu \| \) (J)</th>
    <th>\( \| \mu \|^{2} \) (J)</th>
    <th>F (%)</th>
    <th>\( \rho \) (%)</th>
    <th>\( \kappa \) (%)</th>
  </tr>
  <tr>
    <td rowspan="2">5B</td>
    <td rowspan="2">125M</td>
    <td>SGD (300)</td>
    <td>3.17 (0)</td>
    <td>0.333 (3)</td>
    <td><b>0.99 (0)</b></td>
    <td>0.00 (0)</td>
    <td><b>0.01 (0)</b></td>
    <td>45 (1)</td>
    <td>71 (1)</td>
    <td>15.5 (4)</td>
  </tr>
  <tr>
    <td>Coupled Adam</td>
    <td>3.12 (1)</td>
    <td><b>0.339 (2)</b></td>
    <td>0.94 (1)</td>
    <td>0.02 (0)</td>
    <td>0.01 (0)</td>
    <td>55 (0)</td>
    <td>87 (1)</td>
    <td>4.8 (2)</td>
  </tr>
  <tr>
    <td rowspan="2">10B</td>
    <td rowspan="2">125M</td>
    <td>SGD (300)</td>
    <td>3.07 (1)</td>
    <td>0.341 (4)</td>
    <td><b>0.99 (0)</b></td>
    <td>0.01 (0)</td>
    <td><b>0.01 (0)</b></td>
    <td>49 (1)</td>
    <td>71 (4)</td>
    <td>12.2 (2)</td>
  </tr>
  <tr>
    <td>Coupled Adam</td>
    <td>3.03 (0)</td>
    <td>0.343 (1)</td>
    <td>0.91 (1)</td>
    <td>0.05 (0)</td>
    <td>0.02 (0)</td>
    <td>57 (2)</td>
    <td>82 (0)</td>
    <td>3.6 (8)</td>
  </tr>
  <tr>
    <td rowspan="2">20B</td>
    <td rowspan="2">125M</td>
    <td>SGD (400)</td>
    <td>3.00 (0)</td>
    <td>0.346 (5)</td>
    <td><b>0.98 (1)</b></td>
    <td>0.01 (0)</td>
    <td><b>0.02 (0)</b></td>
    <td>54 (0)</td>
    <td>76 (5)</td>
    <td>7.4 (1.1)</td>
  </tr>
  <tr>
    <td>Coupled Adam</td>
    <td>2.97 (0)</td>
    <td>0.350 (1)</td>
    <td>0.83 (1)</td>
    <td>0.11 (0)</td>
    <td>0.03 (0)</td>
    <td>57 (2)</td>
    <td>77 (0)</td>
    <td>1.7 (5)</td>
  </tr>
</table>

Table 4: Comparison of models whose embeddings were trained with SGD and Coupled Adam. The SGD models were obtained after hyperparameter search for the learning rate. The associated factor \( f \) is specified in parentheses in the Optimizer column. Bold values indicate better results with statistical significance, see the caption of Tab. 1 for more details.

ments in Adam that causes a collective shift of the embedding vectors away from the origin. To mitigate the problem, we have introduced Coupled Adam, which enforces the same effective learning rate for every embedding vector, and thus suppresses the collective shift of the embeddings. We have found that Coupled Adam consistently improves embedding-specific metrics across all experiments, while also achieving better downstream and upstream performance for large datasets, as they are typically used in LLM training. The code to reproduce our results is available at github.com/flxst/coupled-adam .

9 Limitations

Although our method is generally applicable to all common LLM architectures, as they share the same language modeling head and embeddings, only dense decoders were used in our experiments. In addition, only models with up to \( N = 2.6B \) parameters have been tested. Our experiments involve pre-training and few-shot downstream evaluation, yet fine-tuning tasks have not been included. The cosine decay learning rate schedule was applied throughout all experiments (App. E.2). Alternatives such as an infinite learning rate schedule are not incorporated in our study. It would also be interesting to extend our work to optimizers other than SGD and Adam. Furthermore, as mentioned at the end of Sec. 5, we have not explicitly verified that the slight residual shift of the mean embedding, which is observed even for Coupled Adam, is caused by weight tying. Finally, we have used a straightforward implementation of Coupled Adam, closely following Algorithm 1. More sophisticated implementations might lead to increased efficiency and further improvements; we leave it for future work to investigate this.

Acknowledgements

Our computational experiments used around 20000 GPU hours. They were partly run on the EuroHPC supercomputers MeluXina and MareNostrum5 in conjunction with the grants EHPC-DEV-2023D10-032 and EHPC-EXT-2023E02-038.

References

Sanjeev Arora, Yuanzhi Li, Yingyu Liang, Tengyu Ma, and Andrej Risteski. 2016. A latent variable model approach to PMI-based word embeddings. Transactions of the Association for Computational Linguistics, 4:385–399.

Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O’Brien, Eric Halahan, Mohammad Afsh Khan, Shivanshu Purohit, Usvsn Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar Van Der Wal. 2023. Pythia: A suite for analyzing large language models across training and scaling. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 2397–2430. PMLR.

Daniel Biš, Maksim Podkorytov, and Xiwen Liu. 2021. Too much in common: Shifting of embeddings in transformer language models and its implications. In North American Chapter of the Association for Computational Linguistics (NAACL).

Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. 2020. Language models are few-shot learners.

Elia Bruni, Nam Khanh Tran, and Marco Baroni. 2014. Multimodal distributional semantics. J. Artif. Int. Res., 49(1):1–47.

Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind
Tafjord. 2018. Think you have solved question answering? try arc, the ai2 reasoning challenge.

Kawin Ethayarajh. 2019. How contextual are contextualized word representations? Comparing the geometry of BERT, ELMo, and GPT-2 embeddings. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 55–65, Hong Kong, China. Association for Computational Linguistics.

Lev Finkelstein, Evgeniy Gabrilovich, Yossi Matias, Ehud Rivlin, Zach Solan, Gadi Wolfman, and Eytan Ruppin. 2001. Placing context in context: The concept revisited. volume 20, pages 406–414.

Jun Gao, Di He, Xu Tan, Tao Qin, Liwei Wang, and Tie-Yan Liu. 2019. Representation degeneration problem in training natural language generation models.

Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noach, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. 2023. A framework for few-shot language model evaluation.

Aaron Gokaslan and Vanya Cohen. 2019. Openwebtext corpus. http://Skylion007.github.io/OpenWebTextCorpus.

Felix Hill, Roi Reichart, and Anna Korhonen. 2015. SimLex-999: Evaluating semantic models with (genuine) similarity estimation. Computational Linguistics, 41(4):665–695.

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. 2022. Training compute-optimal large language models.

Andrej Karpathy. 2022. NanoGPT. https://github.com/karpathy/nanoGPT.

Diederik Kingma and Jimmy Ba. 2014. Adam: A method for stochastic optimization. International Conference on Learning Representations.

Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy. 2017. RACE: Large-scale ReAding comprehension dataset from examinations. In Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing, pages 785–794, Copenhagen, Denmark. Association for Computational Linguistics.

Stephanie Lin, Jacob Hilton, and Owain Evans. 2022. TruthfulQA: Measuring how models mimic human falsehoods. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 3214–3252, Dublin, Ireland. Association for Computational Linguistics.

Thang Luong, Richard Socher, and Christopher Manning. 2013. Better word representations with recursive neural networks for morphology. In Proceedings of the Seventeenth Conference on Computational Natural Language Learning, pages 104–113, Sofia, Bulgaria. Association for Computational Linguistics.

Max Lübbering, Mehdi Ali, Felix Stollenwerk, Michael Fromm, Alexander Arno Weber, and Richard Ruttmann. 2024. Modalities: A pytorch-native framework for distributed and reproducible foundation model training. https://github.com/Modalities/modalities.

Anemily Machina and Robert Mercer. 2024. Anisotropy is not inherent to transformers. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pages 4892–4907, Mexico City, Mexico. Association for Computational Linguistics.

Jiaqi Mu, Suma Bhat, and Pramod Viswanath. 2018. All-but-the-top: Simple and effective postprocessing for word representations.

Denis Paperno, Germán Kruszewski, Angeliki Lazari-dou, Ngoc Quan Pham, Raffaela Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. 2016. The LAMBADA dataset: Word prediction requiring a broad discourse context. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1525–1534, Berlin, Germany. Association for Computational Linguistics.

Ofir Press and Lior Wolf. 2017. Using the output embedding to improve language models. In Proceedings of the 15th Conference of the European Chapter of the Association for Computational Linguistics: Volume 2, Short Papers, pages 157–163, Valencia, Spain. Association for Computational Linguistics.

Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. 2019. Language models are unsupervised multitask learners.

Sebastian Ruder. 2017. An overview of gradient descent optimization algorithms.

William Rudman and Carsten Eickhoff. 2024. Stable anisotropic regularization.

William Rudman, Nate Gillman, Taylor Rayne, and Carsten Eickhoff. 2022. IsoScore: Measuring the uniformity of embedding space utilization. In Findings of the Association for Computational Linguistics: ACL 2022, pages 3325–3339, Dublin, Ireland. Association for Computational Linguistics.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. 2020. Winogrande: An adversarial winograd schema challenge at scale. Proceedings of the AAAI Conference on Artificial Intelligence, 34(05):8732–8740.

Noam Shazeer. 2020. Glu variants improve transformer.

Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. 2023. SlimPajama: A 627B token cleaned and deduplicated version of RedPajama.

Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. 2023. Roformer: Enhanced transformer with rotary position embedding.

Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajijwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esioibu, Jude Fernandes, Jeremy Fu, Wenying Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jena Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rash Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadar, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. 2023. Llama 2: Open foundation and fine-tuned chat models.

Dilin Wang, Chengyue Gong, and Qiang Liu. 2019. Improving neural language modeling via adversarial training. In Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 6555–6565. PMLR.

Lingxiao Wang, Jing Huang, Kevin Huang, Ziniu Hu, Guangtao Wang, and Quanquan Gu. 2020. Improving neural language generation with spectrum control. In International Conference on Learning Representations.

Sangwon Yu, Jongyoon Song, Heeseung Kim, Seongmin Lee, Woo-Jong Ryu, and Sungroh Yoon. 2022. Rare tokens degenerate all tokens: Improving neural text generation via adaptive gradient gating for rare token embeddings. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 29–45, Dublin, Ireland. Association for Computational Linguistics.

Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. 2019. HellaSwag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 4791–4800, Florence, Italy. Association for Computational Linguistics.

Zhong Zhang, Chongming Gao, Cong Xu, Rui Miao, Qinli Yang, and Junming Shao. 2020. Revisiting representation degeneration problem in language modeling. In Findings of the Association for Computational Linguistics: EMNLP 2020, pages 518–527, Online. Association for Computational Linguistics.

Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, and Sham Kakade. 2024. Deconstructing what makes a good optimizer for language models.

A Unigram Probability Distribution

Fig. 4 shows the unigram probability distribution for the example of the OpenWebText Corpus dataset and the GPT-2 tokenizer.

![Logarithm log(\hat{p}_i) of the unigram probability distribution for the OpenWebText Corpus and the GPT-2 tokenizer. The maximum probability is max_i \hat{p}_i \approx 0.037 or max_i \log(\hat{p}_i) \approx -3.30. The minimum probability (not shown) is min_i \hat{p}_i = 0 or min_i \log(\hat{p}_i) = -\infty.](page_1012_1342_388_388.png)

Figure 4: Logarithm \( \log(\hat{p}_i) \) of the unigram probability distribution for the OpenWebText Corpus and the GPT-2 tokenizer. The maximum probability is \( \max_i \hat{p}_i \approx 0.037 \) or \( \max_i \log(\hat{p}_i) \approx -3.30 \). The minimum probability (not shown) is \( \min_i \hat{p}_i = 0 \) or \( \min_i \log(\hat{p}_i) = -\infty \).

B Embedding Gradients

We explicitly derive Eq. (5), which we recall here for convenience:

\[
g_i := \frac{\partial \mathcal{L}}{\partial e_i} = - (\delta_{it} - p_i) \cdot h
\]

The chain rule yields

\[
\frac{\partial \mathcal{L}}{\partial e_i} = \sum_{k=1}^V \frac{\partial \mathcal{L}}{\partial p_t} \cdot \frac{\partial p_t}{\partial l_k} \cdot \frac{\partial l_k}{\partial e_i},
\]
where the individual factors can directly be obtained from Eqs. (2)-(4):

\[
\frac{\partial \mathcal{L}}{\partial p_t} = -\frac{1}{p_t} \tag{31}
\]
\[
\frac{\partial p_t}{\partial l_k} = \frac{\delta_{kt} \exp(l_t) \cdot \Sigma - \exp(l_t) \exp(l_k)}{\Sigma^2}
= \delta_{kt} p_t - p_t p_k
= p_t (\delta_{kt} - p_k) \tag{32}
\]
\[
\frac{\partial l_k}{\partial e_i} = \delta_{ki} h \tag{33}
\]

Note that in the first line of Eq. (32), we use the abbreviation \( \Sigma = \left( \sum_{j=1}^V \exp(l_j) \right) \). Inserting Eqs. (31), (32) and (33) into Eq. (30) directly leads to Eq. (5):

\[
\frac{\partial \mathcal{L}}{\partial e_i} = - \sum_{k=1}^V \frac{1}{p_t} \cdot p_t (\delta_{kt} - p_k) \cdot \delta_{ki} h
= - \sum_{k=1}^V (\delta_{kt} - p_k) \cdot \delta_{ki} h
= - (\delta_{it} - p_i) \cdot h
\]

C SGD Algorithm

For completeness and comparison to (Coupled) Adam as displayed in Algorithm 1, we summarize the SGD algorithm in Algorithm 2.

<table>
  <tr>
    <th colspan="2">Input:</th>
    <td>\( \eta \) (lr), \( e_i^{(0)} \) (initial embeddings), \( \mathcal{L}(e_i) \) (objective), \( \gamma \) (momentum), \( T \) (number of time steps)</td>
  </tr>
  <tr>
    <th colspan="2">Output:</th>
    <td>\( e^{(T)} \) (final embeddings)</td>
  </tr>
  <tr>
    <td colspan="3">
      1: for \( \tau = 1 \ldots T \) do<br>
      2:   for \( i = 1 \ldots V \) do<br>
      3:     \( g_i^{(\tau)} \gets \nabla_{e_i} \mathcal{L}^{(\tau)}(e_i^{(\tau-1)}) \)<br>
      4:     if \( t > 1 \) then<br>
      5:       \( b_i^{(\tau)} \gets \gamma b_i^{(\tau-1)} + g_i^{(\tau)} \)<br>
      6:     else<br>
      7:       \( b_i^{(\tau)} \gets g_i^{(\tau)} \)<br>
      8:       \( e_i^{(\tau)} \gets e_i^{(\tau-1)} - \eta b_i^{(\tau)} \)<br>
      9:   return \( e^{(T)} \)
    </td>
  </tr>
</table>

Algorithm 2: Pseudocode for the SGD algorithm with optional momentum, applied to the embedding vectors \( e_i \).

D Magnitude of the Second Moment in Adam

In this appendix, the validity of

\[
\mathbb{E} \left[ \tilde{g}_i^2 \right] \propto \tilde{p}_i
\]

is verified. Due to the linearity of lines 5 and 7 in Algorithm 2, it suffices to show that the squared gradient has the property in question:

\[
\mathbb{E} \left[ g_i^2 \right] \propto \tilde{p}_i \tag{34}
\]

We do this in two different ways. First, we derive Eq. (34) using a semi-theoretical approach with minimal experimental input. Afterwards, we confirm the relationship in a purely experimental manner.

D.1 Semi-theoretical Derivation

Here, we derive an expression for the expectation value of the squared gradient in terms of simple observables (Theorem 2). Subsequently, the dependency of those observables on \( \tilde{p}_i \) is determined experimentally. Together, this will yield the proportionality expressed by Eq. (34). We begin our reasoning with a lemma.

Lemma 1 (Expectation Value Decomposition).
The expectation value of the squared gradient can be decomposed into conditional expectation values as follows:

\[
\mathbb{E} \left[ g_i^2 \right] = \tilde{p}_i \cdot \mathbb{E} \left[ g_i^2 \mid i = t \right]
+ (1 - \tilde{p}_i) \cdot \mathbb{E} \left[ g_i^2 \mid i \neq t \right] \tag{35}
\]

Proof. Our starting point is the definition of the expectation value for the continuous random variable \( g_i^2 \):

\[
\mathbb{E} \left[ g_i^2 \right] = \int g_i^2 \, p(g_i) \, dg_i , \tag{36}
\]

where \( p \) denotes the probability distribution of \( g_i \). Since the vocabulary item \( i \) can only be either the true token \( t \) or not, we can decompose \( p \) into a sum of joint probability distributions (using the law of total probabilities), each of which can be expressed in terms of conditional probabilities like so:

\[
\begin{align*}
p(g_i) &= p(g_i, i = t) + p(g_i, i \neq t) \\
&= p(g_i \mid i = t) \cdot p(i = t) \\
&\quad + p(g_i \mid i \neq t) \cdot p(i \neq t) \tag{37}
\end{align*}
\]

Using the unigram probability \( \tilde{p}_i = p(i = t) \), this can also be written as

\[
p(g_i) = \tilde{p}_i \cdot p(g_i \mid i = t)
+ (1 - \tilde{p}_i) \cdot p(g_i \mid i \neq t) \tag{38}
\]

If we insert Eq. (38) back into Eq. (36), the expectation value becomes

\[
\mathbb{E} \left[ g_i^2 \right] = \tilde{p}_i \cdot \int g_i^2 \, p(g_i \mid i = t) \, dg_i
+ (1 - \tilde{p}_i) \cdot \int g_i^2 \, p(g_i \mid i \neq t) \, dg_i , \tag{39}
\]

which by definition of the (conditional) expectation value, Eq. (36), is equivalent to Eq. (35). □
Theorem 2 (Expectation Value Squared Gradient).
Given that the squared hidden state vector \( h^2 \) is independent of \( p_i \) and whether \( i \) is the true token or not, the expectation value of the squared gradient \( g_i^2 \) is given by
\[
\mathbb{E}\left[ g_i^2 \right] = S \cdot \left[ \tilde{p}_i \cdot X_i^{(i=t)} + (1 - \tilde{p}_i) \cdot X_i^{(i \neq t)} \right],
\]
with
\[
S := \mathbb{E}\left[ h^2 \right] \tag{41} \\
X_i^{(i=t)} := \mathbb{E}\left[ (1 - p_i)^2 \mid i = t \right] \tag{42} \\
X_i^{(i \neq t)} := \mathbb{E}\left[ p_i^2 \mid i \neq t \right] \tag{43}
\]
Proof. We start from Lemma 1 and the square of the gradient,
\[
g_i^2 \overset{(5)}{=} (\delta_{it} - p_i)^2 h^2 \tag{44}
\]
Note that squared variables of vectors in \( \mathbb{R}^H \) always denote the elementwise (Hadamard) product, e.g.
\[
g_i^2 \equiv g_i \odot g_i \in \mathbb{R}_{\geq 0}^H, \tag{45}
\]
with strictly non-negative elements. Using Eq. (44), the expectation values on the right side of Eq. (35) can be expressed as
\[
\mathbb{E}\left[ g_i^2 \mid i = t \right] = \mathbb{E}\left[ (1 - p_i)^2 \cdot h^2 \mid i = t \right] \tag{46} \\
\mathbb{E}\left[ g_i^2 \mid i \neq t \right] = \mathbb{E}\left[ p_i^2 \cdot h^2 \mid i \neq t \right] \tag{47}
\]
Given our assumptions regarding \( h^2 \), its expectation value can be factored out:
\[
\mathbb{E}\left[ g_i^2 \mid i = t \right] = S \cdot X_i^{(i=t)} \tag{48} \\
\mathbb{E}\left[ g_i^2 \mid i \neq t \right] = S \cdot X_i^{(i \neq t)} \tag{49}
\]
Inserting Eqs. (48) and (49) into Eq. (35) yields Eq. (40). \( \square \)

Note that Eq. (40) is a vector equation, with \( \mathbb{E}\left[ g_i^2 \right], S \in \mathbb{R}_{\geq 0}^H \) and \( \tilde{p}_i, X_i^{(i=t)}, X_i^{(i \neq t)} \in \mathbb{R}_{>0} \). It states that the expectation value of \( g_i^2 \) factorizes into a global constant \( S \) that is \( i \)-independent, and a factor that is \( i \)-dependent. The latter is a specific combination of the unigram probability \( \tilde{p}_i \), determined by the data, and the conditional expectation values \( X_i^{(i=t)} \) and \( X_i^{(i \neq t)} \), determined by the model.

Experimental Input Regarding the unigram probability, we know that
1. \( \tilde{p}_i \ll 1 \).
   This is the case for virtually all natural language datasets with a common vocabulary size of \( V > 10000 \), according to Zipf’s law.

The conditional expectation values \( X_i^{(i=t)} \) and \( X_i^{(i \neq t)} \) can be empirically estimated by applying training data to different checkpoints. We consider the three small-scale experiments of Sec. 4.1 with \( N \in \{125M, 355M, 760M\} \) and \( D = 20B \), and take ten equidistant checkpoints after \( D' \in \{2B, 4B, \ldots, 20B\} \) seen tokens for each of them. We then continue pseudo-training on 20 batches (\( \approx 2k \) samples or 2M tokens, see Tab. 5) of data using a zero learning rate, and measure the conditional probabilities in Eqs. (42, 43) from which our target quantities can be estimated. Subsequently, linear fits of the form
\[
X_i^{(i=t)} = A^{(i=t)} \cdot \tilde{p}_i \tag{50} \\
X_i^{(i \neq t)} = A^{(i \neq t)} \cdot \tilde{p}_i , \tag{51}
\]
with fit parameters \( A^{(i=t)} \) and \( A^{(i \neq t)} \) are performed. \( R^2 \) is used to assess the quality of the fits. In addition, the mutual information I between the response and the explanatory variable is computed. Since we observe only a very weak dependence of the results for \( R^2 \) and I on \( N \) and \( D' \), we specify the mean and standard deviation over all experiments for them. Our findings are:
2. \( X_i^{(i=t)} \) is independent of \( \tilde{p}_i \).
   The linear fits yield \( R^2 = 0.003(1) \), and the mutual information is \( I\left( X_i^{(i=t)} ; \tilde{p}_i \right) = 0.14(2) \).
3. \( X_i^{(i \neq t)} \) is proportional to \( \tilde{p}_i \).
   The linear fits yield \( R^2 = 0.92(1) \), and the mutual information is \( I\left( X_i^{(i \neq t)} ; \tilde{p}_i \right) = 0.50(2) \).

The three empirical results above, together with Theorem 2, immediately lead to Eq. (34).

D.2 Experimental Confirmation
We reuse the experiments from the previous section to measure the second moment \( \hat{\nu}_i \) directly, in order to estimate \( \mathbb{E}\left[ \hat{\nu}_i \right] \). Again, linear fits of the form
\[
\mathbb{E}\left[ \hat{\nu}_i \right] = A \cdot \tilde{p}_i \tag{52}
\]
are performed and the mutual information is computed. We find

4. \( \mathbb{E}[\hat{\sigma}_i] \) is proportional to \( \tilde{p}_i \).
   The linear fits yield \( R^2 = 0.85(7) \), and the mutual information is \( I(\mathbb{E}[\hat{\sigma}_i]; \tilde{p}_i) = 1.18(9) \).

The results for \( N = 125\mathrm{M} \) and \( D = D' = 20\mathrm{B} \) are depicted in Fig. 5, as an example.

![Experimental results for \( \mathbb{E}[\hat{\sigma}_i] \) (vertical axis) vs. \( \tilde{p}_i \) (horizontal axis) for \( N = 125\mathrm{M} \) and \( D = D' = 20\mathrm{B} \). The blue line shows the linear fit with \( R^2 = 0.91 \).](page_186_370_489_370.png)

Figure 5: Experimental results for \( \mathbb{E}[\hat{\sigma}_i] \) (vertical axis) vs. \( \tilde{p}_i \) (horizontal axis) for \( N = 125\mathrm{M} \) and \( D = D' = 20\mathrm{B} \). The blue line shows the linear fit with \( R^2 = 0.91 \).

Note that while \( R^2 \) and I are again virtually independent of \( N \) and \( D' \), the fit parameter \( A \) is not. Instead, it seems to increase with \( D' \), as shown in Fig. 6. However, as stated in Eq. (19), the order of magnitude is \( A \approx 10^{-4} \) throughout our experiments.

![Experimental results for the linear fit parameter \( A \) as a function of \( N \) and \( D' \).](page_186_820_489_370.png)

Figure 6: Experimental results for the linear fit parameter \( A \) as a function of \( N \) and \( D' \).

E Experimental Details

E.1 Model and Dataset Sizes

The model sizes \( N \) and dataset sizes \( D \) employed in our experiments are depicted in Fig. 7.

![Overview of the dataset (horizontal axis) and model sizes (vertical axis) involved in our small-scale (blue, green and orange circles) and large-scale (red squares) experiments. The dashed, black line shows \( N = D/20 \), which is approximately the compute-optimal trajectory according to Hoffmann et al. (2022).](page_749_104_489_370.png)

Figure 7: Overview of the dataset (horizontal axis) and model sizes (vertical axis) involved in our small-scale (blue, green and orange circles) and large-scale (red squares) experiments. The dashed, black line shows \( N = D/20 \), which is approximately the compute-optimal trajectory according to Hoffmann et al. (2022).

E.2 Training Hyperparameters

In Tab. 5, we list the general hyperparameters used in our small-scale (Sec. 4.1) and large-scale (Sec. 4.2) experiments. During warm-up, the learning rate is increased from zero to the maximum learning rate. This is followed by a cosine decay which reduces the learning rate to 10% of the maximum at the end of training. Note that weight decay is applied only to linear layers, not layer norms or embeddings. Tab. 6 shows the hyperparameters related to model size, following GPT-3 (Brown et al., 2020).

<table>
  <tr>
    <th>Description</th>
    <th>Small-scale</th>
    <th>Large-scale</th>
  </tr>
  <tr>
    <td>optimizer</td>
    <td>AdamW</td>
    <td></td>
  </tr>
  <tr>
    <td>\( \beta_1 \)</td>
    <td>0.9</td>
    <td></td>
  </tr>
  <tr>
    <td>\( \beta_2 \)</td>
    <td>0.95</td>
    <td></td>
  </tr>
  <tr>
    <td>\( \epsilon \)</td>
    <td>1e-8</td>
    <td></td>
  </tr>
  <tr>
    <td>weight decay</td>
    <td>0.1</td>
    <td></td>
  </tr>
  <tr>
    <td>gradient clipping</td>
    <td>1.0</td>
    <td></td>
  </tr>
  <tr>
    <td>dropout</td>
    <td>0.0</td>
    <td></td>
  </tr>
  <tr>
    <td>weight tying</td>
    <td>true</td>
    <td></td>
  </tr>
  <tr>
    <td>vocab size</td>
    <td>502k</td>
    <td></td>
  </tr>
  <tr>
    <td>learning rate schedule</td>
    <td>cosine decay</td>
    <td></td>
  </tr>
  <tr>
    <td>layer normalization</td>
    <td>LayerNorm</td>
    <td></td>
  </tr>
  <tr>
    <td>precision</td>
    <td>BF16</td>
    <td></td>
  </tr>
  <tr>
    <td>hidden activation</td>
    <td>GeLU</td>
    <td>SwiGLU</td>
  </tr>
  <tr>
    <td>positional embedding</td>
    <td>absolute (learned)</td>
    <td>RoPE</td>
  </tr>
  <tr>
    <td>sequence length</td>
    <td>1024</td>
    <td>2048</td>
  </tr>
  <tr>
    <td>batch size (samples)</td>
    <td>96</td>
    <td>256</td>
  </tr>
  <tr>
    <td>batch size (tokens)</td>
    <td>\(~100k\)</td>
    <td>\(~500k\)</td>
  </tr>
  <tr>
    <td>warmup</td>
    <td>100 steps</td>
    <td>1% of steps</td>
  </tr>
  <tr>
    <td>training framework</td>
    <td>nanoGPT</td>
    <td>Modalities</td>
  </tr>
  <tr>
    <td>training parallelism</td>
    <td>DDP</td>
    <td>FSDP</td>
  </tr>
</table>

Table 5: General hyperparameters used in our two sets of experiments.

F Error Analysis and Statistical Significance

For the error analysis, two separate random variables, \( X_0 \) and \( X_1 \), are considered. The symbol \( X \)
<table>
  <tr>
    <th>N</th>
    <th>lr</th>
    <th>heads</th>
    <th>layers</th>
    <th>emb. dim.</th>
  </tr>
  <tr>
    <td>125M</td>
    <td>6.0e-4</td>
    <td>12</td>
    <td>12</td>
    <td>768</td>
  </tr>
  <tr>
    <td>350M</td>
    <td>3.0e-4</td>
    <td>16</td>
    <td>24</td>
    <td>1024</td>
  </tr>
  <tr>
    <td>760M</td>
    <td>2.5e-4</td>
    <td>16</td>
    <td>24</td>
    <td>1536</td>
  </tr>
  <tr>
    <td>1.3B</td>
    <td>2.0e-4</td>
    <td>32</td>
    <td>24</td>
    <td>2048</td>
  </tr>
  <tr>
    <td>2.6B</td>
    <td>1.6e-4</td>
    <td>32</td>
    <td>32</td>
    <td>2560</td>
  </tr>
</table>

Table 6: Model-size dependent hyperparameters used in our experiments. N denotes the model size in terms of parameters, while lr corresponds to the maximum learning rate.

represents one of the metrics discussed in Sec. 4.3, while 0 and 1 stand for two approaches that are to be compared, like standard Adam and Coupled Adam, for instance. For each of the two random variables \( i = \{0, 1\} \), we conduct and evaluate \( S \) training runs with different seeds, yielding results

\[
\{ X_i^{(1)}, \ldots, X_i^{(S)} \}
\]

While it is desirable to have a large sample size \( S \), it is prohibitively expensive for large model and dataset sizes to repeat training runs. We use

\[
S = 3 \tag{54}
\]

except for the large-scale experiments (Sec. 4.2), where we restrict ourselves to

\[
S = 1 \tag{55}
\]

We are interested in the difference

\[
d = X_1 - X_0 \tag{56}
\]

For \( S = 1 \), it can be computed straight forwardly. However, no statement about the statistical uncertainty or significance of \( d \) can be made. In the case of \( S = 3 \), we apply a one-sided Student’s t-test with a confidence level of

\[
\alpha = 95\% \tag{57}
\]

First, the sample means

\[
\bar{X}_i = \frac{1}{S} \sum_{s=1}^S X_i^{(s)} \tag{58}
\]

and the corrected sample standard deviations

\[
\hat{\sigma}_i^2 = \frac{1}{S-1} \sum_{s=1}^S \left( X_i^{(s)} - \bar{X}_i \right)^2 \tag{59}
\]

for the two samples \( i \in \{0, 1\} \) are estimated. The sample means from Eq. (58) are combined to an estimate for their difference,

\[
\bar{d} = \bar{X}_1 - \bar{X}_0 \tag{60}
\]

and the sample standard deviations from Eq. (59) are propagated to the sample standard deviation of \( d \) via Gaussian error propagation:

\[
\hat{\sigma}_d = \sqrt{ \left( \frac{\partial d}{\partial X_0} \cdot \hat{\sigma}_0 \right)^2 + \left( \frac{\partial d}{\partial X_1} \cdot \hat{\sigma}_1 \right)^2 } \tag{56}
\]
\[
= \sqrt{ \hat{\sigma}_0^2 + \hat{\sigma}_1^2 } \tag{61}
\]

Student’s t-distribution for the chosen confidence level \( \alpha \) (see Eq. (57)) and the

\[
\nu = S - 1 \overset{(54)}{=} 2 \tag{62}
\]

degrees of freedom yields

\[
t_{\alpha, \nu} = 2.92 \tag{63}
\]

With \( S \), \( \sigma_d \) and \( t_{\alpha, \nu} \) from Eqs. (54), (61) and (63) as ingredients, the one-sided confidence threshold for the difference can be computed as

\[
d_{\text{significance}} = t_{\alpha, \nu} \cdot \frac{\hat{\sigma}_d}{\sqrt{S}} \tag{64}
\]

Hence, the estimate \( \bar{d} \) from Eq. (60) is considered a statistically significant improvement of approach \( i = 1 \) over approach \( i = 0 \) if

\[
\bar{d} < -d_{\text{significance}} \tag{65}
\]

for metrics where smaller values are desirable (e.g. \( \mathcal{L} \)), and

\[
\bar{d} > d_{\text{significance}} \tag{66}
\]

for metrics where larger values are better (e.g. Acc).

G Additional Results

G.1 Scaled Coupled Adam

Tab. 3 of Sec. 6.1 shows the results of varying the scaling exponent \( n \) (see Eq. (29)) for \( D = 20B \). The dependency of the loss is visualized in Fig. 3. Here, in Fig. 8, we extend the visualization of the results to \( D \in \{5B, 10B, 20B\} \) and the other evaluation metrics.

G.2 SGD

In Tab. 4 of Sec. 6.2, we showed results for SGD using the best hyperparameter \( f \). Detailed results of the corresponding hyperparameter searches can be found in Tab. 7.

G.3 Individual Downstream Task Performance

In Tab. 8-11, we list the individual downstream task performance for all our experiments (Sec. 4-6).
<table>
  <tr>
    <th>D</th>
    <th>N</th>
    <th>Optimizer</th>
    <th>\( \mathcal{L} (\downarrow) \)</th>
    <th>Acc (\(\uparrow\))</th>
    <th>Iso (\(\uparrow\))</th>
    <th>\( \| \mu \| (\downarrow) \)</th>
    <th>\( \| \mu \|^\perp (\downarrow) \)</th>
    <th>\( \tau (\uparrow) \)</th>
    <th>\( \rho (\uparrow) \)</th>
    <th>\( \kappa (\uparrow) \)</th>
  </tr>
  <tr>
    <td rowspan="7">5B</td>
    <td rowspan="7">125M</td>
    <td>SGD (100)</td>
    <td>3.23 (0)</td>
    <td>0.325 (1)</td>
    <td>1.00 (0)</td>
    <td><b>0.00 (0)</b></td>
    <td><b>0.01 (0)</b></td>
    <td>33 (1)</td>
    <td>59 (1)</td>
    <td><b>25.0 (1.1)</b></td>
  </tr>
  <tr>
    <td>SGD (200)</td>
    <td>3.19 (1)</td>
    <td>0.327 (2)</td>
    <td>1.00 (0)</td>
    <td>0.00 (0)</td>
    <td>0.01 (0)</td>
    <td>41 (1)</td>
    <td>66 (1)</td>
    <td>19.2 (0.9)</td>
  </tr>
  <tr>
    <td>SGD (300)</td>
    <td>3.17 (0)</td>
    <td>0.333 (3)</td>
    <td>0.99 (0)</td>
    <td>0.00 (0)</td>
    <td>0.01 (0)</td>
    <td>45 (1)</td>
    <td>71 (1)</td>
    <td>15.5 (1.4)</td>
  </tr>
  <tr>
    <td>SGD (400)</td>
    <td>3.18 (2)</td>
    <td>0.332 (3)</td>
    <td>0.99 (0)</td>
    <td>0.01 (0)</td>
    <td>0.01 (0)</td>
    <td>48 (3)</td>
    <td>75 (1)</td>
    <td>14.2 (5.4)</td>
  </tr>
  <tr>
    <td>SGD (500)</td>
    <td>3.19 (2)</td>
    <td>0.330 (6)</td>
    <td>0.99 (0)</td>
    <td>0.01 (0)</td>
    <td>0.01 (0)</td>
    <td>49 (1)</td>
    <td>77 (1)</td>
    <td>13.2 (2)</td>
  </tr>
  <tr>
    <td>SGD (600)</td>
    <td>3.24 (2)</td>
    <td>0.326 (3)</td>
    <td>0.99 (0)</td>
    <td>0.01 (0)</td>
    <td>0.01 (0)</td>
    <td>48 (2)</td>
    <td>79 (1)</td>
    <td>11.6 (2)</td>
  </tr>
  <tr>
    <td>Standard Adam</td>
    <td>3.14 (0)</td>
    <td>0.340 (2)</td>
    <td>0.31 (2)</td>
    <td>1.10 (6)</td>
    <td>0.67 (3)</td>
    <td>15 (3)</td>
    <td>-54 (3)</td>
    <td>0.6 (1)</td>
  </tr>
  <tr>
    <td>Coupled Adam</td>
    <td><b>3.12 (1)</b></td>
    <td>0.339 (2)</td>
    <td>0.94 (1)</td>
    <td>0.02 (0)</td>
    <td>0.01 (0)</td>
    <td><b>55 (0)</b></td>
    <td><b>87 (1)</b></td>
    <td>4.8 (2)</td>
  </tr>
</table>

<table>
  <tr>
    <th>D</th>
    <th>N</th>
    <th>Optimizer</th>
    <th>\( \mathcal{L} (\downarrow) \)</th>
    <th>Acc (\(\uparrow\))</th>
    <th>Iso (\(\uparrow\))</th>
    <th>\( \| \mu \| (\downarrow) \)</th>
    <th>\( \| \mu \|^\perp (\downarrow) \)</th>
    <th>\( \tau (\uparrow) \)</th>
    <th>\( \rho (\uparrow) \)</th>
    <th>\( \kappa (\uparrow) \)</th>
  </tr>
  <tr>
    <td rowspan="7">10B</td>
    <td rowspan="7">125M</td>
    <td>SGD (100)</td>
    <td>3.13 (1)</td>
    <td>0.337 (4)</td>
    <td>1.00 (0)</td>
    <td><b>0.00 (0)</b></td>
    <td><b>0.01 (0)</b></td>
    <td>41 (1)</td>
    <td>58 (4)</td>
    <td><b>22.2 (7)</b></td>
  </tr>
  <tr>
    <td>SGD (200)</td>
    <td>3.09 (0)</td>
    <td>0.339 (4)</td>
    <td>0.99 (0)</td>
    <td>0.01 (0)</td>
    <td>0.01 (0)</td>
    <td>48 (1)</td>
    <td>65 (4)</td>
    <td>16.1 (3)</td>
  </tr>
  <tr>
    <td>SGD (300)</td>
    <td>3.07 (1)</td>
    <td>0.341 (4)</td>
    <td>0.99 (0)</td>
    <td>0.01 (0)</td>
    <td>0.01 (0)</td>
    <td>49 (1)</td>
    <td>71 (4)</td>
    <td>12.2 (2)</td>
  </tr>
  <tr>
    <td>SGD (400)</td>
    <td>3.07 (0)</td>
    <td>0.340 (3)</td>
    <td>0.99 (1)</td>
    <td>0.01 (0)</td>
    <td>0.01 (0)</td>
    <td>51 (2)</td>
    <td>76 (3)</td>
    <td>10.7 (8)</td>
  </tr>
  <tr>
    <td>SGD (500)</td>
    <td>3.09 (0)</td>
    <td>0.338 (3)</td>
    <td>0.99 (0)</td>
    <td>0.01 (0)</td>
    <td>0.01 (0)</td>
    <td>52 (1)</td>
    <td>79 (3)</td>
    <td>10.3 (2)</td>
  </tr>
  <tr>
    <td>SGD (600)</td>
    <td>3.11 (1)</td>
    <td>0.341 (4)</td>
    <td>0.98 (1)</td>
    <td>0.01 (0)</td>
    <td>0.01 (0)</td>
    <td>53 (0)</td>
    <td>81 (1)</td>
    <td>9.4 (7)</td>
  </tr>
  <tr>
    <td>Standard Adam</td>
    <td>3.07 (0)</td>
    <td>0.343 (3)</td>
    <td>0.21 (3)</td>
    <td>1.58 (5)</td>
    <td>0.75 (0)</td>
    <td>9 (2)</td>
    <td>-64 (5)</td>
    <td>0.4 (0)</td>
  </tr>
  <tr>
    <td>Coupled Adam</td>
    <td><b>3.03 (0)</b></td>
    <td>0.343 (1)</td>
    <td>0.91 (1)</td>
    <td>0.05 (0)</td>
    <td>0.02 (0)</td>
    <td><b>57 (2)</b></td>
    <td><b>82 (0)</b></td>
    <td>3.6 (8)</td>
  </tr>
</table>

<table>
  <tr>
    <th>D</th>
    <th>N</th>
    <th>Optimizer</th>
    <th>\( \mathcal{L} (\downarrow) \)</th>
    <th>Acc (\(\uparrow\))</th>
    <th>Iso (\(\uparrow\))</th>
    <th>\( \| \mu \| (\downarrow) \)</th>
    <th>\( \| \mu \|^\perp (\downarrow) \)</th>
    <th>\( \tau (\uparrow) \)</th>
    <th>\( \rho (\uparrow) \)</th>
    <th>\( \kappa (\uparrow) \)</th>
  </tr>
  <tr>
    <td rowspan="7">20B</td>
    <td rowspan="7">125M</td>
    <td>SGD (100)</td>
    <td>3.05 (1)</td>
    <td>0.343 (2)</td>
    <td>0.99 (0)</td>
    <td><b>0.01 (0)</b></td>
    <td><b>0.01 (0)</b></td>
    <td>50 (1)</td>
    <td>57 (7)</td>
    <td><b>18.4 (4)</b></td>
  </tr>
  <tr>
    <td>SGD (200)</td>
    <td>3.02 (0)</td>
    <td>0.345 (1)</td>
    <td>0.99 (0)</td>
    <td>0.01 (0)</td>
    <td>0.01 (0)</td>
    <td>52 (1)</td>
    <td>64 (7)</td>
    <td>12.7 (7)</td>
  </tr>
  <tr>
    <td>SGD (300)</td>
    <td>3.01 (1)</td>
    <td>0.350 (1)</td>
    <td>0.98 (1)</td>
    <td>0.01 (0)</td>
    <td>0.02 (0)</td>
    <td>53 (2)</td>
    <td>70 (6)</td>
    <td>9.0 (2)</td>
  </tr>
  <tr>
    <td>SGD (400)</td>
    <td>3.00 (0)</td>
    <td>0.346 (5)</td>
    <td>0.98 (1)</td>
    <td>0.01 (0)</td>
    <td>0.02 (0)</td>
    <td>54 (0)</td>
    <td>76 (5)</td>
    <td>7.4 (1.1)</td>
  </tr>
  <tr>
    <td>SGD (500)</td>
    <td>3.01 (1)</td>
    <td>0.348 (4)</td>
    <td>0.98 (1)</td>
    <td>0.02 (0)</td>
    <td>0.02 (0)</td>
    <td>53 (1)</td>
    <td>78 (5)</td>
    <td>7.5 (1.6)</td>
  </tr>
  <tr>
    <td>SGD (600)</td>
    <td>3.04 (3)</td>
    <td>0.348 (5)</td>
    <td>0.95 (5)</td>
    <td>0.02 (0)</td>
    <td>0.02 (0)</td>
    <td>52 (5)</td>
    <td>81 (5)</td>
    <td>6.8 (2.8)</td>
  </tr>
  <tr>
    <td>Standard Adam</td>
    <td>3.03 (0)</td>
    <td>0.346 (1)</td>
    <td>0.10 (3)</td>
    <td>2.14 (7)</td>
    <td>0.82 (2)</td>
    <td>5 (1)</td>
    <td>-66 (2)</td>
    <td>0.3 (0)</td>
  </tr>
  <tr>
    <td>Coupled Adam</td>
    <td><b>2.97 (0)</b></td>
    <td>0.350 (1)</td>
    <td>0.83 (1)</td>
    <td>0.11 (0)</td>
    <td>0.03 (0)</td>
    <td><b>57 (2)</b></td>
    <td><b>77 (0)</b></td>
    <td>1.7 (5)</td>
  </tr>
</table>

Table 7: Results of our experiments with SGD. Values are highlighted in bold if they are significantly better than all the other values in the same column.
Figure 8: Dependency of different metrics on the scaling exponent \( n \), see Eq. (29). From top to bottom: loss (upstream performance), average accuracy (downstream performance), isotropy, mean embedding norm ratio and \( \tau \). Each plot shows the difference to the respective metric obtained for \( n = 0 \). The arrows indicate whether larger (\( \uparrow \)) or smaller (\( \downarrow \)) values are desirable.
<table>
  <tr>
    <th>D</th>
    <th>N</th>
    <th>Adam</th>
    <th>ARC easy</th>
    <th>ARC challenge</th>
    <th>HellaSwag</th>
    <th>LAMBADA</th>
    <th>RACE</th>
    <th>TruthfulQA</th>
    <th>WinoGrande</th>
    <th>Acc (↑)</th>
  </tr>
  <tr>
    <td rowspan="4">5B</td>
    <td rowspan="2">125M</td>
    <td>Standard</td>
    <td>0.41 (0)</td>
    <td>0.19 (1)</td>
    <td>0.27 (0)</td>
    <td>0.26 (0)</td>
    <td>0.28 (1)</td>
    <td>0.45 (0)</td>
    <td>0.52 (1)</td>
    <td>0.340 (2)</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>0.40 (0)</td>
    <td>0.18 (1)</td>
    <td>0.28 (0)</td>
    <td><b>0.27 (0)</b></td>
    <td>0.28 (0)</td>
    <td>0.45 (1)</td>
    <td>0.52 (1)</td>
    <td>0.339 (2)</td>
  </tr>
  <tr>
    <td>355M</td>
    <td>Standard</td>
    <td>0.43 (1)</td>
    <td>0.19 (1)</td>
    <td>0.29 (0)</td>
    <td>0.31 (0)</td>
    <td>0.29 (0)</td>
    <td>0.43 (0)</td>
    <td>0.52 (1)</td>
    <td>0.352 (3)</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>0.43 (2)</td>
    <td>0.19 (1)</td>
    <td>0.29 (0)</td>
    <td>0.31 (1)</td>
    <td>0.29 (0)</td>
    <td>0.43 (1)</td>
    <td>0.51 (0)</td>
    <td>0.350 (4)</td>
  </tr>
  <tr>
    <td rowspan="4">10B</td>
    <td>760M</td>
    <td>Standard</td>
    <td>0.46 (1)</td>
    <td>0.20 (1)</td>
    <td>0.30 (0)</td>
    <td>0.34 (1)</td>
    <td>0.29 (1)</td>
    <td>0.43 (1)</td>
    <td>0.50 (1)</td>
    <td>0.360 (3)</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>0.45 (0)</td>
    <td>0.20 (1)</td>
    <td>0.30 (0)</td>
    <td>0.34 (1)</td>
    <td>0.29 (1)</td>
    <td>0.42 (1)</td>
    <td>0.50 (0)</td>
    <td>0.357 (3)</td>
  </tr>
  <tr>
    <td>125M</td>
    <td>Standard</td>
    <td>0.42 (1)</td>
    <td>0.19 (0)</td>
    <td>0.28 (0)</td>
    <td>0.29 (1)</td>
    <td>0.28 (1)</td>
    <td>0.44 (1)</td>
    <td>0.51 (1)</td>
    <td>0.343 (3)</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>0.42 (1)</td>
    <td>0.19 (1)</td>
    <td>0.28 (0)</td>
    <td>0.29 (0)</td>
    <td>0.28 (1)</td>
    <td>0.44 (0)</td>
    <td>0.50 (1)</td>
    <td>0.343 (1)</td>
  </tr>
  <tr>
    <td rowspan="4">20B</td>
    <td>355M</td>
    <td>Standard</td>
    <td>0.45 (1)</td>
    <td>0.19 (1)</td>
    <td>0.30 (0)</td>
    <td>0.35 (0)</td>
    <td>0.29 (1)</td>
    <td>0.41 (0)</td>
    <td>0.51 (1)</td>
    <td>0.359 (2)</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>0.46 (1)</td>
    <td>0.19 (1)</td>
    <td>0.30 (0)</td>
    <td>0.35 (1)</td>
    <td>0.30 (1)</td>
    <td>0.42 (0)</td>
    <td>0.52 (1)</td>
    <td>0.365 (2)</td>
  </tr>
  <tr>
    <td>760M</td>
    <td>Standard</td>
    <td>0.47 (0)</td>
    <td>0.21 (0)</td>
    <td>0.32 (0)</td>
    <td><b>0.39 (1)</b></td>
    <td>0.30 (1)</td>
    <td>0.41 (0)</td>
    <td>0.53 (1)</td>
    <td>0.375 (2)</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>0.48 (1)</td>
    <td>0.21 (1)</td>
    <td>0.32 (0)</td>
    <td>0.38 (0)</td>
    <td>0.30 (1)</td>
    <td>0.41 (1)</td>
    <td>0.51 (1)</td>
    <td>0.372 (3)</td>
  </tr>
  <tr>
    <td rowspan="4">20B</td>
    <td>125M</td>
    <td>Standard</td>
    <td>0.43 (0)</td>
    <td>0.18 (1)</td>
    <td>0.28 (0)</td>
    <td>0.29 (1)</td>
    <td>0.28 (1)</td>
    <td>0.44 (1)</td>
    <td>0.51 (1)</td>
    <td>0.346 (1)</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>0.44 (1)</td>
    <td>0.19 (1)</td>
    <td>0.29 (0)</td>
    <td>0.31 (0)</td>
    <td>0.29 (1)</td>
    <td>0.43 (1)</td>
    <td>0.51 (0)</td>
    <td>0.350 (1)</td>
  </tr>
  <tr>
    <td>355M</td>
    <td>Standard</td>
    <td>0.46 (0)</td>
    <td>0.21 (1)</td>
    <td>0.31 (0)</td>
    <td>0.37 (2)</td>
    <td>0.29 (0)</td>
    <td>0.42 (1)</td>
    <td>0.51 (1)</td>
    <td>0.366 (4)</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>0.47 (1)</td>
    <td>0.21 (1)</td>
    <td>0.33 (0)</td>
    <td>0.38 (1)</td>
    <td>0.30 (0)</td>
    <td>0.42 (1)</td>
    <td>0.51 (0)</td>
    <td>0.372 (6)</td>
  </tr>
  <tr>
    <td rowspan="4">20B</td>
    <td>760M</td>
    <td>Standard</td>
    <td>0.49 (0)</td>
    <td>0.21 (1)</td>
    <td>0.33 (0)</td>
    <td><b>0.42 (0)</b></td>
    <td>0.30 (1)</td>
    <td>0.41 (1)</td>
    <td>0.53 (0)</td>
    <td>0.385 (3)</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>0.51 (1)</td>
    <td>0.22 (1)</td>
    <td>0.33 (0)</td>
    <td>0.41 (0)</td>
    <td>0.31 (1)</td>
    <td>0.41 (1)</td>
    <td><b>0.54 (0)</b></td>
    <td>0.392 (2)</td>
  </tr>
</table>

Table 8: Detailed downstream task results of our small-scale experiments from Sec. 4.1 and 5.1. Compare to Tab. 1.

<table>
  <tr>
    <th>D</th>
    <th>N</th>
    <th>Adam</th>
    <th>ARC easy</th>
    <th>ARC challenge</th>
    <th>HellaSwag</th>
    <th>LAMBADA</th>
    <th>RACE</th>
    <th>TruthfulQA</th>
    <th>WinoGrande</th>
    <th>Acc (↑)</th>
  </tr>
  <tr>
    <td rowspan="2">26B</td>
    <td>1.3B</td>
    <td>Standard</td>
    <td><b>0.55</b></td>
    <td><b>0.23</b></td>
    <td><b>0.36</b></td>
    <td><b>0.45</b></td>
    <td><b>0.32</b></td>
    <td>0.37</td>
    <td><b>0.54</b></td>
    <td>0.402</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>0.52</td>
    <td>0.23</td>
    <td>0.36</td>
    <td>0.43</td>
    <td>0.32</td>
    <td>0.38</td>
    <td>0.53</td>
    <td>0.396</td>
  </tr>
  <tr>
    <td rowspan="2">52B</td>
    <td>2.6B</td>
    <td>Standard</td>
    <td><b>0.61</b></td>
    <td><b>0.27</b></td>
    <td><b>0.43</b></td>
    <td><b>0.55</b></td>
    <td><b>0.35</b></td>
    <td><b>0.38</b></td>
    <td><b>0.58</b></td>
    <td>0.451</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td>0.60</td>
    <td>0.25</td>
    <td>0.42</td>
    <td>0.54</td>
    <td>0.34</td>
    <td>0.37</td>
    <td>0.56</td>
    <td>0.441</td>
  </tr>
  <tr>
    <td rowspan="2">105B</td>
    <td>1.3B</td>
    <td>Standard</td>
    <td>0.58</td>
    <td>0.27</td>
    <td>0.42</td>
    <td>0.55</td>
    <td>0.35</td>
    <td>0.38</td>
    <td>0.57</td>
    <td>0.446</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td><b>0.60</b></td>
    <td>0.26</td>
    <td>0.42</td>
    <td>0.54</td>
    <td>0.35</td>
    <td>0.39</td>
    <td>0.57</td>
    <td>0.447</td>
  </tr>
  <tr>
    <td rowspan="2">210B</td>
    <td>2.6B</td>
    <td>Standard</td>
    <td>0.65</td>
    <td>0.29</td>
    <td>0.48</td>
    <td>0.63</td>
    <td>0.37</td>
    <td>0.37</td>
    <td>0.63</td>
    <td>0.490</td>
  </tr>
  <tr>
    <td>Coupled</td>
    <td><b>0.67</b></td>
    <td><b>0.32</b></td>
    <td><b>0.48</b></td>
    <td><b>0.63</b></td>
    <td><b>0.37</b></td>
    <td><b>0.39</b></td>
    <td><b>0.61</b></td>
    <td><b>0.492</b></td>
  </tr>
</table>

Table 9: Detailed downstream task results of our large-scale experiments from Sec. 4.2 and 5.2. Compare to Tab. 2.

<table>
  <tr>
    <th>r</th>
    <th>ARC easy</th>
    <th>ARC challenge</th>
    <th>HellaSwag</th>
    <th>LAMBADA</th>
    <th>RACE</th>
    <th>TruthfulQA</th>
    <th>WinoGrande</th>
    <th>Acc (↑)</th>
  </tr>
  <tr>
    <td>-5</td>
    <td>0.43 (0)</td>
    <td>0.20 (0)</td>
    <td>0.28 (0)</td>
    <td>0.30 (1)</td>
    <td>0.29 (1)</td>
    <td>0.43 (1)</td>
    <td>0.51 (0)</td>
    <td>0.349 (2)</td>
  </tr>
  <tr>
    <td>-4</td>
    <td>0.43 (1)</td>
    <td>0.19 (1)</td>
    <td>0.28 (0)</td>
    <td>0.31 (0)</td>
    <td>0.28 (0)</td>
    <td>0.44 (1)</td>
    <td>0.51 (1)</td>
    <td>0.348 (5)</td>
  </tr>
  <tr>
    <td>-3</td>
    <td>0.43 (0)</td>
    <td>0.20 (1)</td>
    <td>0.28 (0)</td>
    <td>0.31 (0)</td>
    <td>0.28 (0)</td>
    <td>0.44 (2)</td>
    <td>0.52 (1)</td>
    <td>0.352 (5)</td>
  </tr>
  <tr>
    <td>-2</td>
    <td>0.43 (1)</td>
    <td>0.19 (1)</td>
    <td>0.29 (0)</td>
    <td>0.31 (1)</td>
    <td>0.29 (1)</td>
    <td>0.44 (0)</td>
    <td>0.52 (1)</td>
    <td>0.352 (1)</td>
  </tr>
  <tr>
    <td>-1</td>
    <td>0.43 (0)</td>
    <td>0.19 (1)</td>
    <td>0.29 (0)</td>
    <td>0.31 (1)</td>
    <td>0.28 (1)</td>
    <td>0.43 (1)</td>
    <td>0.50 (1)</td>
    <td>0.348 (3)</td>
  </tr>
  <tr>
    <td>0</td>
    <td>0.44 (1)</td>
    <td>0.19 (1)</td>
    <td>0.29 (0)</td>
    <td>0.31 (0)</td>
    <td>0.29 (1)</td>
    <td>0.43 (1)</td>
    <td>0.51 (0)</td>
    <td>0.350 (1)</td>
  </tr>
  <tr>
    <td>1</td>
    <td>0.43 (1)</td>
    <td>0.20 (1)</td>
    <td>0.29 (0)</td>
    <td>0.32 (1)</td>
    <td>0.29 (0)</td>
    <td>0.43 (2)</td>
    <td>0.51 (0)</td>
    <td>0.351 (4)</td>
  </tr>
  <tr>
    <td>2</td>
    <td>0.44 (0)</td>
    <td>0.19 (0)</td>
    <td>0.29 (0)</td>
    <td>0.31 (1)</td>
    <td>0.29 (1)</td>
    <td>0.43 (1)</td>
    <td>0.52 (1)</td>
    <td>0.353 (2)</td>
  </tr>
  <tr>
    <td>3</td>
    <td>0.45 (1)</td>
    <td>0.19 (0)</td>
    <td>0.29 (0)</td>
    <td>0.31 (0)</td>
    <td>0.29 (1)</td>
    <td>0.43 (1)</td>
    <td>0.51 (0)</td>
    <td>0.352 (0)</td>
  </tr>
  <tr>
    <td>4</td>
    <td>0.43 (0)</td>
    <td>0.19 (1)</td>
    <td>0.28 (0)</td>
    <td>0.31 (0)</td>
    <td>0.30 (1)</td>
    <td>0.43 (1)</td>
    <td>0.52 (1)</td>
    <td>0.352 (1)</td>
  </tr>
  <tr>
    <td>5</td>
    <td>0.44 (1)</td>
    <td>0.19 (1)</td>
    <td>0.29 (0)</td>
    <td>0.30 (0)</td>
    <td>0.28 (0)</td>
    <td>0.43 (1)</td>
    <td>0.51 (0)</td>
    <td>0.349 (4)</td>
  </tr>
</table>

Table 10: Detailed downstream task results of our ablations on Scaled Coupled Adam from Sec. 6.1. Compare to Tab. 3.

<table>
  <tr>
    <th>D</th>
    <th>N</th>
    <th>Optimizer</th>
    <th>ARC easy</th>
    <th>ARC challenge</th>
    <th>HellaSwag</th>
    <th>LAMBADA</th>
    <th>RACE</th>
    <th>TruthfulQA</th>
    <th>WinoGrande</th>
    <th>Acc (↑)</th>
  </tr>
  <tr>
    <td rowspan="2">5B</td>
    <td>125M</td>
    <td>SGD (300)</td>
    <td>0.40 (1)</td>
    <td>0.18 (1)</td>
    <td>0.27 (0)</td>
    <td>0.26 (0)</td>
    <td>0.28 (0)</td>
    <td>0.44 (1)</td>
    <td>0.50 (2)</td>
    <td>0.333 (3)</td>
  </tr>
  <tr>
    <td>Coupled Adam</td>
    <td>0.40 (0)</td>
    <td>0.18 (1)</td>
    <td>0.28 (0)</td>
    <td><b>0.27 (0)</b></td>
    <td>0.28 (0)</td>
    <td>0.45 (1)</td>
    <td>0.52 (1)</td>
    <td><b>0.339 (2)</b></td>
  </tr>
  <tr>
    <td rowspan="2">10B</td>
    <td>125M</td>
    <td>SGD (300)</td>
    <td>0.41 (0)</td>
    <td>0.19 (1)</td>
    <td>0.28 (0)</td>
    <td>0.28 (1)</td>
    <td>0.28 (1)</td>
    <td>0.44 (1)</td>
    <td>0.50 (1)</td>
    <td>0.341 (4)</td>
  </tr>
  <tr>
    <td>Coupled Adam</td>
    <td>0.42 (1)</td>
    <td>0.19 (1)</td>
    <td>0.28 (0)</td>
    <td>0.29 (0)</td>
    <td>0.28 (1)</td>
    <td>0.44 (0)</td>
    <td>0.50 (1)</td>
    <td>0.343 (1)</td>
  </tr>
  <tr>
    <td rowspan="2">20B</td>
    <td>125M</td>
    <td>SGD (400)</td>
    <td>0.42 (1)</td>
    <td><b>0.20 (0)</b></td>
    <td>0.28 (0)</td>
    <td>0.30 (1)</td>
    <td>0.28 (1)</td>
    <td>0.43 (1)</td>
    <td>0.51 (2)</td>
    <td>0.346 (5)</td>
  </tr>
  <tr>
    <td>Coupled Adam</td>
    <td>0.44 (1)</td>
    <td>0.19 (1)</td>
    <td>0.29 (0)</td>
    <td>0.31 (0)</td>
    <td>0.29 (1)</td>
    <td>0.43 (1)</td>
    <td>0.51 (0)</td>
    <td>0.350 (1)</td>
  </tr>
</table>

Table 11: Detailed downstream task results of our ablations on SGD from Sec. 6.2. Compare to Tab. 4.

Reviewer paper summaries: This paper investigates the issue of anisotropic word embeddings in large language models and identifies the second moment of the Adam optimizer as a key contributing factor. To address this, the authors propose Coupled Adam, a modified optimizer that normalizes embedding updates to mitigate collective shift of Adam. Experimental results demonstrate that Coupled Adam improves embedding quality and enhances both upstream and downstream model performance across both small- and large-scale datasets.

This paper addresses the anisotropy problem in LLM embeddings by identifying the second moment in Adam as a contributing factor. The authors propose a modified optimizer, Coupled Adam, to mitigate this issue. Their experiments demonstrate that Coupled Adam improves embedding quality and enhances both upstream and downstream performance, particularly on large datasets.
Summary of strengths: -	This paper provides valuable insights about the relationship between the Adam optimizer’s second moment and anisotropy problem.
-	This paper introduces Coupled Adam, a modified optimization algorithm designed to mitigate the anisotropy problem by preventing collective shift of Adam while maintaining computational efficiency.
-	The approach is validated across multiple model architectures and datasets, including both small-scale and large-scale language models.

- The paper is well-written and clearly structured, making it easy to follow.
- The empirical evaluation is thorough, including appropriate metrics and an ablation study.
- The proposed optimizer, Coupled Adam, is a well-motivated and practical solution to the problem.
- The findings contribute to a deeper theoretical understanding of anisotropic embeddings.
Summary of weaknesses: In this paper, the authors proposed the individual scales of the second moments as the root cause of the anisotropy problem. While their study provides valuable insights, it primarily focuses on theoretical and experimental comparison using Adam and SGD as baselines, leaving room for further exploration to fully substantiate their claim. For instance, it could be beneficial to examine (1) whether optimizers that rely on first-order momentum[1] avoid this issue and (2) whether the same problem persists in other recent optimizers that utilize second-order momentum[2].


[1] Chen, Xiangning, et al. "Symbolic discovery of optimization algorithms." Advances in neural information processing systems 36 (2023): 49205-49233.

[2] Gupta, Vineet, Tomer Koren, and Yoram Singer. "Shampoo: Preconditioned stochastic tensor optimization." International Conference on Machine Learning. PMLR, 2018.

- While the theoretical analysis appears sound, I do not have the expertise to assess the mathematical derivations rigorously.
Comments, suggestions and typos: Suggestion

-	Could the authors provide an analysis of (1) and (2) mentioned in the Weakness section?

Typo

(L214) unigram probabilitity => unigram probability

Figures should be vectorized
Overall assessment scores: 3.5

4
Metareview score: 4
Metareview: This paper presents _Coupled Adam_, a variation on the Adam optimizer, that is designed to avoid learning anisotropic word embeddings in large language model training. The paper shows analytically that while SGD leads to updates of the output embedding table with mean zero over the vocabulary, the effective learning rate for each embedding in Adam (based on the second moment estimate, $\hat{v}$) breaks this property. Coupled Adam preserves this property by setting the effective learning rate for each embedding to be the mean over the vocabulary.

Coupled Adam is evaluated empirically by pretraining relatively small (125M to 2.6B parameters) decoder-only transformers from scratch. While intrinsic and extrinsic performance of Coupled Adam-trained models is typically improved relative to vanilla Adam, the embedding tables learned by Coupled Adam are far closer to isotropic than the vanilla Adam embeddings.
Summary of reasons to publish: All reviewers agreed that the evaluation of Coupled Adam was thorough, with the finding that Coupled Adam improves embedding isotropy holding across variation in random seed, number of training tokens, and model size. The motivating analysis of mean-invariance in embedding updates with respect to SGD and Adam is clear and compelling.

Concerns raised by reviewers, including the lack of pertinent references and individual task performance, are relatively minor, and have been addressed by the authors during the discussion period.