We can't find the internet
Attempting to reconnect
Something went wrong!
Attempting to reconnect
← nn.Module Mechanics step 4 of 7
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 tiedLinearweights 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.
Stuck?
PyTorch reference solution
Sign in to attempt this problem and reveal the reference solution.