Skip to content

← nn.Module Mechanics step 4 of 7

Medium Primitives

Share one weight parameter between two layers

Jump to the task

Weight tying lets two parts of a model use one Parameter object. For example, a language model can share its token embedding matrix with the output projection. Both weights have shape (vocabulary_size, embedding_dim). This saves one matrix; the fraction of the whole model saved depends on its other parameters, and the effect on accuracy depends on the model.

Here is a smaller example with two bias-free Linear layers:

import torch
import torch.nn as nn

model = nn.Module()
model.a = nn.Linear(2, 2, bias=False)
model.b = nn.Linear(2, 2, bias=False)
model.b.weight = model.a.weight
print(model.a.weight is model.b.weight)  # True
print(len(list(model.parameters())))    # 1

Equal values are not a shared parameter

The assignment above shares the Parameter itself. Copying values into a separate Parameter leaves two independent objects. They may start equal, but different gradients or optimizer updates can make them diverge.

Use object identity (a.weight is b.weight) to check the tie. Matching values alone are not enough. Avoid using .data to implement sharing; it bypasses normal autograd handling and does not establish Parameter identity.

When the shared weight is used twice in a calculation, both uses contribute to its gradient:

x = torch.ones(1, 2)
loss = (model.a(x) + model.b(x)).sum()
loss.backward()
print(model.a.weight.grad.tolist())  # [[2.0, 2.0], [2.0, 2.0]]

The default parameters() iterator returns the shared Parameter once. If the two Linear layers also had separate biases, there would be three parameter tensors: the shared weight and two biases. This exercise has no biases.

Preserve the tie when loading

The state_dict contains both names, a.weight and b.weight. With the default keep_vars=False, the values are detached tensors sharing storage; do not expect the two dictionary values to be the same Python object.

The default load_state_dict(..., assign=False) copies values into the receiving model's existing parameters. Construct the receiving model with the tie already in place. Loading equal values into an untied model does not make its two Parameter objects identical.

This lesson uses the default copy-based load. If you use assign=True or replace a Parameter later, check the sharing again and construct the optimizer after the final parameter setup so it references the right objects.

Your task

Complete Tied so its two bias-free Linear(dim, dim) layers share one weight Parameter. Then implement:

def tie_facts(dim: int) -> dict

Build a Tied model and load its state_dict into a fresh Tied model using the default load_state_dict call. Return:

  • n_params: the number of parameter tensors in the original model.
  • shared: whether its two weight attributes are the same object.
  • shared_after_load: the same identity check on the fresh model after loading.

For dim=4, return {"n_params": 1, "shared": True, "shared_after_load": True}. The round trip transfers the state dictionary in memory; no file is needed.