Better Embeddings with Coupled Adam

• Anisotropic embeddings cluster in a small subspace away from the origin, limiting semantic usefulness and model expressiveness and generalizability.Prior work traced the issue largely to a shifted mean embedding, while the optimizer’s role remained uninvestigated.Source: §1 · p. 1
• Adam scales rare-word embedding updates up relative to frequent-word updates, relevant because training-token frequencies can differ by several orders of magnitude.Source: §1 · p. 1
• The authors identify Adam as playing a crucial role in causing anisotropic embeddings.Source: §1 · p. 1
• Coupled Adam adjusts Adam for embedding parameters to alleviate anisotropy.The adjustment is described as easy to implement and efficient.Source: §1 · p. 2
• For sufficiently large datasets, Coupled Adam significantly improves embedding quality and benefits upstream and downstream performance.Source: §1 · p. 2

2.1 Language Modeling Head

• The language-modeling analysis assumes weight tying and considers only output-layer contributions, applying to architectures with a standard language-modeling head.The head maps hidden states to token logits through embedding vectors.Source: §2 · p. 2, §2.1 · p. 2

2.2 Vanishing Sum of Embedding Gradients

• The true-token embedding update points along h, while other embeddings update oppositely; their gradients sum to zero after accounting for predicted probabilities.The common-enemy explanation based only on update directions is incomplete because it omits probability-dependent gradient scaling.Source: §2.2 · p. 2

2.3 Invariant Mean Embedding with SGD

• SGD’s embedding updates sum to zero at every step, keeping the mean embedding invariant, including with momentum.Individual embeddings can still move in different directions and by different amounts.Source: §2.3 · p. 3

2.4 Shifted Mean Embedding with Adam

• With Adam, embedding-dependent second moments make update sums generally nonzero, changing the mean embedding and shifting embeddings away from the origin.The paper argues that the expected second moment is proportional to unigram probability.Source: §2.4 · pp. 3–4
• Across models, linear fits of measured second moments against unigram probabilities averaged R² = 0.85(7).The authors report experimental confirmation of the relationship.Source: §2.4 · p. 4

3 Coupled Adam

• Averaging the second moments over vocabulary items removes their individual scaling while retaining a global scale for embedding updates.The method replaces the individual exponentially averaged second moments with their average.Source: §3 · p. 4
• The resulting algorithm is named Coupled Adam.Source: §3 · p. 4
• Coupled Adam couples embedding second moments, making the effective learning rate independent of token; its updates sum to zero while retaining second-moment normalization.This combines the zero-sum update property of SGD with Adam’s second-moment normalization.Source: §3 · p. 4

4 Experiments

• Across experiments, standard Adam is retained for non-embedding parameters while standard or Coupled Adam is used for embeddings, across different datasets, frameworks, and architectures.The study includes small-scale runs and large-scale runs with models exceeding 1B parameters.Source: §4 · pp. 4–5

4.1 Small-scale Experiments

• Small-scale runs use GPT-2-style models on OpenWebText, with dataset sizes of 5B, 10B, and 20B tokens, model sizes of 125M, 355M, and 760M, and three seeds per setting.Source: §4.1 · p. 5

4.2 Large-scale Experiments

• Large-scale runs use 1.3B- and 2.6B-parameter models with near compute-optimal dataset sizes and datasets four times larger.The near compute-optimal sizes are 26B and 52B tokens; the larger sizes are 105B and 210B tokens.Source: §4.2 · p. 5

4.3 Evaluation

• Evaluation tracks test loss, average downstream accuracy, isotropy, embedding norms, semantic similarity, vector-length/unigram-probability correlation, and condition number.Source: §4.3 · p. 5
• Table 1 reports rounded means and standard deviations across three seeds, highlighting significant better values; Figure 2 plots loss and accuracy differences by dataset and model size.Source: §4.1 · p. 5, Table 1 · p. 6, Figure 2 · p. 6

5.1 Small-scale Experiments

• With sufficiently large datasets, Coupled Adam improves small-scale upstream and downstream performance, with gains increasing monotonically with dataset size; isotropy exceeds 0.90 in all but one case.The mean embedding is close to the origin, while semantic correlation r, condition number κ, and the positive embedding-length/unigram-probability correlation also improve.Source: §5.1 · p. 6

5.2 Large-scale Experiments

• In large-scale experiments, Coupled Adam performs worse on upstream and downstream metrics at compute-optimal data sizes but better with four times more data.Embedding-specific metrics improve consistently, although the mean embedding shifts farther from the origin as model and dataset sizes increase, reducing isotropy; the authors attribute the small/large-scale performance difference to larger batches and fewer optimization steps.Source: §5.2 · pp. 6–7

6 Ablations

• The additional ablations use a 125M model, dataset sizes of 5B, 10B, and 20B tokens, and three random seeds per setting.Source: §6 · p. 7

6.1 Scaled Coupled Adam

• For scaled Coupled Adam, loss is minimized near n = 0 and worsens significantly at the smallest and largest tested exponents; n ≠ 0 changes the embeddings’ effective learning rate.Smaller n improves isotropy and mean embedding position but harms model performance, while downstream-performance evidence is statistically inconclusive and semantic usefulness appears to suffer for n < 0.Source: §6.1 · p. 7, Figure 3 · p. 8

6.2 SGD

• Coupled Adam consistently outperforms momentum SGD on upstream and downstream task metrics while requiring one fewer hyperparameter to tune, although SGD is better on isotropy, mean shift, and condition number.The SGD learning-rate multiplier is tuned to 300 for 5B and 10B tokens, and 400 for 20B tokens.Source: §6.2 · pp. 7–8

• Gao et al. (2019) describe representation degeneration and cosine regularization; later mitigation approaches include adversarial noise, spectrum control, and Laplacian regularization.Source: §7 · p. 8
• Bís et al. (2021) trace anisotropy largely to a common shift in a dominant direction, while Yu et al. (2022) propose Adaptive Gradient Gating based on rare-token embedding gradients.The paper says its analysis attributes the rare-embedding effect to Adam’s massive gradient up-scaling.Source: §7 · p. 8
• Rudman et al. (2022) define mean-agnostic IsoScore, whereas concurrent work finds simplified Adam variants outperform SGD but remain slightly worse than Adam.Adalayer* averages the second moment over hidden space rather than vocabulary space, unlike Coupled Adam.Source: §7 · p. 8

8 Conclusions

• Table 4 compares models trained with SGD and Coupled Adam for embeddings, with SGD learning rates selected by hyperparameter search.Bold values indicate statistically significant better results.Source: Table 4 · p. 9
• Adam’s common enemy effect and individual second moments cause a collective shift of embedding vectors away from the origin.The paper presents this combination as advancing the theoretical understanding of embedding anisotropy.Source: §8 · pp. 8–9
• Coupled Adam gives every embedding vector the same effective learning rate, suppressing the collective shift while retaining second-moment normalization.The paper reports consistent improvements in embedding-specific metrics across all experiments.Source: §8 · p. 9
• Coupled Adam improves upstream and downstream performance for large datasets, as typically used in LLM training.The paper also reports consistent improvements in embedding-specific metrics across all experiments.Source: §8 · p. 9

9 Limitations

• Experiments use dense decoders up to 2.6B parameters, omit fine-tuning and alternative learning-rate schedules, and do not verify the cause of Coupled Adam’s residual mean-embedding shift.The study also uses cosine decay, tests pre-training and few-shot downstream evaluation, and leaves other optimizers and more efficient implementations for future work.Source: §9 · p. 9

A Unigram Probability Distribution

• Figure 4 illustrates the GPT-2-tokenizer unigram probability distribution for OpenWebText, whose maximum probability is approximately 0.037.The minimum probability is zero and is not shown.Source: A Unigram Probability Distribution · p. 11, Figure 4 · p. 11

B Embedding Gradients

• The passages explicitly derive Eq. (5), whose individual factors come directly from Eqs. (2)–(4); inserting Eqs. (31), (32), and (33) into Eq. (30) leads to Eq. (5).Source: B Embedding Gradients · pp. 11–12

C SGD Algorithm

• The appendix presents SGD pseudocode for updating embedding vectors, with optional momentum and final embeddings as output.It provides the algorithm for comparison with Coupled Adam.Source: C SGD Algorithm · p. 12

D Magnitude of the Second Moment in Adam

• The second-moment analysis combines a semi-theoretical derivation with experimental confirmation of a squared-gradient relationship.The derivation expresses expected squared gradients using observables, then determines their dependence on token frequency experimentally.Source: D Magnitude of the Second Moment in Adam · p. 12, §D.1 · p. 12

D.1 Semi-theoretical Derivation

• Under a stated independence assumption, the semi-theoretical analysis finds one conditional expectation independent of token frequency and another proportional to it.For natural-language datasets with vocabularies larger than 10,000, token probabilities are generally much less than one under Zipf’s law; the empirical fits report R2 = 0.003(1) and 0.92(1), respectively.Source: §D.1 · p. 13

D.2 Experimental Confirmation

• E[bv_i] is proportional to e_i, with linear-fit R2 = 0.85(7) and mutual information I(E[bv_i]; e_i) = 1.18(9).Figure 5 depicts E[bv_i] against e_i for N = 125M and D = D′ = 20B, with a linear fit having R2 = 0.91.Source: §D.2 · p. 14, Figure 5 · p. 14
• The fitted parameter A varies with D′, while R2 and mutual information are virtually independent of N and D′.A appears to increase with D′, and its order of magnitude is approximately 10^-4 throughout the experiments.Source: §D.2 · p. 14, Figure 6 · p. 14

E.1 Model and Dataset Sizes

• Figure 7 maps dataset size against model size for small- and large-scale experiments and marks N = D/20 as an approximately compute-optimal trajectory.The figure distinguishes small-scale experiments with circles from large-scale experiments with squares.Source: Figure 7 · p. 14

E.2 Training Hyperparameters

• Training increases the learning rate from zero to its maximum during warm-up, then cosine-decays it to 10% of that maximum by training’s end.Weight decay is applied to linear layers, not layer norms or embeddings.Source: §E.2 · p. 14, Table 5 · p. 14

F Error Analysis and Statistical Significance

• For S = 1, the difference can be computed, but statistical uncertainty or significance cannot be stated; for S = 3, the appendix uses a one-sided Student’s t-test.It describes combining sample means and propagating standard deviations to estimate the difference and its uncertainty.Source: F Error Analysis and Statistical… · p. 15

G.1 Scaled Coupled Adam

• Figure 8 extends the scaling-exponent results to D ∈ {5B, 10B, 20B} and plots loss, average accuracy, isotropy, mean embedding norm ratio, and r relative to n = 0.The arrows indicate whether larger or smaller values are desirable for each metric.Source: §G.1 · p. 15, Figure 8 · p. 17

G.2 SGD

• The SGD results use the best hyperparameter f from a learning-rate search, with detailed results listed in Table 7.Table 7 highlights values significantly better than all other values in the same column.Source: §G.2 · p. 15, Table 7 · p. 16

G.3 Individual Downstream Task Performance

• Tables 8–11 provide individual downstream-task results for the experiments and ablations, complementing the aggregate comparisons.Tables 8–9 cover the small- and large-scale experiments; Tables 10–11 cover scaled Coupled Adam and SGD ablations.Source: §G.3 · p. 15, Table 8 · p. 18, Table 9 · p. 18