Learn AI Create the Future — with Pranay
Main site

Inside a Transformer Block

Tied input and output embeddings

Use the same table to turn words into numbers on the way in and numbers back into words on the way out, which saves a large slice of a small model's parameters.

Read these first

On this page 5
  1. What the two tables are
  2. Why sharing them makes sense
  3. Why big models often stop doing it
  4. Remember this
  5. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

Tied embeddings mean one table is used twice. Once to look words up on the way in, and once on the way out.

Think of a two-way dictionary — the kind with English and Hindi in one volume. You can go from a word to its meaning. You can go from a meaning back to a word. Same book. Nobody prints two separate volumes for the two directions.

A language model has exactly these two lookups. On the way in it needs "what numbers represent the word chai?". On the way out it needs "given these numbers, how much do they point at chai?".

Tying is the decision to use one table for both.

What the two tables are

Coming in, the model holds a big table with one row per word in its vocabulary. Look up your word, take that row, and those numbers are what enters the model.

Going out, the model has to produce a score for every word in the vocabulary. It compares its final answer against a row for each word. The closest match scores highest.

Both tables are the same shape. One row per word, one column per slot. So the question is whether they should hold the same numbers.

Why sharing them makes sense

There is an argument from meaning, and it is a good one. Two words that mean similar things should have similar rows on the way in. The model then treats them alike. They should also have similar rows on the way out, so the model considers both when guessing. The two tables want the same structure.

There is a stronger argument from size. For a small model, these tables are enormous.

Take a vocabulary of fifty thousand words and a width of 768 slots. That is about thirty-eight million numbers, per table. Say the rest of the model is around eighty-five million. Two separate tables push the total to a hundred and sixty-three million. Tying brings it to a hundred and twenty-four million.

That is a quarter of the model, saved by using one book instead of two.

Why big models often stop doing it

Here is the part that explains why practice changed.

The vocabulary table's size depends on the vocabulary and the width. Nothing else. It does not grow when you add layers.

The rest of the model grows with depth and with the square of the width. So as models get bigger, the table shrinks as a share of the total.

At around a hundred million parameters, tying saves a quarter of everything. At eight billion, it saves under a tenth. The saving stops being worth the constraint, and larger models increasingly leave the two tables separate.

So the rule of thumb is: small models tie, large models often do not. And that is exactly what you find in real configurations.

Remember this

  • One table serves both directions: word to numbers going in, numbers to word coming out.
  • For small models this saves a large share of the total size.
  • The saving shrinks as the model grows, which is why large models frequently untie them.

What to learn next

Developer — Code and libraries.

Setup

bash
pip install torch

Tying, and the arithmetic of when it is worth it

tied.py
import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)
V, d = 1000, 64

class TinyLM(nn.Module):
    def __init__(self, tie):
        super().__init__()
        self.embed = nn.Embedding(V, d)
        self.body  = nn.Linear(d, d)                 # stands in for the transformer blocks
        self.head  = nn.Linear(d, V, bias=False)
        if tie:
            self.head.weight = self.embed.weight     # one tensor, used twice
    def forward(self, idx):
        return self.head(torch.tanh(self.body(self.embed(idx))))

for tie in (False, True):
    m = TinyLM(tie)
    unique = {id(p): p for p in m.parameters()}      # dedupe: a shared tensor counts once
    print(f"tied={str(tie):5s} total={sum(p.numel() for p in unique.values()):>7,} "
          f"same object? {m.head.weight is m.embed.weight}")

m = TinyLM(True)
loss = F.cross_entropy(m(torch.tensor([[3, 7]])).reshape(-1, V), torch.tensor([7, 3]))
loss.backward()
print("\nwith tying, one gradient tensor collects updates from BOTH uses:")
print("  number of distinct parameter tensors:", len({id(p) for p in m.parameters()}))
print("  grad on the shared matrix, shape:", tuple(m.embed.weight.grad.shape))
print("  rows that received gradient:", int((m.embed.weight.grad.abs().sum(1) > 0).sum()),
      "out of", V)
print("  -> rows 3 and 7 got gradient from the input side; every row got it from the output side")

print("\nhow much does tying save? exact numbers from two real configurations.")
print(f"{'model':>12} {'everything but the vocab':>26} {'V x d':>15} "
      f"{'untied total':>15} {'tied total':>15} {'saved':>7}")
for name, body, V_, d_ in (("GPT-2 small",   85_842_432,  50_257,   768),
                           ("Llama 3 8B", 6_979_588_096, 128_256, 4_096)):
    vocab = V_ * d_
    untied, tied = body + 2 * vocab, body + vocab
    print(f"{name:>12} {body:>26,} {vocab:>15,} {untied:>15,} {tied:>15,} "
          f"{1 - tied/untied:>6.1%}")
print("  (GPT-2 ships tied: 124,439,808 parameters. Llama 3 8B ships untied: 8,030,261,248.)")
print("  -> the body grows with depth and width. the vocabulary side does not.")
Output
tied=False total=132,160 same object? False
tied=True  total= 68,160 same object? True

with tying, one gradient tensor collects updates from BOTH uses:
  number of distinct parameter tensors: 3
  grad on the shared matrix, shape: (1000, 64)
  rows that received gradient: 1000 out of 1000
  -> rows 3 and 7 got gradient from the input side; every row got it from the output side

how much does tying save? exact numbers from two real configurations.
       model   everything but the vocab           V x d    untied total      tied total   saved
 GPT-2 small                 85,842,432      38,597,376     163,037,184     124,439,808  23.7%
  Llama 3 8B              6,979,588,096     525,336,576   8,030,261,248   7,504,924,672   6.5%
  (GPT-2 ships tied: 124,439,808 parameters. Llama 3 8B ships untied: 8,030,261,248.)
  -> the body grows with depth and width. the vocabulary side does not.

Three things worth noticing

The tied total matches the published GPT-2 count exactly. 124,439,808, which is what you get from sum(p.numel() for p in GPT2LMHeadModel(GPT2Config()).parameters()). The untied count in the same table, 8,030,261,248, is Llama 3 8B's published size. Neither number is an estimate.

23.7 percent against 6.5 percent. That is the entire argument, in two numbers. Same technique, same reasoning, and the value falls by nearly four times as the model grows. For a 350M-parameter model with a 256k vocabulary — the Gemma shape — the share is larger still.

Every row of the shared matrix received gradient, not only rows 3 and 7. The input lookup touches two rows. The output projection touches all thousand, because cross-entropy assigns a gradient to every logit. A tied matrix therefore gets a dense output-side gradient added to a sparse input-side one, every step. That difference in gradient statistics is a real consideration. It interacts with weight decay and optimiser state in ways that are easy to get wrong.

Doing it correctly

python
self.head.weight = self.embed.weight     # correct: same tensor object

Not this:

python
self.head.weight.data = self.embed.weight.data.clone()   # wrong: a copy that drifts apart

The first line makes both modules point at one nn.Parameter. The optimiser sees it once, and updates apply to both uses. The second creates a copy that is equal at step zero and diverges immediately.

Check with model.head.weight is model.embed.weight. In Hugging Face models the setting is tie_word_embeddings in the config, applied by tie_weights().

Details people trip over

nn.Embedding stores (vocab, d) and nn.Linear(d, vocab) stores (vocab, d) too. Their shapes match without a transpose, because nn.Linear holds its weight transposed relative to how it is applied. That is why the assignment above works directly.

Saving and loading. A tied checkpoint contains one copy of the matrix. Loading it into an untied model leaves the head randomly initialised. The warning is easy to miss. Loading an untied checkpoint into a tied model silently discards one of the two matrices.

Quantization. Input lookup and output projection have different sensitivity to precision loss. Tied weights must satisfy both, so quantizing a tied model can be harder than quantizing an untied one. See quantization in practice.

Fine-tuning that grows the vocabulary. Adding tokens means resizing the embedding. With tying, the head must be re-tied afterwards. resize_token_embeddings handles this; hand-written code often does not.

When to tie

Tie when the vocabulary times the width is a large share of the total. That means small models, large vocabularies, and on-device deployments.

Do not tie when that share is small and you have the budget. Also skip it when input and output want different geometry. Note that a large final logit matrix is also where several practical tricks live, and tying constrains them.

Measure rather than assume. Compute V * d and compare it against the body size for your configuration. That one division answers the question.

Try it yourself

Set V = 32000 and d = 512, with a body of about twenty million parameters. Compute the tied and untied totals. Then raise the width to 4096 and keep the vocabulary. Recompute the body as roughly twelve times depth times width squared. Watch the saving collapse. That collapse is the whole design story.

What to learn next

Researcher — Mathematics and papers.

The construction

Let $E \in \mathbb{R}^{V \times d}$ be the input embedding matrix and $U \in \mathbb{R}^{V \times d}$ the output projection. Logits for final hidden state $h \in \mathbb{R}^d$ are $z = U h$. Tying sets $U = E$, so

$$ z_v = \langle e_v, h \rangle $$

The next-token distribution becomes a softmax over inner products between the hidden state and the input embeddings. The output head is then a similarity search in embedding space.

The original arguments

Press and Wolf (2017), EACL 2017, arXiv:1608.05859, make a representational observation. The output embedding already learns better word similarity structure than the input embedding. Tying forces the input embedding to inherit that structure. They report perplexity improvements and note the regularising effect of halving the parameters in the largest single component.

Inan et al. (2017), arXiv:1611.01462, derive the same construction from an augmented loss. They frame the language model as estimating a distribution. Tying then arises naturally when the target is smoothed toward words similar to the true one.

Both papers were written when models were small enough that $Vd$ dominated. That context matters for how the results transfer.

Why the benefit is guaranteed to shrink

Non-embedding parameters scale approximately as $12 L d^2$. The vocabulary side scales as $Vd$ per table. The tied fraction of the total is

$$ \frac{Vd}{12 L d^2 + 2Vd} $$

$V$ is roughly fixed by tokenizer design, at 32k to 256k across current models. $L$ and $d$ both grow with scale. The ratio therefore falls monotonically. The measured values above are that formula evaluated at two points. They are 23.7 percent for GPT-2 small and 6.5 percent for Llama 3 8B.

There is no scale at which tying becomes harmful on parameter-count grounds. It becomes less useful, and at some point the constraint on representation is no longer worth the diminishing saving.

The constraint it imposes

Tying forces one matrix to serve two functions with different requirements.

  • The input side needs $e_v$ to be a good starting representation of token $v$ in the residual stream.
  • The output side needs $e_v$ to be a good discriminative direction against all other tokens. That is judged given a final hidden state.

Chung et al. (2021), arXiv:2010.12821, decouple the two in an encoder setting. Untied input embeddings can be made much smaller than $d$ via a factorised projection. Output embeddings stay full-rank, and parameter efficiency beats tying. ALBERT (Lan et al., 2020, arXiv:1909.11942) uses the same factorisation argument.

There is also a documented representational effect. Tying pulls token embeddings toward the output space. That is not self-evidently desirable, and it is an active subject of study.

The softmax bottleneck

Here $z = Eh$, with $E$ fixed across contexts. The matrix of log-probabilities over all contexts has rank at most $d + 1$. Yang et al. (2018), Breaking the Softmax Bottleneck, arXiv:1711.03953, argue the true matrix has higher rank. Any model of this form is therefore misspecified. They propose a mixture of softmaxes.

Tying does not create this bottleneck — an untied $U$ has the same rank constraint. But it removes the freedom to choose $U$ independently. That is one of the few available levers on the problem.

Current practice

ModelVocabularyWidthTied
GPT-2 small50,257768yes
Llama 3.2 1B128,2562,048yes
Llama 3 8B128,2564,096no

The split follows scale, and the reported reasoning follows the parameter-share argument rather than any measured quality difference. Controlled ablations at large scale are essentially absent from public technical reports. The honest position is that a clear budget calculation drives the practice, on a thin quality evidence base.

Papers

What to learn next