Building on Word Embeddings, this note covers the input embedding layer specifically as it functions inside a large-scale LLM — including a common weight-sharing trick worth knowing.
The Embedding Layer, at LLM Scale
For a model with a 50,000-token vocabulary and \(d_{\text{model}}=4096\), the embedding matrix alone has over 200 million parameters — often a genuinely significant fraction of a smaller model's total size, though proportionally less significant for the very largest models where the transformer blocks themselves dominate.
Weight Tying — Sharing Input and Output Embeddings
Recall from Vocabulary that both the input embedding layer and the final output projection layer have matching shapes (\(\text{vocab\_size}\times d_{\text{model}}\), just transposed relative to each other). Many LLM implementations deliberately tie these two weight matrices together — using the exact same learned matrix for both converting tokens into vectors and converting the final hidden state back into vocabulary logits. This roughly halves the parameter count contributed by these two layers, and is motivated by the intuition that "how a token should be represented" and "what vector should predict this token" are closely related concepts worth sharing.
Code
import torch.nn as nn
class TiedEmbeddingModel(nn.Module):
def __init__(self, vocab_size, d_model):
super().__init__()
self.embedding = nn.Embedding(vocab_size, d_model)
self.output_layer = nn.Linear(d_model, vocab_size, bias=False)
self.output_layer.weight = self.embedding.weight # WEIGHT TYING: same parameters, shared
model = TiedEmbeddingModel(vocab_size=50000, d_model=768)
print(model.embedding.weight.shape) # (50000, 768)
print(model.output_layer.weight.shape) # (50000, 768) -- literally the SAME tensor, not a copyCommon Mistakes
- Assuming weight tying is always beneficial regardless of model scale — for very large models, the relative parameter savings from tying become proportionally less significant, and some architectures intentionally choose not to tie for other performance reasons.
- Forgetting that tied weights mean a gradient update to one layer automatically affects the other — since they're literally the same underlying parameter tensor, not two separately-initialized tensors that happen to start out equal.
Interview Relevance
Q: "What is weight tying in an LLM, and why is it used?" Sharing the exact same weight matrix between the input token embedding layer and the final output projection layer (which converts hidden states into vocabulary logits) — since both have matching shapes and arguably related roles. This roughly halves the parameter count these two layers would otherwise contribute, at little to no cost in model quality for many architectures.
Practice Question
For a vocabulary of 32,000 tokens and \(d_{\text{model}}=1024\), how many parameters does the embedding layer alone contain? How many would the (separate) output layer add without weight tying?
