REVIEW 3 major objections 4 minor 38 references
Better Embeddings with Coupled Adam
T0 review · 3 major / 4 minor · reviewed 2026-08-08 · deepseek-v4-flash
Pith's one-line read Adam's per-embedding second moment, not the common-enemy effect alone, makes LLM word embeddings anisotropic; replacing it with the vocabulary average restores isotropy and improves performance on large datasets.
desk verdict Coupled Adam is a simple, effective fix for anisotropic embeddings, but the paper's causal story has a real gap around weight tying and the abstract overclaims on large-scale gains. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The mechanism is Eq. (16): the sum of embedding update vectors is nonzero because Adam's effective learning rate $\eta_i = \eta / \sqrt{\hat v_i + \epsilon}$ depends on the token index. Coupled Adam is the fix: in Algorithm 1 it sets $\hat v_i^{(\tau)}$ to $\hat\nu^{(\tau)} = \frac{1}{V}\sum_j \hat v_j^{(\tau)}$ before the update, making the effective learning rate the same for every embedding. The two supporting pillars are the softmax identity $\sum_i g_i = 0$ (Eq. 8), which guarantees that a vanishing sum of updates is the right target, and the empirical proportionality $E[\hat v_i] \propto p_i$ with fitted constant $A \approx 10^{-4}$, which turns the nonzero weighted sum into a frequency-driven shift dominated by rare tokens.
What would settle it
Log the per-embedding effective learning rates $\eta_i$ and first moments $m_i$ during a few steps of standard Adam training and evaluate the sum $\sum_i \eta_i m_i$; if it is consistently zero despite the per-token second moments, or if replacing every $\hat v_i$ with the vocabulary mean leaves the mean-embedding trajectory essentially unchanged, the proposed mechanism is falsified.
Extended reading notes
Core claim
The central claim is that anisotropy has an optimizer-level cause on top of the common-enemy effect, in which all non-target tokens are pushed in the opposite direction from the hidden state. For the softmax language-modeling head, the embedding gradients sum to zero, so with SGD the mean embedding stays exactly where it started. Adam breaks this because its per-parameter second moments make the effective learning rate depend on the token index: the first moments sum to zero, but the weighted sum $\sum_i \eta_i m_i$ does not, and the mean embedding shifts. The paper measures that a token's second moment is proportional to its corpus frequency, $E[\hat v_i] \propto p_i$, so rare tokens receive disproportionately large normalized updates and drive the collective drift. Coupled Adam removes the index dependence by replacing $\hat v_i$ with the vocabulary-mean second moment $\hat\nu$, which makes the sum of updates vanish while preserving Adam's normalization; the experiments show isotropy values around 0.8 to 0.98, mean embedding norms near zero, and word-similarity correlations rising from roughly 3 to 16 up to 55 to 67.
Load-bearing premise
The whole argument leans on the empirical proportionality between a token's Adam second moment and its corpus frequency (linear fits with $R^2 \approx 0.85$) together with the assumption that the squared hidden state is independent of the token identity; if those relations break on other models or data, the rare-token version of the explanation does, though the general i-dependent-second-moment argument would survive.
Editorial extensions
If this is right
- With Coupled Adam, embedding isotropy reaches 0.9 or above in nearly all small-scale runs and the mean-embedding norm drops by one to two orders of magnitude, so the geometric degeneracy described in the literature is largely removed at the optimizer level.
- Word-similarity correlation rises from single digits to 55 to 58 in small-scale runs and from about 52 to 66 to 67 in large-scale runs, indicating that the usable semantic signal in the embeddings is substantially larger once the shared drift is suppressed.
- On sufficiently large datasets, Coupled Adam also lowers test loss and raises average downstream accuracy; the gains grow with dataset size in the 125M to 760M experiments, while at compute-optimal scale the effect is mixed.
- The method is a single substitution inside Adam for the embedding parameters, so it can be added to existing dense transformer training pipelines without new hyperparameters, and the ablation shows the canonical coupling is close to the optimum.
Reading between the lines
- Going beyond the paper: any adaptive optimizer that normalizes each embedding by its own gradient statistics, such as AdamW variants with different beta schedules or other per-parameter adaptive methods, should produce the same frequency-dependent mean drift, so coupling the second moment is a general design principle rather than an Adam-specific patch.
- The residual mean shift under Coupled Adam at 26B to 210B tokens offers a direct test of the authors' weight-tying hypothesis: training the same architecture with untied input and output embeddings should remove most of that residual drift if the hypothesis is right.
- If the mechanism is correct, anisotropy is in part a training artifact of the optimizer rather than an intrinsic property of learned word meaning; this reframes post-hoc isotropy fixes as corrections for an optimization bias and suggests that optimizer choices deserve a place alongside data and architecture in studies of embedding geometry.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper argues that Adam's per-parameter second moments make embedding update vectors depend on token identity, so that the sum of embedding updates is nonzero even though the sum of output-side gradients vanishes; this, the authors claim, is a root cause of anisotropic embeddings in LLMs. They propose Coupled Adam, which replaces each embedding's Adam second moment by the average second moment over the vocabulary, making the effective learning rate i-independent and restoring a zero update sum. The paper evaluates Coupled Adam against standard Adam on GPT-2-scale models (125M to 760M parameters) and larger models (1.3B and 2.6B parameters) trained on up to 210B tokens, reporting large improvements in isotropy, mean-embedding norm, embedding similarity benchmarks, and condition number, together with mixed but often positive upstream and downstream results at larger dataset sizes. The authors additionally provide ablations on scaling the coupled second moment and on using SGD for embeddings.
Significance. If the causal story is correct, this is a valuable and counterintuitive connection between optimizer design and representation quality, with a simple and practical fix that is cheap to implement and backed by an unusually broad set of experiments. The paper ships reproducible code, reports repeated-seed small-scale runs with significance testing, and includes honest limitations. The clean derivation of the vanishing output-side gradient sum and the observation that Adam's i-dependent effective learning rate breaks that property are useful contributions in themselves. However, the theoretical claim is currently proved only for output-layer gradients, while the experiments use tied embeddings for which input-side gradients also contribute; the residual mean shift observed at scale under Coupled Adam is exactly the regime where this gap matters. The empirical proportionality E[bv_i] proportional to p_i is also fit on the same models that are later used as evidence, so the quantitative statements in Section 2.4 should be treated as motivation rather than exact law.
major comments (3)
- [Sec. 2.2-2.4 and Sec. 3, Eqs. (8), (16)] The vanishing-sum argument in Eq. (8) uses only the output-layer gradient (5). Because all experiments use weight tying (Table 5), the true gradient with respect to an embedding also includes input-side contributions from the embedding lookup, and the sum of those input-side gradients over the vocabulary is not shown to vanish. Therefore the conclusion in Eq. (16) that the Adam update sum is nonzero, and the claim in Sec. 3 that Coupled Adam 'Evidently' makes the update sum vanish, hold only for the output-layer contribution. This matters empirically: Table 2 shows Coupled Adam leaves a residual mean shift at scale (||mu|| up to 0.49), which the paper attributes to weight tying without verification. The central causal claim needs either a proof that the total (input plus output) gradient sum vanishes under tied embeddings, or a reformulation that explicitly separates and measures the input-side contribution.
- [Sec. 5.2 and Sec. 9] The residual shift of the mean embedding under Coupled Adam at scale is not explained within the paper's framework. The text hypothesizes that weight tying causes it and cites Machina and Mercer (2024), but also states in the Limitations that this is not explicitly verified. Since the paper's headline theoretical result is that Coupled Adam removes the collective shift, this residual shift is a load-bearing gap. Please add an experiment or analysis that isolates the input-side gradient contribution, for example by training a small untied model or by measuring the sum of input-side gradients during training, and report how much of the residual ||mu|| it accounts for.
- [Sec. 2.4 and App. D, Eqs. (18)-(23)] The derivation of Eq. (18) relies on Theorem 2's assumption that h^2 is independent of p_i and of token identity, and on linear fits with R^2 = 0.85(7) performed on the same models that are later used as evidence. This is not circular for the design of Coupled Adam, since removing i-dependence works regardless of the exact proportionality, but it does mean the quantitative statements in Eqs. (19)-(23) are empirical scaling relations rather than exact identities. The fitted constant A also varies with D' as shown in Fig. 6. Please soften Eqs. (19)-(21) to approximate relations and state explicitly which parts of the causal claim depend on the fitted proportionality rather than on the weaker property that the second moments are i-dependent.
minor comments (4)
- [Eq. (5) and surrounding text] The notation dL/d e_i in Eq. (5) is misleading under weight tying because it omits input-side gradient terms; consider writing g_i^{out} or explicitly labeling Eq. (5) as the output-side contribution.
- [Eq. (21) and Fig. 6] The text around Eq. (21) presents A/V as a concrete number, but Fig. 6 shows A changing with dataset size D'; the 'order of magnitude' phrasing in Eq. (19) should be carried through consistently to Eqs. (20)-(23).
- [Algorithm 1] The pseudocode line 3 writes L(e_i) as the objective for each embedding; clarifying that the loss is a global function L^{(tau)} evaluated at the current parameters would avoid the impression that each embedding is updated against its own separate loss.
- [App. G.3] There is a stray period directly before the caption of Table 11; please remove it.
Circularity Check
No significant circularity; the mean-shift improvement of Coupled Adam is by construction, but the semantic similarity and loss/downstream gains are independent evidence.
-
self definitional
[Section 3, Eq. (24) and the 'Evidently' paragraph]
"In practice, the exponentially averaged second moments bv(τ ) i as they appear in Eq. (14) are replaced by their average: bν(τ ) := 1 V PV i=1 bv(τ ) i ... 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."
Coupled Adam is defined by Eq. (24), which by construction makes the effective learning rate in Eq. (14) independent of i. The subsequent result that the sum of embedding updates, and therefore the mean-embedding shift, vanishes is an algebraic consequence of this definition rather than an empirical prediction. Reporting ||mu|| values near zero in Tables 1 and 2 as an embedding-quality improvement is therefore a consistency check on the construction. This is a minor by-construction element: the central causal claim that Adam's i-dependent second moment causes the shift is not itself reduced to this definition, and the semantic-similarity r, loss, and downstream-accuracy results are independent evidence.
full rationale
The paper's main theoretical chain is not circular: Eq. (8) follows from the output-layer gradient of cross-entropy, Eq. (5) is a rederivation of an external result (Bi´s et al., 2021), and Eqs. (13)-(17) apply the standard Adam update rule to show that i-dependent second moments generally produce a nonzero weighted sum of first moments. The proportionality E[bv_i] proportional to p_i is established by fits on the same models (App. D.2, R^2=0.85), and the constant A is then used in Eqs. (19)-(21) to interpret the rare-token upscaling; this is a fitted explanatory input rather than a held-out prediction, so it introduces self-reference but not a circular derivation. The only true by-construction element is the mean-shift improvement of Coupled Adam, which follows immediately from the definition in Eq. (24). The paper does not rely on load-bearing self-citations: the only author-overlapping citation is the Modalities training framework, which is not used as evidence for any scientific claim. I also note the correctness risk flagged in Section 2.2: weight tying is assumed but only output-layer gradients are summed, and the residual mean shift at scale is attributed to weight tying without verification (Sections 5.2 and 9); this undermines completeness but is not circularity. Overall, the semantic similarity improvements (r from about 10 to about 57) and upstream/downstream results provide independent support, so the central mechanism is not forced by definition.
Assumptions & free parameters
free parameters (3)
- A =
about 1e-4
- A^(i=t) =
not reported numerically
- A^(i!=t) =
not reported numerically
assumptions (4)
- domain assumption The squared hidden state vector h^2 is independent of p_i and of whether i is the true token (Theorem 2).
- domain assumption Weight tying is used and only output-layer contributions are considered.
- standard math The sum of embedding gradients vanishes at each step (Eq. 8).
- domain assumption Unigram probability p_i equals p(i = t) for the true token during training.
Cite this review
Pith. "Pith review of Better Embeddings with Coupled Adam." pith.science (2026). https://pith.science/paper/CYY6RE3Q
@misc{pith2026250208441,
author = {Pith},
title = {Pith review of: Better Embeddings with Coupled Adam},
year = {2026},
howpublished = {\url{https://pith.science/paper/CYY6RE3Q}},
note = {Machine review of arXiv:2502.08441}
}
read the original 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.
Figures
Figures from the paper (5 more)
Reference graph
Works this paper leans on
-
[1]
Sanjeev Arora, Yuanzhi Li, Yingyu Liang, Tengyu Ma, and Andrej Risteski. 2016. https://doi.org/10.1162/tacl_a_00106 A latent variable model approach to PMI -based word embeddings . Transactions of the Association for Computational Linguistics, 4:385--399
-
[2]
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, Usvsn Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar Van Der Wal. 2023. https://proceedings.mlr.press/v202/biderman23a.html Pythia: A suite for analyzing large language models across trai...
2023
-
[3]
Daniel Biś, Maksim Podkorytov, and Xiuwen 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)
work page 2021
-
[4]
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...
arXiv 2020
-
[5]
Elia Bruni, Nam Khanh Tran, and Marco Baroni. 2014. Multimodal distributional semantics. J. Artif. Int. Res., 49(1):1–47
2014
-
[6]
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. 2018. http://arxiv.org/abs/1803.05457 Think you have solved question answering? try arc, the ai2 reasoning challenge
arXiv 2018
-
[7]
Kawin Ethayarajh. 2019. https://doi.org/10.18653/v1/D19-1006 How contextual are contextualized word representations? C omparing the geometry of BERT , ELM o, 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-IJCN...
-
[8]
Lev Finkelstein, Evgeniy Gabrilovich, Yossi Matias, Ehud Rivlin, Zach Solan, Gadi Wolfman, and Eytan Ruppin. 2001. https://doi.org/10.1145/503104.503110 Placing search in context: The concept revisited . volume 20, pages 406--414
arXiv 2001
Show all 38 references
-
[9]
Jun Gao, Di He, Xu Tan, Tao Qin, Liwei Wang, and Tie-Yan Liu. 2019. http://arxiv.org/abs/1907.12009 Representation degeneration problem in training natural language generation models
2019 arXiv
-
[10]
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang S...
2023 doi
-
[11]
Aaron Gokaslan and Vanya Cohen. 2019. Openwebtext corpus. http://Skylion007.github.io/OpenWebTextCorpus
2019
-
[12]
Felix Hill, Roi Reichart, and Anna Korhonen. 2015. https://doi.org/10.1162/COLI_a_00237 S im L ex-999: Evaluating semantic models with (genuine) similarity estimation . Computational Linguistics, 41(4):665--695
2015 doi
-
[13]
Rae, Oriol Vinyals, and Laurent Sifre
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 Osin...
2022 arXiv
-
[14]
Andrej Karpathy. 2022. NanoGPT . https://github.com/karpathy/nanoGPT
2022
-
[15]
Diederik Kingma and Jimmy Ba. 2014. Adam: A method for stochastic optimization. International Conference on Learning Representations
2014
-
[16]
Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy. 2017. https://doi.org/10.18653/v1/D17-1082 RACE : Large-scale R e A ding comprehension dataset from examinations . In Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing, pages...
2017 doi
-
[17]
Stephanie Lin, Jacob Hilton, and Owain Evans. 2022. https://doi.org/10.18653/v1/2022.acl-long.229 T ruthful QA : Measuring how models mimic human falsehoods . In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pa...
2022 doi
-
[18]
Thang Luong, Richard Socher, and Christopher Manning. 2013. https://aclanthology.org/W13-3512 Better word representations with recursive neural networks for morphology . In Proceedings of the Seventeenth Conference on Computational Natural Language Learning, pages 104--113, So...
2013
-
[19]
Max Lübbering, Mehdi Ali, Felix Stollenwerk, Michael Fromm, Alexander Arno Weber, and Richard Rutmann. 2024. https://github.com/Modalities/modalities Modalities: A pytorch-native framework for distributed and reproducible foundation model training. https://github.com/Modalitie...
2024
-
[20]
Anemily Machina and Robert Mercer. 2024. https://doi.org/10.18653/v1/2024.naacl-long.274 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 Technologi...
2024 doi
-
[21]
Jiaqi Mu, Suma Bhat, and Pramod Viswanath. 2018. http://arxiv.org/abs/1702.01417 All-but-the-top: Simple and effective postprocessing for word representations
2018 arXiv
-
[22]
Denis Paperno, Germ \'a n Kruszewski, Angeliki Lazaridou, Ngoc Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fern \'a ndez. 2016. https://doi.org/10.18653/v1/P16-1144 The LAMBADA dataset: Word prediction requiring a broad discourse cont...
2016 doi
-
[23]
Ofir Press and Lior Wolf. 2017. https://aclanthology.org/E17-2025/ Using the output embedding to improve language models . In Proceedings of the 15th Conference of the E uropean Chapter of the Association for Computational Linguistics: Volume 2, Short Papers , pages 157--163, ...
2017
-
[24]
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. 2019. Language models are unsupervised multitask learners
2019
-
[25]
Sebastian Ruder. 2017. http://arxiv.org/abs/1609.04747 An overview of gradient descent optimization algorithms
2017 arXiv
-
[26]
William Rudman and Carsten Eickhoff. 2024. http://arxiv.org/abs/2305.19358 Stable anisotropic regularization
2024 arXiv
-
[27]
William Rudman, Nate Gillman, Taylor Rayne, and Carsten Eickhoff. 2022. https://doi.org/10.18653/v1/2022.findings-acl.262 I so S core: Measuring the uniformity of embedding space utilization . In Findings of the Association for Computational Linguistics: ACL 2022, pages 3325--...
2022 doi
-
[28]
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. 2020. https://doi.org/10.1609/aaai.v34i05.6399 Winogrande: An adversarial winograd schema challenge at scale . Proceedings of the AAAI Conference on Artificial Intelligence, 34(05):8732--8740
2020 doi
-
[29]
Noam Shazeer. 2020. http://arxiv.org/abs/2002.05202 Glu variants improve transformer
2020 arXiv
-
[30]
Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. 2023. https://huggingface.co/datasets/cerebras/SlimPajama-627B SlimPajama: A 627B token cleaned and deduplicated version of RedPajama
2023
-
[31]
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. 2023. http://arxiv.org/abs/2104.09864 Roformer: Enhanced transformer with rotary position embedding
2023 arXiv
-
[32]
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, W...
2023 arXiv
-
[33]
Dilin Wang, Chengyue Gong, and Qiang Liu. 2019. https://proceedings.mlr.press/v97/wang19f.html Improving neural language modeling via adversarial training . In Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Re...
2019
-
[34]
Lingxiao Wang, Jing Huang, Kevin Huang, Ziniu Hu, Guangtao Wang, and Quanquan Gu. 2020. https://api.semanticscholar.org/CorpusID:211145667 Improving neural language generation with spectrum control . In International Conference on Learning Representations
2020
-
[35]
Sangwon Yu, Jongyoon Song, Heeseung Kim, Seongmin Lee, Woo-Jong Ryu, and Sungroh Yoon. 2022. https://doi.org/10.18653/v1/2022.acl-long.3 Rare tokens degenerate all tokens: Improving neural text generation via adaptive gradient gating for rare token embeddings . In Proceedings ...
2022 doi
-
[36]
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. 2019. https://doi.org/10.18653/v1/P19-1472 H ella S wag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 4791--4...
2019 doi
-
[37]
Zhong Zhang, Chongming Gao, Cong Xu, Rui Miao, Qinli Yang, and Junming Shao. 2020. https://doi.org/10.18653/v1/2020.findings-emnlp.46 Revisiting representation degeneration problem in language modeling . In Findings of the Association for Computational Linguistics: EMNLP 2020,...
2020 doi
-
[38]
Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, and Sham Kakade. 2024. http://arxiv.org/abs/2407.07972 Deconstructing what makes a good optimizer for language models
2024 arXiv
Reviewed August 8, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.