Already used CRF? Time to get acquainted with the faster MEMM

CRF用过了，不妨再了解下更快的MEMM？ 苏剑林 · 2020-02-24 · Machine Learning model probabilistic graphical model CRF

HMM, MEMM, and CRF are known as the three classic probabilistic graphical models. Before the deep learning era, they were widely used in machine learning for all sorts of sequence-labeling tasks. An interesting phenomenon is that, in the deep learning era, HMM and MEMM both seem to have "fallen out of favor," leaving only CRF on the stage. I believe that NLP practitioners, even if they haven't personally used it, have at least heard of BiLSTM+CRF for Chinese word segmentation, named entity recognition, and similar tasks — yet almost never hear about BiLSTM+HMM or BiLSTM+MEMM. Why is that?

Today, let's take a closer look at MEMM, and by comparing it with CRF, gain a deeper understanding of the ideas and design philosophy behind probabilistic graphical models.

Model derivation

MEMM stands for Maximum Entropy Markov Model. I have to say, this name might scare off 80% of beginners: if you haven't quite grasped maximum entropy, and you don't recognize Markov either, surely combining the two must be gibberish? But in fact, both MEMM and CRF are, as models, far simpler than their names suggest — their concepts and designs are quite plain and natural, and not hard to understand at all. more

Revisiting CRF

For comparison, let's first revisit CRF. I say "revisit" because I've already written about CRF before; readers not yet familiar with CRF can first read my earlier post A Concise Introduction to Conditional Random Fields (CRF), with a Pure Keras Implementation. For simplicity, both the CRF and MEMM discussed in this post are the simplest "linear-chain" versions.

Throughout this post we'll use sequence labeling as our running example: given an input sequence $\boldsymbol{x}=(x_1,x_2,\dots,x_n)$, we want to output a label sequence $\boldsymbol{y}=(y_1,y_2,\dots,y_n)$ of the same length, so what we're modeling is the probability distribution

\begin{equation}P(\boldsymbol{y}|\boldsymbol{x})=P(y_1,y_2,\dots,y_n|\boldsymbol{x})\label{eq:target}\end{equation}

CRF treats $\boldsymbol{y}$ as a whole and computes a single overall score, using the formula

\begin{equation}\begin{aligned}f(y_1,y_2,\dots,y_n;\boldsymbol{x})=&\,f(y_1;\boldsymbol{x})+g(y_1,y_2)+\dots+g(y_{n-1},y_n)+f(y_n;\boldsymbol{x})\\ =&\,f(y_1;\boldsymbol{x}) + \sum_{k=2}^n \big(g(y_{k-1},y_k)+f(y_k;\boldsymbol{x})\big)\end{aligned}\end{equation}

The distinguishing feature of this scoring function is that it explicitly accounts for the correlations between adjacent labels; $g(y_{k-1},y_k)$ here is called the transition matrix. Now that we've computed the score, the probability is just the softmax of the score, so the final probability distribution takes the form

\begin{equation}P(\boldsymbol{y}|\boldsymbol{x})=\frac{e^{f(y_1,y_2,\dots,y_n;\boldsymbol{x})}}{\sum\limits_{y_1,y_2,\dots,y_n}e^{f(y_1,y_2,\dots,y_n;\boldsymbol{x})}}\label{eq:crf-p}\end{equation}

If we stick purely to the concept, that's really all there is to CRF. In short: treat the target sequence as a whole, first design a scoring function for the target, then take an overall softmax of the scoring function — this modeling philosophy is exactly the same as an ordinary classification problem. Where CRF gets tricky is in the implementation, because the denominator in the equation above involves a sum over all possible paths, which is not trivial to compute. But conceptually, I don't think there's anything particularly difficult about it.

The more straightforward MEMM

Now let's introduce MEMM, which can be seen as an extremely simplified seq2seq model. For the target $\eqref{eq:target}$, it considers the factorization

\begin{equation}P(y_1,y_2,\dots,y_n|\boldsymbol{x})=P(y_1|\boldsymbol{x})P(y_2|\boldsymbol{x},y_1)P(y_3|\boldsymbol{x},y_1,y_2)\dots P(y_n|\boldsymbol{x},y_1,y_2,\dots,y_{n-1})\end{equation}

and then assumes that label dependencies only occur between adjacent positions, so

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

Next, following the design of the linear-chain CRF, we can set

\begin{equation}P(y_1|\boldsymbol{x})=\frac{e^{f(y_1;\boldsymbol{x})}}{\sum\limits_{y_1}e^{f(y_k;\boldsymbol{x})}},\quad P(y_k|\boldsymbol{x},y_{k-1})=\frac{e^{g(y_{k-1},y_k)+f(y_k;\boldsymbol{x})}}{\sum\limits_{y_k}e^{g(y_{k-1},y_k)+f(y_k;\boldsymbol{x})}}\label{eq:memm}\end{equation}

And that's it — this gives us MEMM. Since MEMM has already decomposed the overall probability distribution into a product of step-by-step distributions, computing the loss just requires summing the cross-entropy at each step.

The relationship between the two

Substituting equation $\eqref{eq:memm}$ back into equation $\eqref{eq:p-f}$, we get

\begin{equation}P(\boldsymbol{y}|\boldsymbol{x})=\frac{e^{f(y_1;\boldsymbol{x})+g(y_1,y_2)+\dots+g(y_{n-1},y_n)+f(y_n;\boldsymbol{x})}}{\left(\sum\limits_{y_1}e^{f(y_1;\boldsymbol{x})}\right)\left(\sum\limits_{y_2}e^{g(y_1,y_2)+f(y_2;\boldsymbol{x})}\right)\dots\left(\sum\limits_{y_n}e^{g(y_{n-1},y_n)+f(y_n;\boldsymbol{x})}\right)}\label{eq:memm-p}\end{equation}

Comparing equations $\eqref{eq:memm-p}$ and $\eqref{eq:crf-p}$, we can see that the only difference between MEMM and CRF lies in how the denominator (i.e., the normalizing factor) is computed. We call CRF's equation $\eqref{eq:crf-p}$ globally normalized, and MEMM's equation $\eqref{eq:memm-p}$ locally normalized.

Model analysis

In this section we'll analyze the pros and cons of MEMM, along with improvements and experimental results.

Pros and cons of MEMM

An obvious characteristic of MEMM is that it's simple to implement and fast, because it only needs to run a softmax independently at each step, which means MEMM is fully parallelizable — its speed is basically the same as just doing a plain step-by-step softmax. CRF, on the other hand, has a denominator in equation $\eqref{eq:crf-p}$ that isn't so easy to compute; it eventually reduces to a recursive computation, which can be evaluated in $\mathcal{O}(n)$ time (for details, please refer to A Concise Introduction to Conditional Random Fields (CRF), with a Pure Keras Implementation). Being recursive means it's inherently sequential, so when the main body of our model is a highly parallelizable architecture (such as a pure CNN or pure Attention architecture), CRF will seriously slow down training. Later we'll compare the training speeds of MEMM and CRF (of course, it's only training that's slower — at inference time, MEMM and CRF run at the same speed).

As for downsides, naturally there are some. As mentioned earlier, MEMM can be seen as an extremely simplified seq2seq model. Given that, it inherits all the usual drawbacks of ordinary seq2seq models. One well-known issue in seq2seq is exposure bias; in MEMM, this is called label bias. Roughly speaking: when training MEMM, the prediction at the current step always assumes that the true label of the previous step is known. As a result, if some label $A$ can only be followed by label $B$, the model can achieve this purely by optimizing the transition matrix, without needing to optimize the influence of the input $\boldsymbol{x}$ on $B$ (i.e., $f(B;\boldsymbol{x})$ doesn't get properly optimized). However, at inference time the true label is unknown, and we may not be able to predict the previous step's label $A$ with high confidence; and since $f(B;\boldsymbol{x})$ at the current step wasn't reinforced during training, the current step's $B$ also can't be predicted accurately — which can lead to incorrect predictions.

Bidirectional MEMM

Label bias might be a bit hard to grasp intuitively, but we can look at MEMM's shortcomings from another angle: compared with CRF, one obviously less elegant aspect of MEMM is its asymmetry — it factorizes the probability from left to right. My experiments show that fixing this asymmetry can slightly improve MEMM's performance. My approach is: also run MEMM from right to left, in which case the corresponding probability distribution is

\begin{equation}P(\boldsymbol{y}|\boldsymbol{x})=\frac{e^{f(y_1;\boldsymbol{x})+g(y_1,y_2)+\dots+g(y_{n-1},y_n)+f(y_n;\boldsymbol{x})}}{\left(\sum\limits_{y_n}e^{f(y_n;\boldsymbol{x})}\right)\left(\sum\limits_{y_{n-1}}e^{g(y_n,y_{n-1})+f(y_{n-1};\boldsymbol{x})}\right)\dots\left(\sum\limits_{y_1}e^{g(y_2,y_1)+f(y_1;\boldsymbol{x})}\right)}\end{equation}

We then compute a cross-entropy for this as well, and average it with the cross-entropy from the left-to-right equation $\eqref{eq:memm-p}$ to get the final loss. This way, the model takes both left-to-right and right-to-left directions into account, without adding any extra parameters, remedying the asymmetry defect. To distinguish it, borrowing the naming convention of Bi-LSTM, I call this Bi-MEMM.

Note: the term Bi-MEMM did not first appear here. As far as I can tell, the concept of Bi-MEMM was first proposed in the paper
Bidirectional Inference with the Easiest-First Strategy for Tagging Sequence Data
, where Bi-MEMM refers to a bidirectional decoding strategy for MEMM, which is not the same as what Bi-MEMM means in this blog post.

Experimental results

To verify and compare the performance of MEMM, I implemented both CRF and MEMM in bert4keras, and wrote two scripts, one for Chinese word segmentation (task_sequence_labeling_cws_crf.py) and one for Chinese named entity recognition (task_sequence_labeling_ner_crf.py). In these two scripts, switching from CRF to MEMM is very simple — you just need to replace ConditionalRandomField with MaximumEntropyMarkovModel.

I won't post the detailed experimental data here — it's just a bunch of numbers anyway — but here are some relative comparison results:

1. Under the same experimental settings, Bi-MEMM always outperforms MEMM, and MEMM always outperforms plain Softmax;
2. Under the same experimental settings, CRF is basically never worse than Bi-MEMM;
3. When the encoder model is fairly powerful, CRF and Bi-MEMM perform about the same; when the encoder model is weaker, CRF outperforms Bi-MEMM, by roughly 0.5%;
4. Using a 12-layer BERT base model as the encoder, Bi-MEMM is 25% faster than CRF; using a 2-layer BERT base model as the encoder, Bi-MEMM is 1.5x faster than CRF.

(Note: since I found that Bi-MEMM's performance is consistently slightly better than MEMM's, and the training times of the two are essentially the same, the MaximumEntropyMarkovModel in bert4keras defaults to Bi-MEMM.)

Reflections and extensions

Based on the conclusions above, MEMM's "decline" in the deep learning era seems understandable — aside from faster training, MEMM doesn't seem to offer much advantage over CRF. The two have identical inference speed, and in many cases what we mainly care about is inference speed and performance, so a somewhat slower training speed doesn't matter much. This comparison between the two models is representative — one could say it captures precisely the difference between all globally normalized and locally normalized models in general: globally normalized models usually perform somewhat better, but are usually relatively harder to implement; locally normalized models usually don't surpass globally normalized ones in performance, but win out in ease of implementation and ease of extension.

How exactly are they easy to extend? Let me give two examples.

The first example: suppose the number of labels is very large — for instance, when doing text correction or text generation via sequence labeling (see the paper Fast Structured Decoding for Sequence Models for a related example) — the number of labels equals the vocabulary size $|V|$. Even with subword tokenization, the vocabulary size is easily tens of thousands, at which point the transition matrix would have hundreds of millions of parameters ($|V|^2$), making it impractical to train. You might think of low-rank factorization — indeed, low-rank factorization can bring the transition matrix's parameter count down to $2d|V|$, where $d$ is the dimension of the intermediate factorized layer. Unfortunately, for CRF, low-rank factorization doesn't change the fact that computing the normalizing factor is expensive, because CRF's normalizing factor still needs the full $|V|\times|V|$ transition matrix to be reconstructed before it can be computed. So for scenarios with a huge number of labels, CRF simply can't be used directly. Fortunately, for MEMM, low-rank factorization can effectively reduce the computational cost during training, so it can still be used efficiently. The MaximumEntropyMarkovModel included in bert4keras already has low-rank factorization built in; interested readers can check the source code for details.

The second example: the CRF and MEMM introduced above only consider correlations between adjacent labels. What if we want to consider more complex neighboring correlations — say, jointly considering the correlation between $y_k$ and $y_{k-1},y_{k-2}$? In that case, CRF's global-normalization approach becomes very hard to work with, again because the normalizing factor is hard to compute; but with MEMM's local-normalization approach, this is easy to do. In fact, the hierarchical tagging scheme for information extraction I designed previously can also be described as a locally normalized probabilistic graphical model similar to MEMM, one that handles even more complex correlations.

Summary

This post introduced and gave a brief extension of MEMM, another classic case of probabilistic graphical models alongside CRF. The main difference between the two lies in their normalization scheme. I then made a simple experimental comparison between the two, arriving at the conclusion that MEMM trains faster but doesn't outperform CRF. Even so, I believe MEMM still has its merits, so in closing I sketched out some extensions to MEMM.