Embedding layers and tied weights

4.2.5 · AI-ML › Tokenization & Language Modeling

Overview

Embedding layers discrete token IDs ko continuous vector representations mein convert karte hain, jabki weight tying input embedding aur output projection layers ke beech parameters share karta hai, jisse model size kam hoti hai aur generalization improve hoti hai.

[!intuition] Core Idea

Embeddings ko ek lookup table ki tarah socho jahan har token ID ek row number hai, aur tum us row ka vector retrieve karte ho. Token "cat" (ID=42) hamesha usi ek learned vector ${\mathbf{e}}_{42}\in {\mathbb{R}}^d$ se map hota hai.

Weights tie kyun karein? Input embedding poochta hai "is token ka matlab kya hai?" aur output projection poochta hai "agla token kaunsa hona chahiye?" Ye inverse operations hain—isliye same weight matrix (transposed) use karna model ko ek consistent representation space seekhne par majboor karta hai.

Analogy: Agar English→French aur French→English translation seekh rahe ho, toh dono directions ke liye same dictionary use karna consistency ensure karta hai.

[!definition] Embedding Layer

Ek embedding layer ek learned matrix $\mathbf{E}\in {\mathbb{R}}^{V\times d}$ hai jahan:

• $V$ = vocabulary size
• $d$ = embedding dimension
• Har row $\mathbf{E}[i]$ token $i$ ka embedding vector hai

Operation: Input token ID $x\in \{0,1,\dots ,V-1\}$ ke liye, embedding yeh hai: $\text{Embed}(x)=\mathbf{E}[x]={\mathbf{e}}_x\in {\mathbb{R}}^d$

Yeh ek matrix multiplication nahi hai—yeh ek indexing operation (retrieval) hai.

[!formula] Embedding Operation Derive Karna

Goal: Token ID $x$ ko vector ${\mathbf{e}}_x$ mein convert karo.

Step 1: One-Hot Encoding (Conceptual)

Token $x$ ko one-hot vector ${\mathbf{o}}_x\in \{0,1{\}}^V$ ke roop mein represent karo jahan: ${\mathbf{o}}_x[i]=\{\begin{array}{ll}1& \text{if }i=x\\ 0& \text{otherwise}\end{array}$

Step 2: Matrix Multiplication Interpretation

Agar hum ${\mathbf{o}}_x^T\mathbf{E}$ compute karein: ${\mathbf{o}}_x^T\mathbf{E}=[0,\dots ,0,1,0,\dots ,0]\left[\begin{array}{c}—{\mathbf{e}}_0—\\ vdots\\ —{\mathbf{e}}_x—\\ ⋮\\ —{\mathbf{e}}_{V-1}—\end{array}\right]={\mathbf{e}}_x$

Yeh step kyun? Yeh dikhata hai ki embeddings mathematically equivalent hain one-hot vectors ke saath matrix multiplication ke, lekin hum one-hot actually kabhi banate nahi (memory inefficient hota hai).

Step 3: Efficient Implementation

Practice mein:

# Conceptual: e_x = one_hot(x) @ E  # O(V*d) - slow!
e_x = E[x]  # O(1) indexing - fast!

Yeh kyun important hai: $V=5000$, $d=512$ ke liye, one-hot multiplication $5000\times 512=2.56\text{M}$ multiplications waste karta (zyaadatar zeros ke against). Direct indexing instant hai.

[!formula] Weight Tying: Poori Derivation

Standard Architecture (No Tying)

Forward pass:

1. Input embedding: ${\mathbf{h}}_0={\mathbf{E}}_{\text{in}}[x]$ jahan ${\mathbf{E}}_{\text{in}}\in {\mathbb{R}}^{V\times d}$
2. Transformer layers: ${\mathbf{h}}_L=\text{Transformer}({\mathbf{h}}_0)$
3. Output projection: $\mathbf{z}={\mathbf{h}}_L{\mathbf{W}}_{\text{out}}+\mathbf{b}$ jahan ${\mathbf{W}}_{\text{out}}\in {\mathbb{R}}^{d\times V}$
4. Softmax: $P(x_{t+1}|x_{\le t})=\text{softmax}(\mathbf{z})$

Total parameters: $V\cdot d$ (input) $+V\cdot d$ (output) $=2Vd$

Weight Tying Constraint

${\mathbf{W}}_{\text{out}}={\mathbf{E}}_{\text{in}}^T$ (transpose) set karo. Ab: $\mathbf{z}={\mathbf{h}}_L{\mathbf{E}}_{\text{in}}^T+\mathbf{b}$

Component-wise: Token $j$ ka logit hai: $z_j={\mathbf{h}}_L\cdot {\mathbf{e}}_j+b_j$

Interpretation: Token $j$ ka score = context representation ${\mathbf{h}}_L$ aur token embedding ${\mathbf{e}}_j$ ke beech similarity.

Yeh step kyun? Dot product ${\mathbf{h}}_L\cdot {\mathbf{e}}_j$ measure karta hai ki context token $j$ ke meaning ke saath kitna "aligned" hai—high similarity → high probability.

Parameter reduction: Ab sirf $Vd+V$ parameters hain (embeddings + bias), $Vd-V\approx Vd$ parameters ki saving hoti hai.

Gradient Flow with Tying

Backpropagation ke dauran: $\frac{\partial \mathcal{L}}{\partial \mathbf{E}}=\frac{\partial \mathcal{L}}{\partial {\mathbf{E}}_{\text{in}}}+{\left(\frac{\partial \mathcal{L}}{\partial {\mathbf{W}}_{\text{out}}}\right)}^T$

Yeh step kyun? Kyunki $\mathbf{E}$ do jagah use hota hai (input aur output), dono paths se gradients accumulate hote hain. Yeh richer supervision deta hai—embeddings dono se seekhte hain: "tokens context mein kya mean karte hain" aur "agla kaunsa token likely hai."

[!example] Example 1: Small Vocabulary Embedding

Setup: $V=4$ tokens ["cat", "dog", "the", "ran"], $d=3$ dimensions.

Embedding matrix: $\mathbf{E}=\left[\begin{array}{ccc}0.2& -0.5& 0.8\\ 0.3& -0.4& 0.7\\ -0.9& 0.1& 0.2\\ 0.6& 0.3& -0.5\end{array}\right]\begin{array}{c}\text{(cat, ID=0)}\\ \text{(dog, ID=1)}\\ \text{(the, ID=2)}\\ \text{(ran, ID=3)}\end{array}$

Query: "dog" (ID=1) embed karo.

Solution: $\text{Embed}(1)=\mathbf{E}[1]=[0.3,-0.4,0.7]$

Yeh step kyun? Hum bas row 1 index karte hain. Koi computation nahi chahiye.

[!example] Example 2: Output Projection with Weight Tying

Context: Transformer ke baad, hamare paas ${\mathbf{h}}_L=[0.5,0.2,-0.3]$ hai.

Logits compute karo (tied weights use karke, ${\mathbf{W}}_{\text{out}}={\mathbf{E}}^T$): $\mathbf{z}={\mathbf{h}}_L{\mathbf{E}}^T=[0.5,0.2,-0.3]\left[\begin{array}{cccc}0.2& 0.3& -0.9& 0.6\\ -0.5& -0.4& 0.1& 0.3\\ 0.8& 0.7& 0.2& -0.5\end{array}\right]$

Har logit calculate karo:

• $z_0=0.5(0.2)+0.2(-0.5)+(-0.3)(0.8)=0.1-0.1-0.24=-0.24$
• $z_1=0.5(0.3)+0.2(-0.4)+(-0.3)(0.7)=0.15-0.08-0.21=-0.14$
• $z_2=0.5(-0.9)+0.2(0.1)+(-0.3)(0.2)=-0.45+0.02-0.06=-0.49$
• $z_3=0.5(0.6)+0.2(0.3)+(-0.3)(-0.5)=0.3+0.06+0.15=0.51$

Probabilities (softmax ke baad): $P(\text{next token})=\text{softmax}([-0.24,-0.14,-0.49,0.51])\approx [0.19,0.21,0.15,0.40]$

Interpretation: "ran" (ID=3) ki highest probability hai kyunki uska embedding ${\mathbf{h}}_L$ ke saath sabse zyaada aligned hai.

Yeh step kyun? Weight tying model se yeh poochhta hai: "Meri hidden state ke sabse close kaunsa token embedding hai?" Agar embeddings achhe se seekhe gaye hain, toh yeh exactly wahi hai jo hum chahte hain.

[!example] Example 3: Parameter Count Comparison

Scenario: GPT-2 small with $V=50257$, $d=768$.

Weight tying ke bina:

• Input embeddings: $50257\times 768=38,597,376$
• Output projection: $768\times 50257=38,597,376$
• Total: $77,194,752$ parameters

Weight tying ke saath:

• Shared embeddings: $50257\times 768=38,597,376$
• Output bias: $50257$
• Total: $38,647,633$ parameters

Savings: $77.2M-38.6M=38.6M$ parameters (~50% reduction embedding/projection mein).

Yeh step kyun? Bade vocabularies ke liye, embeddings parameter count mein dominate karte hain. Tying ise almost aadha kar deti hai bina performance hurt kiye (aksar improve karke!).

[!mistake] Common Mistake 1: "Embeddings sirf random vectors hote hain"

Galat intuition: Embeddings arbitrary hote hain—koi bhi random initialization kaam karta hai.

Kyun sahi lagta hai: Initially, embeddings hote hain random. Aur random embeddings bhi kuch signal dete hain.

Fix: Embeddings learned representations hain. Gradient descent ke through:

• Similar tokens (cat/dog) similar vectors paate hain
• Embeddings semantic aur syntactic relationships capture karte hain
• Model ek aise geometry seekhta hai jahan "king - man + woman ≈ queen"

Evidence: Pre-trained embeddings (Word2Vec, GloVe) tasks ke across transfer hote hain kyunki yeh real linguistic structure capture karte hain.

[!mistake] Common Mistake 2: "Weight tying force karta hai input=output exactly"

Galat intuition: Weight tying ke saath, ${\mathbf{E}}_{\text{in}}={\mathbf{W}}_{\text{out}}$, toh input aur output identical hain.

Kyun sahi lagta hai: Hum same matrix use karte hain, toh woh same hone chahiye.

Fix: Hum ${\mathbf{W}}_{\text{out}}={\mathbf{E}}_{\text{in}}^T$ (transpose!) tie karte hain. Operations hain:

• Input: ${\mathbf{h}}_0=\mathbf{E}[x]$ (row $x$ select karo)
• Output: $\mathbf{z}={\mathbf{h}}_L{\mathbf{E}}^T$ (saare dot products compute karo)

Yeh geometrically inverse operations hain:

• Input: "Is token ko kaunsa vector represent karta hai?"
• Output: "Kaunse token vectors is context ke sabse close hain?"

Yeh kyun kaam karta hai: Output layer pooch raha hai "meri hidden state ke saath kaunsa embedding most similar hai?" Agar embeddings achhe se seekhe gaye hain, toh yeh exactly wahi hai jo hum chahte hain.

[!mistake] Common Mistake 3: "Embeddings per-position hote hain"

Galat intuition: Position 5 par token "cat" ko position 10 par "cat" se alag embedding milta hai.

Kyun sahi lagta hai: Transformers positional encodings use karte hain, toh position matter karta hai.

Fix: Token embedding position-independent hota hai: "cat" hamesha ${\mathbf{e}}_{\text{cat}}$ paata hai. Position alag se add hoti hai: ${\mathbf{h}}_0=\mathbf{E}[x]+\text{PositionalEncoding}(\text{pos})$

Yeh step kyun? Token aur position information ko alag karna cleaner hai—embedding capture karta hai ki token kya hai, positional encoding capture karta hai ki woh kahan hai.

[!recall]- Kisi 12-Saal-Ke Bachhe Ko Samjhao

Imagine karo tumhare paas ek magic dictionary hai. Har word ka ek secret code hai (numbers ki ek list). Jab tum "dog" dekhte ho, tumhe hamesha wahi code milta hai, jaise [5, -2, 8].

Ab, tum ek story likh rahe ho. Tumne "The dog ran" likha hai, aur tum agla word predict karna chahte ho. Tumhara brain (transformer) story ke baare mein sochta hai aur apna khud ka code produce karta hai, jaise [6, -1, 7].

Agla word choose karne ke liye, tum apne brain ke code ko dictionary mein har word ke code se compare karte ho. Jis word ka code sabse close hoga woh jeet jaata hai! Shayad "fast" ka code [6, -1.5, 7.2] hai—bahut close—toh tum "fast" predict karte ho.

Weight tying ka matlab hai same dictionary use karna words ko shuru mein lookup karne ke liye AUR end mein compare karne ke liye. Kyun? Kyunki agar "dog" aur "puppy" ke dictionary mein similar codes hain, toh tum chahte ho ki tumhara brain predict karte waqt bhi unhe similarly treat kare. Yeh ek jagah jaane aur wapas aane ke liye ek hi map use karne jaisa hai—zyaada efficient aur consistent!

[!mnemonic] Weight Tying Yaad Karo

"Shared Shoes": Input aur output ek hi shoes (weight matrix) pehnte hain, lekin ek aage chalta hai (tokens embed karna), doosra peeche chalta hai (tokens predict karna). Same shoes, opposite directions. Paison ki bachat (parameters), phir bhi destination tak pahunch jaate ho!

Key Formulas

Embedding Lookup

${\mathbf{e}}_x=\mathbf{E}[x]\text{where }\mathbf{E}\in {\mathbb{R}}^{V\times d}$

Output Logits (Weight Tying)

$z_j={\mathbf{h}}_L\cdot {\mathbf{e}}_j+b_j={\sum }_{k=1}^d{h}_{L,k}\cdot e_{j,k}+b_j$

Probability Distribution

$P(x_{t+1}=j\mid x_{\le t})=\frac{\exp (z_j)}{{\sum }_{i=1}^V\exp (z_i)}$

Connections

• 4.2.01-Tokenization-fundamentals — Tokens IDs mein convert hote hain, phir embed hote hain
• 4.2.04-Subword-tokenization-BPE-WordPiece — Vocabulary size $V$ embedding matrix size determine karta hai
• 4.3.02-Positional-encoding — Token embeddings mein add hota hai position info inject karne ke liye
• 4.4.03-Self-attention-mechanism — Embedded token vectors par operate karta hai
• 5.1.03-Cross-entropy-loss — Loss function jo embeddings ke through backprop karta hai
• 6.2.04-Parameter-efficient-fine-tuning — Adapter methods aksar embeddings freeze karte hain

Practical Notes

Weight tying kab use karein:

• ✅ Standard language models (GPT, BERT)
• ✅ Chhote se medium models jahan parameter count matter karta hai
• ✅ Jab vocabulary badi ho ($V>10k$)

Kab use NA karein:

• ❌ Encoder-decoder models jahan alag source/target vocabularies hon
• ❌ Models jahan input/output dimensions differ karein ($d_{\text{in}}\ne d_{\text{out}}$)
• ❌ Kuch multimodal models (vision+language) jahan output token prediction nahi hai

Implementation tip: Zyaadatar frameworks (PyTorch, TensorFlow) weight tying ko trivial banate hain:

self.output_proj.weight = self.embedding.weight  # Share reference

#flashcards/ai-ml

Concept Map