PyTorch Weight Tying

Weight tying makes two layers share one weight tensor. The shared tensor is stored once, so the model holds fewer parameters.

Two layers, two tables

An output head is the last layer, and it turns vectors into class scores. The embedding table and the head start separate, and each holds its own weight tensor.

torch.manual_seed(0)
embedding = nn.Embedding(20, 8)
head = nn.Linear(8, 20)
print("Embedding:", sum(p.numel() for p in embedding.parameters()))
print("Head:", sum(p.numel() for p in head.parameters()))
print("Same tensor:", head.weight is embedding.weight)

Tie the head to the table

Tying is one assignment of a weight from one layer to another. The two shapes print equal, which the assignment requires.

head.weight = embedding.weight
print("Same tensor:", head.weight is embedding.weight)
print("Head weight shape:", tuple(head.weight.shape))
print("Embedding weight shape:", tuple(embedding.weight.shape))

One tensor, two jobs

Editing the shared tensor changes both layers at once. Row 3 of the table is row 3 of the head weight.

print("Row 3 equal:", torch.equal(embedding.weight[3], head.weight[3]))
with torch.no_grad():
    embedding.weight[3].fill_(1.0)
print("Head row 3 sum after the edit:", head.weight[3].sum().item())

Example

import torch
from torch import nn


class WordModel(nn.Module):
    def __init__(self, tie):
        super().__init__()
        self.embedding = nn.Embedding(20, 8)
        self.head = nn.Linear(8, 20)
        if tie:
            self.head.weight = self.embedding.weight

    def forward(self, index):
        return self.head(self.embedding(index))


torch.manual_seed(0)
untied = WordModel(tie=False)
torch.manual_seed(0)
tied = WordModel(tie=True)

untied_total = sum(p.numel() for p in untied.parameters())
tied_total = sum(p.numel() for p in tied.parameters())

print("Untied total:", untied_total)
print("Tied total:", tied_total)
print("Saved by tying:", untied_total - tied_total)
print("Same tensor:", tied.head.weight is tied.embedding.weight)

Key points

Weight tying: two layers sharing one weight tensor
Use head.weight = embedding.weight to tie a head to an embedding table
Shape rule: the two weights must have the same shape before tying
A shared tensor is stored once, so parameters() counts it once
One tensor does two jobs, as a lookup table and as a scoring weight

In this example, two WordModel models are built with the same seed, one tied and one not. Each model holds an embedding table of 20 words at width 8 and a head of the same shape. The untied model counts 340 parameters across both layers. The tied model counts 180, because the shared weight is stored once. Tying saves 160 parameters, the full size of one weight tensor. The last line confirms the head weight and the embedding weight are the same tensor. The Build an LLM from Scratch course ties its language model head to its embedding table in this same way.

Note

The head bias is not tied, because the embedding table has no bias to share

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch
from torch import nn


class WordModel(nn.Module):
    def __init__(self, tie):
        super().__init__()
        self.embedding = nn.Embedding(20, 8)
        self.head = nn.Linear(8, 20)
        if tie:
            self.head.weight = self.embedding.weight

    def forward(self, index):
        return self.head(self.embedding(index))


torch.manual_seed(0)
untied = WordModel(tie=False)
torch.manual_seed(0)
tied = WordModel(tie=True)

untied_total = sum(p.numel() for p in untied.parameters())
tied_total = sum(p.numel() for p in tied.parameters())

print("Untied total:", untied_total)
print("Tied total:", tied_total)
print("Saved by tying:", untied_total - tied_total)
print("Same tensor:", tied.head.weight is tied.embedding.weight)

Exercise

A model ties a (20, 8) embedding table to its head. How many parameters does tying save?

20 8 160 180