Next Word Prediction Using Recurrent Neural Network
Sequence models and memory
Overview
We’ve modeled multilayer perceptrons both by hand and with frameworks. That worked for our binary classification, but it breaks down the moment inputs form a sequence, where order carries meaning. If you’re like me and have a love-hate codependent relationship with LLMs, you should know that they’re essentially a really sophisticated version of next-word predictors, meaning each prediction depends on the words that came before it.
Take this NYT headline: the rise of the machine
Given part of the sentence, we want the model to predict what comes next, such as:
input: the rise of the → output: machine
We’ll build this sequential model using a recurrent neural network (RNN). Recurrent means each step passes a hidden state forward into the next step, so the network carries a learned summary of everything it has seen so far.
Sequences and Time Steps
In our predictor, [ ] marks the word being predicted, and it slides forward one word at a time:
the [ ] → the rise [ ] → the rise of [ ] → the rise of the [ ]
The beauty is that every position of every sentence is a 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\) |
|---|---|---|---|---|
| input \(x_t\) | the | rise | of | the |
| target \(y_t\) | rise | of | the | machine |
It’s worth noting that \(t\) is not a timestamp, is not a timestamp, 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 the words “of” ([6F 66]) and “machine” ([6D 61 63 68 69 6E 65]): it would produce vectors of different sizes. Moreover, “machine” sits closer to “machete” than to “robot”, 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 size is \(V = 4\), so every word will become a vector with 4 numbers:
| word | the | rise | of | machine |
|---|---|---|---|---|
| index | 0 | 1 | 2 | 3 |
A particular word is selected with a one-hot vector. It turns the number at the word’s index to 1, and the rest to 0. For example, “rise” would be [0, 1, 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, training becomes difficult, and 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 \end{bmatrix}}_{\text{one-hot ("rise")}} \cdot \underbrace{\begin{bmatrix} 0.6 & -0.2 \\ 0.4 & 0.7 \\ -0.5 & 0.3 \\ 0.1 & -0.8 \end{bmatrix}}_{\text{embedding } E} = \underbrace{\begin{bmatrix} 0.4 & 0.7 \end{bmatrix}}_{\text{input vector } x_t} \]
Above, “rise” 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.
\(\tanh\) and Softmax Activation Function
It’s time to introduce another activation function:
\[ \tanh(x) = \frac{e^{2x}-1}{e^{2x}+1} \]
\(\tanh\) takes an input and outputs a number between \(-1\) and \(1\). This is useful in modeling sequential data, because it can either increase or decrease the hidden state. Since the output stays bounded no matter how large the input gets, the values can’t grow arbitrarily large, which would otherwise result in enormous step sizes for gradient descent.
Softmax turns a whole vector of scores into a probability distribution whose entries sum to 1. Intuitively, for each output vector \(\hat{y}_t\), we want the entry corresponding to the next word to carry the highest probability. The equation can be expressed as:
\[ \hat{y}_{t,k} = \frac{e^{z_{t,k}}}{\sum_{j=1}^{V} e^{z_{t,j}}} \]
where \(t\) is the time step, \(z_t\) is the vector of raw scores before softmax, often referred to as logits, \(V\) is the vocabulary size, \(k\) is the index of the entry being computed, and \(j\) runs over every index in the vector.
Adding Memory
Forming memory means the previous output must affect the current output in some way, and researchers already figured this out. Recall the perceptron: \(\hat{y} = \sigma(z)\), where \(z = b + \sum_i x_i w_i\). The recurrent version adds one term to that sum — 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.
- \(h_{t-1}\) is the same quantity one step earlier. At step 0, \(h_0\) is set to zero.
- \(w_h\) is a new learnable weight that decides how much of the past to carry forward.
If we stack many of these neurons side by side, the whole layer becomes a single matrix operation:
\[ 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 prediction at step \(t\) is read off from the hidden state:
\[ z_t = h_t W_{hy} \]
where \(z_t\) is the vector of logits that feeds the softmax.
Calculation by Hand
We will blaze through the forward pass from \(t=1\) to \(t=4\). Since we already have grasp on matrix operations, we’ll skip some of the tedious details.
Step 1 - Set Up Weights
Embedding
\[ E = \begin{bmatrix} 0.6 & -0.2 \\ 0.4 & 0.7 \\ -0.5 & 0.3 \\ 0.1 & -0.8 \end{bmatrix} \]
Input Weights
\[ W_{xh} = \begin{bmatrix} 0.5 & -0.3 \\ 0.2 & 0.4 \end{bmatrix} \]
Output Weights
\[ W_{hy} = \begin{bmatrix} 0.7 & -0.1 & 0.2 & -0.5 \\ -0.3 & 0.6 & -0.4 & 0.2 \end{bmatrix} \]
This toy network omits bias vectors so the arithmetic stays small. That is why the parameter count is exactly \(8 + 4 + 4 + 8 = 24\): 8 values in \(E\), 4 in \(W_{xh}\), 4 in \(W_{hh}\), and 8 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 “rise”
1. Embedding lookup. One-hot vector for “the” is [1, 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{rise}} &= (0.25430)(-0.1) + (-0.25430)(0.6) = -0.17801 \\ z_{1,\text{of}} &= (0.25430)(0.2) + (-0.25430)(-0.4) = 0.15258 \\ z_{1,\text{machine}} &= (0.25430)(-0.5) + (-0.25430)(0.2) = -0.17801 \end{aligned} \]
\[ z_1 = \begin{bmatrix} 0.25430 & -0.17801 & 0.15258 & -0.17801 \end{bmatrix} \]
6. Softmax.
\[ e^{z_1} \approx \begin{bmatrix} 1.28956 & 0.83693 & 1.16484 & 0.83693 \end{bmatrix}, \qquad \sum_j e^{z_{1,j}} \approx 4.12826 \]
\[ \hat{y}_1 = \frac{e^{z_1}}{4.12826} \approx \begin{bmatrix} 0.31237 & \mathbf{0.20273} & 0.28216 & 0.20273 \end{bmatrix} \]
The network gives the word “rise” a probability of 0.20273.
\(t = 2\) — input “rise”, target “of”
1. Embedding lookup. One-hot vector for “rise” is [0, 1, 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{rise}} &= (0.39380)(-0.1) + (0.18333)(0.6) = 0.07062 \\ z_{2,\text{of}} &= (0.39380)(0.2) + (0.18333)(-0.4) = 0.00543 \\ z_{2,\text{machine}} &= (0.39380)(-0.5) + (0.18333)(0.2) = -0.16023 \end{aligned} \]
\[ z_2 = \begin{bmatrix} 0.22066 & 0.07062 & 0.00543 & -0.16023 \end{bmatrix} \]
6. Softmax.
\[ e^{z_2} \approx \begin{bmatrix} 1.24690 & 1.07317 & 1.00544 & 0.85195 \end{bmatrix}, \qquad \sum_j e^{z_{2,j}} \approx 4.17746 \]
\[ \hat{y}_2 = \frac{e^{z_2}}{4.17746} \approx \begin{bmatrix} 0.29848 & 0.25690 & \mathbf{0.24068} & 0.20394 \end{bmatrix} \]
The correct word “of” gets 0.24068.
\(t = 3\) — input “of”, target “the”
1. Embedding lookup. “of” 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 “rise”:
\[ 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{rise}} &= (-0.18513)(-0.1) + (0.44826)(0.6) = 0.28747 \\ z_{3,\text{of}} &= (-0.18513)(0.2) + (0.44826)(-0.4) = -0.21633 \\ z_{3,\text{machine}} &= (-0.18513)(-0.5) + (0.44826)(0.2) = 0.18222 \end{aligned} \]
\[ z_3 = \begin{bmatrix} -0.26407 & 0.28747 & -0.21633 & 0.18222 \end{bmatrix} \]
6. Softmax.
\[ e^{z_3} \approx \begin{bmatrix} 0.76792 & 1.33305 & 0.80547 & 1.19988 \end{bmatrix}, \qquad \sum_j e^{z_{3,j}} \approx 4.10632 \]
\[ \hat{y}_3 = \frac{e^{z_3}}{4.10632} \approx \begin{bmatrix} \mathbf{0.18701} & 0.32463 & 0.19615 & 0.29220 \end{bmatrix} \]
The correct word “the” gets 0.18701.
\(t = 4\) — input “the”, target “machine”
1. Embedding lookup. “the” is index 0 again, so \(x_4\) is the same row of \(E\) we already used at \(t=1\):
\[ x_4 = \begin{bmatrix} 0.6 & -0.2 \end{bmatrix}, \qquad h_3 = \begin{bmatrix} -0.18513 & 0.44826 \end{bmatrix} \]
2. Transform the input. Identical input word means an identical input transform:
\[ x_4 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. This is the whole point of the recurrence. At \(t=1\) this term was zero; now it carries three words of context:
\[ h_3 W_{hh} = \begin{bmatrix} -0.18513 & 0.44826 \end{bmatrix} \begin{bmatrix} 0.1 & 0.4 \\ -0.2 & 0.3 \end{bmatrix} = \begin{bmatrix} -0.10816 & 0.06043 \end{bmatrix} \]
4. Add and squash.
\[ a_4 = \begin{bmatrix} 0.26 - 0.10816 & -0.26 + 0.06043 \end{bmatrix} = \begin{bmatrix} 0.15184 & -0.19957 \end{bmatrix} \]
\[ h_4 = \begin{bmatrix} \tanh(0.15184) & \tanh(-0.19957) \end{bmatrix} = \begin{bmatrix} 0.15068 & -0.19697 \end{bmatrix} \]
Same word in, different state out. \(h_4 \neq h_1\) purely because of memory.
5. Project onto the vocabulary.
\[ \begin{aligned} z_{4,\text{the}} &= (0.15068)(0.7) + (-0.19697)(-0.3) = 0.16457 \\ z_{4,\text{rise}} &= (0.15068)(-0.1) + (-0.19697)(0.6) = -0.13325 \\ z_{4,\text{of}} &= (0.15068)(0.2) + (-0.19697)(-0.4) = 0.10892 \\ z_{4,\text{machine}} &= (0.15068)(-0.5) + (-0.19697)(0.2) = -0.11473 \end{aligned} \]
\[ z_4 = \begin{bmatrix} 0.16457 & -0.13325 & 0.10892 & -0.11473 \end{bmatrix} \]
6. Softmax.
\[ e^{z_4} \approx \begin{bmatrix} 1.17889 & 0.87525 & 1.11507 & 0.89161 \end{bmatrix}, \qquad \sum_j e^{z_{4,j}} \approx 4.06082 \]
\[ \hat{y}_4 = \frac{e^{z_4}}{4.06082} \approx \begin{bmatrix} 0.29031 & 0.21554 & 0.27459 & \mathbf{0.21956} \end{bmatrix} \]
The correct word “machine” gets 0.21956.
Step 3 - Calculate Loss
[TODO: refine by introduce probility techniques]
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_{j=1}^{V} y_{t,j} \ln \hat{y}_{t,j} = -\ln \hat{y}_t[\text{target}] \]
\[ \begin{aligned} L_1 = -\ln(0.20273) = 1.59586 \\ L_2 = -\ln(0.24068) = 1.42428 \\ L_3 = -\ln(0.18701) = 1.67659 \\ L_4 = -\ln(0.21956) = 1.51611 \end{aligned} \]
Average over the sequence:
\[ L = \frac{1}{4}(1.59586 + 1.42428 + 1.67659 + 1.51611) = 1.55321 \]
Backpropagation Through Time
Backpropagation through time (BPTT) is the same idea as backpropagation in a simple perceptron, except the hidden state dependency now unrolls across time steps.
Establish Common Factors
Since the hidden state at \(t=4\) is affected by the hidden state at \(t=3\), the chain rule can get quite long. It helps to isolate common factors so we don’t rewrite them for every time step.
For example, every weight gradient at \(t=4\) begins with the same two factors:
\[ \frac{\partial L_4}{\partial z_4} = \frac{\partial L_4}{\partial \hat{y}_4} \cdot \frac{\partial \hat{y}_4}{\partial z_4} \]
The same applies at any timestep, so we call this 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. The expression for its derivative is:
\[ tanh'(x) = \left( \frac{{e^{2x}}-1}{{e^{2x}}+1} \right)' = 1 - \tanh^2(x), \]
However, in the case of RNN, since \(h_t = \tanh(\dots)\), we can replace \(\tanh^2\) with \(h^2_t\):
\[ tanh'(a_t) = \frac{\partial h_t}{\partial a_t} = 1 - h_t^2, \]
Multivariable Chain Rule Review
Before we start to compute the gradient, a slight detour is required to understand the math behind the variation. 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 four ways:
- Nudging \(W_{hy}\) will change \(L_1\)
- Nudging \(W_{hy}\) will change \(L_2\)
- Nudging \(W_{hy}\) will change \(L_3\)
- Nudging \(W_{hy}\) will change \(L_4\)
All the \(L_t\) can change the 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}} + \frac{\partial L_4}{\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\) together with \(\frac{\partial z_t}{\partial W_{hy}}=(h_tW_{hy})'=h_t\), 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 four time steps, \(L = \frac{1}{4}\sum_{t=1}^{4} L_t\), each per-timestep gradient picks up a factor of \(\frac{1}{4}\). So \(\delta^z_4\) would be \(\frac{1}{4}(\hat{y}_4 - y_4)\).
\[ \begin{aligned} \delta^z_4 = \frac{1}{4} \left( \begin{bmatrix} 0.29031 & 0.21554 & 0.27459 & \mathbf{0.21956} \end{bmatrix} - \begin{bmatrix} 0 & 0 & 0 & 1 \end{bmatrix} \right) \\ h_4^\top \delta^z_4 = \begin{bmatrix} 0.15068 \\ -0.19697 \end{bmatrix} \begin{bmatrix} 0.07258 & 0.05389 & 0.06865 & \mathbf{−0.19511} \end{bmatrix} \\ \frac{\partial L_4}{\partial W_{hy}} = \begin{bmatrix} 0.01094 & 0.00812 & 0.01034 & −0.02940\\ −0.01430 & −0.01061 & −0.01352 & 0.03843 \end{bmatrix} \end{aligned} \]
\[ \begin{aligned} \delta^z_3 = \frac{1}{4} \left( \begin{bmatrix} \mathbf{0.18701} & 0.32463 & 0.19615 & 0.29220 \end{bmatrix} - \begin{bmatrix} 1 & 0 & 0 & 0 \end{bmatrix} \right) \\ h_3^\top \delta^z_3 = \begin{bmatrix} -0.18513 \\ 0.44826 \end{bmatrix} \begin{bmatrix} \mathbf{−0.20325} & 0.08116 & 0.04904 & 0.07305 \end{bmatrix} \\ \frac{\partial L_3}{\partial W_{hy}} = \begin{bmatrix} 0.03763 & −0.01502 & −0.00908 & −0.01352\\ −0.09111 & 0.03638 & 0.02198 & 0.03275 \end{bmatrix} \end{aligned} \]
\[ \begin{aligned} \delta^z_2 = \frac{1}{4} \left( \begin{bmatrix} 0.29848 & 0.25690 & \mathbf{0.24068} & 0.20394 \end{bmatrix} - \begin{bmatrix} 0 & 0 & 1 & 0 \end{bmatrix} \right) \\ h_2^\top \delta^z_2 = \begin{bmatrix} 0.39380 \\ 0.18333 \end{bmatrix} \begin{bmatrix} 0.07462 & 0.06422 & \mathbf{−0.18983} & 0.05098 \end{bmatrix} \\ \frac{\partial L_2}{\partial W_{hy}} = \begin{bmatrix} 0.02939 & 0.02529 & −0.07475 & 0.02008\\ 0.01368 & 0.01177 & −0.03480 & 0.00935 \end{bmatrix} \end{aligned} \]
\[ \begin{aligned} \delta^z_1 = \frac{1}{4} \left( \begin{bmatrix} 0.31237 & \mathbf{0.20273} & 0.28216 & 0.20273 \end{bmatrix} - \begin{bmatrix} 0 & 1 & 0 & 0 \end{bmatrix} \right) \\ h_1^\top \delta^z_1 = \begin{bmatrix} 0.25430 \\ -0.25430 \end{bmatrix} \begin{bmatrix} 0.07809 & \mathbf{−0.19932} & 0.07054 & 0.05068 \end{bmatrix} \\ \frac{\partial L_1}{\partial W_{hy}} = \begin{bmatrix} 0.01986 & −0.05069 & 0.01794 & 0.01289\\ −0.01986 & 0.05069 & −0.01794 & −0.01289 \end{bmatrix} \end{aligned} \]
After summing up the gradients, we have:
\[ \frac{\partial L}{\partial W_{hy}} = \begin{bmatrix} 0.09781 & −0.03230 & −0.05555 & −0.00996\\ −0.11158 & 0.08823 & −0.04428 & 0.06763 \end{bmatrix} \]
Jacobian Matrix Review
Before we get to hidden gradient computation, we need another intermission in our quest. So far, we’ve only taken derivatives of functions that map a scalar to a scalar. However, for the hidden state, both input and output are vectors. To gain some intuition, it helps to start from the 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 tells us how the output responds to nudging each input. To demonstrate this, say we want to increase \(x_1\) by 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 functions \(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 a crucial component for understanding \(W_{hh}\) gradient calculation.
Calculate Input Gradient \(W_{xh}\)
The total gradient decomposes over the four losses, exactly as before:
\[ \frac{\partial L}{\partial W_{xh}} = \frac{\partial L_4}{\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_4}/{\partial W_{xh}}\) term by term. From the forward pass, \(W_{xh}\) reaches \(h_4\) through the same four 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=4} \rightarrow h_4 \\ W_{xh} \space \text{used at t=3} \rightarrow h_3 \rightarrow h_4 \\ W_{xh} \space \text{used at t=2} \rightarrow h_2 \rightarrow h_3 \rightarrow h_4 \\ W_{xh} \space \text{used at t=1} \rightarrow h_1 \rightarrow h_2 \rightarrow h_3 \rightarrow h_4 \\ \end{aligned} \]
So the hidden state changes with respect to \(W_{xh}\) as:
\[ \frac{\partial h_4}{\partial W_{xh}} = \frac{\partial^+ h_4}{\partial W_{xh}} + \frac{\partial h_4}{\partial h_3}\frac{\partial^+ h_3}{\partial W_{xh}} + \frac{\partial h_4}{\partial h_3}\frac{\partial h_3}{\partial h_2}\frac{\partial^+ h_2}{\partial W_{xh}} + \frac{\partial h_4}{\partial h_3}\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=4\) carries the same bracket:
\[ \frac{\partial L_4}{\partial W_{xh}} = \frac{\partial L_4}{\partial h_4}\left[\underbrace{\frac{\partial^+ h_4}{\partial W_{xh}}}_{\text{used at } t=4} + \underbrace{\frac{\partial h_4}{\partial h_3}\frac{\partial^+ h_3}{\partial W_{xh}}}_{\text{used at } t=3} + \underbrace{\frac{\partial h_4}{\partial h_3}\frac{\partial h_3}{\partial h_2}\frac{\partial^+ h_2}{\partial W_{xh}}}_{\text{used at } t=2} + \underbrace{\frac{\partial h_4}{\partial h_3}\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=4\):
\[ \left.\frac{\partial L_4}{\partial W_{xh}}\right|_{\text{used at } t=4} = \underbrace{ \frac{\partial L_4}{\partial \hat{y}_4} \cdot \frac{\partial \hat{y}_4}{\partial z_4} }_{\delta^z_4} \cdot \underbrace{ \frac{\partial z_4}{\partial h_4} }_{W_{hy}^\top} \odot \underbrace{ \frac{\partial h_4}{\partial a_4} }_{1-h_4^2} \cdot \underbrace{ \frac{\partial a_4}{\partial W_{xh}} }_{x_4} \]
Now compare that against the \(W_{hh}\) version of the same slice:
\[ \begin{aligned} \left.\frac{\partial L_4}{\partial W_{hh}}\right|_{\text{used at } t=4} = \delta^z_4 \cdot W_{hy}^\top \odot (1-h_4^2) \cdot \underbrace{h_3}_{\partial^+ a_4 / \partial W_{hh}} \\ \left.\frac{\partial L_4}{\partial W_{xh}}\right|_{\text{used at } t=4} = \delta^z_4 \cdot W_{hy}^\top \odot (1-h_4^2) \cdot \underbrace{x_4}_{\partial^+ a_4 / \partial W_{xh}} \end{aligned} \]
Notice how every factor is identical except the last one. We’ve already seen:
\[ \delta^a_4 = \left(\delta^z_4\, W_{hy}^\top\right) \odot \left(1-h_4^2\right) \]
which is the same \(\delta^a_4\) we already computed during the \(W_{hh}\) pass. So the slice shortens to a single outer product, with \(x_4\) swapped in for \(h_3\):
\[ \left.\frac{\partial L_4}{\partial W_{xh}}\right|_{\text{used at } t=4} = x_4^\top\, \delta^{a(4)}_4 \]
Now since the pattern is obvious, we’ll fast-track to the final gradient calculation:
\[ \begin{aligned} \left.\frac{\partial L_4}{\partial W_{xh}}\right|_{\text{used at } t=4} = x_4^\top\, \delta^4_4 &= \begin{bmatrix} 0.6 \\ -0.2 \end{bmatrix} \begin{bmatrix} 0.15314 & -0.05375 \end{bmatrix} \\ &= \begin{bmatrix} 0.09189 & -0.03225 \\ -0.03063 & 0.01075 \end{bmatrix} \end{aligned} \]
\[ \begin{aligned} \left.\frac{\partial L_4}{\partial W_{xh}}\right|_{\text{used at } t=3} = x_3^\top\, \delta^4_3 &= \begin{bmatrix} -0.5 \\ 0.3 \end{bmatrix} \begin{bmatrix} -0.00598 & -0.03736 \end{bmatrix} \\ &= \begin{bmatrix} 0.00299 & 0.01868 \\ -0.00179 & -0.01121 \end{bmatrix} \end{aligned} \]
\[ \begin{aligned} \left.\frac{\partial L_4}{\partial W_{xh}}\right|_{\text{used at } t=2} = x_2^\top\, \delta^4_2 &= \begin{bmatrix} 0.4 \\ 0.7 \end{bmatrix} \begin{bmatrix} -0.01313 & -0.00968 \end{bmatrix} \\ &= \begin{bmatrix} -0.00525 & -0.00387 \\ -0.00919 & -0.00677 \end{bmatrix} \end{aligned} \]
Unlike the \(W_{hh}\) pass, the \(t=1\) slice does not vanish here, because \(W_{xh}\) pairs with \(x_1\) rather than \(h_0\):
\[ \begin{aligned} \left.\frac{\partial L_4}{\partial W_{xh}}\right|_{\text{used at } t=1} = x_1^\top\, \delta^4_1 &= \begin{bmatrix} 0.6 \\ -0.2 \end{bmatrix} \begin{bmatrix} -0.00485 & -0.00026 \end{bmatrix} \\ &= \begin{bmatrix} -0.00291 & -0.00016 \\ 0.00097 & 0.00005 \end{bmatrix} \end{aligned} \]
Summing the four slices gives \(L_4\)’s full contribution:
\[ \frac{\partial L_4}{\partial W_{xh}} = \begin{bmatrix} 0.08671 & -0.01760 \\ -0.04064 & -0.00718 \end{bmatrix} \]
\(L_3\) only reaches back three steps, so it has three slices:
\[ \frac{\partial L_3}{\partial W_{xh}} = \begin{bmatrix} 0.10468 & -0.01080 \\ -0.04619 & 0.06250 \end{bmatrix} \]
\(L_2\) only reaches back two steps, so it has two slices. \(L_1\) has only one slice. Therefore we get the gradient:
\[ \frac{\partial L_2}{\partial W_{xh}} = \begin{bmatrix} 0.01539 & 0.05785 \\ -0.01755 & 0.06308 \end{bmatrix} , \quad \frac{\partial L_1}{\partial W_{xh}} = \begin{bmatrix} 0.03556 & -0.09041 \\ -0.01185 & 0.03014 \end{bmatrix} \]
Summing the four per-loss gradients gives the input weight gradient.
\[ \begin{aligned} \frac{\partial L}{\partial W_{xh}} &= \begin{bmatrix} 0.08671 & -0.01760 \\ -0.04064 & -0.00718 \end{bmatrix} + \begin{bmatrix} 0.10468 & -0.01080 \\ -0.04619 & 0.06250 \end{bmatrix} + \begin{bmatrix} 0.01539 & 0.05785 \\ -0.01755 & 0.06308 \end{bmatrix} + \begin{bmatrix} 0.03556 & -0.09041 \\ -0.01185 & 0.03014 \end{bmatrix} \\ &= \begin{bmatrix} 0.24234 & -0.06095 \\ -0.11624 & 0.14854 \end{bmatrix} \end{aligned} \]
Calculate Embedding Gradient \(E\)
The embedding matrix is a parameter like any other, so it decomposes over the four losses the same way:
\[ \frac{\partial L}{\partial E} = \frac{\partial L_4}{\partial E} + \frac{\partial L_3}{\partial E} + \frac{\partial L_2}{\partial E} + \frac{\partial L_1}{\partial E} \]
Comparing term with \(W_{xh}\) gradient calculation, we can see a similar pattern:
\[ \begin{aligned} \left.\frac{\partial L_4}{\partial W_{xh}}\right|_{\text{used at } t=4} = \delta^z_4 \cdot W_{hy}^\top \odot (1-h_4^2) \cdot \underbrace{x_4}_{\partial^+ a_4 / \partial W_{xh}} \\ \left.\frac{\partial L_4}{\partial E}\right|_{\text{used at } t=4} = \delta^z_4 \cdot W_{hy}^\top \odot (1-h_4^2) \cdot \underbrace{W_{xh}^\top \cdot \text{one-hot}}_{\text{two more steps back}} \end{aligned} \]
We fast-track again to final gradient:
\[ \begin{aligned} \delta^a_1 = \begin{bmatrix} 0.11273 & -0.10692 \end{bmatrix}, \quad \delta^a_1 W_{xh}^\top = \begin{bmatrix} 0.08844 & -0.02022 \end{bmatrix} \\ \delta^a_2 = \begin{bmatrix} -0.01423 & 0.14647 \end{bmatrix}, \quad \delta^a_2 W_{xh}^\top = \begin{bmatrix} -0.05105 & 0.05574 \end{bmatrix} \\ \delta^a_3 = \begin{bmatrix} -0.17701 & 0.04627 \end{bmatrix}, \quad \delta^a_3 W_{xh}^\top = \begin{bmatrix} -0.10239 & -0.01689 \end{bmatrix} \\ \delta^a_4 = \begin{bmatrix} 0.15314 & -0.05375 \end{bmatrix}, \quad \delta^a_4 W_{xh}^\top = \begin{bmatrix} 0.09270 & 0.00913 \end{bmatrix} \end{aligned} \]
“the” is fed in at both \(t=1\) and \(t=4\), so row 0 collects two contributions and they simply add: \(\begin{bmatrix} 0.08844 & -0.02022 \end{bmatrix} + \begin{bmatrix} 0.09270 & 0.00913 \end{bmatrix} = \begin{bmatrix} 0.18114 & -0.01109 \end{bmatrix}\). The word “machine” never appears as an input across \(t=1\ldots4\), so its row receives no gradient this pass:
\[ \frac{\partial L}{\partial E} = \begin{bmatrix} 0.18114 & -0.01109 \\ -0.05105 & 0.05574 \\ -0.10239 & -0.01689 \\ 0 & 0 \end{bmatrix} \]
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.24234 & -0.06095 \\ -0.11624 & 0.14854 \end{bmatrix} = \begin{bmatrix} 0.37883 & -0.26952 \\ 0.25812 & 0.32573 \end{bmatrix} \]
\[ W_{hh} \leftarrow \begin{bmatrix} 0.1 & 0.4 \\ -0.2 & 0.3 \end{bmatrix} - 0.5\begin{bmatrix} -0.10168 & 0.06542 \\ 0.03981 & -0.05286 \end{bmatrix} = \begin{bmatrix} 0.15084 & 0.36729 \\ -0.21991 & 0.32643 \end{bmatrix} \]
\[ \begin{aligned} W_{hy} \leftarrow \begin{bmatrix} 0.7 & -0.1 & 0.2 & -0.5 \\ -0.3 & 0.6 & -0.4 & 0.2 \end{bmatrix} - 0.5\begin{bmatrix} 0.09781 & -0.03230 & -0.05555 & -0.00996 \\ -0.11158 & 0.08823 & -0.04428 & 0.06763 \end{bmatrix} \\ = \begin{bmatrix} 0.65110 & -0.08385 & 0.22778 & -0.49502 \\ -0.24421 & 0.55589 & -0.37786 & 0.16618 \end{bmatrix} \end{aligned} \]
\[ E \leftarrow \begin{bmatrix} 0.6 & -0.2 \\ 0.4 & 0.7 \\ -0.5 & 0.3 \\ 0.1 & -0.8 \end{bmatrix} - 0.5\begin{bmatrix} 0.18114 & -0.01109 \\ -0.05105 & 0.05574 \\ -0.10239 & -0.01689 \\ 0 & 0 \end{bmatrix} = \begin{bmatrix} 0.50943 & -0.19445 \\ 0.42553 & 0.67213 \\ -0.44881 & 0.30845 \\ 0.10000 & -0.80000 \end{bmatrix} \]
The “machine” embedding row is untouched, since it received no gradient this pass. Every other parameter has shifted a little in the direction that lowers the loss.
Verify Effectiveness
Run the same four words through the updated weights:
| before | after | |
|---|---|---|
| \(\hat{y}_1[\text{rise}]\) | 0.20273 | 0.21849 |
| \(\hat{y}_2[\text{of}]\) | 0.24068 | 0.25283 |
| \(\hat{y}_3[\text{the}]\) | 0.18701 | 0.21684 |
| \(\hat{y}_4[\text{machine}]\) | 0.21956 | 0.23889 |
| mean loss | 1.55321 | 1.46410 |
One pass of BPTT nudged every correct word upward and pulled the mean loss down by roughly 0.09 nats.
Exploding and Vanishing Gradients
We’ve seen how sigmoid can cause vanishing gradient with large number of hidden layers. RNN has same issue during backpropagation:
\[ \frac{\partial h_t}{\partial h_{t-k}} = \prod_{i=t-k+1}^{t} W_{hh}\, \mathrm{diag}\!\left(1 - h_i^2\right), \quad k = \text{number of steps 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 offset each other. Repeated multiplication by one matrix compounds relentlessly. Drop the factor below 1 and gradients vanish exponentially with distance; push it above 1 and they explode:
| steps back \(k\) | 1 | 5 | 10 | 20 | |
|---|---|---|---|---|---|
| Vanishing | \(w = 0.5\): \(\;0.5^k\) | 0.5 | 0.031 | 0.00098 | 0.00000095 |
| Exploding | \(w = 1.5\): \(\;1.5^k\) | 1.5 | 7.6 | 57.7 | 3,325 |
The result is a weight that either stops moving, ot takes steps large enough to overshoot the minimum, or blows up to NaN.
Remedy 1 - Gradient Clipping
Exploding gradients cause the loss to spikes, parameters leap to huge values, and training collapses into NaN. The standard fix is 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} \]
Remember that 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.
Remedy 2 - Use Architecture Variance
Vanishing gradients are harder to spot. 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. Hence 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.
Prep data for embedding
sentence = "the rise of the machine"
words = sentence.split()
# Extract unique vocabulary to build embedding
vocab = list(dict.fromkeys(words))
# vocab: {'the': 0, 'rise': 1, 'of': 2, 'machine': 3}
word_to_index = {w: i for i, w in enumerate(vocab)}
embedding_class = len(vocab)Training
We’ll train for 100 epochs and print the updated parameters after the first epoch, so we can check them against the calculation we did by hand.
For PyTorch, we’ll use the built-in class (nn.RNN) for convenience. It abstracts away much of the forward pass. In the next article, you’ll see what’s happening under the hood by using PyTorch’s stateless functions.
NOTE: In Python, if you don’t know what a function does, you can call the
helpfunction. It’ll print the signature and documentation for that function. For example,help(model.fit).
import torch
import torch.nn as nn
torch.set_printoptions(precision=5, sci_mode=False)
# Create input and target vector
ids = [word_to_index[w] for w in words]
inputs = torch.tensor([ids[:-1]])
# Inputs: ['the', 'rise', 'of', 'the']
targets = torch.tensor([ids[1:]])
# Targets: ['rise', 'of', 'the', 'machine']
# Create parameters
embedding = torch.tensor([[0.6, -0.2], [0.4, 0.7], [-0.5, 0.3], [0.1, -0.8]])
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.3, 0.6, -0.4, 0.2]])
# Initialize RNN 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)
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)
# Training
loss_fn = nn.CrossEntropyLoss()
# Backprop config
optimizer = torch.optim.SGD(model.parameters(), lr=0.5)
epochs = 100
for epoch in range(epochs):
optimizer.zero_grad()
logits = model(inputs)
loss = loss_fn(logits.view(-1, embedding_class), targets.view(-1))
# Backpropagation
loss.backward()
# Gradient clipping
grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
# Update parameters
optimizer.step()
# Print loss every 5 epochs
if epoch % 5 == 0:
print(f"epoch {epoch+1:3d} loss {loss.item():.5f} grad norm {grad_norm.item():.5f}")
# Print first update
if epoch == 0:
updated = {
"E": model.embed.weight.data,
"W_xh": model.rnn.weight_ih_l0.data.T,
"W_hh": model.rnn.weight_hh_l0.data.T,
"W_hy": model.head.weight.data.T,
}
# This should be identical to our manual calculation
for name in updated:
print(f"\n{name}:\n{updated[name]}")import numpy as np
import tensorflow as tf
np.set_printoptions(precision=5, suppress=True)
# Create parameters
ids = [word_to_index[w] for w in words]
inputs = tf.constant([ids[:-1]])
targets = tf.constant([ids[1:]])
T = len(ids) - 1
E = np.array([[0.6, -0.2], [0.4, 0.7], [-0.5, 0.3], [0.1, -0.8]], dtype='float32')
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.3, 0.6, -0.4, 0.2]], dtype='float32')
# Create model architecture
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),
])
# Use our parameters
embed, rnn, head = model.layers
embed.set_weights([E])
rnn.set_weights([W_xh, W_hh])
head.set_weights([W_hy])
# Config backprop
model.compile(
optimizer=tf.keras.optimizers.SGD(learning_rate=0.5, global_clipnorm=1.0),
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
)
# Training
history = model.fit(inputs, targets, epochs=100, verbose=0)
for i, l in enumerate(history.history["loss"], 1):
if i % 5 == 0:
print(f"epoch {i:3d} loss {l:.5f}")
# This should be identical to our manual calculation
print("W_xh:\n", rnn.weights[0].numpy())
print("W_hh:\n", rnn.weights[1].numpy())
print("W_hy:\n", head.weights[0].numpy())
print("W_hy:\n", embed.weights[0].numpy())Wrap Up
RNN is a useful deep learning architecture, and it’s fun to train on small sequential data: building a JSON object from a sentence, or picking out people and companies from text. Though it’s also important to recognize the bottlenecks:
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. It’s 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. (vanishing & exploding gradient)
Next we’ll explore a slightly different architecture (LSTM and GRU) inspired by RNN that attempts to address some of the bottle necks.