Neural Language Modeling With Implicit Cache Pointers

Let us denote the hidden state of an RNNLM at time step $t$ as $\mathbf{h}_{t}$. Conventionally, the RNNLM output $\mathbf{y}_{t}$ is determined as

$$\mathbf{y}_{t}=\text{softmax}(\mathbf{W}\mathbf{h}_{t}+\mathbf{b}),$$

(1)

where $\mathbf{W}\in\mathbb{R}^{V\times H}$, $\mathbf{b}\in\mathbb{R}^{V}$, and $\mathbf{y}_{t}\in\mathbb{R}^{V}$, with $V$ and $H$ being the vocabulary size and hidden state dimension, respectively.

We extend the output dimension by a predefined size $L$. The extended part represents the $L$ immediately preceding words in history. Activation of these $L$ extended units, denoted as $\mathbf{p}_{t}$ in Figure 1, is computed via a linear projection of $\mathbf{h}_{t}$ from the last hidden layer of the RNNLM. We thus have

		$\displaystyle\mathbf{p}_{t}=\mathbf{W}_{p}\mathbf{h}_{t},$		(2)
		$\displaystyle\mathbf{z}_{t}=\text{concat}(\mathbf{W}\mathbf{h}_{t}+\mathbf{b},\mathbf{p}_{t}),$		(3)
		$\displaystyle\mathbf{y}_{t}=\text{softmax}\left(\mathbf{z}_{t}\right),$		(4)

where $\mathbf{W}_{p}\in\mathbb{R}^{L\times H}$, $\mathbf{z}_{t}\in\mathbb{R}^{V+L}$, and $\mathbf{p}_{t}\in\mathbb{R}^{L}$. Applying softmax on $\mathbf{z}_{t}$ generates the extended output $\mathbf{y}_{t}\in\mathbb{R}^{V+L}$.

Since the $L$ extended outputs indicate where to copy from the history, we call $\mathbf{p}_{t}$ the pointer component of our model. It only introduces $L$ $\times$ $H$ additional parameters.

The objective for training a neural LM is to maximize the log likelihood of training data. The loss function is written as

$$L(\theta)=-\frac{1}{T}\sum_{t=1}^{T}\log(\mathbf{y}_{t}\bm{\cdot}\mathbf{s}_{t}),$$

(5)

where $T$ is the total number of words in the training data, $\mathbf{s}_{t}$ is the supervision vector, and $\bm{\cdot}$ is vector dot product operation. In conventional training, $\mathbf{s}_{t}$ is a one-hot vector with 1 in the index of the target word. To train the proposed neural LM with the pointer component, the supervision vector $\mathbf{s}_{t}$ is set to have additional ones in history positions where the the target was previously seen. So $\mathbf{s}_{t}$ is an at-least-one-hot vector.

The pointer component and the modified training supervision makes an RNNLM be aware of where to copy from the history. However, it may still be challenging for the model to memorize which words are particularly likely to reoccur, i.e. are ”bursty”. To learn the burstiness of words, one additional unit, denoted by $m_{t}$ is introduced alongside the pointer $\mathbf{p}_{t}$, as shown in Figure 1.

This additional unit is computed by a vector dot product of the hidden state $\mathbf{h}_{t}$ and a parameter vector with the same dimension as $\mathbf{h}_{t}$, but is not used in computing the softmax. It influences the probability that a word may repeat through $\mathbf{p}_{t}$. Specifically, these additional units from the $L$ immediately preceding word positions are concatenated to form

$$\mathbf{m}_{t}=\text{concat}(m_{t-(L-1)},...,m_{t}),$$

(6)

where $\mathbf{m}_{t}\in\mathbb{R}^{L}$, and $\mathbf{m}_{t}$ is element-wisely added to the pointer component $\mathbf{p}_{t}$, i.e.

$$\mathbf{p}_{t}:=\mathbf{p}_{t}+\mathbf{m}_{t}.$$

(7)

Thus $m_{t}$ influences the output $\mathbf{y}_{t}$ indirectly by modifying $\mathbf{p}_{t}$ in (3), which in turn is a part of $\mathbf{z}_{t}$ in (4).

Compared with an RNNLM, the memory augmented pointer component only has $(L+1)$ $\times$ $H$ total additional parameters while a PSMM has extra $H^{2}+2H$ parameters. Without attention and gating mechanism, the proposed model is simpler than PSMM and has fewer extra parameters when $L$ $\leq$ $H$.

This pointer mechanism described above is for RNNLMs, while it can also be easily incorporated into Transformer-based LMs. In the latter, the context vector from the last Transformer block is treated as the hidden state in RNNLMs.

We first compare the proposed model with PSMM and neural cache under the plain LSTM setup. Perplexities on PTB and WikiText-2 are shown in Table 3. The performance gap between the PSMM in the paper [19] and ours is mainly caused by different implementations of truncated back-propagation through time (BPTT). They use an explicit truncated BPTT while we follow the normal way discussed in [19] considering efficiency and convenience of data preprocessing. We first concatenate all text words and then chunk them with fixed size $L$. So, if the truncated BPTT length is $L$, each training word on average experiences $L$/2 instead of $L$ time-steps for back-propagation. This means each training word sees $L$/2 history words on average.

In Table 3 “Memory Aug” refers to the memory augmented pointer. We set history length as 100, equal to the truncated BPTT length. Setting $L$ = 50 for neural cache is a fair comparison with others. Results on both datasets show that memory augmentation provides further improvement on top of the pointer component. And with memory augmentation, the proposed model outperforms the rest on both datasets. In subsequent tables, “Proposed” refers to LMs with the memory augmented pointer.

To verify whether the proposed approach is robust, we conduct experiments on a stronger baseline Frage-AWD-LSTM setup [23]. We reproduced their results and implemented the proposed approach on top of theirs, without tuning meta-parameters. Perplexity results in Table 4 shows that the proposed model on both datasets achieves better results than Frage-AWD-LSTM. Further improvements are observed with increasing the history length from 50 to 100, as expected. We also observe complementary effects of the proposed model and neural cache model.

We experiment with both LSTM- and Transformer-based LMs on SWBD. Perplexity on dev set from Kaldi RNNLM is 50. The history length as well as BPTT length is set to 100 for the proposed model and PSMM. Results of Pytorch trained models are in Table 5.

For LSTM-based models, the proposed approach outperforms the baseline LSTM, but performs slightly worse than the PSMM and neural cache models. For Transformer LMs, the proposed approach also achieves better perplexity than the two Transformer baselines.

We notice the performance gains of the proposed model over both LSTM and Transformer baselines on SWBD are smaller than on PTB and WikiText-2. To check whether this may relate to style (only SWBD is spoken style) and average sentence length (sentences on Switchboard are the shortest on average), we experiment with Transformer-based LMs on WSJ. The proposed approach reduces perplexity from 71.5 to 65.6 on test set (eval92), compared with a baseline Transformer LM. The improvement is relative 8.2% and in similar range with gains on WikiText-2. And this results in 0.1 absolute WER reduction from 1.5 to 1.4 on the same test set by $N$-best rescoring. While no conclusions can be drawn yet, experiments show that the proposed model perform better on written-style text which usually has longer average sentence length.

One desired feature of the proposed model is that it may predict rare words better than LSTMs. To verify this, we further examine LM performance on the test set of each dataset. We split the vocabulary of each corpus into 10 buckets based on word frequencies in training data such that we get a roughly equal number of test tokens in each bucket. We then compute the differences between the test cross entropy of the proposed model and the LSTM baseline on words in each bucket for each dataset. Figure 2 shows the results on WikiText-2. As expected, larger reductions in cross-entropy are observed from the proposed model on rare words. Similar trends are seen on PTB and SWBD, though the overall perplexity improvement on SWBD is marginal. Figures for them are omitted due to page limits.

As Transformers show similar perplexity gains as LSTMs on SWBD, we only experiment with LSTMs for $N$-best rescoring. WERs on the full HUB5’00 evaluation set (Eval’00), the SWBD subset (SWB), and Callhome subset (CH) are in Table 6. “State-carry” means when scoring a hypothesis for current utterance, the initial hidden state is copied from the last hidden state of the best hypothesis for the previous utterance, instead of being zero initialized. WER improvements by the LSTM with “state-carry” in Table 6 indicate that cross-sentence context is useful. Similar observations are presented in [30]. If not specified with “w/o state-carry”, models are evaluated in the state-carry way.

To investigate the effect on WERs of rare words, we conduct a similar analysis as Section 5.3 does. Words with errors on Eval’00 (<5000) are divided into 5 buckets based on frequency. Relative WER reductions on words in these buckets in Figure 3 show that the proposed model improves performance on both relatively rare words and very frequent ones. As expected, some rare words occur within the context window, for example, masters and offered, are recognized correctly. Decoded output also shows that frequent words such as train, short, and were are correctly recognized. Though the overall WER improvement by the proposed model is marginal, correctly recognizing relatively rare words plays an important role in user experience of ASR-based products or service.

We also notice the proposed model sometimes introduces errors on words that are wrongly recognized in first pass decoding. A possible reason is that the supervision vectors for pointer components are from decoded hypotheses and hence may contain errors. We verify this by using test transcription in rescoring and observe a further 0.1 absolute WER reduction on Eval’00 of SWBD. The mismatched condition between training and evaluation is a common issue for the proposed approach, PSMM, and neural cache model. To alleviate the mismatch, word level confidence scores and error adaptive training approaches could be considered.