Elamur Research⌕
☰
home / papers / 2026 / Weight Tying Biases Token Embeddings Towards the Output Space

Weight Tying Biases Token Embeddings Towards the Output Space

Paper · 2026 · Antonio Lopardo, Avyukth Harish, Catherine Arnett et al.

Summary

Many language models share one matrix between the input embedding and the output (unembedding) layer, but it has been unclear what this sharing does to the learned representations. The authors examine this using matched tied/untied open-weight pairs: OLMo-1B vs OLMo-1B-0724, Qwen3-4B vs Qwen3-8B, and GPT-Neo-2.7B vs Pythia-2.8B. They apply four kinds of analysis:

  • linear alignment of embedding matrices with cosine similarity,
  • tuned lens probing of intermediate layers,
  • tracking of gradient flow,
  • a training intervention that scales up input-side gradients.

They find that a shared matrix ends up resembling an untied output matrix far more than an untied input matrix. The cause is that output-layer gradients dominate early in training. As a result, early layers of tied models show higher tuned-lens KL divergence, which suggests they contribute less to the residual stream. Boosting input gradients reduces this bias, but it costs output quality and does not improve downstream scores.

The work offers a mechanistic reason why large models increasingly untie embeddings. It also points to training fixes that could help small models, where embedding parameters make up much of the budget.

Key points

  • In the per-token analysis on OLMo-1B, the tied matrix reaches a mean cosine similarity of 0.719 with the untied output matrix, compared with 0.525 with the untied input matrix. KNN overlap and spectral distance checks support this across all three model families.
  • In tuned lens analysis, tied models show systematically elevated KL divergence in early layers compared with their untied counterparts.
  • Scaling input gradients during training shifts the tied matrix toward input-like structure, which is causal evidence for the gradient imbalance. However, this trades off against output compatibility and brings no consistent downstream gains.
  • Embedding matrices make up about 73.4% of parameters in Pythia 70M but only about 9.2% at 2.8B scale. This is consistent with Qwen3 tying embeddings only at 0.6B–4B.
  • The study has several limitations:
    • Each setting relies on a single training run.
    • The model pairs carry confounds: different Dolma snapshots for OLMo, model size for Qwen3, and architecture differences for GPT-Neo/Pythia.
    • It does not establish a precise threshold above which untying becomes preferable.
    • Code is available at github.com/AntonioLopardo/weight-tying-bias.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓University of California, BerkeleyBerkeley · United States175 works · 227,440 citations

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 works

Cite

Cite as Antonio Lopardo, Avyukth Harish, Catherine Arnett and Akshat Gupta. 2026. Weight Tying Biases Token Embeddings Towards the Output Space. Findings of the Association for Computational Linguistics: ACL 2026.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2026.findings-acl.2027

@inproceedings{lopardo2026weight,
  title = {{Weight Tying Biases Token Embeddings Towards the Output Space}},
  author = {Antonio Lopardo, Avyukth Harish, Catherine Arnett and Akshat Gupta},
  booktitle = {Findings of the Association for Computational Linguistics: ACL 2026},
  year = {2026},
  doi = {10.18653/v1/2026.findings-acl.2027},
  url = {https://doi.org/10.18653/v1/2026.findings-acl.2027}
}

References

46 references, 21 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Julien Abadji, Pedro Ortiz Suarez, Laurent Romary, and Benoît Sagot. 2022. Towards a cleaner document-oriented multilingual crawled corpus. In Proceedings of the Thirteenth Language Resources and Evaluation Conference, pages 4344–4355, Marseille, France. European Language Resources Association.
  2. Mikel Artetxe, Gorka Labaka, and Eneko Agirre. 2016. Learning principled bilingual mappings of word embeddings while preserving monolingual invariance. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pages 2289–2294. Association for Computational Linguistic…
  3. Nora Belrose, Zach Furman, Logan Smith, Danny Halawi, Igor Ostrovsky, Lev McKinney, Stella Biderman, and Jacob Steinhardt. 2023. Eliciting latent predictions from transformers with the tuned lens. arXiv preprint arXiv:2303.08112.
  4. Francesco Bertolotti and Walter Cazzola. 2024. By tying embeddings you are assuming the distributional hypothesis. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 3584–3610. PMLR. find
  5. Stella Biderman, Hailey Schoelkopf, Quentin Anthony, Herbie Bradley, Kyle O’Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, and 1 others. 2023. Pythia: A suite for analyzing large language models across training and scaling. arXiv preprint arXiv:2304.0… find
  6. Jeffrey L Elman. 1990. Finding structure in time. Cognitive Science, 14(2):179–211. find
  7. Daniel Biś, Maksim Podkorytov, and Xiuwen Liu. 2021. Too much in common: Shifting of embeddings in transformer language models and its implications. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pag… find
  8. Jun Gao, Di He, Xu Tan, Tao Qin, Liwei Wang, and Tie-Yan Liu. 2019. Representation degeneration problem in training natural language generation models. In Proceedings of ICLR.
  9. Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. 2020. PIQA: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 7432–7439.
  10. Sid Black, Leo Gao, Phil Wang, Connor Leahy, and Stella Biderman. 2021. GPT-Neo: Large Scale Autoregressive Language Modeling with Mesh-Tensorflow. find
  11. Hyung Won Chung, Thibault Févry, Henry Tsai, Melvin Johnson, and Sebastian Ruder. 2020. Rethinking embedding coupling in pre-trained language models. arXiv preprint arXiv:2010.12821.
  12. Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. 2019. BoolQ: Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics:…
  13. 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. arXiv preprint arXiv:1803.05457.
  14. Alexis Conneau, Shijie Wu, Haoran Li, Luke Zettlemoyer, and Veselin Stoyanov. 2020. Emerging cross-lingual structure in pretrained language models. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 6022– 6034. Association for Computational Linguistics.
  15. DeepSeek-AI. 2024. DeepSeek-V3 technical report. arXiv preprint arXiv:2412.19437. find
  16. Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACL-HLT, pages 4171–4186.
  17. Brandon Duderstadt, Hayden S. Helm, and Carey E. Priebe. 2023. Comparing foundation models using data kernels. arXiv preprint arXiv:2305.05126. — [arXiv:2305.05126](https://arxiv.org/abs/2305.05126) find
  18. Kawin Ethayarajh. 2019. How contextual are contextualized word representations? comparing the geometry of BERT, ELMo, and GPT-2 embeddings. In Proceedings of EMNLP-IJCNLP, pages 55–65. find
  19. Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. 2020. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027. find
  20. 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 Sutawika, and 5 others.… find
  21. Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, and 1 others. 2024. Gemma: Open models based on Gemini research and technology. arXiv preprint arXiv:2403.08295. find
  22. Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, and 1 others. 2024. The Llama 3 herd of models. arXiv preprint arXiv:2407.21783.
  23. Dirk Groeneveld, Iz Beltagy, Pete Walsh, Akshita Bhagia, Rodney Kinney, Oyvind Tafjord, Ananya Harsh Jha, Hamish Ivison, Ian Magnusson, Yizhong Wang, and 1 others. 2024. Olmo: Accelerating the science of language models. arXiv preprint arXiv:2402.00838. find
  24. Junjie Hu, Sebastian Ruder, Aditya Siddhant, Graham Neubig, Orhan Firat, and Melvin Johnson. 2020. XTREME: A massively multilingual multi-task benchmark for evaluating cross-lingual generalization. arXiv preprint arXiv:2003.11080.
  25. Hakan Inan, Khashayar Khosravi, and Richard Socher. 2017. Tying word vectors and word classifiers: A loss framework for language modeling. In Proceedings of ICLR. find
  26. Ganesh Jawahar, Benoît Sagot, and Djamé Seddah. 2019. What does BERT learn about the structure of language? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 3651–3657, Florence, Italy. Association for Computational Linguistics.
  27. Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. 2019. RoBERTa: A robustly optimized BERT pretraining approach. arXiv preprint arXiv:1907.11692.
  28. 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). Association for Computational Linguistics. find
  29. Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. 2017. Pointer sentinel mixture models. In International Conference on Learning Representations.
  30. Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. 2018. Can a suit of armor conduct electricity? A new dataset for open book question answering. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pages 2381–2391. find
  31. Tomas Mikolov, Quoc V Le, and Ilya Sutskever. 2013. Exploiting similarities among languages for machine translation. arXiv preprint arXiv:1309.4168. nostalgebraist. 2020. interpreting GPT: the logit lens. LessWrong. find
  32. 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. Association for Computational Linguistics.
  33. Qwen Team. 2025. Qwen3 technical report. Preprint, arXiv:2505.09388.
  34. Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. 2019. Language models are unsupervised multitask learners. Technical report, OpenAI. find
  35. Sara Rajaee and Mohammad Taher Pilehvar. 2022. An isotropy analysis in the multilingual BERT embedding space. In Findings of the Association for Computational Linguistics: ACL 2022. Association for Computational Linguistics. find
  36. Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. 2020. WinoGrande: An adversarial winograd schema challenge at scale. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 8732–8740. and 1 others. 2024. Dolma: An open corpus of three trillion tokens… find
  37. Felix Stollenwerk and Tobias Stollenwerk. 2025. Better embeddings with coupled adam. arXiv preprint arXiv:2502.08441. find
  38. Team Olmo, Allyson Ettinger, Amanda Bertsch, Bailey Kuehl, David Graham, David Heineman, Dirk Groeneveld, Faeze Brahman, Finbarr Timbers, Hamish Ivison, and 1 others. 2025. Olmo 3. arXiv preprint arXiv:2512.13961.
  39. Team OLMo, Pete Walsh, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Shane Arora, Akshita Bhagia, Yuling Gu, Shengyi Huang, Matt Jordan, Nathan Lambert, Dustin Schwenk, Oyvind Tafjord, and 1 others. 2024. Olmo 2: The best fully open language model to date. arXiv preprint arXiv:2501.00656. find
  40. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30.
  41. 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. find
  42. 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. find
  43. Alex Warstadt, Alicia Parrish, Haokun Liu, Anhad Mohananey, Wei Peng, Sheng-Fu Wang, and Samuel R. Bowman. 2020. BLiMP: The benchmark of linguistic minimal pairs for English. Transactions of the Association for Computational Linguistics, 8:377– 392.
  44. Chao Xing, Dong Wang, Chao Liu, and Yiye Lin. 2015. Normalized word embedding and orthogonal transform for bilingual word translation. In Proceedings of the 2015 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 1006–1011. A…
  45. 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. find
  46. Peter H Schönemann. 1966. A generalized solution of the orthogonal procrustes problem. Psychometrika, 31(1):1–10. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/lopardo-2026-weight-tying-biases-token-embeddings/