Teaching BERT to Wait: Balancing Accuracy and Latency for
Streaming Disfluency Detection

Streaming settings force models to make trade-offs between accuracy and latency. More accurate predictions can be obtained by providing longer right context at the cost of incurring additional latency. However, since most tokens are not disfluent, a model may not require the right context in order to accurately classify fluent tokens. Rather than using a fixed lookahead window, we train a BERT model to jointly classify tokens and simultaneously choose the lookahead window size dynamically at each token.

Our proposed model architecture consists of a pre-trained BERT model with two separate token classification heads added on top, as shown in Figure 2. Each classification head consists of a linear layer applied to the hidden layer outputs of the BERT model. The first classification head, the disfluency classifier, is trained to classify whether each token is disfluent or not. The second classification head (the wait classifier) is trained to classify whether the model should wait for further input (temporarily abstain from predicting) or immediately output a prediction for the given token. In effect, the wait classifier decides how large of a lookahead window the model needs to make its prediction on the current input. At inference time, we only output tokens that lie to the left of the first token for which the model outputs a wait prediction and that are predicted to be fluent. This avoids potentially producing disjoint output in the case where the model produces predictions of wait followed by predict, making the output more clear for the user’s display.

We also adapt the training objective such that it accounts for both the accuracy and latency of the model’s outputs on each successive prefix. Let $(x,y)$ be the pair of input sequence and target output sequence with prefixes $x_{1},x_{2},\cdots,x_{|x|}$ and $y_{1},y_{2},\cdots,y_{|x|}$, respectively where $|x|$ is the length of the full sentence. We also denote $f(x)$ as the output logits of the disfluency classifier, $g(x)$ as the output logits of the wait classifier, $\sigma(\cdot)$ as the softmax function, and $H(\cdot,\cdot)$ as the cross-entropy loss. Then the traditional cross-entropy loss on the full input and target sequences is

$$\ell_{\textsc{Full}}(x,y)=H(\sigma(f(x)),y)$$

(1)

However, for each prefix we wish to only compute the cross-entropy loss on the tokens to the left of the first token for which the model outputs a wait prediction. To accomplish this, we devise a binary mask that zeros out the loss on the tokens to the right of and including the first wait prediction:

$$\mathbf{m}(\sigma(g(x_{i}))=(m_{1},\cdots,m_{|x_{i}|})$$

(2)

where

	$\displaystyle m_{j}$	$\displaystyle=\begin{cases}1&\mbox{if }j<k\\ 0&\mbox{otherwise}\end{cases}$		(3)
	$\displaystyle k$	$\displaystyle=\min\,\{j|\sigma(g(x_{i})_{j})>0.5\}$		(4)

where $g(x_{i})_{j}$ is the $j$-th element of vector $g(x_{i})$. We then apply this mask to the cross-entropy loss for each prefix of example $x$ to obtain prefix loss:

$$\ell_{\textsc{Prefix}}(x,y)=\sum_{i=1}^{|x|-1}\mathbf{m}(\sigma(g(x_{i})))\circ H(\sigma(f(x_{i})),y_{i})$$

(5)

where we abuse notation here by denoting $H(\sigma(f(x_{i})),y_{i})$ as the vector for which the $j$-th element corresponds to the cross-entropy of $(\sigma(f(x_{i}))_{j},y_{i,j})$ and $\circ$ is element-wise multiplication. Lastly, we define a latency cost that scales with both the probability of abstaining from classifying the $j$-th token in the $i$-th prefix ($\sigma(g(x_{i})_{j})$) and with the expected wait time, as measured by number of tokens, incurred by abstaining starting from token $j$ in prefix $x_{i}$:

$$\ell_{\textsc{Latency}}(x)=\sum_{i=1}^{|x|-1}\sum_{j=1}^{i}(i-j)\sigma(g(x_{i})_{j})$$

(6)

Putting these together, the total loss for a single example $(x,y)$ is:

	$\displaystyle\ell(x,y)$	$\displaystyle=\ell_{\textsc{Full}}(x,y)+\gamma\ell_{\textsc{Prefix}}(x,y)$
		$\displaystyle+\lambda\ell_{\textsc{Latency}}(x)$		(7)

with hyperparameters $\gamma$ and $\lambda$ controlling the relative strengths of the prefix and latency costs, respectively. We also include the cross-entropy loss on the full sequences ($\ell_{\textsc{Full}}$) in addition to the prefix losses ($\ell_{\textsc{Prefix}}$) because we wish for the model to maintain its ability to make predictions on full sequences. Since $g(x)$ does not appear anywhere in $\ell_{\textsc{Full}}$, the model is effectively forced to make predictions once it receives the full utterance.

Similarly, the $\ell_{\textsc{Latency}}$ term is essential because without it, the model could achieve minimal loss by always waiting (e.g. $\sigma(g(x_{i})_{j})=1$ for all prefixes $i$ and time steps $j$), and only learning to classify disfluent tokens after receiving the full sequence. This is equivalent to the non-incremental classification loss. If we instead set $\sigma(g(x_{i})_{j})=0$ for all $i,j$ (the case where the model never waits), the resulting loss is equivalent to the learning objective for strongly incremental training (see Section 3.2). In essence, our training objective is a generalization of the strongly incremental objective.

Although BERT models are typically fine-tuned using complete pre-segmented sequences, an incremental model must process partial inputs at inference time, resulting in a distributional shift between the complete utterances typically seen in training datasets and the partial utterances seen at inference time. A simple solution is to fine-tune BERT both on complete and partial inputs, a training scheme known as restart incrementality (Kahardipraja et al., 2021). By providing successively extended prefixes of a given utterance to the model and computing the loss on the model outputs for each prefix, we can mimic the streaming data that the model would encounter in real time. In all of our experiments, each successive prefix adds a single word to the previous prefix, a setting known as strongly incremental (Shalyminov et al., 2018). Although this approach requires re-computation of the model outputs for each successive prefix, this also enables the model to correct its previous predictions, or to switch between waiting and predicting when it receives helpful right context. Incorporating prefixes during training in incremental disfluency detection has been previously explored by Rohanian and Hough (2021). This serves as a strong baseline in our experiments.

All of our experiments use small distilled BERT models, specifically a small vocabulary BERT model (Zhao et al., 2021) ($\text{BERT}_{\text{SV}}$) with 12 hidden layers of size 128 that is pre-trained on English Wikipedia and BookCorpus (Zhu et al., 2015). The details of our hyperparameter tuning can be found in the Appendix (Section A.2).

We use small models for two reasons: 1) fine-tuning a model on all given prefixes of each training example is resource intensive, and 2) many streaming natural language understanding applications run entirely on mobile devices which precludes the use of large models. Previous work on small non-incremental BERT-based models used for disfluency detection Rocholl et al. (2021) showed significant improvement in memory and latency without compromising task performance. The core $\text{BERT}_{\text{SV}}$ model is a distilled version of $\text{BERT}_{\text{BASE}}$ with smaller vocabulary and reduced hidden layer dimensions Zhao et al. (2021). Due to its smaller vocabulary size (5K versus 30K tokens), the model has only about 3.1M parameters, as compared to $\text{BERT}_{\text{BASE}}$’s approximately 108.9M parameters, achieving around 80% latency reduction.

In order to isolate the effects of training with restart incrementality (Section 3.2) versus the improvements derived directly from incorporating our new training objective, we also evaluate two other models: 1) a non-incremental $\text{BERT}_{\text{SV}}$ model trained in the usual way, on full sequences; and 2) a $\text{BERT}_{\text{SV}}$ model trained with restart incrementality - i.e., on all prefixes of every training example (which we will refer to as “all prefixes” in following tables). Setup (1) is equivalent to ablating both $\ell_{\textsc{Prefix}}$ and $\ell_{\textsc{Latency}}$ whereas setup (2) is equivalent to ablating only $\ell_{\textsc{Latency}}$. We do not ablate $\ell_{\textsc{Prefix}}$ in isolation since this leaves only the $\ell_{\textsc{Full}}$ and $\ell_{\textsc{Latency}}$ terms, and there does not exist a meaningful measure of latency when the model never needs to wait for more input (since it is always given the full utterance as input). For each of these baseline models we also follow Rohanian and Hough (2021) and Kahardipraja et al. (2021) by evaluating with different fixed lookahead (LA) window sizes of $\text{LA}=0,1,2$.

Accuracy alone is not a sufficient measure of success to robustly evaluate a streaming model. Since a streaming model is meant to operate in real time, it should return output as soon as possible after it receives new input. As such, we also need to evaluate it with respect to latency – i.e. the number of new tokens a model must consume before producing a prediction for the current token. Furthermore, streaming models are often designed to be capable of retroactively changing their predictions on previous tokens as new input arrives. This introduces the risk of output “jitter” or “flicker,” where the output changes dramatically as new input is consumed, necessitating evaluation of stability. To capture all these important dimensions of streaming model performance, we evaluate the models using the following diachronic metrics (with formulas and further details in the Appendix):

• •

  Streaming $F_{1}$: An accuracy metric scored in the same way as the typical $F_{1}$ score, albeit we score the predictions for a single token over the course of multiple time steps separately as if they were predictions for separate tokens.

• •

  Edit Overhead (EO) (Buß and Schlangen, 2011): A stability metric that measures the average number of unnecessary edits, normalized by utterance length.

• •

  Time-to-detection (TTD) (Hough and Schlangen, 2017): A latency metric that is only computed on disfluent tokens that are classified correctly. It is the average amount of time (in number of tokens consumed) that the model requires before first detecting a disfluency. As mentioned earlier, we include both reparanda and interregna as disfluencies.

• •

  Average waiting time (AWT): The average amount of time (in number of tokens consumed) that the model waits for further input before making a prediction on a given token. For models with a fixed lookahead window, this is equivalent to the lookahead window size. For the streaming model, this is equivalent to the average lookahead window size.

• •

  First time to detection (FTD) (Zwarts et al., 2010; Rohanian and Hough, 2021): Similar to to the TDD metric described above with the main difference being that the latency (in number of words) is calculated starting from the onset of a gold standard repair.

Table 1 shows both the incremental and non-incremental evaluation metrics. Our proposed streaming $\text{BERT}_{\text{SV}}$ model achieved a $9\%$ increase¹¹1All percentages mentioned in this section are computed as a percentage of the original number, rather than as a difference in percentage points. in streaming $F_{1}$ over both of the baselines (with lookahead = 0), as well as a $71\%$ and $72\%$ decrease in edit overhead compared to the non-streaming models trained on full sequences and all prefixes, respectively. Despite being trained with a different architecture and loss objective, the streaming model does not sacrifice its non-incremental performance, yielding a final output $F_{1}$ score that is only one point less than its non-streaming counterparts. Generally speaking, when the streaming model does output a prediction, it classifies tokens as disfluent less often than the non-streaming models with zero LA window, achieving much higher precision (P) and marginally lower recall (R), resulting in a model that “flickers" less frequently. However, this does contribute to a slightly higher time-to-detection score compared to the baselines with zero lookahead, since the streaming model is generally less aggressive but more precise with outputting disfluent predictions. When compared to the models with fixed lookahead (the lower half of Table 1), however, the streaming model always achieves lower TTD while achieving significantly lower waiting time and comparable accuracy and stability.

Figure 3 shows an error analysis on the models’ predictions. We computed the percentage of the time that the streaming model misclassified a token, counting each incidence of the token across each time step separately. All models achieved the lowest misclassification rates on fluent and edit tokens, and the highest misclassification rates on reparanda tokens. For tokens that were fluent, edit terms, or part of a repair or repair onset, the streaming model achieved significantly lower misclassification rates than the baselines. However, all three models performed comparably on interregna and reparanda. Since we measure misclassification rate on every token at every time step, the high misclassification rates on reparanda are expected as it is often not feasible to detect a reparandum until an interregnum or repair onset has been seen.

As shown in Table 3, in comparison with the $\text{BERT}_{\text{LARGE}}$-based prophecy decoding model proposed in Rohanian and Hough (2021), our streaming model achieves state-of-the-art stability (85% decrease in EO) and latency (65% decrease in FTD)²²2In addition to shorter word-level latency metrics presented in the results, the runtime latency of the $\text{BERT}_{\text{SV}}$ model is 80% lower than that of $\text{BERT}_{\text{BASE}}$ Rocholl et al. (2021)., despite having far fewer parameters.