Understanding how to implement a character-based RNN language model

What RNNs do

I expect readers to have a basic idea of what RNN do and why they work well for some problems. RNN are well-suited for problem domains where the input (and/or output) is some sort of a sequence - time-series financial data, words or sentences in natural language, speech, etc.

There is a lot of material about this online, and the basics are easy to understand for anyone with even a bit of machine learning background. However, there is not enough coherent material online about how RNNs are implemented and trained - this is the goal of this post.

Character-based RNN language model

The basic structure of min-char-rnn is represented by this recurrent diagram, where x is the input vector (at time step t), y is the output vector and h is the state vector kept inside the model.

The line leaving and returning to the cell represents that the state is retained between invocations of the network. When a new time step arrives, some things are still the same (the weights inherent to the network, as we shall soon see) but some things are different - h may have changed. Therefore, unlike stateless NNs, y is not simply a function of x; in RNNs, identical xs can produce different ys, because y is a function of x and h, and h can change between steps.

The character-based part of the model's name means that every input vector represents a single character (as opposed to, say, a word or part of an image). min-char-rnn uses one-hot vectors to represent different characters.

A language model is a particular kind of machine learning algorithm that learns the statistical structure of language by "reading" a large corpus of text. This model can then reproduce authentic language segments - by predicting the next character (or word, for word-based models) based on past characters.

Internal-structure of the RNN cell

Let's proceed by looking into the internal structure of the RNN cell in min-char-rnn:

• Bold-faced symbols in reddish color are the model's parameters, weights for matrix multiplication and biases.
• The state vector h is shown twice - once for its past value, and once for its currently computed value. Whenever the RNN cell is invoked in sequence, the last computed state h is passed in from the left.
• In this diagram y is not the final answer of the cell - we compute a softmax function on it to obtain p - the probabilities for output characters [1]. I'm using these symbols for consistency with the code of min-char-rnn, though it would probably be more readable to flip the uses of p and y (making y the actual output of the cell).

Mathematically, this cell computes:

Learning model parameters with backpropagation

This section will examine how we can learn the parameters W and b for this model. Mostly it's standard neural-network fare; we'll compute the derivatives of all the steps involved and will then employ backpropagation to find a parameter update based on some computed loss.

There's one serious issue we'll have to address first. Backpropagation is usually defined on acyclic graphs, so it's not entirely clear how to apply it to our RNN. Is h an input? An output? Both? In the original high-level diagram of the RNN cell, h is both an input and an output - how can we compute the gradient for it when we don't know its value yet? [2]

The way out of this conundrum is to unroll the RNN for a few steps. Note that we're already doing this in the detailed diagram by distinguishing between and . This makes every RNN cell locally acyclic, which makes it possible to use backpropagation on it. This approach has a cool-sounding name - Backpropagation Through Time (BPTT) - although it's really the same as regular backpropagation.

Note that the architecture used here is called "synced many-to-many" in Karpathy's Unreasonable Effectiveness of RNNs post, and it's useful for training a simple char-based language model - we immediately observe the output sequence produced by the model while reading the input. Similar unrolling can be applied to other architectures, like encoder-decoder.

Here's our RNN again, unrolled for 3 steps:

Now the same diagram, with the gradient flows depicted with orange-ish arrows:

With this unrolling, we have everything we need to compute the actual weight updates during learning, because when we want to compute the gradients through step 2, we already have the incoming gradient , and so on.

Do you now wonder what is for the final step at time t?

In some models, sequence lengths are fairly limited. For example, when we translate a single sentence, the sequence length is rarely over a couple dozen words; for such models we can fully unroll the RNN. The h state output of the final step doesn't really "go anywhere", and we assume its gradient is zero. Similarly, the incoming state h for the first step is zero.

Other models work on potentially infinite sequence lengths, or sequences much too long for unrolling. The language model in min-char-rnn is a good example, because it can theoretically ingest and emit text of any length. For these models we'll perform truncated BPTT, by just assuming that the influence of the current state extends only N steps into the future. We'll then unroll the model N times and assume that is zero. Although it really isn't, for a large enough N this is a fairly safe assumption. RNNs are hard to train on very long sequences for other reasons, anyway (we'll touch upon this point again towards the end of the post).

Finally, it's important to remember that although we unroll the RNN cells, all parameters (weights, biases) are shared. This plays an important part in ensuring translation invariance for the models - patterns learned in one place apply to another place [3]. It leaves the question of how to update the weights, since we compute gradients for them separately in each step. The answer is very simple - just add them up. This is similar to other cases where the output of a cell branches off in two directions - when gradients are computed, their values are added up along the branches - this is just the basic chain rule in action.

We now have all the necessary background to understand how an RNN learns. What remains before looking at the code is figuring out how the gradients propagate inside the cell; in other words, the derivatives of each operation comprising the cell.

Flowing the gradient inside an RNN cell

As we saw above, the formulae for computing the cell's output from its inputs are:

To be able to learn weights, we have to find the derivatives of the cell's output w.r.t. the weights. The full backpropagation process was explained in this post, so here is only a brief refresher.

Recall that is the predicted output; we compare it with the "real" output () during learning, to find the loss (error):

To perform a gradient descent update, we'll need to find , for every weight value w. To do this, we'll have to:

1. Find the "local" gradients for every mathematical operation leading from w to L.
2. Use the chain rule to propagate the error backwards through these local gradients until we find .

We start by formulating the chain rule to compute :

Next comes:

Let's say the weight w we're interested in is part of , so we have to propagate some more:

We'll then proceed to propagate through the tanh function, bias addition and finally the multiplication by , for which the derivative by w is computed directly without further chaining.

Let's now see how to compute all the relevant local gradients.

Cross-entropy loss gradient

We'll start with the derivative of the loss function, which is cross-entropy in the min-char-rnn model. I went through a detailed derivation of the gradient of softmax followed by cross-entropy in this post; here is only a brief recap:

Re-formulating this for our specific case, the loss is a function of , assuming the "real" class r is constant for every training example:

Since inputs and outputs to the cell are 1-hot encoded, let's just use r to denote the index where is non-zero. Then the Jacobian of L is only non-zero at index r and its value there is .

Softmax gradient

A detailed computation of the gradient for the softmax function was also presented in this post. For being the softmax of y, the Jacobian is:

Fully-connected layer gradient

Next on our path backwards is:

From my earlier post on backpropagating through a fully-connected layer, we know that . But that's not all; note that on the forward pass splits in two - one edge goes into the fully-connected layer, another goes to the next RNN cell as the state. When we backpropagate the loss gradient to , we have to take both edges into account; more specifically, we have to add the gradients along the two edges. This leads to the following backpropagation equation:

In addition, note that this layer already has model parameters that need to be learned - and - a "final" destination for backpropagation. Please refer to my fully-connected layer backpropagation post to see how the gradients for these are computed.

Gradient of tanh

The vector is produced by applying a hyperbolic tangent nonlinearity to another fully connected layer.

To get to the model parameters , and , we have to first backpropagate the loss gradient through tanh. tanh is a scalar function; when it's applied to a vector we apply it in element-wise fashion to every element in the vector independently, and collect the results in a similarly-shaped result vector.

Its mathematical definition is:

To find the derivative of this function, we'll use the formula for deriving a ratio:

So:

Just like for softmax, it turns out that there's a convenient way to express the derivative of tanh in terms of tanh itself. When we apply the chain rule to derivatives of tanh, for example: where k is a function of w. We get:

In our case k(w) is a fully-connected layer; to find its derivatives w.r.t. the weight matrices and bias, please refer to the backpropagation through a fully-connected layer post.

Learning model parameters with Adagrad

We've just went through all the major parts of the RNN cell and computed local gradients. Armed with these formulae and the chain rule, it should be possible to understand how the min-char-rnn code flows the loss gradient backwards. But that's not the end of the story; once we have the loss derivatives w.r.t. to some model parameter, how do we update this parameter?

The most straightforward way to do this would be using the gradient descent algorithm, with some constant learning rate. I've written about gradient descent in the past - please take a look for a refresher.

Most real-world learning is done with more advanced algorithms these days, however. One such algorithm is called Adagrad, proposed in 2011 by some experts in mathematical optimization. min-char-rnn happens to use Adagrad, so here is a simplified explanation of how it works.

The main idea is to adjust the learning rate separately per parameter, because in practice some parameters change much more often than others. This could be due to rare examples in the training data set that affect a parameter that's not often affected; we'd like to amplify these changes because they are rare, and dampen changes to parameters that change often.

Therefore the Adagrad algorithm works as follows:

# Same shape as the parameter array x
memory = 0
while True:
  dx = compute_grad(x)

  # Elementwise: each memory element gets the corresponding dx^2 added to it.
  memory += dx * dx

  # The actual parameter update for this step. Note how the learning rate is
  # modified by the memory. epsilon is some very small number to avoid dividing
  # by 0.
  x -= learning_rate * dx / (np.sqrt(memory) + epsilon)

If a given element in dx was updated significantly in the past, its corresponding memory element will grow and thus the learning rate is effectively decreased.

Gradient clipping

If we unroll the RNN cell 10 times, the gradient will be multiplied by ten times on its way from the last cell to the first. For some structures of , this may lead to an "exploding gradient" effect where the value keeps growing [4].

To mitigate this, min-char-rnn uses the gradient clipping trick. Whenever the gradients are updated, they are "clipped" to some reasonable range (like -5 to 5) so they will never get out of this range. This method is crude, but it works reasonably well for training RNNs.

The flip side problem of vanishing gradient (wherein the gradients keep getting smaller with each step) is much harder to solve, and usually requires more advanced recurrent NN architectures.

min-char-rnn model quality

While min-char-rnn is a complete RNN implementation that manages to learn, it's not really good enough for learning a reasonable model for the English language. The model is too simple for this, and suffers seriously from the vanishing gradient problem.

For example, when training a 16-step unrolled model on a corpus of Sherlock Holmes books, it produces the following text after 60,000 iterations (learning on about a MiB of text):

one, my dred, roriny. qued bamp gond hilves non froange saws, to mold his a work, you shirs larcs anverver strepule thunboler muste, thum and cormed sightourd so was rewa her besee pilman

It's not complete gibberish, but not really English either. Just for fun, I wrote a simple Markov chain generator and trained it on the same text with a 4-character state. Here's a sample of its output:

though throughted with to taken as when it diabolice, and intered the stairhead, the stood initions of indeed, as burst, his mr. holmes' room, and now i fellows. the stable. he retails arm

Which, you'll admit, is quite a bit better than our "fancy" deep learning approach! And it was much faster to train too...

To have a better chance of learning a good model, we'll need a more advanced architecture like LSTM. LSTMs employ a bunch of tricks to preserve long-term dependencies through the cells and can learn much better language models. For example, Andrej Karpathy's char-rnn model from the Unreasonable Effectiveness of RNNs post is a multi-layer LSTM, and it can learn fairly nice models for a varied set of domains, ranging from Shakespeare sonnets to C code snippets in the Linux kernel.

Conclusion

The goal of this post wasn't to develop a very good RNN model; rather, it was to explain in detail the math behind training a simple RNN. More advanced RNN architerctures like LSTM are somewhat more complicated, but all the core ideas are very similar and this post should be helpful in nailing the basics.

Update: An extension of this post to LSTMs.