ALBERT › Why Factorization Works
Language AI HandbookAll chapters ALBERT › Why Factorization WorksOne section · Sections in this chapter4 / 18 ↓
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
At first glance, reducing embedding dimension seems like it would hurt representation quality. Tokens have less room to encode meaning. But the key insight is that embedding vectors and hidden states serve different purposes.
Token embeddings capture context-independent information. The embedding for "bank" is the same whether it appears in "river bank" or "savings bank." This relatively simple lookup doesn't require the full expressiveness of the hidden dimension, which must capture context-dependent representations built through multiple layers of attention.
def demonstrate_projection_capacity():
"""Show that small embeddings project well to larger spaces."""
torch.manual_seed(42)
vocab_size = 1000
embedding_dim = 64
hidden_dim = 256
# Create factorized embedding
factorized = FactorizedEmbedding(vocab_size, embedding_dim, hidden_dim)
# Sample some token IDs
token_ids = torch.tensor([[0, 1, 2, 3, 4]])
# Get intermediate and final representations
intermediate = factorized.token_embedding(token_ids)
final = factorized(token_ids)
return intermediate, finalRepresentation Dimensions: Intermediate (E=64): torch.Size([1, 5, 64]) After projection (H=256): torch.Size([1, 5, 256]) Intermediate variance: 0.9986 Final variance: 0.3171
The intermediate representation is 4x smaller than the final output (64 vs 256 dimensions), yet it preserves the information needed for the projection layer to expand it. The variance change shows how the projection redistributes activation magnitudes across the larger hidden space.
The projection layer learns to expand the compressed representations into the full hidden space. During pretraining, this projection adapts to whatever patterns the model needs, effectively learning a task-appropriate decompression of the initial token representations.
We can visualize this by examining how different tokens map through the factorized embedding space:
