Weight Tying › Effects on Training Dynamics
Language AI HandbookAll chapters Weight Tying › Effects on Training DynamicsOne section · Sections in this chapter9 / 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
Weight tying does not just reduce parameters. It changes the fundamental dynamics of how the model learns, and understanding these dynamics helps explain why tied models often outperform untied models with the same number of free parameters.
The most important change is in how gradients flow. With separate embedding matrices, the input embedding receives gradients only through the transformer layers, while the output projection receives gradients directly from the cross-entropy loss. With tied weights, the single shared matrix receives gradients from both sources simultaneously. This is a form of implicit multi-task learning: the embedding must simultaneously satisfy two objectives, and each objective provides additional signal about what good representations look like.
Consider a rare token like a technical term that appears infrequently in the training corpus. With separate embeddings, the input embedding for this token receives sparse gradient updates (only when it appears in the input), and the output embedding receives sparse updates (only when it should be predicted). With tied weights, both sources of gradient contribute to the same parameters. This provides more learning signal per occurrence. This is particularly useful for vocabulary items that are infrequent as either input or output but not both.
Gradient Flow
With tied weights, the embedding matrix receives gradients from two sources:
-
Input gradients: Backpropagated through the transformer from the loss. When a token appears in the input, its embedding affects all downstream computations, and gradients flow backward through all the transformer layers to update the embedding.
-
Output gradients: Direct gradients from the output projection. When a token should be predicted, the dot product between the hidden state and that token's embedding directly contributes to the cross-entropy loss. This provides a gradient that says "move this embedding closer to the current hidden state."
This double gradient flow can be viewed as implicit multi-task learning. The embedding must simultaneously satisfy two objectives: representing tokens well for input processing so the transformer can reason about them, and providing good targets for output prediction so the model can accurately identify which token to generate next.
class GradientTracker(nn.Module):
"""Track gradient magnitudes from different sources."""
def __init__(self, vocab_size, hidden_dim):
super().__init__()
self.embedding = nn.Embedding(vocab_size, hidden_dim)
self.fc = nn.Linear(hidden_dim, hidden_dim) # Simulates transformer
def forward(self, token_ids, return_hidden=False):
# Input path
x = self.embedding(token_ids)
hidden = self.fc(x)
# Output path (tied)
logits = torch.matmul(hidden, self.embedding.weight.T)
if return_hidden:
return logits, hidden
return logits
# Track gradients
grad_model = GradientTracker(1000, 64)
tokens = torch.randint(0, 1000, (4, 16))
labels = torch.randint(0, 1000, (4, 16))
logits = grad_model(tokens)
loss = F.cross_entropy(logits.view(-1, 1000), labels.view(-1))
loss.backward()
grad_magnitude = grad_model.embedding.weight.grad.abs().mean()Mean gradient magnitude on tied embedding: 0.001368 This gradient combines contributions from both: - Forward pass through the input embedding - Backward pass through the output projection Gradient norm: 0.9875 With separate embeddings, each matrix would receive only half this signal.
The gradient magnitude shows how much the embedding weights would change in a single training step (before learning rate scaling). With tied weights, this gradient is typically larger than it would be with separate embeddings because it aggregates signals from both the input and output paths. This can speed up learning for rare tokens that might otherwise receive sparse gradient updates, since each appearance provides double the learning signal.
The combined gradient also has a regularizing effect. The input gradient pushes the embedding toward representations that make downstream processing easier. The output gradient pushes it toward representations that align well with the hidden states the transformer produces for contexts where this token is appropriate. When both gradients point in compatible directions, the embedding is in a region of representation space that satisfies both constraints simultaneously, which is exactly where you want it.
Embedding Scale
A subtle issue arises with weight tying: the optimal scale for input embeddings may differ from the optimal scale for output projections. This affects real training dynamics and model performance.
Input embeddings are often scaled by before being added to positional encodings, following the original transformer paper by Vaswani et al. (2017). The scaling formula is:
where:
- : the scaled input embedding that enters the transformer
- : the embedding dimension
- : the raw embedding lookup for token
- : the scaling factor, which counterbalances the variance reduction that occurs when embeddings are initialized with small standard-deviation values
Without this scaling, embeddings initialized with small weights (standard deviation around 0.02) would be numerically negligible compared to the positional encodings, which are also on a small scale. The positional encodings would dominate the input signal, preventing the model from distinguishing token identities. The scaling amplifies the token embeddings to a magnitude that balances both sources of information.
However, if we applied this scaling directly to the weight matrix, it would also affect the output logits when using tied weights. Scaled embeddings as output targets would produce much larger logits than expected, potentially destabilizing softmax and making the output distribution unreasonably peaked.
Modern implementations handle this by applying scaling at input time rather than modifying the embedding matrix itself:
class ScaledTiedEmbeddings(nn.Module):
"""
Tied embeddings with proper scaling.
Input embeddings are scaled up, but output projection uses raw weights.
"""
def __init__(self, vocab_size, hidden_dim):
super().__init__()
self.hidden_dim = hidden_dim
self.embedding = nn.Embedding(vocab_size, hidden_dim)
self.scale = np.sqrt(hidden_dim)
def embed(self, token_ids):
"""Scaled embedding for input."""
return self.embedding(token_ids) * self.scale
def project(self, hidden_states):
"""Unscaled projection for output."""
return torch.matmul(hidden_states, self.embedding.weight.T)
scaled_model = ScaledTiedEmbeddings(vocab_size, hidden_dim)
test_input = torch.randint(0, vocab_size, (1, 10))
# Forward pass
input_embeds = scaled_model.embed(test_input)
# ... transformer processing would happen here ...
output_logits = scaled_model.project(
input_embeds
) # Using embeds as proxy for hidden stateInput embedding scale factor: 27.71 Input embedding norm (per token): 765.35 Raw (unscaled) embedding norm: 27.62 The scale is applied at input time, not in the embedding matrix itself, so the output projection sees the unscaled embeddings. Scale ratio: 27.71x amplification for input path
The scale factor of approximately 27.7 for hidden dimension 768 significantly amplifies the input embeddings. By applying this scaling at runtime rather than baking it into the weights, we decouple the input and output requirements. The transformer processes scaled embeddings while the output projection uses raw embeddings, allowing each path to operate at its optimal magnitude. This is an example of the broader principle that tying weights does not mean treating both uses identically; you can still apply path-specific transformations.