RNN Architecture › RNN for Language Modeling

A language model assigns a probability to a sequence of tokens. It can also generate new text by sampling one token at a time, each conditioned on the tokens already generated.

The RNN language model works as follows. At each time step $t$, the network receives the embedding of the current token $\mathbf{x}_t$ and updates its hidden state. The output logits $\mathbf{y}_t$ represent the unnormalized log-probabilities for all tokens in the vocabulary, and a softmax converts these to a probability distribution over the next token:

$$P(\text{token}_{t+1} \mid \text{token}_1, \ldots, \text{token}_t) = \text{softmax}\!\left(\mathbf{W}_{hy}\,\mathbf{h}_t + \mathbf{b}_y\right)$$

The training objective is to maximize the log-likelihood of the correct next token at every position, which is equivalent to minimizing the cross-entropy loss:

$$\mathcal{L} = -\frac{1}{T} \sum_{t=1}^{T} \log P(\text{token}_{t+1} \mid \text{token}_1, \ldots, \text{token}_t)$$

This objective is defined at every time step rather than only at the end of the sequence. The gradients flow backward through all $T$ positions, giving the model dense supervision signal. This density of supervision is one of the reasons language modeling is such a powerful pre-training objective: the model receives feedback about every token, rather than a single label for the entire sequence.

In the worked example above, we used one-hot encodings: a $V$-dimensional vector with a single 1 at the token's index. This works for tiny vocabularies but is impractical at scale. For a vocabulary of 50,000 tokens, each one-hot vector has 50,000 dimensions, most of which are zero. Multiplying $\mathbf{W}_{xh} \in \mathbb{R}^{H \times V}$ by a one-hot vector is equivalent to just selecting the corresponding column of $\mathbf{W}_{xh}$, which is computationally wasteful.

An embedding layer solves this by representing each token as a dense vector $\mathbf{e}_t \in \mathbb{R}^{E}$ where $E \ll V$ (typically $E = 64$ to $512$). The embedding lookup is implemented as a matrix $\mathbf{E} \in \mathbb{R}^{V \times E}$, and retrieving the embedding for token $i$ means fetching row $i$: $\mathbf{e}_i = \mathbf{E}[i]$. This is computationally equivalent to multiplying by a one-hot vector but implemented as an efficient index operation.

The embedding matrix is jointly trained with the rest of the model by gradient descent. Tokens that appear in similar contexts learn to have similar embeddings. The embedding layer effectively performs a learned dimensionality reduction, converting the sparse high-dimensional token index into a dense, semantically meaningful vector before feeding it into the RNN.

A natural question during training: at step $t$, should the model use its own predicted token as the next input, or should it use the ground-truth token?

Teacher forcing uses the ground-truth token at each step. During training, regardless of what the model predicts, the next input is the actual next token from the training data. This has two benefits:

• Stability: if the model makes an error early in training, it does not get compounded through the rest of the sequence, keeping gradients informative
• Speed: the model receives correct context at every step, which typically makes training converge faster

The downside is exposure bias: at inference time, the model uses its own predictions as inputs (because there is no ground truth to fall back on). If the model's predictions are imperfect, the distribution of inputs at inference differs from training. This mismatch can cause generation quality to degrade, especially for long sequences.

Exposure bias is most noticeable in text generation tasks. A model that trained with teacher forcing has never, during training, seen what happens when it generates a wrong token and must then recover. At inference time, a single early mistake shifts the distribution of inputs to a regime the model did not encounter during training, potentially causing a cascade of errors.

Several approaches attempt to bridge this gap. Scheduled sampling gradually replaces ground-truth tokens with model-generated tokens as training progresses. Minimum-risk training optimizes a sequence-level loss rather than token-level cross-entropy, but requires efficient algorithms for computing expectations over sequences. Professor forcing and curriculum learning offer other perspectives on the training/inference mismatch. Despite these alternatives, teacher forcing remains the standard approach for initial training, with refinements applied as needed.