Your CRF Layer's Learning Rate Might Not Be High Enough

你的CRF层的学习率可能不够大 苏剑林 · 2020-02-07 · Machine Learning model probabilistic graphical model CRF

CRF is a classic method for sequence labeling. It's theoretically elegant and works well in practice. Readers unfamiliar with CRF are welcome to check out my earlier post A Concise Introduction to Conditional Random Fields (with a Pure Keras Implementation). After BERT came out, quite a few works explored combining BERT with CRF for sequence labeling tasks. However, many experimental results show (e.g., the paper BERT Meets Chinese Word Segmentation) that whether for Chinese word segmentation or named entity recognition tasks, BERT+CRF doesn't seem to bring any improvement over the simpler BERT+Softmax—this is quite different from the behavior of traditional BiLSTM+CRF or CNN+CRF models.

Illustration of the CRF-based 4-tag segmentation model

These past couple of days I added a Chinese word segmentation example using CRF to bert4keras (task_sequence_labeling_cws_crf.py), and during debugging I discovered that the CRF layer might be under-trained. I ran a few more comparison experiments, and the results suggest this might be the main reason CRF doesn't bring much improvement when combined with BERT. So I'm recording my analysis here and sharing it with everyone.more

A Poor Transition Matrix

Since I'm using my own implementation of the CRF layer, in order to confirm that my implementation wasn't buggy, after running the BERT+CRF experiment (using the base version of BERT), I first inspected the transition matrix, whose values roughly looked like this:

$$\begin{array}{c|cccc} & s & b & m & e \\ \hline s & -0.0517 & -0.459 & -0.244 & 0.707 \\ b & -0.564 & -0.142 & 0.314 & 0.613 \\ m & 0.196 & -0.334 & -0.794 & 0.672 \\ e & 0.769 & 0.841 & -0.683 & 0.572 \\ \end{array}$$

Here the value in row $i$, column $j$ represents the score for transitioning from $i$ to $j$ (denoted $S_{i\to j}$); the absolute value of the score is meaningless—only relative comparisons matter. By the way, the Chinese word segmentation in this post uses the $(s,b,m,e)$ character-tagging scheme; if you're not familiar with it, see Chinese Word Segmentation Series 3: Character Tagging and the HMM Model.

Intuitively, though, this is not a well-learned transition matrix—it might even have a negative effect. Look at the first row, for instance: $S_{s\to b} = -0.459$, $S_{s\to e}=0.707$, i.e., $S_{s\to b}$ is noticeably smaller than $S_{s\to e}$. But according to the tagging scheme of $(s,b,m,e)$, $s$ can be followed by $b$, but not by $e$, so $S_{s\to b} < S_{s\to e}$ being what it is is clearly unreasonable—it might lead to invalid tag sequences. Ideally, $S_{s\to e}$ should be $-\infty$.

This unreasonable transition matrix made me suspect, for a while, that there was a bug in my CRF implementation. But after repeated checking and comparison against the official Keras implementation, I eventually confirmed that my implementation was correct. So where was the problem coming from?

Unequal Learning Rates

If we set aside the question of whether the transition matrix makes sense, and simply feed the model's training results directly into the Viterbi algorithm for decoding and prediction, then evaluate with the official script, we find an F1 of around 96.1% (on the PKU task)—already state of the art.

The transition matrix is terrible, yet the final result is still very good—this can only mean that the transition matrix has almost no effect on the final result. Under what circumstances would the transition matrix have almost no effect? A plausible reason is that the per-character label scores output by the model are far larger than the values in the transition matrix, and are already highly discriminative, so the transition matrix simply can't influence the overall outcome—in other words, at this point plain Softmax followed by argmax already works great. To confirm this, I randomly picked a few sentences and looked at the model's output label distribution for each character, and indeed found that the highest label score for each character was generally between 6 and 8, while the remaining label scores were generally more than 3 points lower than the top score—an order of magnitude larger than the values in the transition matrix, so it's clearly very hard for the transition matrix to have any influence. This confirms the conjecture.

A good transition matrix should obviously help with prediction—at the very least it should help rule out unreasonable label transitions, or at least guarantee no negative effect. So it's worth pondering: what exactly is preventing the model from learning a good transition matrix? My guess is: the learning rate.

After BERT is pretrained, when fine-tuning it on a downstream task, you only need a very small learning rate (typically on the order of $10^{-5}$); too large a rate might actually prevent convergence. Despite the small learning rate, convergence is quite fast for most downstream tasks—many tasks converge to their optimum in just 2–3 epochs. On the other hand, BERT's fitting capacity is very strong, so it can fit the training data quite thoroughly.

What does this tell us? First, recall that the per-character label distribution is computed directly by the BERT model, while the transition matrix is an add-on, with no direct connection to BERT. When we fine-tune with a learning rate on the order of $10^{-5}$, the BERT part converges quickly—that is, the per-character label distribution gets fitted rapidly, and because BERT's fitting capacity is so strong, it quickly converges to a fairly optimal state (i.e., the target label gets a high score, with a large gap from non-target labels). But since the transition matrix has little connection to BERT, while the per-character label distribution rapidly converges to a near-optimal value, the transition matrix keeps "sauntering along" at the rate of $10^{-5}$, ending up an order of magnitude smaller than the per-character label scores. Moreover, once the per-character label distribution can already fit the target sequence well, the transition matrix is no longer really needed (its gradient becomes very small, so it hardly gets updated at all).

Thinking along these lines, a natural idea emerges: what if we increase the learning rate of the CRF layer? I tried increasing the CRF layer's learning rate, and after several experiments, found that once the CRF layer's learning rate exceeds 100 times the main learning rate, the transition matrix starts to become sensible. Below is a transition matrix trained with a BERT main learning rate of $10^{-5}$ and a CRF layer learning rate of $10^{-2}$ (i.e., 1000 times larger):

$$\begin{array}{c|cccc} & s & b & m & e \\ \hline s & 3.17 & 2.16 & -3.97 & -2.04 \\ b & -3.89 & -0.451 & 1.67 & 0.874 \\ m & -3.9 & -4.41 & 3.82 & 2.45 \\ e & 1.88 & 0.991 & -2.48 & -0.247 \\ \end{array}$$

This transition matrix makes sense, and the magnitudes are right too—it has learned the correct label transitions, e.g., $s\to s,b$ scores higher than $s\to m,e$, $b\to m,e$ scores higher than $b\to s,b$, and so on. However, even with the boosted CRF layer learning rate, the result is not noticeably better than without the adjustment. Ultimately, this comes down to the fact that BERT's fitting capacity is so strong that even Softmax alone achieves state-of-the-art performance, so the transition matrix naturally can't provide much of a boost.

(Note: for implementation tricks on increasing the learning rate, see "Making Keras Even Cooler!": Layer-wise Learning Rates and Free Gradients.)

More Experimental Analysis

CRF doesn't change BERT's performance much, because BERT's fitting capacity is so strong that good results are achieved even without the transition matrix. So, would lowering BERT's fitting capacity lead to a more noticeable difference?

In the earlier experiments we fine-tuned using the output of the 12th layer of BERT base. Now let's fine-tune using only the output of the 1st layer, to test whether the adjustments described above make a significant difference. The results are as follows:

$$\begin{array}{c|cc|cc} \hline & \text{main learning rate} & \text{CRF learning rate} & \text{epoch 1 test F1} & \text{optimal test set F1}\\ \hline \text{CRF-1} & 10^{-3} & 10^{-3} & 0.914 & 0.925\\ \text{CRF-2} & 10^{-3} & 10^{-2} & \textbf{0.929} & \textbf{0.930}\\ \text{CRF-3} & 10^{-2} & 10^{-2} & 0.673 & 0.747\\ \text{Softmax} & 10^{-3} & \text{-} & 0.899 & 0.907\\ \hline \end{array}$$

Since only 1 layer of BERT is used, the main learning rate is set to $10^{-3}$ (the shallower the model, the larger the learning rate can appropriately be); the main comparison here is the improvement brought by adjusting the CRF layer's learning rate. From the table we can see:

1. Under an appropriate learning rate, CRF outperforms Softmax;
2. Appropriately increasing the CRF layer's learning rate also brings some improvement over the original CRF.

This suggests that for models whose fitting capacity isn't especially strong (e.g., using only the first few layers of BERT, or for certain particularly difficult tasks where even the full BERT's fitting capacity isn't sufficient), CRF and its transition matrix still provide some benefit, and fine-tuning the CRF layer's learning rate can yield an even bigger improvement. Furthermore, all of the experiments above were based on BERT—does the same trick work for traditional BiLSTM+CRF or CNN+CRF as well? I did a simple test, and found that in some cases it does help, so I suspect this is a general trick for the CRF layer.

Summary

Starting from the CRF example added to bert4keras, this post found that when combining BERT with CRF, the CRF layer might be under-trained. I then hypothesized a possible cause, further confirmed the hypothesis through experiments, and finally proposed boosting the CRF layer's learning rate to improve its effectiveness, providing preliminary validation of its usefulness (for certain tasks).