Skip to content

Parameter Sharing & Weight Tying

Overview

Parameter sharing is using the same nn.Parameter object in multiple places of a model. PyTorch's attribute-based discovery means assigning one parameter to several modules (or reassigning it) makes the optimizer update it once — while the forward pass uses it everywhere.

  • One parameter, many forward uses, one gradient accumulation → memory-efficient + enforces structure.
  • Weight tying (e.g., tied embeddings in language models, encoder/decoder sharing) reduces parameter count and can improve generalization.
  • Sharing works because state_dict keys follow one canonical location.

đź’ˇ If you share parameters, serialization saves them once (single key). Load with the same structure or assign manually via strict=False + mapping.

Simple Weight Tying

import torch, torch.nn as nn

class TiedLinear(nn.Module):
    """Y = W @ X computed 'twice' through one parameter."""
    def __init__(self, in_f, out_f):
        super().__init__()
        self.W = nn.Parameter(torch.randn(out_f, in_f) * 0.1)
        self.fc = nn.Linear(out_f, in_f)  # unrelated, for contrast

    def forward(self, x):
        # use W directly twice: memory 1x, updates once
        return torch.mm(x, self.W.t()) * (torch.mm(x, self.W.t()) > 0).float()

m = TiedLinear(8, 8)
print("num params:", sum(p.numel() for p in m.parameters()))

Tied Embeddings (the classic use)

class TiedEmbeddingLM(nn.Module):
    """Embedding and output projection SHARE the weight matrix."""
    def __init__(self, vocab, dim):
        super().__init__()
        self.embed = nn.Embedding(vocab, dim)
        # tie: output projection reuses embed.weight
        self.out_proj = nn.Linear(dim, vocab, bias=False)
        self.out_proj.weight = self.embed.weight       # <-- the tie

    def forward(self, ids):
        return self.out_proj(self.embed(ids))

m = TiedEmbeddingLM(1000, 64)
print("total params:", sum(p.numel() for p in m.parameters()))
print("expected     :", 1000 * 64, "(saved another 64k)")

Verify the tie really is one storage

print("same storage:", m.out_proj.weight.data_ptr() == m.embed.weight.data_ptr())

⚠️ Tying after building the module: set self.out_proj.weight = self.embed.weight — PyTorch re-registers the parameter under out_proj.weight, and state_dict will only store it once under embed.weight? Actually it stores it under BOTH keys pointing at same storage — verify with strict=False on load.

Sharing by Reference (same object, two modules)

class Shared(nn.Module):
    def __init__(self):
        super().__init__()
        w = nn.Parameter(torch.randn(4, 4) * 0.1)
        self.mod_a = nn.Linear(4, 4, bias=False)
        self.mod_b = nn.Linear(4, 4, bias=False)
        self.mod_a.weight = w
        self.mod_b.weight = w     # same object!

s = Shared()
print("ptr equal:", s.mod_a.weight.data_ptr() == s.mod_b.weight.data_ptr())
optim = torch.optim.SGD(s.parameters(), lr=0.1)  # ONE param -> step hits both

Encoder–Decoder Tying (seq2seq)

class Seq2SeqTied(nn.Module):
    def __init__(self, vocab, dim):
        super().__init__()
        self.embed = nn.Embedding(vocab, dim)
        self.encoder = nn.LSTM(dim, dim, batch_first=True)
        self.decoder = nn.LSTM(dim, dim, batch_first=True)
        self.out = nn.Linear(dim, vocab, bias=False)
        self.out.weight = self.embed.weight          # tie the LM head
        # optionally tie decoder input embed too:
        self.dec_embed = self.embed                  # reuse entirely

    def forward(self, src, tgt):
        enc, _ = self.encoder(self.embed(src))
        dec, _ = self.decoder(self.embed(tgt))
        return self.out(dec)

Hooks to See Sharing in Action

seen = {}
for name, p in m.named_parameters():
    seen.setdefault(p.data_ptr(), []).append(name)
for ptr, names in seen.items():
    if len(names) > 1:
        print("SHARED:", names)   # ['embed.weight', 'out_proj.weight']

Pitfalls & Rules

  • Sharing via assignment (a.weight = b.weight) is the safe path; sharing via param.data = ... copies values, not identity.
  • Save/load: strict=False + custom load_state_dict mapping if keys duplicate.
  • Optimizers dedupe automatically, but manual gradient lists (param.grad) will get added twice if you accumulate manually — be aware.
  • torch.compile handles tied weights fine; older exports (TorchScript) may duplicate them — check produced size.

Safe save & load of shared modules

sd = m.state_dict()          # saves shared param under multiple keys w/ same storage
m2 = TiedEmbeddingLM(1000, 64)
m2.load_state_dict(sd, strict=False)   # keys align; tie again after load
m2.out_proj.weight = m2.embed.weight

Comparison: Share vs Clone

Approach Memory Grad behavior Use case
a.weight = b.weight 1Ă— 1 update twice tying
.clone() each 2Ă— independent pretrain-then-finetune
a.load_state_dict(b) 2Ă— independent copies copy weights once

Key Takeaways

  • Sharing = assign the same Parameter object to multiple modules/heads.
  • The optimizer sees one parameter → one update, one memory footprint.
  • Tying embeddings → LM head is the canonical NLP win (saves vocabĂ—dim).
  • Verify with data_ptr() equality; re-tie after load_state_dict.
  • Design sharing before save/load; handle duplicates with strict=False and explicit mapping.