Weight Tying 🔗

Table of Contents

• 0. Intuition
• 1. Motivation: Parameter Dominance at the Vocabulary Boundary
  • 1.1 Where Parameters Live
  • 1.2 The Symmetry Observation
• 2. The Technique: Formal Definition
  • 2.1 Notation
  • 2.2 The Standard (Untied) Forward Pass
  • 2.3 Weight Tying Defined
  • 2.4 Parameter Savings
• 3. Why It Works: Geometry and Regularization
  • 3.1 The Semantic Space Argument
  • 3.2 Implicit Regularization
• 4. Formal Gradient Analysis
  • 4.1 Setup: Loss and Computational Graph
  • 4.2 Gradient from the Output Projection Path
  • 4.3 Gradient from the Embedding Lookup Path
  • 4.4 The Combined Gradient
  • 4.5 Gradient Imbalance and Its Consequences
• 5. When Weight Tying Hurts
  • 5.1 Large Vocabularies at Scale
  • 5.2 Structural Mismatch Between Input and Output
  • 5.3 Multilingual and Heavily Subword-Tokenized Models
  • 5.4 Loss Landscape Perspective
• 6. Variations
  • 6.1 Factored Embeddings (ALBERT)
  • 6.2 Output Projection Scaling
  • 6.3 Cross-Layer Weight Sharing
• References

0. Intuition 🧠

A language model touches the vocabulary at exactly two points in its forward pass: at the very start, when it converts input token IDs into vectors, and at the very end, when it converts its final hidden state back into a probability distribution over tokens. Two big matrices do this work — call them the embedding matrix and the unembedding matrix.

Weight tying is the observation that these two matrices are doing essentially the same job, just in opposite directions — and so there’s no reason they can’t share the same numbers.

Here’s the intuition. The embedding matrix has to answer: *“What does token* dog mean as a context clue?” The unembedding matrix has to answer: ”How strongly does this hidden state point toward predicting dog?” Both questions are about the same underlying thing — the meaning of dog in the distributional sense. A token that tends to appear after golden retriever (as a prediction target) is exactly the kind of token whose embedding should encode”follows pet-related context” (as an input signal). The two representations are different facets of the same semantic fact.

The rest of this note formalizes this picture, analyzes when it works, when it fails, and how the shared weight is updated during training.

1. Motivation: Parameter Dominance at the Vocabulary Boundary 📐

1.1 Where Parameters Live

A standard autoregressive language model over a vocabulary of size \(V\) with hidden dimension \(d\) and \(L\) transformer layers has the following parameter budget (ignoring biases):

The vocabulary-boundary matrices each cost \(Vd\) parameters. For a typical large language model — say \(V = 32{,}000\), \(d = 4{,}096\), \(L = 32\) — the vocabulary matrices together consume \(2 \times 32{,}000 \times 4{,}096 \approx 262\text{M}\) parameters, while a single transformer layer costs roughly \(4 \times 4{,}096^2 \approx 67\text{M}\). For small models (sub-100M parameters), the two vocabulary matrices can constitute the majority of all parameters.

1.2 The Symmetry Observation

Both \(W_E\) and \(W_U\) mediate between the same two spaces: the vocabulary space \(\mathbb{R}^V\) (or equivalently, the simplex over tokens) and the model’s internal representation space \(\mathbb{R}^d\). Specifically:

• \(W_E \in \mathbb{R}^{V \times d}\): maps a one-hot token encoding to a dense vector in \(\mathbb{R}^d\).
• \(W_U \in \mathbb{R}^{d \times V}\): maps a hidden state in \(\mathbb{R}^d\) to unnormalized logits over \(\mathbb{R}^V\).

The \(i\)-th row of \(W_E\), written \(e_i \in \mathbb{R}^d\), is the input embedding of token \(i\). The \(i\)-th column of \(W_U\), written \(u_i \in \mathbb{R}^d\), is the output embedding of token \(i\): the direction in hidden-state space toward which the model must point to assign high probability to token \(i\).

The observation is that \(e_i\) and \(u_i\) both purport to represent token \(i\) in \(\mathbb{R}^d\). This raises a natural question: must they be distinct?

2. The Technique: Formal Definition 🔑

2.1 Notation

Let: - \(V \in \mathbb{N}\) — vocabulary size. - \(d \in \mathbb{N}\) — model (hidden) dimension. - \(W_E \in \mathbb{R}^{V \times d}\) — the input embedding matrix; row \(i\) gives the embedding of token \(i\). - \(W_U \in \mathbb{R}^{d \times V}\) — the unembedding matrix (output projection); column \(j\) gives the output embedding of token \(j\). - \(h_t \in \mathbb{R}^d\) — the hidden state of the model at position \(t\), output of the final transformer layer. - \(\mathbf{x} = (x_1, \ldots, x_T)\) — input token sequence, with \(x_t \in \{1, \ldots, V\}\).

2.2 The Standard (Untied) Forward Pass

The LM forward pass proceeds in three stages.

Stage 1 — Embedding lookup. For each position \(t\), select row \(x_t\) of \(W_E\):

\[\text{emb}_t = W_E[x_t, :] = e_{x_t} \in \mathbb{R}^d.\]

Stage 2 — Transformer. Apply \(L\) transformer layers to the sequence of embeddings to produce hidden states \(h_1, \ldots, h_T \in \mathbb{R}^d\). This stage is agnostic to the vocabulary matrices.

Stage 3 — Output projection. Compute unnormalized logits and a next-token distribution:

\[\ell_t = W_U^\top h_t \in \mathbb{R}^V, \qquad p_t = \text{softmax}(\ell_t) \in \Delta^{V-1}.\]

Note: \(W_U^\top \in \mathbb{R}^{V \times d}\) acts on \(h_t \in \mathbb{R}^d\) to produce logits in \(\mathbb{R}^V\). Equivalently, the logit for token \(j\) at position \(t\) is the dot product \(\langle u_j, h_t \rangle\) where \(u_j\) is the \(j\)-th column of \(W_U\).

The cross-entropy loss over the sequence is

\[\mathcal{L} = -\frac{1}{T}\sum_{t=1}^{T} \log p_t(x_{t+1}) = -\frac{1}{T}\sum_{t=1}^{T} \log \frac{\exp(\langle u_{x_{t+1}}, h_t \rangle)}{\sum_{v=1}^{V} \exp(\langle u_v, h_t \rangle)}.\]

2.3 Weight Tying Defined

Definition (Weight Tying). Weight tying is the constraint

\[W_U = W_E^\top,\]

i.e., column \(j\) of \(W_U\) equals row \(j\) of \(W_E\): \(u_j = e_j\) for all \(j \in \{1, \ldots, V\}\). Under this constraint, the logit for token \(j\) at position \(t\) becomes

\[\ell_t(j) = \langle e_j, h_t \rangle,\]

and the model is parameterized by the single matrix \(W \in \mathbb{R}^{V \times d}\) serving both roles simultaneously.

The tied forward pass is:

flowchart LR
    tokens["Token IDs
x_1, ..., x_T"] -->|"row select W"| emb["Embeddings
e_{x_t} ∈ R^d"]
    emb --> transformer["Transformer
L layers"]
    transformer --> hidden["Hidden states
h_t ∈ R^d"]
    hidden -->|"W^T h_t"| logits["Logits
ℓ_t ∈ R^V"]
    logits --> softmax["Softmax
p_t ∈ Δ^{V-1}"]
    W["Shared W ∈ R^{V x d}"] --> emb
    W --> logits

2.4 Parameter Savings

Without tying, the model maintains \(W_E\) and \(W_U\) separately, at a total cost of \(2Vd\) parameters for the vocabulary matrices. With tying, only \(W_E\) (or equivalently \(W\)) is stored, at cost \(Vd\).

Exact parameter savings: \(Vd\) parameters eliminated by weight tying.

3. Why It Works: Geometry and Regularization 💡

3.1 The Semantic Space Argument

The core justification for weight tying is geometric. Both \(e_i\) and \(u_i\) are learned representations of token \(i\) in \(\mathbb{R}^d\), and both must encode the token’s semantic content relative to other tokens for the model to function.

• The input embedding \(e_i\) must encode token \(i\) such that the transformer can compute useful attention patterns and feed-forward activations. Syntactically similar tokens (in distributional sense) should have nearby embeddings.

• The output embedding \(u_i\) must encode token \(i\) such that the dot product \(\langle u_i, h_t \rangle\) is large exactly when token \(i\) is the appropriate continuation of context \(h_t\). Again, tokens that appear in similar distributional contexts should have similar \(u_i\).

Both criteria derive from the distributional hypothesis: tokens with similar distributional contexts have similar representations. This is the same hypothesis underlying word2vec, GloVe, and every token embedding method. The natural conjecture — which Press and Wolf (2017) validate empirically — is that the optimal \(e_i\) and \(u_i\) should be close (perhaps identical) for most tokens.

Press and Wolf (2017) demonstrate empirically that in untied models, the output embedding \(W_U\) trained from scratch is close to the transpose of the trained input embedding \(W_E\) in cosine similarity. This confirms that tying is not merely a computational shortcut but reflects a real structural property of the optimization landscape.

3.2 Implicit Regularization

Tying adds a hard equality constraint \(W_U = W_E^\top\) to the optimization. This reduces the model’s degrees of freedom by \(Vd\), which has a regularizing effect. Concretely:

1. Reduced expressivity prevents the model from memorizing idiosyncratic input/output pairs that differ only because of noise in training data.

2. Gradient coupling (formalized in Section 4) means that the embedding of a rarely-seen input token still receives gradient signal from its appearances as a prediction target, and vice versa. This improves the representation of tail tokens.

4. Formal Gradient Analysis 📐

4.1 Setup: Loss and Computational Graph

Fix a single training example: a sequence \((x_1, \ldots, x_T, x_{T+1})\) and the cross-entropy loss

\[\mathcal{L} = -\sum_{t=1}^{T} \log p_t(x_{t+1}),\]

where \(p_t(j) = \text{softmax}(W^\top h_t)(j) = \frac{\exp(\langle e_j, h_t \rangle)}{\sum_{v} \exp(\langle e_v, h_t \rangle)}\).

Under weight tying, the single matrix \(W \in \mathbb{R}^{V \times d}\) (rows indexed by token) participates in the computation graph via two distinct paths:

1. Path A (Embedding Lookup): \(W[x_t, :] = e_{x_t}\) is selected and passed into the transformer.
2. Path B (Output Projection): \(W^\top \cdot h_t = \ell_t\) computes logits over all tokens.

We want \(\nabla_W \mathcal{L}\), which by the chain rule accumulates contributions from both paths.

4.2 Gradient from the Output Projection Path

For a fixed position \(t\), the logit vector \(\ell_t = W^\top h_t\) contributes to \(\mathcal{L}\) through \(p_t = \text{softmax}(\ell_t)\). The gradient of the cross-entropy with respect to \(\ell_t\) is the standard softmax gradient:

\[\frac{\partial \mathcal{L}}{\partial \ell_t} = p_t - \mathbf{1}_{x_{t+1}},\]

where \(\mathbf{1}_{x_{t+1}} \in \mathbb{R}^V\) is the one-hot indicator of the target token. Writing this as \(\delta_t^{\text{out}} = p_t - \mathbf{1}_{x_{t+1}} \in \mathbb{R}^V\), the gradient with respect to \(W\) from Path B at position \(t\) is:

\[\frac{\partial \mathcal{L}}{\partial W}\bigg|_{\text{Path B}, t} = \delta_t^{\text{out}} \otimes h_t = \delta_t^{\text{out}} h_t^\top \in \mathbb{R}^{V \times d}.\]

Here \(\otimes\) denotes the outer product. Explicitly, the gradient for row \(j\) of \(W\) from Path B is:

\[\left.\frac{\partial \mathcal{L}}{\partial w_j}\right|_{\text{Path B}, t} = (p_t(j) - \mathbf{1}[j = x_{t+1}]) \cdot h_t.\]

Summing over positions:

\[\left.\frac{\partial \mathcal{L}}{\partial W}\right|_{\text{Path B}} = \sum_{t=1}^{T} \delta_t^{\text{out}} h_t^\top.\]

4.3 Gradient from the Embedding Lookup Path

The embedding lookup selects row \(x_t\) of \(W\) at each position \(t\), so only rows corresponding to tokens that appear in the input sequence receive gradient through Path A. Let \(\bar{h}_t \in \mathbb{R}^d\) be the upstream gradient flowing back from the transformer into the embedding at position \(t\):

\[\bar{h}_t = \frac{\partial \mathcal{L}}{\partial e_{x_t}} = \frac{\partial \mathcal{L}}{\partial \text{emb}_t}.\]

This upstream gradient depends on the transformer’s internal computation and is in general a non-trivial function of all parameters. For row \(j\) of \(W\), the gradient from Path A is:

\[\left.\frac{\partial \mathcal{L}}{\partial w_j}\right|_{\text{Path A}} = \sum_{t : x_t = j} \bar{h}_t.\]

The sum is over all positions at which token \(j\) appears as input. For a rare token \(j\), this sum may be empty (zero gradient from Path A), while Path B always contributes (token \(j\) appears in the denominator of every softmax computation).

4.4 The Combined Gradient

By the chain rule and linearity of differentiation, the total gradient on row \(j\) of the tied weight matrix is the sum of contributions from both paths:

\[\boxed{\frac{\partial \mathcal{L}}{\partial w_j} = \underbrace{\sum_{t : x_t = j} \bar{h}_t}_{\text{Path A: embedding lookup}} + \underbrace{\sum_{t=1}^{T}(p_t(j) - \mathbf{1}[j = x_{t+1}]) \cdot h_t}_{\text{Path B: output projection}}}\]

Key conclusion: The tied weight row \(w_j\) receives gradient from (A) every input position at which token \(j\) appears, and from (B) every position in the sequence via the softmax residual.

4.5 Gradient Imbalance and Its Consequences

The Path B gradient sums over all \(T\) positions for every token \(j\), because token \(j\) contributes to the softmax denominator (and potentially the numerator) at every step. The Path A gradient sums only over the typically much smaller set \(\{t : x_t = j\}\).

Empirically, Path B dominates throughout training. Iyer et al. (2025) measure gradient norms in tied models and find that output-path gradients account for approximately 70% of the total signal during early training. The tied matrix therefore evolves to resemble an output embedding more than an input embedding — a bias baked in by the asymmetric gradient magnitudes.

5. When Weight Tying Hurts ⚠️

5.1 Large Vocabularies at Scale

The efficiency argument for tying is predicated on the vocabulary matrices constituting a non-trivial fraction of total parameters. As model depth and width increase, the transformer layers dominate:

\[\text{fraction of params in embeddings (tied)} = \frac{Vd}{Vd + 12Ld^2} = \frac{1}{1 + 12Ld/V}.\]

For large \(L\) or large \(d/V\), this fraction vanishes. Iyer et al. (2025) report:

Surprisingly, at billion-parameter scale the efficiency gain from tying largely disappears, while the representational cost (forcing \(e_j = u_j\)) persists. At very large scales, the constraint strictly reduces the model’s capacity in a regime where parameters are cheap.

5.2 Structural Mismatch Between Input and Output

The geometric argument for tying assumes that \(e_j\) and \(u_j\) should be “close.” This fails when the input and output distributions are structurally different.

Example — Pointer networks. A pointer network generates outputs by selecting from the input sequence rather than from a fixed vocabulary. The output distribution is over input positions, not over tokens, so there is no meaningful shared semantic space between the embedding of an input token and the output distribution.

Example — Encoder-only models. In models like BERT that do not have a causal language modeling head, weight tying is not naturally applicable since there is no unembedding matrix in the autoregressive sense.

5.3 Multilingual and Heavily Subword-Tokenized Models

In multilingual models, the vocabulary contains subword pieces from many languages. A subword that is a complete word in one language (e.g., “est” in French) appears in radically different left-context distributions (as an input) versus right-context distributions (as a prediction target) compared to monolingual settings. The asymmetry between \(e_j\) and \(u_j\) can be larger, making the tying constraint more lossy.

5.4 Loss Landscape Perspective

6. Variations 🔧

6.1 Factored Embeddings (ALBERT)

When \(d\) is large, even a single embedding matrix \(W \in \mathbb{R}^{V \times d}\) is expensive. ALBERT (Lan et al., 2020) introduces factored embeddings: instead of learning \(W_E \in \mathbb{R}^{V \times d}\) directly, decompose it as

\[W_E = A B, \quad A \in \mathbb{R}^{V \times k},\; B \in \mathbb{R}^{k \times d}, \quad k \ll d.\]

Parameter count: \(Vk + kd\) versus \(Vd\) for the full matrix.

Savings: The factored approach reduces parameters from \(Vd\) to \(k(V + d)\), for a saving of

\[\Delta_{\text{param}} = Vd - k(V + d).\]

Since \(k \ll \min(V, d)\), the saving is approximately \(Vd\): almost the entire original embedding cost is eliminated.

6.2 Output Projection Scaling

The original Transformer (Vaswani et al., 2017) uses weight tying but multiplies the shared embedding matrix by \(\sqrt{d}\) in the embedding lookup stage:

\[\text{emb}_t = \sqrt{d} \cdot W[x_t, :].\]

The output projection is unchanged: \(\ell_t = W^\top h_t\). The rationale is normalization: if \(W\)’s rows have unit norm (a common initialization), then \(e_{x_t}\) has \(\ell_2\) norm 1, which is small relative to the positional encoding vectors (which have entries of order 1). Scaling by \(\sqrt{d}\) brings the embedding magnitude to \(\mathcal{O}(\sqrt{d})\), comparable to sinusoidal positional encodings.

6.3 Cross-Layer Weight Sharing

ALBERT also introduces a distinct form of weight sharing: tying the parameters of all transformer layers. Every layer uses the same attention weight matrices and FFN weights. This reduces the transformer parameter count from \(O(L d^2)\) to \(O(d^2)\), independent of depth.

This is a different technique from embedding weight tying but is often combined with factored embeddings in ALBERT. The motivation is similar: exploit structural redundancy in the learned representations.

References