Skip to content

← nn.Module Mechanics step 4 of 7

Medium Primitives

One matrix in two places

Language models usually share the embedding matrix with the output projection. The same matrix maps tokens to vectors on the way in and vectors to logits on the way out.

self.embed = nn.Embedding(vocab, dim)
self.out = nn.Linear(dim, vocab, bias=False)
self.out.weight = self.embed.weight        # one tensor, two owners

This halves the parameters of a small model and usually improves it: the two matrices are learning the same relationship from opposite directions, so sharing is a genuine prior rather than only a saving.

What sharing actually means

Assignment, not copying. Both attributes point at one Parameter object. So:

  • parameters() deduplicates. A model with two tied Linear weights has three parameter tensors, not four, which is what the optimiser sees.
  • Gradients from both uses accumulate into the one .grad, which is exactly right and is the accumulation rule doing useful work.
  • state_dict() contains both keys, both mapping to the same tensor.

Where it breaks

Copying the values instead of sharing the object:

self.out.weight.data = self.embed.weight.data.clone()     # not tied

Now there are two tensors with equal values that immediately drift apart. Nothing raises, the parameter count is wrong, and the model quietly loses the prior it was supposed to have.

Re-tying after loading is the other half: load_state_dict writes into the existing tensors, so a model constructed with the tie keeps it. A model that ties after loading is fine too. A model that ties in __init__ and then reassigns self.out.weight somewhere later is not.

Your task

Complete Tied so its two linear layers genuinely share one weight, then:

def tie_facts(dim: int) -> dict

reports the number of parameter tensors, whether the two weights are the same object, and whether they are still the same object after a state_dict round trip through a fresh model.