TeaForN: Giving Teacher Forcing a Bit More "Foresight"
Teacher Forcing is the classic training scheme for Seq2Seq models, and Exposure Bias is its classic flaw — this should be familiar territory for anyone working on text generation. I previously wrote a post, Exposure Bias in Seq2Seq: A Brief Analysis and Countermeasures, which offered a preliminary analysis of the Exposure Bias problem.
This post introduces a new scheme from Google called TeaForN for alleviating Exposure Bias, from the paper TeaForN: Teacher-Forcing with N-grams. It uses a nested iteration approach to let the model anticipate the next $N$ tokens in advance (not just the single token currently being predicted). Its approach has some genuinely clever aspects worth learning from.
(Note: to keep things as consistent as possible with earlier posts on this blog, the notation used here differs somewhat from that of the original paper. Please focus on understanding what the symbols mean rather than memorizing their exact form.) more
Teacher Forcing
The post Exposure Bias in Seq2Seq: A Brief Analysis and Countermeasures already covered Teacher Forcing in reasonable detail, so here I'll just give a brief recap. First, Seq2Seq models factorize the joint probability into a product of conditional probabilities — this is the so-called "autoregressive model":
\begin{equation}\begin{aligned}p(\boldsymbol{y}|\boldsymbol{x})=&\,p(y_1,y_2,\dots,y_n|\boldsymbol{x})\\ =&\,p(y_1|\boldsymbol{x})p(y_2|\boldsymbol{x},y_1)\dots p(y_n|\boldsymbol{x},y_1,\dots,y_{n-1}) \end{aligned}\end{equation}
Then, when training the model $\dots p(y_t|\boldsymbol{x},y_1,\dots,y_{t-1})$ at step $t$, we assume that $\boldsymbol{x},y_1,\dots,y_{t-1}$ are all already known, and have the model predict only $y_t$ — this is Teacher Forcing. But at prediction time, the true $y_1,\dots,y_{t-1}$ are all unknown; instead they are predicted recursively, which can lead to error propagation and similar issues. So the problem with Teacher Forcing is this inconsistency between training and prediction, which makes it hard to gauge prediction-time behavior from the training process.
Lacking Foresight
How can we get a more concrete grasp of the problem caused by this inconsistency? We can think of it as the model "lacking foresight." In the decoder, the input $\boldsymbol{x}$ together with the preceding $t-1$ output tokens is encoded into a vector $h_t$. Under Teacher Forcing, this $h_t$ is only ever used to predict $y_t$, with no direct connection to $\boldsymbol{y}_{> t}$ — in other words, its "field of vision" is limited to this single step $t$.
Teacher Forcing diagram
Take the vector $h_3$ in the figure above as an example: Teacher Forcing only uses it to predict "阴" (overcast). But in reality, the prediction of "阴" will also influence the predictions of "晴" (sunny), "圆" (full), and "缺" (waning) — that is, $h_3$ ought to be related to "晴", "圆", and "缺" as well, and Teacher Forcing fails to establish this connection explicitly. As a result, the model, when decoding, is prone to simply outputting the locally highest-probability token at each step, which easily leads to generic "safe" responses or repetitive decoding.
Student Forcing
To improve the model's "foresight," the most thorough approach is naturally to make the training phase mirror the decoding phase — that is, to have $h_1,h_2,\dots,h_t$ also be predicted recursively during training, just as at decoding time, without relying on ground-truth labels. Let's call this approach Student Forcing. However, training this way brings two serious problems:
First, it sacrifices parallelism. With Teacher Forcing, if the decoder is built from a CNN or Transformer-style architecture, all tokens can be trained in parallel during training (only inference remains sequential). With Student Forcing, however, everything remains sequential throughout.
Second, it is extremely hard to get to converge. Student Forcing typically needs Gumbel-Softmax or reinforcement learning to back-propagate gradients, and training with either of these faces serious instability. Usually one has to pretrain with Teacher Forcing before switching to Student Forcing, and even then it's not particularly stable.
To put it vividly, Student Forcing is like a teacher letting a student work through a complicated problem entirely on their own, with no hand-holding, only giving a final verdict on whether the result is good or bad. If the student manages to work it out, that may indicate genuine ability — but the lack of the teacher's careful guidance means the student is far more likely to hit a wall.
Looking a Few Steps Further Ahead
Is there something in between Teacher Forcing and Student Forcing? Yes — TeaForN, the subject of this post, is one such method. The idea is that ordinary Teacher Forcing amounts to looking only 1 step ahead during training, while Student Forcing amounts to looking $L$ steps ahead (where $L$ is the length of the target sentence). If we simply look a few steps further ahead (equivalent to seeing an N-gram), then in principle we should be able to improve foresight without severely sacrificing the model's parallelism. Here's the diagram:
TeaForN diagram
Intuitively, this amounts to feeding the output forward through several more rounds of iteration, so that the first $t-1$ tokens are no longer just predicting the $t$-th token, but also the $t+1,t+2,\cdots$-th. For instance, in the figure above, we ultimately use $h_6^{(3)}$ to predict the character "缺", and we can see that $h_6^{(3)}$ depends only on the three characters "月", "有", and "阴". So we can also think of $h_4^{(1)}$ as a vector that simultaneously predicts the three characters "晴", "圆", and "缺" — hence the improved foresight.
In Mathematical Terms
To put this in mathematical language, we can split the decoder into an embedding layer $E$ and the remaining part $M$. The embedding layer maps the input sentence $s=[w_0, w_1, w_2, \cdots, w_{L-1}]$ to a sequence of vectors $[e_0, e_1, e_2, \cdots, e_{L-1}]$ (where $w_0$ is a fixed decoding start marker — i.e., [S] in the figure above, sometimes written as <bos> in other papers), which is then handed to the model $M$ to produce the vector sequence $[h_1, h_2, h_3, \cdots, h_L]$, i.e.
\begin{equation}[h_1, h_2, h_3, \cdots, h_L] = M(E([w_0, w_1, w_2, \cdots, w_{L-1}]))\end{equation}
We then obtain the token probability distribution at step $t$ via $p_t = softmax(Wh_t + b)$, and finally train with $-\log p_t[w_t]$ as the loss function — this is standard Teacher Forcing.
It's natural to imagine that the output vector sequence $[h_1, h_2, h_3, \cdots, h_{L-1}]$, which is responsible for mapping to token distributions, is in some sense similar to the embedding sequence $[e_1, e_2, e_3, \cdots, e_{L-1}]$. So what if we append an extra $e_0$, and feed $[e_0, h_1, h_2, \cdots, h_{L-1}]$ back into the model $M$ for another round of processing? That is,
\begin{equation}\begin{aligned}[] \left[e_0,e_1,e_2,\cdots,e_{L-1}\right]& = E\left(\left[w_0, w_1,w_2,\cdots,w_{L-1}\right]\right)\\ \left[h_1^{(1)},h_2^{(1)},h_3^{(1)},\cdots,h_L^{(1)}\right]& = M\left(\left[e_0,e_1,e_2,\cdots,e_{L-1}\right]\right)\\ \left[h_1^{(2)},h_2^{(2)},h_3^{(2)},\cdots,h_L^{(2)}\right]& = M\left(\left[e_0, h_1^{(1)},h_2^{(1)},\cdots,h_{L-1}^{(1)}\right]\right)\\ \left[h_1^{(3)},h_2^{(3)},h_3^{(3)},\cdots,h_L^{(3)}\right]& = M\left(\left[e_0, h_1^{(2)},h_2^{(2)},\cdots,h_{L-1}^{(2)}\right]\right)\\ &\,\,\vdots \end{aligned}\end{equation}
We then compute the probability distribution $p_t^{(i)} = softmax(Wh_t^{(i)} + b)$ for every $h$, and finally compute the cross-entropy and take a weighted sum:
\begin{equation}\text{loss} = -\sum_{t=1}^L \sum_{i=1}^N \lambda_i \log p_t^{(i)}[w_t]\end{equation}
Once training is complete, we only use $E$ and $M$ for standard decoding operations (such as beam search) — that is, we only need $h_t^{(1)}$, not $h_t^{(2)},h_t^{(3)},\cdots$ anymore. This whole pipeline is precisely TeaForN, the star of this post.
Results, Reflections, and Discussion
As for experimental results, there is indeed an improvement. Looking at the experiment tables in the original paper, the gain is more pronounced when the beam size is relatively large. This is not hard to understand: in principle, this kind of processing should, at worst, not make things worse — so it counts as a "sure bet" strategy of sorts.
One of TeaForN's experimental results (text summarization)
The original paper discusses a few points worth debating, and we'll take a look at them here too.
First, should the $M$ used at each iteration step share weights? Intuitively, sharing seems better — if the weights aren't shared, then looking $N$ steps ahead would multiply the parameter count by roughly $N$, which doesn't seem great. Of course, it's best to rely on experiments, and the original paper indeed ran this comparison, confirming our intuition.
TeaForN's performance on machine translation, including a comparison of shared vs. unshared weights
Second, perhaps the biggest question is: is it actually reasonable to treat $[h_1, h_2, h_3, \cdots, h_{L-1}]$ as $[e_1, e_2, e_3, \cdots, e_{L-1}]$ during the iteration process? Of course, the experimental results have already shown that it works, which is the most convincing evidence there is. But since the mapping from $h_t$ to $p_t$ is built via an inner product, $h_t$ need not actually resemble $e_t$ — so would things work better if we could make them closer? The original paper considers the following approach:
\begin{equation}\frac{\sum\limits_{w\in \text{Top}_k(p_t)}p_t[w] e_w}{\sum\limits_{w\in \text{Top}_k(p_t)}p_t[w]}\end{equation}
That is, after computing $p_t$ at each step, we take the top $k$ tokens by probability, and use a weighted average of their embedding vectors as the input for the next iteration step. The original paper experiments with $k=4$ and $k=|V|$ (the vocabulary size), with results as shown below. Overall, the effect of Top-k is not particularly stable — the best cases are roughly on par with just using $h_t$ directly, so there's no real need to try anything else.
Effect of replacing h with a Top-k weighted average of embeddings
That said, I do think the paper would have been even more complete if it had also compared using Gumbel-Softmax to simulate sampling.
Closing Summary
This post covered a new training scheme from Google called TeaForN, which sits between Teacher Forcing and Student Forcing. It can alleviate the model's Exposure Bias problem without seriously sacrificing training parallelism — a strategy well worth trying. Beyond that, it actually offers a new way of thinking about this class of problems (preserving both parallelism and foresight through iteration), which has plenty to reflect on.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.