Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Weight Tying › A Worked Example

Michael BrenndoerferUpdated June 17, 20257 min read

Language AI HandbookAll chapters

Part 16 of 67Transformer Architectures

Chapter 115Weight Tying › A Worked Example

Reading section 6 of 15 · A Worked Example · 7 minSection 6 of 15 · 7 min

Section 6 of 15 · 7 min read

Weight Tying › A Worked ExampleOne section · Sections in this chapter6 / 15 ↓

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

The formulas above describe weight tying abstractly, but seeing the mechanism in action makes it concrete. Let's build a tiny language model and trace through exactly how tied embeddings convert a hidden state into token probabilities. We will work with small enough numbers to inspect every value manually.

Setting Up a Toy VocabularyLink Copied

We will work with a vocabulary of just seven words using 4-dimensional embeddings. These small numbers let us inspect every value and understand exactly what is happening at each step. In a real model you might have 50,000 words and 768 dimensions, making manual inspection impractical, but the same mathematical operations apply at every scale.

In[11]:

Code

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

# Create a tiny vocabulary for illustration
tiny_vocab = ["the", "cat", "sat", "on", "mat", "dog", "ran"]
tiny_size = len(tiny_vocab)
tiny_dim = 4

# Initialize with small, interpretable embeddings


# Create tied embedding model
class TinyLM(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(tiny_size, tiny_dim)
        # Initialize with small random values
        nn.init.uniform_(self.embedding.weight, -0.5, 0.5)

    def forward(self, hidden_state):
        """Get probabilities for next token given hidden state."""
        logits = torch.matmul(hidden_state, self.embedding.weight.T)
        return F.softmax(logits, dim=-1)


model = TinyLM()

Out[12]:

Console

Vocabulary: ['the', 'cat', 'sat', 'on', 'mat', 'dog', 'ran']

Embedding matrix (7 tokens x 4 dims):
---------------------------------------------
  the   : [-0.366, -0.499, +0.362, -0.212]
  cat   : [-0.203, -0.113, -0.029, -0.206]
  sat   : [+0.114, -0.049, -0.306, -0.257]
  on    : [+0.264, +0.247, -0.104, +0.128]
  mat   : [+0.276, -0.166, +0.333, -0.484]
  dog   : [-0.399, -0.407, +0.308, +0.145]
  ran   : [-0.390, -0.450, -0.434, +0.312]

Each word has a 4-dimensional embedding vector. These are randomly initialized, as they would be before training. In a trained model, semantically similar words would have similar embeddings. The words "cat" and "dog" would cluster together because they are both animals that appear in similar grammatical contexts. The words "sat" and "ran" would form another cluster as verbs. Weight tying means that these same clusters serve as targets for the output prediction: when the model wants to generate "cat," it aims to produce a hidden state that points in the direction of the "cat" embedding vector.

The Prediction MechanismLink Copied

Now comes the key insight. Suppose the transformer has processed some context and produced a final hidden state. With weight tying, we predict the next token by computing dot products between this hidden state and every embedding in our vocabulary.

Let's simulate this. We'll create a hidden state that is similar to the "cat" embedding, as if the model were about to predict "cat" as the next word. This simulates what would happen if the transformer correctly learned to associate some context with the concept of "cat":

In[13]:

Code

# Suppose the transformer outputs a hidden state similar to "cat"'s embedding
cat_embedding = model.embedding.weight[1].detach()  # Index 1 = "cat"

# Move from the other-token centroid toward "cat" so the intended token wins.
# This fixed construction keeps every render variant identical without reseeding.
other_embeddings = torch.cat(
    [model.embedding.weight[:1], model.embedding.weight[2:]], dim=0
).detach()
other_centroid = other_embeddings.mean(dim=0)
hidden_state = cat_embedding + 2.0 * (cat_embedding - other_centroid)
hidden_state = hidden_state.unsqueeze(0)  # Add batch dimension

# Get next token probabilities
probs = model(hidden_state).squeeze()
assert probs.argmax().item() == 1, (
    "The pedagogical hidden state must predict 'cat'."
)

The hidden state is intentionally constructed to align with "cat"'s embedding while moving away from the average of the other token embeddings. In a real model, the transformer layers would produce this hidden state based on the input context. The contrastive offset represents the fact that the model's output is not a perfect copy of any single embedding; it is a point in the embedding space that aligns more strongly with some tokens than others.

Tracing the Dot Products Step by StepLink Copied

Now let's examine what happens inside the forward pass. For each vocabulary token, we compute the dot product between the hidden state and that token's embedding. This gives us a raw score, which softmax then converts into a probability:

Out[14]:

Console

Hidden state (similar to 'cat' embedding):
  [-0.4412608   0.10209379 -0.14097796 -0.49593773]

Dot products with each token embedding:
--------------------------------------------------
  the   : dot = +0.1648  =>  P = 0.1576
  cat   : dot = +0.1844  =>  P = 0.1607
  sat   : dot = +0.1155  =>  P = 0.1500
  on    : dot = -0.1397  =>  P = 0.1162
  mat   : dot = +0.0545  =>  P = 0.1411
  dog   : dot = +0.0190  =>  P = 0.1362
  ran   : dot = +0.0325  =>  P = 0.1381

Most likely next token: 'cat'
Probability assigned:   0.1607

All probabilities sum to: 1.000000

The dot product reveals the alignment between the hidden state and each embedding. "Cat" receives the highest score because its embedding is most similar to the hidden state we constructed. After softmax normalization, this translates to the highest probability. The other tokens receive lower but non-zero probabilities. This reflects the fact that uncertainty exists about the correct next token.

This is the heart of weight tying: the same embedding that represents "cat" for input processing is also the target pattern the model aims to produce when generating "cat". The transformer's job is to transform the input context into a hidden state that matches the correct next token's embedding. Learning happens both by improving the embedding representations and by improving the transformer's ability to produce hidden states that match those representations.

Numerical Walkthrough of SoftmaxLink Copied

Let us trace the complete computation from logits to probabilities for the first three tokens to make the math fully concrete. Suppose the hidden state produces logits of $z = [0.8, 1.4, 0.3, 0.1, 0.5, 0.2, 0.4]$ for our seven tokens. The softmax computation proceeds:

\begin{aligned} \exp(z_{\text{the}}) &= \exp(0.8) \approx 2.226 \\ \exp(z_{\text{cat}}) &= \exp(1.4) \approx 4.055 \\ \exp(z_{\text{sat}}) &= \exp(0.3) \approx 1.350 \\ &\vdots \end{aligned}

The normalizing denominator sums all seven exponentials, and each token's probability is its exponential divided by this sum. The token with the highest logit, "cat" in this example, receives the highest probability. The key observation is that with tied weights, these logit values are dot products with the input embedding vectors, so improving the input embedding for "cat" simultaneously makes it easier to predict "cat" as an output.

Visualizing Embedding SimilarityLink Copied

To better understand how the dot product works as a similarity measure, let's visualize the pairwise similarities between all embeddings in our tiny vocabulary along with the hidden state:

Out[15]:

Visualization

Heatmap showing dot product similarities between seven token embeddings and a hidden state vector. — Pairwise dot product similarities between token embeddings and the hidden state. The hidden state row shows highest similarity with 'cat', which is expected since we constructed it to be close to the cat embedding. The diagonal tends toward high values because each embedding is maximally aligned with itself. With tied weights, the same geometric structure that determines input representations also determines output predictions.

The heatmap reveals the geometric structure of our embedding space. The bottom row and rightmost column show how similar the hidden state is to each vocabulary embedding. Since we constructed the hidden state to resemble "cat," that cell shows the highest similarity in the hidden row. The diagonal entries reflect each token's self-similarity. Off-diagonal entries show how similar different tokens are to each other in the learned representation space.

With tied weights, this same similarity matrix governs both how tokens are processed as input and how they are selected as output. Making "cat" more recognizable as an input (by moving its embedding to a more distinctive location in the space) simultaneously makes it easier to generate "cat" as an output (by making the target point more distinct from other tokens).