Weight Tying › Implementation Details
Language AI HandbookAll chapters Weight Tying › Implementation DetailsOne section · Sections in this chapter5 / 15 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement a complete language model head with weight tying, handling the practical details that matter in real systems. A production-grade implementation needs to manage positional encodings, dropout, and the subtleties of the embedding scale.
The key engineering insight in a full implementation is that the embedding matrix is shared by reference, not by copying. Both the encode and decode operations point to the same underlying tensor. When gradients update the shared matrix, they simultaneously improve both the input representations and the output targets. This is not a coincidence or a side effect. It is the core mechanism that makes weight tying beneficial.
class LanguageModelHead(nn.Module):
"""
Complete language model embedding layer with weight tying.
Handles token embedding, positional encoding, and output projection
using a single shared vocabulary embedding matrix.
"""
def __init__(
self, vocab_size, hidden_dim, max_seq_len, dropout=0.1, tie_weights=True
):
super().__init__()
self.hidden_dim = hidden_dim
self.tie_weights = tie_weights
# Token embeddings (always needed)
self.token_embedding = nn.Embedding(vocab_size, hidden_dim)
# Position embeddings (not tied)
self.position_embedding = nn.Embedding(max_seq_len, hidden_dim)
# Output projection: tied or separate
if tie_weights:
self.output_projection = None # Will use token_embedding.weight
else:
self.output_projection = nn.Linear(
hidden_dim, vocab_size, bias=False
)
self.dropout = nn.Dropout(dropout)
# Initialize embeddings
self._init_weights()
def _init_weights(self):
"""Initialize embedding weights."""
nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02)
nn.init.normal_(self.position_embedding.weight, mean=0.0, std=0.02)
if self.output_projection is not None:
nn.init.normal_(self.output_projection.weight, mean=0.0, std=0.02)
def embed(self, token_ids):
"""Convert token IDs to embeddings with position information."""
batch_size, seq_len = token_ids.shape
# Token embeddings
tok_emb = self.token_embedding(token_ids)
# Position embeddings
positions = torch.arange(seq_len, device=token_ids.device)
pos_emb = self.position_embedding(positions)
# Combine and apply dropout
return self.dropout(tok_emb + pos_emb)
def project(self, hidden_states):
"""Project hidden states to vocabulary logits."""
if self.tie_weights:
# Use shared embedding matrix
return torch.matmul(hidden_states, self.token_embedding.weight.T)
else:
return self.output_projection(hidden_states)
# Compare tied vs untied
model_tied = LanguageModelHead(
vocab_size, hidden_dim, max_seq_len=512, tie_weights=True
)
model_untied = LanguageModelHead(
vocab_size, hidden_dim, max_seq_len=512, tie_weights=False
)
def count_parameters(model):
return sum(p.numel() for p in model.parameters())Parameter comparison: Tied weights: 38,793,216 parameters Untied weights: 77,193,216 parameters Difference: 38,400,000 parameters The tied model saves 49.7% of embedding-related parameters
The key implementation detail is in the project method. With tied weights, we directly use self.token_embedding.weight.T instead of a separate projection layer. This ensures that gradients flow through the same parameters during both forward and backward passes. There is no additional memory allocated for a second embedding matrix. The model literally shares the same tensor object.
Notice that position embeddings are not tied and do not participate in the output projection. This is by design. Position embeddings encode sequential position information that has no meaning in the output vocabulary space. Only the token embeddings, which capture the semantic identity of each vocabulary item, are shared between input and output.
Verifying the Tying
Let's verify that our implementation shares parameters correctly and that gradients flow as expected:
# Test the tied model
test_tokens = torch.randint(0, vocab_size, (2, 10)) # Batch of 2, length 10
# Get embeddings
embeddings = model_tied.embed(test_tokens)
# Simulate transformer processing (just use embeddings as output for demo)
hidden_output = embeddings
# Get logits
logits = model_tied.project(hidden_output)
# Verify gradient flow
loss = logits.sum()
loss.backward()Input shape: torch.Size([2, 10]) Embedding shape: torch.Size([2, 10, 768]) Output logits shape: torch.Size([2, 10, 50000]) Token embedding gradient shape: torch.Size([50000, 768]) Gradient flows through shared weights: True Gradient norm on shared embedding: 1143.0863 Gradients arrive from both input and output paths simultaneously.
The gradient from the output loss flows directly into the token embedding matrix, confirming that the tying works correctly. With tied weights, every backward pass accumulates gradients from multiple sources into the same underlying parameters. This is the mechanism by which tied models receive more learning signal per parameter.