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.

-