CS546 Learning Notes: From Word Embedding to RNN, Understanding BPTT and Vanishing Gradients

1. Embedding matrix: each word corresponds to one row

First, assume each word is a single token. Real tokenizers don't necessarily segment by complete words; we use words here just for convenience in hand calculation.

The vocabulary order is:

Token IDs start from 0, so the ID for cats is 2. The vocabulary size is $V=5$, and each embedding has $d=3$ numbers:

For example:

Embedding lookup retrieves the corresponding row based on token ID:

This is index 2, which corresponds to the 3rd row written out.

The sentence I like cats has 3 tokens, and looking them up in the table yields a $3\times3$ matrix. The two 3's represent the number of tokens and embedding dimension respectively, which happen to be the same.

Therefore, 2 sentences, each padded to 4 tokens, with each vector of length 3, the input tensor is $2\times4\times3$. This is the common batch-first convention.

2. One-hot times matrix is embedding lookup

Write the one-hot for cats as a row vector:

Then:

All other rows are multiplied by 0, only the cats row is multiplied by 1. Programs typically use direct lookup instead of actually creating a one-hot and performing many multiplications involving zeros.

The embedding matrix is usually a learnable parameter: later predictions generate a loss, gradients flow back, and the corresponding word vectors are updated.

3. Why is the gradient of embedding $x^Tg_e$?

When I saw this transpose, I was confused again:

How is the transpose of $x$ derived? Can you work through it for me?

Start with a single element. The forward computation is $e=xE$, expanding the $j$-th output component:

Here, the subscripts use mathematical notation starting from 1, so cats corresponds to $i=3$; its program index is still 2.

Let the gradient flowing back be a row vector:

where $g_j=\partial L/\partial e_j$. The parameter $E_{ij}$ directly affects only $e_j$, and:

The chain rule gives:

Arranging all elements back into a matrix:

The shapes also match:

This is called the outer product. The justification for the transpose comes from element-wise differentiation yielding $x_i g_j$; shape is a tool for checking the result. When directly writing $xg_e$, the inner dimension $5\neq1$, and matrix multiplication won't even work.

For example, if $g_e=[2,-1,3]$, then:

If the learning rate is $\eta=0.1$, the standard Gradient Descent update is:

I initially wrote the first number as $0.6$. Actually $-0.4-0.2=-0.6$; the second number increases because we subtract a negative gradient: $0.8-0.1(-1)=0.9$.

"Only updating the looked-up row" here has conditions: $E$ is only used for input lookup, uses standard Gradient Descent, and has no other update sources like regularization. If there is weight decay, optimizer history, or shared embedding weights with the output layer, we cannot simply claim other rows remain unchanged.

If the same word appears twice, both lookups use the same row of parameters, so gradients from both positions sum together.

4. Where does the correct answer for next token prediction come from?

My question at that time was:

Previously there were definite labels; how do we provide the answer here?

First, let's distinguish: labels in regression tasks are usually numerical; one-hot is the label representation from earlier classification tasks. Next-token prediction in language models can also be viewed as a classification task over the vocabulary.

The raw text itself provides the answer. For:

Shift the sequence by one position:

So the training input is [I, like, cats], and the target is [like, cats, today]. Although at step 2 we only read in like, the hidden state still carries information from I before it.

This way of constructing supervised signals from the data itself is called self-supervised learning. The target can conceptually be written as one-hot, but programs typically store token IDs directly.

We can also add <BOS> and <EOS> to represent sequence start and end:

The specific usage of special tokens depends on model and data processing conventions; not all implementations must use these two names. When adding new tokens, the vocabulary size must be computed accordingly.

During training, this example feeds the correct token from the source text at each step, called teacher forcing. During generation, the token selected from the previous step is fed back to the next step. The RNN's hidden state still needs to be computed sequentially over time; knowing all labels doesn't mean we can compute the hidden states of all time steps in parallel.

5. Similar embeddings are not humans manually specifying the meaning of each number

Each coordinate in an embedding typically has no human-assigned independent semantics. Vectors gradually adjust by completing training objectives; words that frequently appear in similar contexts may learn similar representations.

For example, cats and dogs are often used similarly, but this is a trend resulting from training, not guaranteed by one or two sentences.

It's especially important to distinguish between input embeddings and output layer weights: in the sample I like → cats, if the two don't share parameters, the input lookup uses I and like, and won't directly look up and update the cats row on the input side just because the target is cats. The output layer learns to predict cats through classification gradients over the vocabulary.

Therefore, "similar context corresponds to similar embedding" cannot be directly inferred from this single sample; we must look at the complete data, training objective, and parameter sharing scheme. Cosine similarity is commonly used when comparing word vectors; it mainly measures whether directions are similar.

6. Why can't we just add the embeddings of all words in a sentence?

Compare:

• dog bites man
• man bites dog

I initially thought their sum representations would be different, but vector addition satisfies commutativity:

For this kind of fixed token embedding, direct summation loses word order. By arranging vectors into a matrix according to word order, we can preserve this difference in the input representation.

If each embedding has length 3, three words flattened have length 9, and five words have length 15. A standard MLP's first layer weight shape is fixed and cannot directly accept these two different lengths.

Padding to a fixed length can be used with MLP and preserve word order through fixed positions, but requires preprocessing length ranges. RNN provides another approach: read one token at a time and pass previous information to the next step.

7. RNN: current input plus previous hidden state

The basic RNN computation is:

Starting from this section, $x_t$ denotes the embedding column vector of the current token, not the one-hot row vector from the lookup derivation earlier. If $E$ stores one word vector per row, then $x_t=E[\mathrm{ID}_t]^T$.

Let the embedding dimension be $d=3$ and hidden dimension be $H=4$:

All three terms yield $4\times1$, so they can be added; $\tanh$ acts element-wise without changing the shape.

Each time step shares the same set of $W_x,W_h,b_h$. Longer sequences increase the number of computation steps, but not the number of these parameters. Hidden state is a learned representation of historical information, with no guarantee of completely remembering all prior context.

8. Manual Verification: Why Can RNN Perceive Order?

Temporarily drop $\tanh$ and simplify both input and hidden state to scalars:

Input sequence is $[1,3]$:

Swap the order, input $[3,1]$:

The same two inputs yield final states 7 and 5 respectively. Earlier inputs pass through the recurrent connection, so their effect on the final state differs from that of the current input, thus sequence order can influence the result.

This shows RNN can express word order differences, but does not mean any parameter setting will necessarily distinguish all different sequences.

9. We Can Already Distinguish Order—Why $\tanh$?

I recall the earlier ReLU: an MLP with all linear layers can be merged into a single linear transformation no matter how many layers. RNN has a similar problem.

Remove activation function and bias:

For a fixed number of steps, it remains a linear combination of inputs and initial state; with bias it becomes an affine transformation. Although it contains order information, this state computation cannot express nonlinear relationships.

After adding $\tanh$, nesting appears:

This typically cannot be merged into a single linear layer anymore.

$\tanh$ also constrains each hidden component to $(-1,1)$. In the linear example $h_t=2h_{t-1}$, $h_0=1$, the state becomes $2,4,8,\ldots$; hidden states after $\tanh$ do not grow unboundedly this way. However, bounded forward state does not guarantee stable backward gradients.

10. Output Next Word from Hidden State

The output layer reuses the Softmax classification from before:

If $H=4$ and vocabulary size $V=5$:

$p_t$ is likewise $5\times1$, with each element being the probability of the corresponding word—all non-negative and sum to 1.

Using the five-word vocabulary, if the target is cats:

Cross-entropy at a single position is:

The gradient is:

Higher probability for the correct word means lower loss. Treating logits as independent variables, the negative gradient direction would increase the correct word's logit and decrease others'. But actual training updates shared parameters, which does not guarantee that after updating the entire model or a batch, every sample's logits strictly change in this direction.

Multiple positions' objectives can be summed:

Or averaged; when averaging, the corresponding total gradient should also be divided by $T$. With padding, loss is typically computed only over valid targets.

11. BPTT: Chain Rule After Unrolling Over Time

Suppose we only look at three time steps, with total loss:

$h_1$ affects $L_1$, and also affects $L_2$ through $h_2$, and affects $L_3$ through $h_2,h_3$. Therefore:

While $h_2$ only receives contributions from $L_2,L_3$. In the forward graph of a basic unidirectional RNN, there is no path from $h_2$ back to affect $L_1$.

This is Backpropagation Through Time (BPTT): unroll the RNN over time, then apply ordinary backpropagation.

In the scalar case, one path from $L_3$ is:

Still the familiar rule: multiply along the path, add across paths.

In the vector case, we cannot freely permute matrix order. Define $J_t=\partial h_t/\partial h_{t-1}\in\mathbb R^{H\times H}$, and write the gradient of loss with respect to state as a column vector, then the contribution flowing back from $L_3$ to $h_1$ is:

12. Why Must Gradients of Shared Parameters Be Summed?

The same $W_h$ is reused across different time steps. During backpropagation, contributions from each use to the total loss must be accumulated to update the same parameter.

To make "each step's contribution" clear, define the pre-activation quantity:

Let $\delta_t=\partial L/\partial a_t\in\mathbb R^{H\times1}$, which already includes contributions from current and subsequent losses, then:

Each term has shape $(H\times1)(1\times H)=H\times H$. As with embeddings, an outer product appears again.

After this complete sequence's backpropagation, ordinary Gradient Descent updates:

We don't treat it as three independently learned $W_h$ just because we unrolled three time steps. Similarly, do not double-count gradients that already include subsequent influences.

13. Vanishing and Exploding Gradients: Consequences of Repeated Multiplication

Start with the simplest scalar RNN:

The local derivative at each step is $w$. Assume the gradient of end-point loss with respect to $h_3$ is 1, then:

When $w=0.5$:

When $w=2$:

The former shows vanishing gradient, the latter exploding gradient. Here we use $h_0$ to observe how the signal changes across time, not claiming the fixed zero initial state must be a parameter to update.

Vanishing gradients mean the learning signal from later losses may be extremely weak for early computations, making long-term dependencies hard to learn. If early positions also have their own loss, they may still receive strong local gradients; what vanishes is the long-distance contribution we care about, not all gradients necessarily becoming zero.

In a real tanh RNN:

Here $h_t^2$ denotes element-wise square, $\operatorname{diag}$ places these activation derivatives on the diagonal. The entire $J_t$ has shape $H\times H$.

Because:

When $h_t$ approaches $1$ or $-1$, the derivative approaches 0, called saturation. Gradients passing through these directions tend to get compressed.

The final manual calculation: end-point gradient is 1, local derivatives at five consecutive steps are all 0.2:

This time it's correct. A local derivative that doesn't seem particularly small, when multiplied repeatedly, can still rapidly weaken learning signals from far away.

In the matrix case, it depends on the product of Jacobians; we cannot just check if some element of $W_h$ is greater than 1. Although each step shares $W_h$, activation derivatives vary with state, so the complete Jacobians at different time steps are usually different.