Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Weight Tying › The Weight Tying Idea

Michael BrenndoerferUpdated June 17, 20256 min read

Language AI HandbookAll chapters

Part 16 of 67Transformer Architectures

Chapter 115Weight Tying › The Weight Tying Idea

Reading section 3 of 15 · The Weight Tying Idea · 6 minSection 3 of 15 · 6 min

Section 3 of 15 · 6 min read

Weight Tying › The Weight Tying IdeaOne section · Sections in this chapter3 / 15 ↓

Sections in this chapter15

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

We have established that language models maintain two large matrices with identical shapes: one for reading tokens and one for generating them. Both have dimensions $V \times d$ , where each row represents one vocabulary token. This structural symmetry hints at a deeper connection worth pursuing carefully.

Think about what information these matrices encode. The input embedding learns "what does this token mean when I read it?" The output projection learns "what hidden state pattern should produce this token?" But at a deeper level, these are really two perspectives on the same question: what is the semantic identity of this token within the model's learned representation space? A token's meaning should not depend on whether you are reading it or generating it.

Think of it this way. If you learned the word "melancholy" in school, you would use the same mental concept whether you are recognizing that word in a sentence you are reading or deciding to write it in a sentence you are composing. You do not maintain one neural representation of "melancholy" for reading and a completely separate representation for writing. You have one underlying concept, and it serves both purposes. Weight tying imposes the same discipline on language models.

a model's representation of any token should be internally consistent: the vector that captures "what this token means as input" should be the same vector that is "the target to match when generating this token as output." Any discrepancy between these two representations is wasted degrees of freedom at best, and incoherence at worst.

The Core FormulaLink Copied

Weight tying formalizes this intuition by collapsing both matrices into one:

\mathbf{W}_{\text{out}} = \mathbf{E}

where:

$\mathbf{W}_{\text{out}}$ : the output projection matrix, now set equal to $\mathbf{E}$ rather than being an independent learned parameter
$\mathbf{E}$ : the input embedding matrix of shape $V \times d$ , which now serves both roles simultaneously

This single equation halves the embedding-related parameters. But what does it mean computationally? Let's trace through the math to understand how the shared matrix serves both roles in a coherent way.

From Intuition to FormulaLink Copied

When a token enters the model, we look up its embedding. This operation is unchanged by tying:

\mathbf{h}_{\text{in}} = \mathbf{E}[t]

The token at index $t$ retrieves row $t$ from the matrix, a simple table lookup. Training updates this row to encode useful information about token $t$ for downstream processing.

When the transformer produces its final hidden state $\mathbf{h}_{\text{out}}$ , we convert this $d$ -dimensional vector into scores for all $V$ vocabulary tokens. With weight tying, we use the same matrix for this conversion:

\mathbf{z} = \mathbf{E} \cdot \mathbf{h}_{\text{out}}

This matrix multiplication produces a vector of $V$ logits. But what is happening geometrically inside this operation? Each logit is a dot product between the hidden state and one row of the embedding matrix.

The Dot Product as SimilarityLink Copied

The logit for token $i$ decomposes as follows. We need a formula that takes the hidden state $\mathbf{h}_{\text{out}}$ and produces a score reflecting how likely token $i$ is to be the correct next word. With tied weights, we compute:

z_i = \mathbf{e}_i \cdot \mathbf{h}_{\text{out}} = \sum_{j=1}^{d} e_{i,j} \cdot h_j

where:

$z_i$ : the logit (unnormalized score) for token $i$ , a single real number
$\mathbf{e}_i$ : the embedding vector for token $i$ (row $i$ of $\mathbf{E}$ ), a $d$ -dimensional vector with components $e_{i,1}, e_{i,2}, \ldots, e_{i,d}$
$\mathbf{h}_{\text{out}}$ : the final hidden state vector, with $d$ components $h_1, h_2, \ldots, h_d$
$d$ : the embedding dimension

Why does this formula make sense? Notice that the dot product measures geometric alignment between two vectors. Two vectors pointing in the same direction in the $d$ -dimensional embedding space yield a large positive value. Orthogonal vectors yield zero. Vectors pointing in opposite directions yield a large negative value. Softmax then converts these alignment scores into probabilities, so the token whose embedding most closely aligns with the hidden state receives the highest probability.

This creates a beautiful geometric interpretation: the transformer's job is to produce a hidden state that points in the direction of the correct next token's embedding. When processing the context "The cat sat on the ___", the transformer should generate a hidden state that is close to the embedding vectors for words like "mat," "floor," or "couch," and far from the embedding vectors for words like "quantum" or "bureaucracy."

With tied weights, the embedding vectors that serve as targets for this similarity search are exactly the same vectors used to represent tokens in the input. The model must learn one coherent embedding for each token that works both as an input representation and as an output target. This constraint is not a restriction; it is a clarification that forces the model to build consistent semantics.

Implementation Note

In PyTorch, the embedding matrix has shape $(V, d)$ and hidden states typically have shape $(B, S, d)$ where $B$ is the batch size and $S$ is the sequence length. To compute logits for a full batch of hidden states, we use hidden_states @ embedding.weight.T, which performs a batched matrix multiplication. The result has shape $(B, S, V)$ , giving one logit per token per position. This single operation replaces what would otherwise require a separate nn.Linear(hidden_dim, vocab_size, bias=False) layer.

ImplementationLink Copied

Translating this math into code is straightforward. We create a single embedding matrix and use it for both encoding (lookup) and decoding (matrix multiplication):

In[5]:

Code

import torch
import torch.nn as nn


class TiedEmbeddings(nn.Module):
    """Language model head with tied input and output embeddings."""

    def __init__(self, vocab_size, hidden_dim):
        super().__init__()
        # Single shared embedding matrix
        self.embedding = nn.Embedding(vocab_size, hidden_dim)

    def encode(self, token_ids):
        """Convert token IDs to embeddings."""
        return self.embedding(token_ids)

    def decode(self, hidden_states):
        """Project hidden states to vocabulary logits."""
        # Use embedding weight matrix transposed
        return torch.matmul(hidden_states, self.embedding.weight.T)


# Create tied embeddings
tied_model = TiedEmbeddings(vocab_size, hidden_dim)
tied_params = vocab_size * hidden_dim

Out[6]:

Console

Parameter count with weight tying:
  Shared embedding:     38,400,000 parameters

Parameter reduction: 38,400,000 parameters (50% savings)
We eliminated one full 50,000 x 768 matrix (38.4M parameters).

The parameter savings are substantial. We eliminated one entire $V \times d$ matrix, cutting embedding-related parameters in half. For models with large vocabularies or smaller hidden dimensions, this reduction represents a significant fraction of total model size. More importantly, as we will see shortly, this reduction often comes with improved performance, not degraded performance.