Week 6 — Day 4: Word Embeddings (nn.Embedding), Sequence Models (RNN, LSTM) & Vanishing Gradients

Mentor's Note (10+ YOE Staff AI Engineer POV): Why do we study nn.Embedding, RNNs, and LSTMs right before building Transformers in Week 7?

1. nn.Embedding(vocab_size, embed_dim) is the very first layer of every LLM (GPT-4, Llama-3)! When a token ID like 42 enters an LLM, nn.Embedding simply looks up Row #42 in a trainable matrix of shape (vocab_size, embed_dim)!
2. Why RNNs Failed & Why Transformers Were Invented: Interviewers love asking: "Why couldn't Recurrent Neural Networks (RNN / LSTM) scale to 100,000-token context windows or train on 10,000 GPUs in parallel?" Today you learn the two mathematical reasons: Sequential Time Dependency ($h_t = f(W_{hh} h_{t-1} + W_{xh} x_t)$) and Vanishing / Exploding Gradients Through Time ($\prod_{k=1}^T \frac{\partial h_k}{\partial h_{k-1}}$)!

Part 1: How nn.Embedding Works Under the Hood

Instead of multiplying a huge one-hot vector [0, 0, ..., 1, ..., 0] (length 100,000) by a matrix, nn.Embedding(num_embeddings=100_000, embedding_dim=768) is an $O(1)$ array index lookup table!

• Input: Integer Tensor of Token IDs of shape (Batch, Seq_Len)
• Output: Float Tensor of Dense Vectors of shape (Batch, Seq_Len, Embedding_Dim)

Part 2: RNN vs LSTM Gates & The Vanishing Gradient Equation

2.1 Vanilla RNN Hidden State Equation

At each time step $t = 1, 2, \dots, T$: $h_t = \tanh(W_{xh} x_t + W_{hh} h_{t-1} + b)$

• Problem #1 (No Parallelism Across Time): You cannot compute step $t = 500$ until step $t = 499$ finishes! GPUs sit idle waiting on sequential loops.
• Problem #2 (Vanishing Gradients in Backpropagation Through Time — BPTT): To compute how token $h_1$ affects the final token $h_T$, the Chain Rule multiplies $T$ Jacobian matrices together: $\frac{\partial h_T}{\partial h_1} = \prod_{t=2}^{T} \frac{\partial h_t}{\partial h_{t-1}} \propto (W_{hh})^T$
  • If the largest eigenvalue of $W_{hh}$ is $< 1.0$ (e.g., 0.8), then over $T = 100$ tokens: $0.8^{100} = 2.0 \times 10^{-10} \approx \mathbf{0.0}$ (Vanishing Gradient — the model forgets the beginning of the sentence!).
  • If the eigenvalue is $> 1.0$ (e.g., 1.2), then $1.2^{100} = 8.2 \times 10^7 \to \mathbf{\text{NaN}}$ (Exploding Gradient!).

2.2 How LSTM (Long Short-Term Memory) Partially Helped

LSTM introduced an additive Cell State highway ($c_t = f_t \odot c_{t-1} + i_t \odot \tilde{c}_t$) controlled by 3 Sigmoid gates (Forget Gate $f_t$, Input Gate $i_t$, Output Gate $o_t$) so gradients could flow through addition (+) instead of repeated matrix multiplication—until Self-Attention (Transformers) eliminated the sequential bottleneck entirely in 2017!

Part 3: Hands-On Lab — nn.Embedding, LSTM & Vanishing Gradient Proof (sequence_models_rnn_lstm.py)

Check out src/sequence_models_rnn_lstm.py (in Tab 2: Codes — click Run in Browser to execute) and run the script and tests:

In Simple English: Reading a Sentence One Word at a Time (RNN vs LSTM & Memory Fade)

Imagine trying to summarize a 500-page novel by writing a 3-sentence sticky note after reading Page 1, erasing and rewriting that same tiny sticky note after Page 2, and continuing to Page 500:

• By Page 500, you have almost completely forgotten who the villain was on Page 3! That is the Vanishing Gradient Problem: multiplying weights step-by-step ($w^{500}$ when $w = 0.8 \implies 0.8^{500} \approx 0.0$) shrinks early learning signals to zero!
• LSTM adds a separate "long-term notebook" (cell state conveyor belt) with Forget/Keep gates, which helps — but still forces you to read word 1 before word 2 before word 3 slowly. Week 7 Transformers solve both problems at once by letting every word look directly at every other word simultaneously!