Navigating the Minefield of MT Beam Search in Cascaded Streaming Speech Translation

To understand our algorithm and its complexity, we need to introduce the core concept of intermediate and final events that flow through the system. ASR is feeding MT two types of events, intermediates and finals. Intermediate events denote an ASR result that is not stable—transcription can change. ASR has not yet made the final determination, and every intermediate will be followed by a subsequent event that can be either another intermediate or a final, to substitute words from the starting point of the last intermediate.

Final events denote ASR results that won’t change anymore, and they advance the “starting point of the last intermediate” to after the final event. The concatenation of all final events is the final transcription of the entire utterance. One observes that frequently, events consist of a single newly-recognized word, or (with word-piece based ASR) sometimes even only a prefix of a word—consider a plural noun reported as two consecutive final events: its singular base form followed by an “s”.

A key goal of our streaming beam search for MT is that it gracefully handles intermediate events and word pieces, by both generating them, but also correctly handling intermediate ASR events.¹¹1How ASR generates intermediates vs. finals is a well-understood problem that is not subject of this paper. An MT intermediate is just like ASR—a preview as to how one or more words would be translated; but future additional input may change or refine the word choices.

When the translation target is a display, displaying intermediate translations early is useful, as they provide a perception of responsiveness. If after receiving more context, the MT model decides that an intermediate was incorrect, the display is rewritten, which some users describe as “intelligent” behavior.

When a display is not available and the translation target is audio/TTS, however, “rewriting” is not possible. Intermediates are not useful, as we must hold back TTS until translated words have been finalized (are stable). Hence, in this scenario, the objective for the MT is to finalize and emit translated words as soon as possible—this is called streaming MT.

One can approximate simultaneous MT with a non-streaming model with a traditional encoder-decoder MT model trained to take a whole sentence as its input, where intermediate translations are generated by repeatedly translating the growing prefix of a sentence as it is being spoken. The updated translations can be updated on the display, sometimes called “screen rewriting.”

One disadvantage is the substantial flicker because the translation will evolve as more context becomes available, often replacing big chunks of the text. The simple solution—simply wait until the end of the sentence–however leads to poor (slow) user experience. We find users value a perception of responsiveness over stability. Another disadvantage of repeated retranslation is its computational inefficiency ($O(T^{2})$)—one ends up running very similar computations over and over again. This is exacerbated with ASR intermediates that cause even more retranslations.

Simultaneous MT addresses both. The idea is that the model keeps a state, and when a new transcription arrives from ASR, it processes it and then decides whether it is ready to emit tokens to the output or whether it requires more input. This is for the case where emitted tokens are final and cannot be changed anymore; we extend it towards emitting intermediates as well in Section 3.3.

The deal-breaking drawback of repeated retranslation is that it does not work with TTS output, as we cannot “rewrite” audio.²²2The time machine has not been invented yet. User-perceived latency is directly determined by the modeling and decoding latency.

In summary, compared to “faking it” via repeated retranslation, simultaneous MT solves two problems: when translating to display, simultaneous MT primarily reduces compute, while when translating to audio, it aims to reduce latency from one sentence at a time to a few words.

An MT model is composed of two parts: encoder and decoder. We used the standard Transformer [5] architecture for both. The simultaneous-MT model, in additional to the probability distribution over tokens, also a returns a READ/WRITE probability at each timestamp: When the probability for WRITE is below a threshold, the model has not seen sufficient input to generate the next token. In this case, we will not generate a token to output, but instead ingest another input token (in a real-time system, this might block until the next ASR event is received).

Several model architectures and training methods for the READ/WRITE classifier have been proposed. The current state of the art approach is Monotonic Infinite Look-back (MILK) [3]. A deeper discussion about WRITE probability can be found in Section 3.5.

Non-streaming beam search is well-known, and relatively straight-forward. The “beam” is the set of the best-scoring $k$ hypotheses during the incremental decoding process. A non-streaming beam decoder first executes the encoder for an entire input sequence. Then it repeatedly runs one step of the decoder for each of the (up to) $k$ hypotheses in the current beam and expands each by all $n$ words in the vocabulary. This yields up to $n\cdot k$ expanded hypotheses from which we select the top $k$ based on their aggregated score. While iterating through the top $k$ hypotheses, we set aside those that ended with an <EOS> or exceed a maximum length. Decoding terminates once $k$ hypotheses have been finished.

Since ASR inputs arrive in real time, a “read” decision may refer to an input token that has not been spoken or recognized yet, and therefore require the decoder to yield the CPU until the next ASR event. This means that the internal decoding state must be carried across incremental invocations of the decoder.

Specifically, we need to maintain the list of the current hypotheses, and for each hypothesis, we need to store the encoder and decoder model state, and couple other variables. We will refer to this as the beam search state. The model state is necessary to store, so that we don’t need to recompute it every time a new token comes in. The encoder and decoder model state caching is implemented using KV-cache mechanism [7] which caches keys and values inside the transformer layer computed in the previous timestamps.

We want to write to the display intermediate outputs while we are decoding input to improve user-perceived latency. We have two types of outputs—intermediate and final. A final output is a prefix of the output that will not change—we won’t rewrite it. Intermediate output can change, so the user might experience a bit of a flicker when words get overwritten on the display.

To determine final events, at each decoding step, we take a look at the top k hypotheses, and determine the common prefix of those. That common prefix is then emitted as a final event since this prefix can no longer change during the following steps of the beam search.

To generate intermediate events, each time ASR sends us an event, we process it until all hypotheses end in a “read”, then we take the top hypothesis, and return that as an intermediate input. While the common prefix of the hypotheses is stable—it won’t change—the rest of the output sequence in the top hypothesis can change, so we treat it as intermediate.

The ASR engine itself is also running beam search similar to ours, and it is therefore generating intermediate results which are not stable and might change. To handle intermediate ASR inputs, we checkpoint the beam search state for the last finalized input, run decoding for the intermediate input, emit the top hypothesis as intermediate output, and then rewind the state to the last finalized/stable state checkpoint. This ensures that the user is getting more feedback about the translation from the system, and we don’t recompute the same state multiple times.

An interesting edge case that we had not anticipated initially was that ASR can return incomplete word prefixes in the stream of inputs. For example, the word token _carpet might be recognized as two word pieces _car followed by pet, and these may be emitted in two consecutive final events. The issue with incomplete words is that the tokenization of the word can change once we get the rest of the word—while the ASR vocabulary may not have a whole-word entry for _carpet, the MT vocabulary may. The way we handle this is that we take the input, and treat the prefix that ends at a word boundary (space or punctuation) as finalized input. We cannot know whether the remainder of the input string has stable tokenization, so we treat it as intermediate input.

To rank beams, we need to compute the score for each beam. It is not completely clear how the score should be computed in the case of streaming NMT model. In case of RNN-T, the score of a beam is the score of the path including the “read” or “write” probabilities. According to pattern-recognition theory, error rate is minimized by selecting the hypothesis with the highest probability of having given rise to the observed input (Maximum-A-Posteriori, or MAP, decoding). This is rigorously reflected by the RNN-T decoding algorithm, and also in the RNN-T training loss. On the other hand, the “write” probability provided by MILK [3] method used to train the streaming-MT model is more of a heuristic. Up to now, MILK models have only been decoded with greedy search, where an accurate probabilistic score for multiple competing hypotheses is not needed.

We experimented with including the heuristic write probability in the path score of a beam. This did not work well. The issue is that the MILK loss does not include the write probability³³3We hope that a future version of the MILK loss model can provide a more rigorous “write” probability estimate that enables the use of a true MAP decoder as shown to be optimal by pattern-recognition theory., just the token probabilities (plus a latency penalty that only indirectly affects the write probability). The approach that we settled for was that the score of a beam is length normalized token sequence probability, excluding the write probability at all. Length normalization is needed because each beam could perform a different number of writes.

Model inference is dominated by memory bandwidth more so than compute cycles. Therefore, batching multiple model evaluations significantly improves the inference time. We don’t need to perform batching on the encoder side because all of the beams read at the same time, which means that we can hold the same encoder state for all of the beams. However, every beam has a different decoder state because each beam could perform a different number of writes. In non-streaming beam search, batching these is easy because all hypotheses share the same length. In the streaming case, it’s not so easy. When we are performing self-attention, the token sequence is different for each beam which makes it hard to store all of the beams in the same tensor.⁴⁴4This problem could be mitigated using LSTM decoder, but that might introduce accuracy regression, so it is about a trade-off.

The model sometimes produces <EOS> token before we get <EOS> token on input. We decided not to allow the model to do that because we still want to finish reading the whole sentence on the input.

Once we read <EOS> on input, we cannot perform any reads. We can only write. In this case, we ignore the model’s write probability, and just keep writing until the model produces <EOS> token. This strategy proved to be better than the strategy where we keep writing until the write probability is above a threshold. Sometimes the model does not produce <EOS> token at all. In that case, we want to stop writing at some point. The heuristic that we settled for is that we have a fixed threshold on the output sequence length that is based on the number of input tokens. The threshold is equal to $ax+b$, where x is the number of input tokens, and a and b coefficients that were tweaked using cross-validation.

The main trade off that we are making when designing a streaming system is between latency and accuracy. In MT, accuracy is commonly measured by the BLEU[8] score, for “bilingual evaluation understudy”. In our case, we only consider final output tokens for BLEU score evaluation. We do not take into account intermediate translations in BLEU score calculation.

We also want to measure the user’s perceived latency. Common metric used in simultaneous translation is word average lag[9] which estimates how many future input words on average we need to read in order to translate an input word.

From the user perspective, it also matters how often we rewrite (change) words on the display. For users, it is better when we don’t rewrite any words so they don’t have to go back and reread the sentence. We came up with a metric called character flicker which quantifies what percentage of characters change on the display on average after each write to the display.

The last metric that we used is total runtime. This measures CPU time running the whole translation pipeline. The computational resources are limited on device and every subsystem running on the device has a certain compute budget, and when MT exceeds the budget, it will get preempted meaning that we don’t be allowed to run MT model for a while on the device. Additionally, the total runtime also contributes to the user’s perceived latency. The total runtime also includes the compute needed for processing intermediate inputs.

In these benchmarks we simulated streaming ASR input by splitting a sentence into words, and then feeding one word after another to the streaming system as an ASR final event. In this case, we only have final inputs. We don’t have intermediate inputs. We used internal English-Spanish dataset for these benchmarks.

In these benchmarks, we transcribed a set of audio recordings from real-life English-Spanish conversations using ASR model and logged the intermediate and final outputs from ASR. Then we fed both intermediates and finals into the translation system in the same way the model would receive them during real-time operation. We compared streaming beam search and fake streaming with repeated retranslations.

Table 3 shows how long it takes to process a single input event (intermediate or final), on average, from ASR, and character flicker on the display.

We can see that CPU time is better for streaming compared to fake streaming for all beam sizes. This is because the fake streaming needs to perform more retranslations when we introduce intermediate input events. The improvement is biggest for the beam size one. Flicker is also much better for all beam sizes for streaming compared to fake streaming.