Weight Tying in Large Language Models: Mathematical Foundations, Geometric Bottlenecks, and Modern Architectural Trade-Offs

In autoregressive language models, the embedding layer at the input and the unembedding projection layer at the output serve as the two bridges between discrete vocabulary tokens and the continuous hidden representation space. In the foundational Transformer architecture (Vaswani et al., 2017) and early generative models like GPT-2 (Radford et al., 2019), the weights of these two matrices were shared, a technique known as weight tying.

While weight tying significantly reduces parameter count—especially as vocabulary sizes expand into hundreds of thousands of tokens—modern frontier architectures increasingly diverge. Models such as LLaMA 3 (Dubey et al., 2024) and Mistral decouple input and output embeddings, whereas compact models like Gemma (Gemma Team, 2024) maintain weight tying. This design choice reflects a fundamental trade-off between parameter efficiency, geometric representation capacity, and gradient optimization dynamics.

TIED EMBEDDINGS (W_in = W_out = W)
Input Tokens:  x_t ──► [ Lookup in W ] ──────────► e_t ──► [ Transformer Layers ] ──► h_L
                                                                                     │
Output Logits: z_t ◄── [ MatMul with W^T ] ◄── RMSNorm(h_L) ◄────────────────────────┘
(Single matrix of size |V| x d_model shared across input lookup and output projection)

UNTIED EMBEDDINGS (W_in ≠ W_out)
Input Tokens:  x_t ──► [ Lookup in W_in ] ───────► e_t ──► [ Transformer Layers ] ──► h_L
                                                                                     │
Output Logits: z_t ◄── [ MatMul with W_out^T ] ◄── RMSNorm(h_L) ◄────────────────────┘
(Two independent matrices: W_in for semantic input mapping, W_out for logit classification)

Mathematical Formulation of Tied vs. Untied Embeddings

Let $V$ denote the vocabulary of discrete tokens with size $|V|$, and let $d_{\text{model}}$ denote the hidden dimension of the transformer.

The Untied Architecture

In an untied autoregressive language model:

1. Input Embedding Matrix: $W_{\text{in}} \in \mathbb{R}^{|V| \times d_{\text{model}}}$. For an input token $x_t \in \{1, \dots, |V|\}$, the input vector $e_t \in \mathbb{R}^{d_{\text{model}}}$ is retrieved via index lookup:

$$e_t = W_{\text{in}}[x_t, :]$$

1. Transformer Backbone: The sequence of input vectors is processed through $L$ transformer layers with residual connections, self-attention, and feed-forward networks to produce the final layer representation $h_t \in \mathbb{R}^{d_{\text{model}}}$.
2. Output Unembedding Matrix: $W_{\text{out}} \in \mathbb{R}^{|V| \times d_{\text{model}}}$. After applying final normalization $\hat{h}_t = \text{RMSNorm}(h_t)$, unnormalized log-probabilities (logits) $z_t \in \mathbb{R}^{|V|}$ are computed via linear projection:

$$z_t = \hat{h}_t W_{\text{out}}^T$$ where the logit for token $k$ is the inner product: $$z_{t, k} = \langle \hat{h}_t, w_{\text{out}, k} \rangle$$

1. Categorical Probability Distribution: The predictive distribution over next tokens is obtained via the softmax function:

$$P(x_{t+1} = k \mid x_{\le t}) = \frac{\exp(z_{t, k})}{\sum_{j=1}^{|V|} \exp(z_{t, j})}$$

The Tied Architecture

In a tied architecture (Press & Wolf, 2016; Inan et al., 2016), the model enforces $W_{\text{out}} = W_{\text{in}} = W$.

In the standard formulation of Vaswani et al. (2017), the input lookup is additionally multiplied by a constant scaling factor $\sqrt{d_{\text{model}}}$ to prevent the embedding magnitude from being dwarfed by subsequent residual additions: $$e_t = \sqrt{d_{\text{model}}} \cdot W[x_t, :]$$ $$z_t = \hat{h}_t W^T$$

This formulation eliminates $|V| \times d_{\text{model}}$ trainable parameters from the network.

Historical Origins: Parameter Scarcity in Recurrent Models

Weight tying was proposed independently by Press & Wolf (2016) and Inan et al. (2016) during the era of recurrent neural network (LSTM/GRU) language modeling.

In 2016, language models operated with hidden dimensions between 256 and 1024, sequence lengths under 100 tokens, and vocabularies of 30,000 to 50,000 tokens. In such setups, two unshared embedding matrices accounted for 30% to 50% of the entire model parameter budget.

Press and Wolf argued on both theoretical and empirical grounds:

1. Duality of Representation: If a word vector captures the semantic meaning of a word when processed as input context, that same vector ought to serve as the ideal target representation when predicting that word as output.
2. Regularization Effect: Forcing parameter sharing acted as a strong regularizer, reducing overfitting on small training corpora such as Penn Treebank and WikiText-2.
3. Data Efficiency for Rare Words: In an untied model, rare tokens only receive gradient updates to their input embedding row when they appear in the input sequence. Under weight tying, rare token embeddings receive gradient updates on every training step via the output softmax denominator.

The Functional Asymmetry Between Input and Output Spaces

Despite the parameter savings, input embedding and output classification perform fundamentally different mathematical operations. Enforcing $W_{\text{in}} = W_{\text{out}}$ forces two contradictory geometric structures into a single matrix.

1. Semantic Manifold Mapping vs. Hyperplane Discrimination

• Input Space ($W_{\text{in}}$): Acts as a continuous coordinate mapping. It projects discrete categorical IDs into a semantic latent space where Euclidean distance and cosine similarity capture distributional properties, syntactic categories, and compositional analogies. These vectors serve as initial keys, queries, and values in self-attention layers.
• Output Space ($W_{\text{out}}$): Acts as a bank of linear classifiers. Each row vector $w_{\text{out}, k}$ represents the normal vector of a decision hyperplane in $\mathbb{R}^{d_{\text{model}}}$. The logit $z_k = \langle \hat{h}_t, w_{\text{out}, k} \rangle$ measures the scalar projection of the final hidden state along direction $w_{\text{out}, k}$.

2. The Frequency-Norm Distortion

Token frequencies in natural language follow a power-law Zipfian distribution. In the output unembedding layer, the unigram marginal probability $P(k)$ directly influences the required magnitude of logits.

To assign higher baseline probability to high-frequency tokens (such as articles, prepositions, and punctuation) across diverse contexts, the cross-entropy loss naturally encourages the output vectors $w_{\text{out}, k}$ of frequent tokens to develop larger $L_2$ norms $\|w_{\text{out}, k}\|_2$.

When weights are tied, these inflated norms propagate directly into the input embeddings. As a result:

• High-frequency tokens enter the Layer 0 residual stream with disproportionately large vector magnitudes.
• Low-frequency tokens enter with small magnitudes.
• This creates an artificial activation imbalance that the initial LayerNorm or RMSNorm blocks must continually rescale.

3. Representation Degeneration and Anisotropy

As analyzed by Gao et al. (2019) in their study of representation degeneration, trained word embeddings often collapse into a narrow convex cone (high anisotropy).

In tied models, this degeneration is magnified: the push-pull dynamics of cross-entropy push output vectors in specific directional clusters to optimize next-token separation, which in turn reduces the angular diversity of the input space, degrading the expressivity of early attention layers.

Gradient Dynamics and Optimization Interference

The core mechanical drawback of weight tying lies in how gradients flow during backpropagation.

Consider the token-level cross-entropy loss $\mathcal{L}_t = -\log P(x_{t+1} = y_t \mid x_{\le t})$, where $y_t$ is the ground-truth target token at step $t$:

$$\mathcal{L}_t = -z_{t, y_t} + \log \sum_{j=1}^{|V|} \exp(z_{t, j})$$

Untied Gradient Flow

In an untied model, the gradient with respect to row $k$ of the output matrix $W_{\text{out}}$ is:

$$\frac{\partial \mathcal{L}_t}{\partial w_{\text{out}, k}} = \left( P(x_{t+1} = k \mid x_{\le t}) - \mathbb{I}(k = y_t) \right) \hat{h}_t$$

Every row $k \in \{1, \dots, |V|\}$ in $W_{\text{out}}$ receives an update at every time step. For the target token ($k = y_t$), the vector is pulled toward $\hat{h}_t$. For all non-target tokens ($k \neq y_t$), the vector is pushed away proportionally to its predicted probability.

Conversely, the gradient with respect to row $k$ of the input matrix $W_{\text{in}}$ is:

$$\frac{\partial \mathcal{L}_t}{\partial w_{\text{in}, k}} = \frac{\partial \mathcal{L}_t}{\partial e_t} \cdot \mathbb{I}(k = x_t)$$

Only the single row corresponding to the active input token $x_t$ is updated, guided solely by the downstream backpropagation through the transformer layers.

Tied Gradient Flow

In a tied model ($W = W_{\text{in}} = W_{\text{out}}$), the total gradient for row $k$ is the sum of both contributions:

$$\frac{\partial \mathcal{L}_t}{\partial w_k} = \underbrace{\left( P(x_{t+1} = k \mid x_{\le t}) - \mathbb{I}(k = y_t) \right) \hat{h}_t}_{\text{Output Classification Gradient}} + \underbrace{\frac{\partial \mathcal{L}_t}{\partial e_t} \cdot \mathbb{I}(k = x_t)}_{\text{Input Feature Gradient}}$$

This summation causes optimization interference:

1. The input representation for token $k$ is continuously modified by negative repulsive forces from output classification steps where $k$ was merely an unselected candidate in the softmax denominator.
2. The directional trajectory needed for optimal linear discrimination at layer $L$ frequently opposes the trajectory needed for optimal key-value retrieval at layer 1.

Modern Scaling and Architectural Choices

In the era of large-scale pre-training, the decision to tie or untie embeddings depends directly on model parameter scale, hidden dimension $d_{\text{model}}$, and vocabulary size $|V|$.

Vocabulary Size vs. Parameter Footprint

Modern tokenizers have grown significantly to accommodate multilingual text, code, and mathematical notation:

• LLaMA 1 / 2: $|V| = 32,000$
• Mistral 7B: $|V| = 32,768$
• LLaMA 3 / 3.1 / 3.2: $|V| = 128,256$
• Gemma 1 / 2: $|V| = 256,000$
• Qwen 2.5: $|V| = 152,064$

The parameter count for a single embedding matrix is $|V| \times d_{\text{model}}$.

| Model | Hidden Dim ($d_{\text{model}}$) | Vocab Size ($|V|$) | Single Matrix Params | Tied / Untied | Embedding % of Total Params | | :--- | :--- | :--- | :--- | :--- | :--- | | Gemma 2B | 2,048 | 256,000 | 524M | Tied | ~20.9% (Tied) vs ~34.6% (if Untied) | | Gemma 2 9B | 3,584 | 256,000 | 917M | Tied | ~9.9% (Tied) vs ~18.0% (if Untied) | | SmolLM 135M | 576 | 49,152 | 28.3M | Tied | ~20.9% (Tied) | | LLaMA 3 8B | 4,096 | 128,256 | 525M | Untied | ~13.1% (Both matrices: 1.05B total) | | LLaMA 3 70B | 8,192 | 128,256 | 1.05B | Untied | ~3.0% (Both matrices: 2.10B total) | | Mistral 7B | 4,096 | 32,768 | 134M | Untied | ~3.7% (Both matrices: 268M total) | | DeepSeek-V3 | 7,168 | 129,280 | 926M | Untied | <0.3% of 671B total |

Why Compact Models Keep Weight Tying

For compact models (under 3 billion parameters) utilizing large vocabularies, untying embeddings introduces severe overhead:

• In Gemma 2B, maintaining untied embeddings would require $2 \times 524\text{M} = 1.048\text{B}$ parameters purely in input and output embeddings.
• Over 50% of the non-embedding parameter budget would be consumed by static projection tables rather than self-attention and feed-forward computation layers.
• Under strict hardware deployment constraints (such as running on mobile edge devices with 4 GB RAM), weight tying preserves critical parameter capacity for transformer depth and width.

Why Frontier Models Decouple Embeddings

In models with 8 billion parameters or more, the proportion of parameters occupied by embeddings drops under 15% (and under 3% at 70B+ scales). At this operating point:

1. Representational Freedom: Decoupled matrices allow $W_{\text{in}}$ to specialize in semantic neighborhood clustering while $W_{\text{out}}$ optimizes logit calibration and output classification hyperplanes.
2. Empirical Performance: As shown by Chung et al. (2020) in Rethinking Embedding Coupling in Pre-trained Language Models, untied models systematically achieve lower validation perplexity and higher downstream evaluation scores when total non-embedding parameter count is held constant.
3. Training Stability at Scale: Decoupled embeddings eliminate gradient interference between input representations and softmax outputs, contributing to more stable pre-training dynamics across trillion-token pre-training runs.

Sources

• Using the Output Embedding to Improve Language Models (Press & Wolf, EACL 2017 / arXiv:1608.05859)
• Tying Word Vectors and Word Classifiers: A Loss Framework for Language Modeling (Inan, Khosravi, & Socher, ICLR 2017 / arXiv:1611.01462)
• Attention Is All You Need (Vaswani et al., NeurIPS 2017 / arXiv:1706.03762)
• Rethinking Embedding Coupling in Pre-trained Language Models (Chung et al., ICLR 2021 / arXiv:2010.12821)
• Representation Degeneration Problem in Language Modeling (Gao et al., ICLR 2019 / arXiv:1907.12009)
• The Llama 3 Herd of Models (Dubey et al., 2024 / arXiv:2407.21783)
• Gemma: Open Models Based on Gemini Research and Technology (Gemma Team, Google DeepMind, 2024 / arXiv:2403.08295)