Discriminative Self-training for Punctuation Prediction

Our model treats punctuation prediction as a sequence labeling task. The input are transcripts without punctuation. For languages without explicit word boundaries such as Chinese, the input is segmented into words. Our model predicts whether there is a specific punctuation mark after each word. The word tokens are processed with WordPiece tokenization and the output are fed into a Transformer encoder [28]. The final hidden state of the encoder corresponding to the last sub-token is used as the input to the softmax classifier for punctuation prediction.

Figure 1 provides an overview of the proposed discriminative self-training approach for punctuation prediction, given large-scale written language data, human-labeled spoken language data with punctuation labels ${(x_{1},y_{1}),(x_{2},y_{2}),\dots,(x_{n},y_{n})}$ and unlabeled spoken language data ${\tilde{x}_{1},\tilde{x}_{2},\dots,\tilde{x}_{m}}$.

First, we pre-train a language model on large-scale well-formed text corpora by self-supervised tasks, such as masked language modeling (MLM) and next sentence prediction (NSP), to obtain deep bidirectional language representations.

We then initialize a teacher model $\theta^{t}$ from the pre-trained language model and train it on spoken language data with human-labeled punctuation marks through minimizing the cross-entropy loss $\ell$:

$\displaystyle\ell=\frac{1}{n}\sum_{i=1}^{n}{\ell(y_{i},f(x_{i},\theta^{t}))}$

(1)

Then, we initialize a student model $\theta^{s}$ from pre-trained language models and train it by minimizing the cross-entropy loss on the combination of human-labeled spoken language data and pseudo-labeled spoken language data.

In this paper, we propose a self-training approach employing weighted loss and label smoothing in a discriminative way when training on the combination of human-labeled data and pseudo-labeled data, denoted discriminative self-training (Disc-ST). Intuitively, the pseudo-labeled data has more noise than human-labeled data. Thus, we use different weights to combine the loss on human-labeled data and pseudo-labeled data (denoted weighted loss), as follows:

$\displaystyle\ell=\sum_{i=1}^{n}{\ell(y_{i},f(x_{i},\theta^{s}))}+\alpha\sum_{i=1}^{m}{\ell(\tilde{y}_{i},f(\tilde{x}_{i},\theta^{s}))}$

(3)

where $\alpha$ is a weight factor. Meanwhile, we explore one of the output regularizers, label smoothing [29], to deal with training with noisy labels. In future work, we plan to explore other output regularizers such as focal loss [30] and bootstrapping loss [31]. The standard cross-entropy loss is calculated as

$\displaystyle\ell$

$\displaystyle=-\sum_{i=1}^{K}y_{i}log(p_{i})=-y^{T}log(p)$

(4)

where $y=\text{One-hot}(i)$, $p=[p_{1},p_{2},\dots,p_{K}]$, $K$ is the number of classes, and $i$ is the label index. For label smoothing, we replace $y$ with $y_{LS}$ as

$\displaystyle y_{LS}$

$\displaystyle=(1-\beta)y+\beta/K$

(5)

Similar to the motivation for weighted loss, considering the different noise levels in human-labeled data and pseudo-labeled data, we use label smoothing discriminatively, i.e., we use different factors ($\beta_{1}$ and $\beta_{2}$) for human-labeled data and pseudo-labeled data. $\alpha$, $\beta_{1}$ and $\beta_{2}$ are hyperparameters optimized on the validation set. After Disc-ST, we put the student back as the teacher, and iterate this process until the performance converges. The student model achieving the best performance on the human-labeled validation set is chosen as our final model.

When evaluating punctuation prediction performance on a test set, it is infeasible to decode a very long utterance without any segmentation all at once, since it will cause unacceptable latency and out-of-memory issues. Previous work studied different decoding strategies to segment long sequences, including overlapping windows [32], streaming input scheme [33], overlapped-chunk split and merging algorithm [34], and a fast decoding strategy [1]. We extend the overlapped-chunk split and merging algorithm [34] to improve punctuation prediction accuracy. As showed in Figure 2, we use a sliding window with both left and right overlapped context, and only keep predictions of the model where there is enough context information on both sides. The step size equals the window size minus the sum of the left overlap size and the right overlap size. In our experiments, we observe that the left context is more important than the right context for punctuation prediction accuracy. The main difference from the original overlapped-chunk split and merging algorithm is that our step size can be tuned by optimizing the left and right overlap sizes (as two hyperparameters) based on the performance on a validation set; whereas their step size is not tuned and is set as half of the window size.

We evaluate the punctuation prediction accuracy on the English IWSLT2011 benchmark dataset and an internal Chinese spoken language dataset. The IWSLT2011 benchmark contains three types of punctuation marks (comma, period, and question mark). We use the same data organization and same tokenized data as Che et al. [10]¹¹1https://github.com/IsaacChanghau/neural_sequence_labeling used. We also compare the proposed approach with previous methods on a large internal Chinese dataset [1]. We use Jieba²²2https://github.com/fxsjy/jieba for Chinese word segmentation. The punctuation annotations consist of four types of punctuation marks (comma, period, question mark, and enumeration comma). For unlabeled spoken language data, we use LibriSpeech [4], Fisher Speech Transcripts Part 1 and Part 2 [35] for the English dataset and use internal speech transcripts without punctuation for the Chinese dataset. Data statistics are summarized in Table 1. We evaluate the punctuation prediction performance using token-based precision (P), recall (R), F₁-score (F₁), following previous works [12].

For all experiments on English IWSLT2011, we use the “BERT-base, Uncased” model (12 layers, 768 hidden units, 12 heads, 110M parameters) [36]³³3https://github.com/google-research/bert and “ELECTRA-large” model (24 layers, 1024 hidden units, 16 heads, 335M parameters) [37]⁴⁴4https://github.com/google-research/electra. For IWSLT2011, we train 3 epochs, and use batch size 32 and initial learning rate 5e-5 for BERT-base; and batch size 16, initial learning rate 2e-5 for ELECTRA-large. The window size is set to 120, the right overlap size is set to 15, and the left overlap size is set to 35 for BERT-base and 40 for ELECTRA-large. For the Chinese dataset, we use the “RoBERTa-wwm-base, Chinese” models (12 layers, 768 hidden units, 12 heads, 110M parameters) [38]⁵⁵5https://github.com/ymcui/Chinese-BERT-wwm. We use batch size 64, initial learning rate 2e-5, and train 2 epochs. The number of self-training iterations is 1 in all of our experiments.

We compare the proposed approach with previous models on English IWSLT2011. The first group of models and results listed in Table 2 is cited from previous works. T-LSTM [12] used uni-directional LSTM and T-BRNN-pre [13] used bidirectional RNN with attention. BLSTM-CRF and Teacher-Ensemble are the best single and ensemble models in [14], respectively. DRNN-LWMA-pre [39] used a deep recurrent neural network architecture with layer-wise multi-head attentions. Self-attention-word-speech [16] used a full sequence Transformer encoder-decoder model with pretrained word2vec and speech2vec embeddings. CT-Transformer [1] used a controllable time-delay Transformer and utilized well-formed text corpora to pretrain. SAPR [40] used a Transformer encoder-decoder which views the punctuation restoration as a translation task. BERT-base+Adversarial [3] used BERT-base with adversarial multi-task learning by adding an extra POS tagging task. BERT-large+Transfer [2] used BERT-large with LSTM-CRF. BERT-base+FocalLoss [20] improved BERT-base with focal loss to alleviate the class imbalance problem. RoBERTa-large+augmentation [21] used RoBERTa-large with data augmentation, including insertion, substitution, and deletion. The current SOTA (RoBERTa-base) [22] used RoBERTa-base with aggregating predictions across multiple context windows.

The second group shows the experimental results from this work (using the decoding strategy in Section 3.3). We finetune BERT-base and ELECTRA-large models with the human-labeled training set and observe that ELECTRA-large outperforms the current SOTA (RoBERTa-base) (F₁ 84.4 versus 83.9) and significantly outperforms BERT-base (F₁ 84.4 versus 77.4). Vanilla self-training (denoted vanilla ST) on BERT-base improves F₁ from 77.4 to 79.0 (+1.6), and on ELECTRA-large improves F₁ from 84.4 to 84.7 (+0.3), demonstrating that vanilla ST obtains significant gains over the strong baselines. Replacing vanilla ST with Discriminative Self-Training (Disc-ST) on BERT-base achieves a further improvement on F₁ from 79.0 to 79.6 (+0.6), a statistically significant improvement with $p<0.01$; for ELECTRA-large, Disc-ST improves F₁ from 84.7 by vanilla ST to 85.2 (+0.5), a statistically significant improvement with $p<0.01$. Our ELECTRA-large+Disc-ST outperforms the current SOTA (RoBERTa-base) by 1.3 absolute F₁ gain.

We also compare the proposed approach with previous models on the Chinese dataset. The first group of models and results in Table 3 is cited from previous works. The second group shows the experimental results from this work. To compare with previous methods in [1], we use the same fast decoding strategy [1] with a low frame rate 3 and the number of look-ahead words after end-of-sentence mark as 6. The RoBERTa-wwm-base model significantly outperforms the current SOTA CT-Transformer (F₁ 59.6 versus 58.8). Furthermore, Discriminative Self-Training improves F₁ from 59.6 to 60.2 (+0.6), a statistically significant improvement with $p<0.01$.

We conduct ablation study to understand the contributions of component algorithms to the overall system performance on the IWSLT2011 test set, as shown in Table 4. When we add vanilla ST to BERT-base, F₁ improves from 77.4 to 79.0 (+1.6). When we further add weighed loss, F₁ improves from 79.0 to 79.3 (+0.3). When we further add equal label smoothing (LS) for both human-labeled data and pseudo-labeled data, F₁ improves from 79.3 to 79.4 (+0.1). When we replace equal LS with discriminative LS, F₁ improves from 79.3 to 79.6 (+0.3). These results demonstrate that both weighted loss and discriminative LS in discriminative self-training contribute to the improvement over vanilla self-training.