"Non-autoregressive" is no slouch either: MLM-based reading comprehension QA

A while back I wrote The all-purpose seq2seq: reading comprehension QA based on seq2seq, exploring how to do reading-comprehension-style QA using the most general seq2seq approach, and got quite good results (0.77 for a single model, surpassing the best fine-tuned model I used in the competition). In this post I continue with the same task but take a different angle, working directly from an MLM model. The final result ends up being basically the same, but with a boost in prediction speed.

Two kinds of generation

Broadly speaking, MLM-style generation also counts as a seq2seq model, except it belongs to "non-autoregressive" generation, whereas what we usually mean (in the narrow sense) by seq2seq refers to autoregressive generation. This section gives a brief introduction to these two concepts.

Autoregressive generation

As the name suggests, autoregressive generation means the decoding stage generates tokens recursively, one at a time. It models the following probability distribution:

\begin{equation}p(y_1,y_2,\dots,y_n|x)=p(y_1|x)p(y_2|x,y_1)\dots p(y_n|x,y_1,\dots,y_{n-1})\label{eq:at}\end{equation}

For a more detailed introduction, see Playing with Keras for seq2seq automatic title generation and From language models to Seq2Seq: Transformer plays it however it likes, all thanks to Mask. I won't go into autoregressive generation in much more detail here.

Non-autoregressive generation

Since autoregressive generation requires recursive decoding and cannot be parallelized, decoding speed is fairly slow. Because of this, in recent years quite a lot of work has gone into researching non-autoregressive generation, with considerable success. Put simply, non-autoregressive generation seeks ways to make the decoding of each token parallelizable, and the simplest non-autoregressive model just directly assumes that each token is independent:

\begin{equation}p(y_1,y_2,\dots,y_n|x)=p(y_1|x)p(y_2|x)\dots p(y_n|x)\label{eq:nat}\end{equation}

This is a very strong assumption, and it's only suitable in some fairly special cases. If you use it directly for ordinary text generation, such as automatic summarization, the results will be poor. For more sophisticated work on non-autoregressive generation, you can find plenty by searching "non-autoregressive text generation" on Arxiv or Google.

As already hinted in the title, the way this post does reading comprehension is "MLM-based". Readers who are familiar with the BERT model will know that MLM (Masked Language Model) is actually a special case of $\eqref{eq:nat}$, so generative models built on MLM fall under the umbrella of non-autoregressive generation.

Model overview

As "overview" suggests, this is going to be really brief, because doing reading comprehension based on MLM is genuinely very simple.

Model diagram

First, fix a maximum length $l_{\max}$, then concatenate the question and passage, insert $l_{\max}$ [MASK] tokens in between, feed all of that into BERT, and finally have the parts corresponding to [MASK] predict the answer (this holds both during training and during prediction). See the diagram below:

Diagram of the MLM-based reading comprehension model (where [M] denotes the [MASK] token)Diagram of the MLM-based reading comprehension model (where [M] denotes the [MASK] token)

Code and results

Code link: task_reading_comprehension_by_mlm.py

Using the evaluation script that comes with SogouQA, the score on the valid set comes out to about 0.77 (Accuracy=0.7282149325820084, F1=0.8207266829447049, Final=0.7744708077633566), which is on par with the The all-purpose seq2seq: reading comprehension QA based on seq2seq model from before. However, the prediction speed is noticeably improved: the previous seq2seq approach could only predict around 2 samples per second, whereas now it reaches 12 samples per second — a 6x speedup without any drop in performance.

Which one should you use?

In principle, seq2seq is the all-purpose tool, and in principle the distribution modeled by seq2seq in equation $\eqref{eq:at}$ should be more reasonable than the one modeled by MLM in $\eqref{eq:nat}$. So why can the MLM approach achieve results on par with seq2seq? When should you use MLM, and when should you use seq2seq?

Training and prediction

First off, the biggest problem with seq2seq is that it's slow — and it gets even slower for long-text generation. So if efficiency is a requirement, you naturally have no choice but to give up on the seq2seq approach.

If we set efficiency aside, is seq2seq then simply the best choice? Not necessarily. Although from a modeling perspective $\eqref{eq:at}$ is more accurate, seq2seq training is done via teacher forcing, so it suffers from the "exposure bias" problem: during training, the input at each time step comes from the ground-truth answer text, whereas during generation, the input at each time step comes from the output of the previous step. So once a single token is generated poorly, the error can propagate forward, causing the generation to get progressively worse.

Put plainly, there is an inconsistency between training and prediction, and this inconsistency can lead to accumulated errors. In contrast, the MLM-based approach behaves consistently between training and prediction, because it never needs the ground-truth label as input (during prediction, the positions corresponding to the answer are also fed [MASK]), so there's no error accumulation. And it's precisely because of this property that decoding no longer needs to be recursive — it can be parallelized, improving decoding speed.

A unique correct answer

Besides that, MLM and other non-autoregressive generation methods are, relatively speaking, better suited to short-text generation, since the shorter the text, the closer it is to satisfying the independence assumption. At the same time, non-autoregressive generation is well suited to scenarios where "there is only one correct answer" — and the reading comprehension task in this post, being primarily extractive, happens to fit exactly this scenario. That's why MLM performs quite well here too.

In fact, sequence labeling models such as per-frame softmax or CRF can also be viewed as non-autoregressive generation models. I think the fundamental reason they work is that "the correct answer sequence is unique," rather than the intuitive belief that "input and output are aligned." In other words, if the condition "there is only one correct answer" holds, then non-autoregressive generation is worth considering.

Note that having a unique answer here doesn't mean each sample only has one human-annotated answer; rather, it means the task is designed in such a way that the answer is made unique. Take word segmentation, for example: once the labeling scheme is fixed, each sentence corresponds to exactly one correct segmentation. Title generation, on the other hand, is different — clearly the same article can have different titles, so the answer for title generation is not unique (even if the training data only has one title per article).

Summary

This post experimented with doing reading-comprehension-style QA via the non-autoregressive generation approach of MLM, and found that the final performance is quite good, with several times the speedup. Along the way, the post also briefly compared the similarities and differences between autoregressive and non-autoregressive generation, and analyzed when the non-autoregressive approach is applicable and why.

English translation of a post from 科学空间 | Scientific Spaces by 苏剑林. Original: https://kexue.fm/archives/7148
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.