Next Word Prediction Using Recurrent Neural Network
Sequence models and memory
Overview
We’ve seen a multilayer perceptron modeled both by hand and with a framework. Each input goes through a weighted sum plus a shift, then an activation function:
\[ z = b + \sum_{i=1}^{m} x_i w_i \]
It worked for our binary classification, but it breaks down the moment inputs form a sequence, where order carries meaning. Take this sentence:
the cat sat on the ___
Most English speakers will fill in the blank with “mat” — because of the context from the words that came before it. A multilayer perceptron has no mechanism for “before”. Its input size is fixed, and every prediction starts from a blank slate.
Researchers created the recurrent neural network (RNN) to handle sequential data. Recurrent means each step passes a hidden state forward into the next step — the network carries a learned summary of its past forward.
To explore this, we’ll build a next-word predictor: given the words so far, output a probability for every word that could come next. It uses the same basic autoregressive cross-entropy idea as modern language models, although large models usually predict tokens rather than whole words and use vastly more parameters. Ours will have 28.
Sequences and Time Steps
The term next-word predictor might already feel intuitive: it’s the same setup as the language models mentioned in the overview, where a prefix such as your question, is what the model continues from.
In our predictor, [ ] marks the word being predicted, and it slides forward one word at a time:
the [ ] → the cat [ ] → the cat sat [ ] → the cat sat on [ ] → the cat sat on the [ ]
The beauty is that every position of every sentence is a free labeled example — the sentence itself fills each bracket, no human annotation required.
The network iterates through the positions in order, one per time step \(t\):
| time step | \(t=1\) | \(t=2\) | \(t=3\) | \(t=4\) | \(t=5\) |
|---|---|---|---|---|---|
| input \(x_t\) | the | cat | sat | on | the |
| target \(y_t\) | cat | sat | on | the | mat |
It’s worth noting that \(t\) is not a timestamp, instead, it’s just a position in the sequence.
Words into Numbers → Indexing and Embedding
An ML model might be arbitrarily complex, but it’s still built on matrix multiplication. Therefore words need to be turned into numbers before we can form an input vector. An initial instinct might be to just use ASCII values. However, think about encoding for word “on” ([6F 6E]) and “mat” ([6D 61 74]), it would produce vectors of different sizes. Moreover, “cat” sits next to “bat” and nowhere near “kitten”, and the network would have to waste capacity unlearning those accidents of spelling. The network needs same-length vectors whose values aren’t hard-coded at all, but are learnable parameters.
Indexing fixes each word’s position. Our vocabulary count is \(V = 5\), so every word will become a vector with 5 numbers:
| word | the | cat | sat | on | mat |
|---|---|---|---|---|---|
| index | 0 | 1 | 2 | 3 | 4 |
A particular word is selected with a one-hot vector. All it does is turn the number at the word’s index to 1, and the rest to 0. For example, “cat” would be [0, 1, 0, 0, 0].
Embedding creates learnable vectors. The embedding matrix holds \(V \times d\) values, where \(d\) is a hyperparameter we choose. A small \(d\) makes words rigid. With \(d = 1\), a word can only be similar to other words in one way at a time. Too big, and training becomes difficult — most of those values would never earn their keep. Classic word vectors used \(d = 50\)–\(300\); large LLMs use thousands. For simplicity, ours is \(d = 2\).
\[ \underbrace{\begin{bmatrix} 0 & 1 & 0 & 0 & 0 \end{bmatrix}}_{\text{one-hot ("cat")}} \cdot \underbrace{\begin{bmatrix} 0.6 & -0.2 \\ 0.4 & 0.7 \\ -0.5 & 0.3 \\ 0.1 & -0.8 \\ -0.3 & -0.6 \end{bmatrix}}_{\text{embedding } E} = \underbrace{\begin{bmatrix} 0.4 & 0.7 \end{bmatrix}}_{\text{input vector } x_t} \]
Above, “cat” is essentially selecting the second row of the embedding matrix. So the whole process is just a word lookup table, except that the table’s values are not hard-coded. They start random and get updated by gradient descent like any other weight, so the network learns where to place each word.
Adding Memory
Forming memory means the previous output must affect the current output in some way, and researchers already figured this out. Remember the perceptron neuron: \(\hat{y} = \sigma(z)\), where \(z = b + \sum_i x_i w_i\). The recurrent version adds one term to that sum, which is the neuron’s own activation from the previous step:
\[ h_t = \tanh\Big(\underbrace{\textstyle\sum_i x_{t,i}\, w_{x,i}}_{\text{current input}} + \underbrace{h_{t-1}\, w_h}_{\text{its own past}}\Big) \]
- \(h_t\) is the neuron’s activation at step \(t\), called the hidden state.
- \(\tanh\) squashes output between \(-1\) and \(1\), often can be expressed as \({{e^{2x}}-1} \over {{e^{2x}}+1}\).
- \(h_{t-1}\) is the same quantity one step earlier. At the first step there is no past, so \(h_0\) is all zeros.
- \(w_h\) is a new learnable weight deciding how much of the past to carry forward.
Stack many of these neurons side by side and the scalars become the matrix form our calculations use:
\[ h_t = \tanh(x_t W_{xh} + h_{t-1} W_{hh}) \]
where \(W_{xh}\) transforms the current word and \(W_{hh}\) carries the old memory into the new one.
The output at step \(t\) is the hidden state projected onto a score for every vocabulary word:
\[ z_t = h_t W_{hy} \]
For turning those \(V\) scores into a prediction, one more activation function is worth mentioning — softmax:
\[ \hat{y}_{t,k} = \frac{e^{z_{t,k}}}{\sum_{j=1}^{V} e^{z_{t,j}}} \]
Sigmoid squashes one number into one probability; softmax squashes a whole vector into a probability distribution that sums to 1 — exactly what a multiclass prediction needs.
rolled unrolled across time
ŷt ŷ1 ŷ2 ŷ3
▲ ▲ ▲ ▲
┌─────┐ ┌─────┐ ┌─────┐ ┌─────┐
h ─▶ RNN ─▶ h h0 ──▶│ RNN │────▶│ RNN │────▶│ RNN │──▶ h3
└─────┘ └─────┘ h1 └─────┘ h2 └─────┘
▲ ▲ ▲ ▲
xt x1 x2 x3
"the" "cat" "sat"Calculation by Hand
We will run RNN from \(t=1\) to \(t=3\). Since we already have grasp on matrix operations, we’ll skip some of the tedious details.
Step 1 - Set Up Weights
\[ E = \begin{bmatrix} 0.6 & -0.2 \\ 0.4 & 0.7 \\ -0.5 & 0.3 \\ 0.1 & -0.8 \\ -0.3 & -0.6 \end{bmatrix} \]
\[ W_{xh} = \begin{bmatrix} 0.5 & -0.3 \\ 0.2 & 0.4 \end{bmatrix} \]
\[ W_{hh} = \begin{bmatrix} 0.1 & 0.4 \\ -0.2 & 0.3 \end{bmatrix} \]
\[ W_{hy} = \begin{bmatrix} 0.7 & -0.1 & 0.2 & -0.5 & 0.1 \\ -0.3 & 0.6 & -0.4 & 0.2 & 0.3 \end{bmatrix} \]
Rows of \(E\) are the word embeddings, top to bottom: the, cat, sat, on, mat.
This toy network omits bias vectors so the arithmetic stays small. That is why the parameter count is exactly \(10 + 4 + 4 + 10 = 28\): 10 values in \(E\), 4 in \(W_{xh}\), 4 in \(W_{hh}\), and 10 in \(W_{hy}\).
Step 2 - Run the Forward Pass
Every time step runs the same six computations: look up the embedding, transform the input (\(x_t W_{xh}\)), transform the memory (\(h_{t-1} W_{hh}\)), add and squash into the new hidden state, project onto the vocabulary, and softmax.
\(t = 1\) — input “the”, target “cat”
1. Embedding lookup. one-hot vector for “the” is [1, 0, 0, 0, 0], so \(x_1\) is row 0 of \(E\):
\[ x_1 = \begin{bmatrix} 0.6 & -0.2 \end{bmatrix}, \qquad h_0 = \begin{bmatrix} 0 & 0 \end{bmatrix} \text{ (memory starts empty)} \]
2. Transform the input. Dot \(x_1\) against each column of \(W_{xh}\):
\[ x_1 W_{xh} = \begin{bmatrix} 0.6 & -0.2 \end{bmatrix} \begin{bmatrix} 0.5 & -0.3 \\ 0.2 & 0.4 \end{bmatrix} = \begin{bmatrix} 0.26 & -0.26 \end{bmatrix} \]
3. Transform the memory. Same recipe with \(h_0\) and \(W_{hh}\)
\[ h_0 W_{hh} = \begin{bmatrix} 0 & 0 \end{bmatrix} \begin{bmatrix} 0.1 & 0.4 \\ -0.2 & 0.3 \end{bmatrix} = \begin{bmatrix} 0 & 0 \end{bmatrix} \]
4. Add and squash.
\[ a_1 = x_1 W_{xh} + h_0 W_{hh} = \begin{bmatrix} 0.26 + 0 & -0.26 + 0 \end{bmatrix} = \begin{bmatrix} 0.26 & -0.26 \end{bmatrix} \]
\[ h_1 = \begin{bmatrix} \tanh(0.26) & \tanh(-0.26) \end{bmatrix} = \begin{bmatrix} 0.25430 & -0.25430 \end{bmatrix} \]
5. Project onto the vocabulary.
\[ \begin{aligned} z_{1,\text{the}} &= (0.25430)(0.7) + (-0.25430)(-0.3) = 0.25430 \\ z_{1,\text{cat}} &= (0.25430)(-0.1) + (-0.25430)(0.6) = -0.17801 \\ z_{1,\text{sat}} &= (0.25430)(0.2) + (-0.25430)(-0.4) = 0.15258 \\ z_{1,\text{on}} &= (0.25430)(-0.5) + (-0.25430)(0.2) = -0.17801 \\ z_{1,\text{mat}} &= (0.25430)(0.1) + (-0.25430)(0.3) = -0.05086 \end{aligned} \]
\[ z_1 = \begin{bmatrix} 0.25430 & -0.17801 & 0.15258 & -0.17801 & -0.05086 \end{bmatrix} \]
6. Softmax.
\[ e^{z_1} \approx \begin{bmatrix} 1.28956 & 0.83693 & 1.16484 & 0.83693 & 0.95041 \end{bmatrix}, \qquad \sum_j e^{z_{1,j}} \approx 5.07867 \]
\[ \hat{y}_1 = \frac{e^{z_1}}{5.07867} \approx \begin{bmatrix} 0.25392 & \mathbf{0.16479} & 0.22936 & 0.16479 & 0.18714 \end{bmatrix} \]
The network gives the word “cat” a probability of 0.16479.
\(t = 2\) — input “cat”, target “sat”
1. Embedding lookup. one-hot vector for “cat” is [0, 1, 0, 0, 0], so \(x_2\) is row 1 of \(E\). we inherit memory from \(t=1\):
\[ x_2 = \begin{bmatrix} 0.4 & 0.7 \end{bmatrix}, \qquad h_1 = \begin{bmatrix} 0.25430 & -0.25430 \end{bmatrix} \]
2. Transform the input.
\[ x_2 W_{xh} = \begin{bmatrix} 0.4 & 0.7 \end{bmatrix} \begin{bmatrix} 0.5 & -0.3 \\ 0.2 & 0.4 \end{bmatrix} = \begin{bmatrix} 0.34 & 0.16 \end{bmatrix} \]
3. Transform the memory. This is where the recurrence wakes up. Both components of \(h_1\) feed both components of the result:
\[ h_1 W_{hh} = \begin{bmatrix} 0.25430 & -0.25430 \end{bmatrix} \begin{bmatrix} 0.1 & 0.4 \\ -0.2 & 0.3 \end{bmatrix} = \begin{bmatrix} 0.07629 & 0.02543 \end{bmatrix} \]
That vector is the memory of “the” leaking into the present. Every later step compounds this.
4. Add and squash.
\[ a_2 = \begin{bmatrix} 0.34 + 0.07629 & 0.16 + 0.02543 \end{bmatrix} = \begin{bmatrix} 0.41629 & 0.18543 \end{bmatrix} \]
\[ h_2 = \begin{bmatrix} \tanh(0.41629) & \tanh(0.18543) \end{bmatrix} = \begin{bmatrix} 0.39380 & 0.18333 \end{bmatrix} \]
5. Project onto the vocabulary.
\[ \begin{aligned} z_{2,\text{the}} &= (0.39380)(0.7) + (0.18333)(-0.3) = 0.22066 \\ z_{2,\text{cat}} &= (0.39380)(-0.1) + (0.18333)(0.6) = 0.07062 \\ z_{2,\text{sat}} &= (0.39380)(0.2) + (0.18333)(-0.4) = 0.00543 \\ z_{2,\text{on}} &= (0.39380)(-0.5) + (0.18333)(0.2) = -0.16023 \\ z_{2,\text{mat}} &= (0.39380)(0.1) + (0.18333)(0.3) = 0.09438 \end{aligned} \]
\[ z_2 = \begin{bmatrix} 0.22066 & 0.07062 & 0.00543 & -0.16023 & 0.09438 \end{bmatrix} \]
6. Softmax.
\[ e^{z_2} \approx \begin{bmatrix} 1.24690 & 1.07317 & 1.00544 & 0.85195 & 1.09898 \end{bmatrix}, \qquad \sum_j e^{z_{2,j}} \approx 5.27644 \]
\[ \hat{y}_2 = \frac{e^{z_2}}{5.27644} \approx \begin{bmatrix} 0.23631 & 0.20339 & \mathbf{0.19055} & 0.16146 & 0.20828 \end{bmatrix} \]
The correct word “sat” gets 0.19055.
\(t = 3\) — input “sat”, target “on”
1. Embedding lookup. “sat” is index 2, so \(x_3\) is row 2 of \(E\):
\[ x_3 = \begin{bmatrix} -0.5 & 0.3 \end{bmatrix}, \qquad h_2 = \begin{bmatrix} 0.39380 & 0.18333 \end{bmatrix} \]
2. Transform the input.
\[ x_3 W_{xh} = \begin{bmatrix} -0.5 & 0.3 \end{bmatrix} \begin{bmatrix} 0.5 & -0.3 \\ 0.2 & 0.4 \end{bmatrix} = \begin{bmatrix} -0.19 & 0.27 \end{bmatrix} \]
3. Transform the memory. Note that \(h_2\) now carries traces of both “the” and “cat”:
\[ h_2 W_{hh} = \begin{bmatrix} 0.39380 & 0.18333 \end{bmatrix} \begin{bmatrix} 0.1 & 0.4 \\ -0.2 & 0.3 \end{bmatrix} = \begin{bmatrix} 0.00271 & 0.21252 \end{bmatrix} \]
4. Add and squash.
\[ a_3 = \begin{bmatrix} -0.19 + 0.00271 & 0.27 + 0.21252 \end{bmatrix} = \begin{bmatrix} -0.18729 & 0.48252 \end{bmatrix} \]
\[ h_3 = \begin{bmatrix} \tanh(-0.18729) & \tanh(0.48252) \end{bmatrix} = \begin{bmatrix} -0.18513 & 0.44826 \end{bmatrix} \]
5. Project onto the vocabulary.
\[ \begin{aligned} z_{3,\text{the}} &= (-0.18513)(0.7) + (0.44826)(-0.3) = -0.26407 \\ z_{3,\text{cat}} &= (-0.18513)(-0.1) + (0.44826)(0.6) = 0.28747 \\ z_{3,\text{sat}} &= (-0.18513)(0.2) + (0.44826)(-0.4) = -0.21633 \\ z_{3,\text{on}} &= (-0.18513)(-0.5) + (0.44826)(0.2) = 0.18222 \\ z_{3,\text{mat}} &= (-0.18513)(0.1) + (0.44826)(0.3) = 0.11597 \end{aligned} \]
\[ z_3 = \begin{bmatrix} -0.26407 & 0.28747 & -0.21633 & 0.18222 & 0.11597 \end{bmatrix} \]
6. Softmax.
\[ e^{z_3} \approx \begin{bmatrix} 0.76792 & 1.33305 & 0.80547 & 1.19988 & 1.12296 \end{bmatrix}, \qquad \sum_j e^{z_{3,j}} \approx 5.22928 \]
\[ \hat{y}_3 = \frac{e^{z_3}}{5.22928} \approx \begin{bmatrix} 0.14685 & 0.25492 & 0.15403 & \mathbf{0.22945} & 0.21474 \end{bmatrix} \]
The correct word “on” gets 0.22945.
Step 3 - Calculate Loss
The output is a full probability distribution, so we use categorical cross-entropy. Instead of binary classification, we use it for \(V\) classes.
\[ L_t = -\sum_{k=1}^{V} y_{t,k} \ln \hat{y}_{t,k} = -\ln \hat{y}_t[\text{target}] \]
\[ \begin{aligned} L_1 = -\ln(0.16479) = 1.80306 \\ L_2 = -\ln(0.19055) = 1.65782 \\ L_3 = -\ln(0.22945) = 1.47206 \end{aligned} \]
Average over the sequence:
\[ L = \frac{1}{3}(1.80306 + 1.65782 + 1.47206) = 1.64431 \]
Step 4 - Backpropagation Through Time
Backpropagation through time (BPTT) has the same concept compared to backpropagation with simple perceptron, except now it’s unrolling hidden state dependency through time step and applying the chain rule.
Establish Common Factors
Since the hidden state at \(t=3\) is affected by the hidden state at \(t=2\), the chain rule can get quite long. It helps to isolate common factors so we don’t rewrite them for every timestep.
For example, every weight gradient at \(t=3\) begins with the same two factors:
\[ \frac{\partial L_3}{\partial z_3} = \frac{\partial L_3}{\partial \hat{y}_3} \cdot \frac{\partial \hat{y}_3}{\partial z_3} \]
The same applies at any timestep, so we call this quantity the output error: \[ output \space error = \delta^z_t = \frac{\partial L}{\partial z_t} = \left( -z_{t, target} +\ln \left( {\sum_{j=1}^{V} e^{z_{t,j}}} \right) \right)' = \hat{y}_t-y_t. \]
It reads as “how much the loss at timestep \(t\) changes when the logits \(z_t\) change.” Every gradient formula that follows will reuse \(\delta^z_t\) instead of re-deriving these factors.
Another common factor is our tanh activation function, which can be expressed as:
\[ tanh'(a_t) = \frac{\partial h_t}{\partial a_t} = \left( {{e^{2a_t}}-1} \over {{e^{2a_t}}+1} \right)' = 1-h_t^2, \]
Multivariable Chain Rule Review
We’ve seen the ordinary chain rule in the perceptron. Given equations:
\[ x=g(w)=w^2, \quad y=f(x)=x^3 \rightarrow y=f(g(w)) \]
The output path can be described as \(w \rightarrow x \rightarrow y\). When we want to determine how a change of \(w\) affects the output of \(y\), we can calculate the derivative using the chain rule:
\[ \left(f(g(w))\right)'=\frac{dy}{dw}=\frac{dy}{dx}\cdot\frac{dx}{dw} \]
However, in RNN, a weight can influence the total loss through multiple paths. Take \(\frac{\partial L}{\partial W_{hy}}\) as an example, from forward pass, we know \(L\) can be reached in three ways:
- Nudging \(W_{hy}\) will change \(L_1\), which will affect total loss \(L\)
- Nudging \(W_{hy}\) will change \(L_2\), which will affect total loss \(L\)
- Nudging \(W_{hy}\) will change \(L_3\), which will affect total loss \(L\)
When an output depends on an input through multiple paths, the derivative along each path is calculated separately, and the result are summed:
\[ \frac{\partial L}{\partial W_{hy}} = \frac{\partial L_1}{\partial W_{hy}} + \frac{\partial L_2}{\partial W_{hy}} + \frac{\partial L_3}{\partial W_{hy}} \]
Calculate Output Gradient \(W_{hy}\)
Gradient at each time step \(t\) can be expressed as:
\[ \frac{\partial L_t}{\partial W_{hy}} = \frac{\partial L_t}{\partial \hat{y}_t} \cdot \frac{\partial \hat{y}_t}{\partial z_t} \cdot \frac{\partial z_t}{\partial W_{hy}} \]
Using common factor \(\delta^z_t\), and \(\frac{\partial z_t}{\partial W_{hy}}=(h_tW_{hy})'=h_t\). Now the gradient is simply:
\[ \frac{\partial L_t}{\partial W_{hy}} = h_t^\top \delta^z_t \]
\(\top\) means transpose. It rotates the hidden state vector to match the dimension of the output error to make matrix operation valid.
Since our total loss is the average over three timesteps, \(L = \frac{1}{3}\sum_{t=1}^{3} L_t\), each per-timestep gradient picks up a factor of \(\frac{1}{3}\). So \(\delta^z_3\) would be \(\frac{1}{3}(\hat{y}_3 - y_3)\).
\[ \begin{aligned} \delta^z_3 = \frac{1}{3} \left( \begin{bmatrix} 0.14685 & 0.25492 & 0.15403 & \mathbf{0.22945} & 0.21474 \end{bmatrix} - \begin{bmatrix} 0 & 0 & 0 & 1 & 0 \end{bmatrix} \right) \\ h_3^\top \delta^z_3 = \begin{bmatrix} -0.18513 \\ 0.44826 \end{bmatrix} \begin{bmatrix} 0.04895 & 0.08497 & 0.05134 & \mathbf{−0.25685} & 0.07158 \end{bmatrix} \\ \frac{\partial L_3}{\partial W_{hy}} = \begin{bmatrix} -0.00906 & -0.01573 & -0.00950 & 0.04755 & -0.01325\\ 0.02194 & 0.03809 & 0.02301 & -0.11514 & 0.03209 \end{bmatrix} \end{aligned} \]
\[ \begin{aligned} \delta^z_2 = \frac{1}{3} \left( \begin{bmatrix} 0.23631 & 0.20339 & \mathbf{0.19055} & 0.16146 & 0.20828 \end{bmatrix} - \begin{bmatrix} 0 & 0 & 1 & 0 & 0 \end{bmatrix} \right) \\ h_2^\top \delta^z_2 = \begin{bmatrix} 0.39380 \\ 0.18333 \end{bmatrix} \begin{bmatrix} 0.07877 & 0.06780 & \mathbf{−0.26982} & 0.05382 & 0.06943 \end{bmatrix} \\ \frac{\partial L_2}{\partial W_{hy}} = \begin{bmatrix} 0.03102 & 0.02670 & −0.10626 & 0.02119 & 0.02734\\ 0.01444 & 0.01243 & −0.04947 & 0.00987 & 0.01273 \end{bmatrix} \end{aligned} \]
\[ \begin{aligned} \delta^z_1 = \frac{1}{3} \left( \begin{bmatrix} 0.25392 & \mathbf{0.16479} & 0.22936 & 0.16479 & 0.18714 \end{bmatrix} - \begin{bmatrix} 0 & 1 & 0 & 0 & 0 \end{bmatrix} \right) \\ h_1^\top \delta^z_1 = \begin{bmatrix} 0.25430 \\ -0.25430 \end{bmatrix} \begin{bmatrix} 0.08464 & \mathbf{−0.27840} & 0.07645 & 0.05493 & 0.06238 \end{bmatrix} \\ \frac{\partial L_1}{\partial W_{hy}} = \begin{bmatrix} 0.02152 & −0.07080 & 0.01944 & 0.01397 & 0.01586\\ −0.02152 & 0.07080 & −0.01944 & −0.01397 & −0.01586 \end{bmatrix} \end{aligned} \]
After summing up the gradients, we have:
\[ \frac{\partial L}{\partial W_{hy}} = \begin{bmatrix} 0.04348 & −0.05983 & −0.09632 & 0.08271 & 0.02995\\ 0.01486 & 0.12132 & −0.04590 & −0.11924 & 0.02896 \end{bmatrix} \]
Jacobian Matrix Review
So far, we’ve only solved derivative that takes in a scalar and produces scalar. However, for hidden state, we have matrix for input, and matrix for output. To gain some intuition, it helps to start from simplest example. Consider this equation:
\[ \begin{bmatrix} f(x_1, x_2) \\[8pt] g(x_1, x_2) \end{bmatrix} = \begin{bmatrix} x_1^2+x_2 \\[8pt] 3x_1+2x_2 \end{bmatrix} , \quad x_1=2 , \quad x_2=1 \]
If we want to determine how changes to either input \(x_1\) or \(x_2\) affect the output, we can construct a matrix with all the partial derivatives:
\[ \begin{bmatrix} {\partial f / \partial x_1} & {\partial f / \partial x_2} \\[8pt] {\partial g / \partial x_1} & {\partial g / \partial x_2} \end{bmatrix} = \begin{bmatrix} 4 & 1 \\[8pt] 3 & 2 \end{bmatrix} \]
This matrix is called the Jacobian. It would tell us that how would nudging each input affect the output? To demonstrate this, say we want to increase \(x_1\) by a 0.01. The \(\delta\) (delta - change) in output can be expressed as:
\[ \Delta \mathrm{output} \approx J \, \Delta x = \begin{bmatrix} 4 & 1 \\ 3 & 2 \end{bmatrix} \begin{bmatrix} \Delta x_1 \\ \Delta x_2 \end{bmatrix} = \begin{bmatrix} 4 & 1 \\ 3 & 2 \end{bmatrix} \begin{bmatrix} 0.01 \\ 0 \end{bmatrix} = \begin{bmatrix} 0.04 \\ 0.03 \end{bmatrix} \]
To verify this, we can plug this increment into the our function \(f\) and \(g\):
\[ f(2.01, 1) - f(2, 1) \approx 0.04, \qquad g(2.01, 1) - g(2, 1) \approx 0.03 \]
This is crucial component for understanding \(W_{hh}\) gradient calculation.
Calculate Input Gradient \(W_{xh}\)
The total gradient decomposes over the three losses, exactly as before:
\[ \frac{\partial L}{\partial W_{xh}} = \frac{\partial L_3}{\partial W_{xh}} + \frac{\partial L_2}{\partial W_{xh}} + \frac{\partial L_1}{\partial W_{xh}} \]
We’ll again walk through \({\partial L_3}/{\partial W_{xh}}\) term by term. From the forward pass, \(W_{xh}\) reaches \(h_3\) through the same three routes \(W_{hh}\) did, because both weights sit inside the same pre-activation \(a_t = x_t W_{xh} + h_{t-1} W_{hh}\):
\[ \begin{aligned} W_{xh} \space \text{used at t=3} \rightarrow h_3 \\ W_{xh} \space \text{used at t=2} \rightarrow h_2 \rightarrow h_3 \\ W_{xh} \space \text{used at t=1} \rightarrow h_1 \rightarrow h_2 \rightarrow h_3 \\ \end{aligned} \]
So the hidden state changes with respect to \(W_{xh}\) as:
\[ \frac{\partial h_3}{\partial W_{xh}} = \frac{\partial^+ h_3}{\partial W_{xh}} + \frac{\partial h_3}{\partial h_2}\frac{\partial^+ h_2}{\partial W_{xh}} + \frac{\partial h_3}{\partial h_2}\frac{\partial h_2}{\partial h_1}\frac{\partial^+ h_1}{\partial W_{xh}} \]
and the gradient at \(t=3\) carries the same bracket:
\[ \frac{\partial L_3}{\partial W_{xh}} = \frac{\partial L_3}{\partial h_3}\left[\underbrace{\frac{\partial^+ h_3}{\partial W_{xh}}}_{\text{used at } t=3} + \underbrace{\frac{\partial h_3}{\partial h_2}\frac{\partial^+ h_2}{\partial W_{xh}}}_{\text{used at } t=2} + \underbrace{\frac{\partial h_3}{\partial h_2}\frac{\partial h_2}{\partial h_1}\frac{\partial^+ h_1}{\partial W_{xh}}}_{\text{used at } t=1}\right] \]
Expanding the term inside the bracket used at \(t=3\):
\[ \left.\frac{\partial L_3}{\partial W_{xh}}\right|_{\text{used at } t=3} = \underbrace{ \frac{\partial L_3}{\partial \hat{y}_3} \cdot \frac{\partial \hat{y}_3}{\partial z_3} }_{\delta^z_3} \cdot \underbrace{ \frac{\partial z_3}{\partial h_3} }_{W_{hy}^\top} \odot \underbrace{ \frac{\partial h_3}{\partial a_3} }_{1-h_3^2} \cdot \underbrace{ \frac{\partial a_3}{\partial W_{xh}} }_{x_3} \]
Now compare that against the \(W_{hh}\) version of the same slice:
\[ \begin{aligned} \left.\frac{\partial L_3}{\partial W_{hh}}\right|_{\text{used at } t=3} = \delta^z_3 \cdot W_{hy}^\top \odot (1-h_3^2) \cdot \underbrace{h_2}_{\partial^+ a_3 / \partial W_{hh}} \\ \left.\frac{\partial L_3}{\partial W_{xh}}\right|_{\text{used at } t=3} = \delta^z_3 \cdot W_{hy}^\top \odot (1-h_3^2) \cdot \underbrace{x_3}_{\partial^+ a_3 / \partial W_{xh}} \end{aligned} \]
Notice how every factor is identical except the last one. We’ve already seen:
\[ \delta^a_3 = \left(\delta^z_3\, W_{hy}^\top\right) \odot \left(1-h_3^2\right) \]
which is the same \(\delta^a_3\) we already computed during the \(W_{hh}\) pass. So the slice shortens to a single outer product, with \(x_3\) swapped in for \(h_2\):
\[ \left.\frac{\partial L_3}{\partial W_{xh}}\right|_{\text{used at } t=3} = x_3^\top\, \delta^{a(3)}_3 \]
Now since the pattern is obvious, we’ll fast track to the final gradient calculation:
\[ \begin{aligned} \left.\frac{\partial L_3}{\partial W_{xh}}\right|_{\text{used at } t=3} = x_3^\top\, \delta^3_3 = \begin{bmatrix} -0.5 \\ 0.3 \end{bmatrix} \begin{bmatrix} 0.16574 & -0.01129 \end{bmatrix} = \begin{bmatrix} -0.08287 & 0.00565 \\ 0.04972 & -0.00339 \end{bmatrix} \end{aligned} \]
\[ \left.\frac{\partial L_3}{\partial W_{xh}}\right|_{\text{used at } t=2} = x_2^\top\, \delta^3_2 = \begin{bmatrix} 0.4 \\ 0.7 \end{bmatrix} \begin{bmatrix} 0.01019 & -0.03531 \end{bmatrix} = \begin{bmatrix} 0.00408 & -0.01412 \\ 0.00713 & -0.02472 \end{bmatrix} \]
\[ \left.\frac{\partial L_3}{\partial W_{xh}}\right|_{\text{used at } t=1} = x_1^\top\, \delta^3_1 = \begin{bmatrix} 0.6 \\ -0.2 \end{bmatrix} \begin{bmatrix} -0.01226 & -0.01181 \end{bmatrix} = \begin{bmatrix} -0.00736 & -0.00709 \\ 0.00245 & 0.00236 \end{bmatrix} \]
Summing the three slices gives \(L_3\)’s full contribution:
\[ \frac{\partial L_3}{\partial W_{xh}} = \begin{bmatrix} -0.08615 & -0.01556 \\ 0.05930 & -0.02574 \end{bmatrix} \]
\(L_2\) only reaches back two steps, so it has two slices. \(L_1\) has only one slice. Therefore we get gradient:
\[ \frac{\partial L_2}{\partial W_{xh}} = \begin{bmatrix} 0.02411 & 0.08842 \\ -0.02604 & 0.09662 \end{bmatrix} , \quad \frac{\partial L_1}{\partial W_{xh}} = \begin{bmatrix} 0.04554 & -0.10849 \\ -0.01518 & 0.03616 \end{bmatrix} \]
Summing the three per-loss gradients gives the input weight gradient.
\[ \begin{aligned} \frac{\partial L}{\partial W_{xh}} = \begin{bmatrix} -0.08615 & -0.01556 \\ 0.05930 & -0.02574 \end{bmatrix} + \begin{bmatrix} 0.02411 & 0.08842 \\ -0.02604 & 0.09662 \end{bmatrix} + \begin{bmatrix} 0.04554 & -0.10849 \\ -0.01518 & 0.03616 \end{bmatrix} \\[8pt] = \begin{bmatrix} -0.01650 & -0.03563 \\ 0.01808 & 0.10704 \end{bmatrix} \end{aligned} \]
Calculate Embedding Gradient \(E\)
The embedding matrix is a parameter like any other, so it decomposes over the three losses the same way:
\[ \frac{\partial L}{\partial E} = \frac{\partial L_3}{\partial E} + \frac{\partial L_2}{\partial E} + \frac{\partial L_1}{\partial E} \]
When we compare term with \(W_{xh}\) gradient calculation, we can observe similar pattern:
\[ \begin{aligned} \left.\frac{\partial L_3}{\partial W_{xh}}\right|_{\text{used at } t=3} = \delta^z_3 \cdot W_{hy}^\top \odot (1-h_3^2) \cdot \underbrace{x_3}_{\partial^+ a_3 / \partial W_{xh}} \\ \left.\frac{\partial L_3}{\partial E}\right|_{\text{used at } t=3} = \delta^z_3 \cdot W_{hy}^\top \odot (1-h_3^2) \cdot \underbrace{W_{xh}^\top \cdot \text{one-hot}}_{\text{two more steps back}} \end{aligned} \]
We fastrack again to final gradient:
\[ \delta^a_1 = \begin{bmatrix} 0.11823 & -0.14612 \end{bmatrix}, \qquad \delta^a_2 = \begin{bmatrix} -0.01142 & 0.11600 \end{bmatrix}, \qquad \delta^a_3 = \begin{bmatrix} 0.16574 & -0.01129 \end{bmatrix} \]
\[ \delta^a_1 W_{xh}^\top = \begin{bmatrix} 0.10295 & -0.03480 \end{bmatrix}, \quad \delta^a_2 W_{xh}^\top = \begin{bmatrix} -0.04051 & 0.04412 \end{bmatrix}, \quad \delta^a_3 W_{xh}^\top = \begin{bmatrix} 0.08626 & 0.02863 \end{bmatrix} \]
The words “on” and “mat” never appear as inputs across \(t=1\ldots3\), so their rows receive no gradient this pass — a direct consequence of the lookup-table view of embeddings. Placing each result in its row:
\[ \frac{\partial L}{\partial E} = \begin{bmatrix} 0.10295 & -0.03480 \\ -0.04051 & 0.04412 \\ 0.08626 & 0.02863 \\ 0 & 0 \\ 0 & 0 \end{bmatrix} \]
Step 5 - Update the Parameters
With every gradient in hand, one step of gradient descent nudges each parameter opposite its gradient, scaled by the learning rate \(\eta = 0.5\):
\[ \theta \leftarrow \theta - \eta\, \frac{\partial L}{\partial \theta} \]
All values are rounded to 5 decimals.
\[ W_{xh} \leftarrow \begin{bmatrix} 0.5 & -0.3 \\ 0.2 & 0.4 \end{bmatrix} - 0.5\begin{bmatrix} -0.01650 & -0.03563 \\ 0.01808 & 0.10704 \end{bmatrix} = \begin{bmatrix} 0.50825 & -0.28219 \\ 0.19096 & 0.34648 \end{bmatrix} \]
\[ W_{hh} \leftarrow \begin{bmatrix} 0.1 & 0.4 \\ -0.2 & 0.3 \end{bmatrix} - 0.5\begin{bmatrix} 0.06236 & 0.02505 \\ 0.03329 & -0.03157 \end{bmatrix} = \begin{bmatrix} 0.06882 & 0.38747 \\ -0.21664 & 0.31578 \end{bmatrix} \]
\[ \begin{aligned} W_{hy} \leftarrow \begin{bmatrix} 0.7 & -0.1 & 0.2 & -0.5 & 0.1 \\ -0.3 & 0.6 & -0.4 & 0.2 & 0.3 \end{bmatrix} - 0.5\begin{bmatrix} 0.04348 & -0.05983 & -0.09632 & 0.08271 & 0.02995 \\ 0.01486 & 0.12132 & -0.04590 & -0.11924 & 0.02896 \end{bmatrix} \\ = \begin{bmatrix} 0.67826 & -0.07009 & 0.24816 & -0.54136 & 0.08502 \\ -0.30743 & 0.53934 & -0.37705 & 0.25962 & 0.28552 \end{bmatrix} \end{aligned} \]
\[ E \leftarrow \begin{bmatrix} 0.6 & -0.2 \\ 0.4 & 0.7 \\ -0.5 & 0.3 \\ 0.1 & -0.8 \\ -0.3 & -0.6 \end{bmatrix} - 0.5\begin{bmatrix} 0.10295 & -0.03480 \\ -0.04051 & 0.04412 \\ 0.08626 & 0.02863 \\ 0 & 0 \\ 0 & 0 \end{bmatrix} = \begin{bmatrix} 0.54852 & -0.18260 \\ 0.42026 & 0.67794 \\ -0.54313 & 0.28569 \\ 0.10000 & -0.80000 \\ -0.30000 & -0.60000 \end{bmatrix} \]
The “on” and “mat” embedding rows are untouched, since they received no gradient this pass. Every other parameter has shifted a little in the direction that lowers the loss.
Verify Effectiveness
Run the same three words through the updated weights:
| before | after | |
|---|---|---|
| \(\hat{y}_1[\text{cat}]\) | 0.16479 | 0.17263 |
| \(\hat{y}_2[\text{sat}]\) | 0.19055 | 0.19861 |
| \(\hat{y}_3[\text{on}]\) | 0.22945 | 0.24168 |
| mean loss | 1.64431 | 1.59772 |
One pass of BPTT nudged every correct word upward and dropped the loss below the randomly guessed result.
Exploding and Vanishing Gradients
In earlier multi-layer perceptron example, we saw that sigmoid derivatives of ~0.07 shrink to \(0.07^3 = 0.000343\) across three layers. It makes weight update extremely slow, and we refer to it as the vanishing gradient problem. RNNs have a worse version. First, an unrolled RNN over a 100-word sentence is effectively a 100-layer network. Second, each hop back in time multiplies by the same matrix \(W_{hh}\) (times the tanh slope), as we saw in the BPTT section:
\[ \frac{\partial h_t}{\partial h_{t-k}} = \prod_{i=t-k+1}^{t} (1 - h_i^2)\, W_{hh}, \quad k=\text{number of step back} \]
NOTE: \(\prod\) has similar function as \(\sum\). Instead of adding, it multiplies. One caveat: unlike \(\sum\), the order matters here, because matrix multiplication is not commutative.
An MLP multiplies by different matrices per layer, so their effects can partially cancel. Repeated multiplication by one matrix compounds relentlessly in whichever direction that matrix pushes. The scalar version makes the danger obvious — suppose the per-step factor were a single number \(w\):
| steps back \(k\) | 1 | 5 | 10 | 20 |
|---|---|---|---|---|
| \(w = 1.5\): \(\;1.5^k\) | 1.5 | 7.6 | 57.7 | 3,325 |
| \(w = 0.5\): \(\;0.5^k\) | 0.5 | 0.031 | 0.00098 | 0.00000095 |
Nudge the factor above 1 and gradients explode exponentially with distance; below 1 and they vanish. There is no comfortable middle to sit in for every step of every sequence.
Gradient Clipping
Exploding gradients announce themselves loudly: the loss spikes, parameters leap to huge values, and training collapses into NaN. The standard fix is blunt and effective — gradient clipping. Before the update, measure the length of the full gradient vector \(g\); if it exceeds a threshold \(\tau\), rescale it to length \(\tau\) while keeping its direction:
\[ \text{if } \lVert g \rVert > \tau: \quad g \leftarrow \tau \, \frac{g}{\lVert g \rVert} \]
NOTE: \(\lVert g \rVert\) means getting the length of \(g\) vector, also called norm. If \(g=\begin{bmatrix}x & y\end{bmatrix}\), then \(\lVert g \rVert = \sqrt{x^2+y^2}\)
To watch clipping actually fire, suppose a longer sequence had produced a gradient many times larger:
\[ \lVert g \rVert = 6.26718, \qquad \frac{\tau}{\lVert g \rVert} = \frac{1g}{6.26718} = 0.15956g \]
Every one of the parameter entries is multiplied by that single factor \(0.15956\). For the hidden weights:
\[ \begin{bmatrix} 1.24720 & 0.50100 \\ 0.66580 & -0.63140 \end{bmatrix} \longrightarrow \begin{bmatrix} 0.19900 & 0.07994 \\ 0.10624 & -0.10075 \end{bmatrix} \]
Now notice what did not change. Before clipping, the first two entries stood in ratio \(1.24720 / 0.50100 = 2.48942\). After clipping, \(0.19900 / 0.07994 = 2.48942\). Every such ratio is preserved across all parameters.
That is also why the threshold has to be global. Rescaling each weight matrix by its own factor would shrink some blocks more than others, quietly rotating the update into a direction no gradient ever pointed.
The Vanishing
Vanishing gradients fail silently, which makes them the more dangerous sibling. The gradient from the loss at step 50 back to the input at step 1 is a product of ~50 factors below 1, which is effectively zero. The early steps stop receiving credit or blame, so the network never learns long-range dependencies.
Clipping can’t help here, because zero can’t be rescaled into a signal. Researchers solve this problem by using LSTM and GRU. Which will be explored in the next article.
Train in PyTorch and TensorFlow
Now we have solid understanding of the process, it’s time to put into practice. Even though we’ll walk through the steps in both PyTorch and TensorFlow, it’s worth mentioning that PyTorch is more widely used in research as of the time of writing.
Step 1 - Initialize Training Data
To create training data, we need to extract unique vocabulary, so we can build out the embeddings.
sentence = "the cat sat on the mat"
words = sentence.split()
# Extract unique vocabulary to build embedding
vocab = list(dict.fromkeys(words)) # {'the': 0, 'cat': 1, 'sat': 2, 'on': 3, 'mat': 4}
word_to_index = {w: i for i, w in enumerate(vocab)}
embedding_class = len(vocab)Step 2 - Initialize Model
We need to initialize parameters to be identical to our hand calculation:
PyTorch:
import torch
import torch.nn as nn
import torch.nn.functional as F
# Create input and target vector
ids = [word_to_index[w] for w in words]
inputs = torch.tensor([ids[:-1]]) # ['the', 'cat', 'sat', 'on', 'the']
targets = torch.tensor([ids[1:]]) # ['cat', 'sat', 'on', 'the', 'mat']
# One-hot and the embedding lookup
embedding = torch.tensor([[0.6, -0.2], [0.4, 0.7], [-0.5, 0.3], [0.1, -0.8], [-0.3, -0.6]])
one_hot = F.one_hot(inputs[0], num_classes=embedding_class).float()
# Initialize Model
class NextWordRNN(nn.Module):
def __init__(self, vocab_size, embed_dim=2, hidden_dim=2):
super().__init__()
self.embed = nn.Embedding(vocab_size, embed_dim)
self.rnn = nn.RNN(embed_dim, hidden_dim, batch_first=True, bias=False)
self.head = nn.Linear(hidden_dim, vocab_size, bias=False)
def forward(self, idx):
x = self.embed(idx)
h, _ = self.rnn(x)
return self.head(h)
model = NextWordRNN(embedding_class)
# Load previous weights
# PyTorch computes h_t = tanh(x_t W_ih^T + h_{t-1} W_hh^T), so its stored
# matrices are the transpose of ours. The embedding table needs no transpose.
W_xh = torch.tensor([[0.5, -0.3], [0.2, 0.4]])
W_hh = torch.tensor([[0.1, 0.4], [-0.2, 0.3]])
W_hy = torch.tensor([[0.7, -0.1, 0.2, -0.5, 0.1], [-0.3, 0.6, -0.4, 0.2, 0.3]])
with torch.no_grad():
model.embed.weight.copy_(embedding)
model.rnn.weight_ih_l0.copy_(W_xh.T)
model.rnn.weight_hh_l0.copy_(W_hh.T)
model.head.weight.copy_(W_hy.T)TensorFlow:
import numpy as np
import tensorflow as tf
# Create input and target vector
ids = [word_to_index[w] for w in words]
inputs = tf.constant([ids[:-1]]) # the cat sat on the
targets = tf.constant([ids[1:]]) # cat sat on the mat
T = len(ids) - 1
# One-hot and the embedding lookup
E = np.array([[0.6, -0.2], [0.4, 0.7], [-0.5, 0.3], [0.1, -0.8], [-0.3, -0.6]], dtype='float32')
one_hot = tf.one_hot(inputs[0], depth=embedding_class)
# Load previous weights
W_xh = np.array([[0.5, -0.3], [0.2, 0.4]], dtype='float32')
W_hh = np.array([[0.1, 0.4], [-0.2, 0.3]], dtype='float32')
W_hy = np.array([[0.7, -0.1, 0.2, -0.5, 0.1], [-0.3, 0.6, -0.4, 0.2, 0.3]], dtype='float32')
model = tf.keras.Sequential([
tf.keras.layers.Input(shape=(T,)),
tf.keras.layers.Embedding(embedding_class, 2),
tf.keras.layers.SimpleRNN(2, return_sequences=True, use_bias=False),
tf.keras.layers.Dense(embedding_class, use_bias=False),
])
embed, rnn, head = model.layers
embed.set_weights([E])
rnn.set_weights([W_xh, W_hh])
head.set_weights([W_hy])Step 3 - Train Model
We will train for 1 epoch to verify our result against the previous calculation:
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.5)
EPOCHS = 1
for epoch in range(EPOCHS):
optimizer.zero_grad()
logits = model(inputs)
loss = loss_fn(logits.view(-1, embedding_class), targets.view(-1))
loss.backward()
grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
print(f"epoch {epoch+1:3d} loss {loss.item():.5f} grad norm {grad_norm.item():.5f}")
per_step = F.cross_entropy(logits.view(-1, embedding_class), targets.view(-1), reduction='none')
print("per-step loss:", [round(v, 5) for v in per_step.tolist()])
# [1.80306, 1.65782, 1.47206, 1.63924, 1.72186]embed.set_weights([E]); rnn.set_weights([W_xh, W_hh]); head.set_weights([W_hy])
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
with tf.GradientTape() as tape:
logits = model(inputs, training=True)
loss = loss_fn(targets, logits)
grads = tape.gradient(loss, model.trainable_variables)
per_step = tf.keras.losses.sparse_categorical_crossentropy(targets, logits, from_logits=True)
print("per-step loss:", [round(float(v), 5) for v in per_step[0]])
# [1.80306, 1.65782, 1.47206, 1.63924, 1.72186]Wrap Up
We’ve ran successfully constructed a small RNN model. However, there’re few things to know before designing the next model.
The size: 28 parameters, embedding dimension of 2, one sentence, no bias vectors. Real models are the same equations with more parameters.
The sequential bottleneck: \(h_t\) cannot be computed until \(h_{t-1}\) exists. Where a feedforward network processes a whole batch at once, an RNN must walk the sequence one step at a time — a hard limit on how much a GPU can help.
The memory itself: Repeated multiplication by \(W_{hh}\) means the gradient’s reach backward is governed by a product of similar factors. Clipping catches the exploding half. The vanishing half is untreated, and it is the reason a vanilla RNN quietly fails to connect a pronoun to a noun forty words back.
Nothing in the architecture lets the network choose to keep a piece of information and discard another; the hidden state is overwritten at every step. LSTM and GRU address exactly this by adding gates — small learned switches that decide what to retain, what to forget, and what to expose.