Weight Tying Biases Token Embeddings Towards the Output Space

Antonio Lopardo; Avyukth Harish; Catherine Arnett; Akshat Gupta

cs.CLMar 27, 2026

Weight Tying Biases Token Embeddings Towards the Output Space

Authors: Antonio Lopardo, Avyukth Harish, Catherine Arnett, Akshat Gupta

Organizations: EleutherAI · University of California Berkeley

Abstract

Weight tying, i.e. sharing parameters between input and output embedding matrices, is common practice in language model design, yet its impact on the learned embedding space remains poorly understood. In this paper, we show that tied embedding matrices align more closely with output (unembedding) matrices than with input embeddings of comparable untied models, indicating that the shared matrix is shaped primarily for output prediction rather than input representation. This unembedding bias arises because output gradients dominate early in training. Using tuned lens analysis, we show this negatively affects early-layer computations, which contribute less effectively to the residual stream. Scaling input gradients during training reduces this bias, providing causal evidence for the role of gradient imbalance. This is mechanistic evidence that weight tying optimizes the embedding matrix for output prediction, compromising its role in input representation. These results help explain why weight tying can harm performance at scale and have implications for training smaller LLMs, where the embedding matrix contributes substantially to total parameter count.

Figures & tables

Figure 1: Per-token cosine similarity (after linear alignment) between the tied embedding matrix and the untied input (blue) and output (orange) matrices for two OLMo-1B training runs (tied and untied). Dashed lines indicate means. The tied matrix is substantially more aligned with the untied output (mean: 0.719) than the untied input (mean: 0.525).

OLMo-1B (tied) vs OLMo-1B-0724 (untied)
Comparison	Identity	Orthogonal	Linear
Input (U) $\rightarrow$ Output (U)	0.012	0.440	0.574
Output (U) $\rightarrow$ Tied	0.014	0.669	0.719
Input (U) $\rightarrow$ Tied	0.001	0.420	0.525
GPT-Neo-2.7B (tied) vs Pythia-2.8B (untied)
Input (U) $\rightarrow$ Output (U)	-0.001	0.456	0.518

Table 1: Geometric alignment between embedding spaces under identity, orthogonal (rotation), and linear (unconstrained least squares) transformations. We report mean cosine similarity between corresponding token embeddings from tied and untied models.

Figure 2: Tuned lens KL divergence for tied vs untied OLMo-1B. Lower values indicate better alignment between a layer’s representations and the output space. The clearest separation appears in the early layers, where the tied model shows higher KL divergence.

Figure 3: For OLMo-1B-0724 (untied). Top: cosine similarity to initial embeddings (step 0), measuring cumulative drift. Bottom: cosine similarity between consecutive checkpoints, measuring between checkpoints change rate.

Figure 4: Gradient flow to the shared embedding matrix in tied OLMo-1B during the first 1000 training steps. Left: L2 norm of gradients from the input embedding (blue) and output projection (orange) pathways on a log scale. Right: relative contribution of each pathway as a percentage of total gradient norm.

Figure 5: Norm-frequency relationship for OLMo-1B after 10k steps (20B tokens). Left: untied model’s input (blue) and output (orange) matrices. Right: tied model’s shared matrix (pink).

Model	vs Untied Input	vs Untied Output
Tied (no scaling)	0.216	0.384
Tied (input $\times$ 5)	0.222	0.369

Table 2: Cosine similarity (after alignment) between tied OLMo-1B embedding matrices and the input/output matrices of OLMo-1B-0724 (untied) at step 10K. Higher gradient scaling shifts mean more input-like structure.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Total (Untied)	Total (Tied)	Saved
Pythia-70M	70.4M	44.7M	25.8M
Pythia-160M	162.3M	123.7M	38.6M
Pythia-410M	405.3M	353.8M	51.5M
Pythia-1B	1011.8M	908.8M	103.0M
Pythia-2.8B	2775.2M	2646.4M	128.8M

Appendix

Table 3: Total parameter counts with and without weight tying for Pythia models. For an untied model, tying removes one embedding matrix, saving

Vd

parameters (half of the untied embedding total

2Vd

). Totals are computed by summing checkpoint parameter counts.

Model	Hidden	Embed Params	% of Total
Pythia-70M	512	51.4M	73.4%
Pythia-160M	768	77.2M	48.3%
Pythia-410M	1024	103.0M	25.1%
Pythia-1B	2048	206.0M	20.6%
Pythia-2.8B	2560	257.4M	9.2%

Appendix

Table 4: Embedding parameters for the Pythia suite (untied). Embed Params shows the combined input and output matrices (

2\times V\times d

where

V=50{,}257

). Weight tying would halve these values, saving over 35% of parameters at small scale but less than 5% at billion-scale.

Comparison	KNN@10 ( $\uparrow$ )	Spectral Dist ( $\downarrow$ )
OLMo-1B (tied) vs OLMo-1B-0724 (untied)
Input (U) $\leftrightarrow$ Output (U)	0.455	1.349
Output (U) $\leftrightarrow$ Tied	0.733	0.588
Input (U) $\leftrightarrow$ Tied	0.496	1.373
Qwen3-4B (tied) vs Qwen3-8B (untied)
Input (U) $\leftrightarrow$ Output (U)	0.366	1.421

Appendix

Table 5: KNN@10 overlap (higher = more similar) and spectral distance (lower = more similar) between embedding spaces. Bold indicates the most similar pair within each family. GPT-Neo/Pythia uses aligned vocabulary (36,938 tokens).

Figure 6: Tuned lens KL divergence across layers for Pythia-2.8B (untied) vs GPT-Neo-2.7B (tied). Both models have 32 layers. GPT-Neo-2.7B shows elevated KL divergence in the early layers.

Figure 7: Tuned lens KL divergence for Qwen3-4B (tied) vs Qwen3-8B (untied). Qwen3-4B (tied) displays significantly higher early-layer KL divergence before starting to converge with Qwen3-8B (untied).

Figure 8: Embedding matrix dynamics for OLMo-7B (untied). Top: cosine similarity to initial embeddings (step 0), measuring cumulative drift. Bottom: cosine similarity between consecutive checkpoints, measuring per-step change rate.

Figure 9: Embedding matrix dynamics for Pythia-1B (untied). Top: cosine similarity to initial embeddings (step 0), measuring cumulative drift. Bottom: cosine similarity between consecutive checkpoints, measuring per-step change rate.

Model	vs Untied Input	vs Untied Output
Tied (no scaling)	0.172	0.197
Tied (input $\times$ 2)	0.173	0.197
Tied (input $\times$ 10)	0.173	0.190

Appendix

Table 6: Cosine similarity after Procrustes alignment of tied OLMo-1B models trained with scaled input gradients at step 1000. Even early in training, higher input gradient scaling slightly increases alignment with the untied input matrix and decreases alignment with the untied output matrix.

Benchmark	Tied	Tied (input $\times$ 5)
WikiText-2 PPL $\downarrow$	35.71	36.64
PIQA	0.620	0.631
HellaSwag	0.329	0.326
Winogrande	0.514	0.509
ARC-Easy	0.396	0.402
ARC-Challenge	0.233	0.241

Appendix

Table 7: Downstream evaluation of tied OLMo-1B models with and without input gradient scaling at step 10K (20B tokens). Higher is better for all benchmarks except WikiText-2 (perplexity, lower is better). No consistent performance difference is observed.

	# PT	# FT	XNLI	NER	PAWS-X	XQuAD	TyDi	Avg
	params	params	Acc	F1	Acc	EM/F1	EM/F1
Coupled	177M	177M	70.7	69.2	85.3	46.2/63.2	40.7/56.7	62.3
Decoupled	269M	177M	71.3	68.9	85.0	46.9/63.8	42.8/58.1	62.7

Appendix

Table 8: Effect of decoupling input and output embedding matrices on XTREME benchmark ( Hu et al., 2020 ) performance ( Chung et al., 2020 ) . PT: Pre-training. FT: Fine-tuning. Decoupling adds parameters during pre-training but improves average downstream performance.

Language Modeling

Language modeling develops systems that learn patterns in text and other sequences to understand, predict, and generate language.

See all 519 papers

Output Embedding Centering for Stable LLM Pretraining
Jan 5, 2026Felix Stollenwerk, Anna Lokrantz, Niclas HertzbergLanguage Model Pretraining Neural Representation Geometry
AI Sweden
Lost in Backpropagation: The LM Head is a Gradient Bottleneck
Mar 10, 2026Nathan Godey, Yoav ArtziLLM Training Efficient Language Model Training
Cornell University, New York, USA
Context Staircase: Signature-Aligned Dynamics of Token Embeddings under Small Initialization
Aug 31, 2026Junjie Yao, Liangkai Hang, Zhi-Qin John XuLanguage Model Pretraining Word Embeddings
School of Mathematical Sciences Shanghai Jiao Tong University Shanghai 200240, China · Institute of Natural Sciences, School of Mathematical Sciences Shanghai Jiao Tong University Shanghai 200240, China

Jan 5, 2026J/K move · Enter open · S save

Output Embedding Centering for Stable LLM Pretraining

Felix Stollenwerk, Anna Lokrantz, Niclas Hertzberg

AI Sweden

Language Model Pretraining Neural Representation Geometry

Open paper AskPDF arXiv

Pretraining of large language models is not only expensive but also prone to certain training instabilities. A specific instability that often occurs at the end of training is output logit divergence. The most widely used mitigation strategies, z-loss and logit soft-capping, merely address the symptoms rather than the underlying cause of the problem. In this paper, we analyze the instability from the perspective of the output embeddings' geometry and identify anisotropic embeddings as its source. Based on this, we propose output embedding centering (OEC) as a new mitigation strategy, and demonstrate that it suppresses output logit divergence. OEC can be implemented in two different ways: as a deterministic operation called

μ

-centering, or a regularization method called

μ

-loss. Our experiments show that both variants outperform z-loss in terms of training stability, while being on par with logit soft-capping. This holds true both in the presence and the absence of weight tying. As a secondary result, we find that

μ

-loss is significantly less sensitive to regularization hyperparameter tuning than z-loss.

Abstract

Language Modeling

Explore similar work

Output Embedding Centering for Stable LLM Pretraining

Lost in Backpropagation: The LM Head is a Gradient Bottleneck

Context Staircase: Signature-Aligned Dynamics of Token Embeddings under Small Initialization