Efficient Encoders for Streaming Sequence Tagging

State-of-the-art text encoding methods assume the offline setting, where the entire input text is available when encoding it. This differs from the streaming setting where the input text grows over time (such as transcribed speech or a typed query) Cho and Esipova (2016); Gu et al. (2017); Chang et al. (2022). Processing streaming input incrementally can enable suggestions on partial inputs Iranzo Sanchez et al. (2022), reduced final latency Zhou et al. (2022) and lead to more interactive NLU agents Cai and de Rijke (2016).

Existing state-of-the-art bidirectional encoders (such as Devlin et al. (2019)) are computationally expensive for streaming processing. When a new token is received, these models require a restart, i.e., re-computation of representation of each token by re-running the bidirectional layer Kahardipraja et al. (2021). This adds to the computational cost (i.e., Higher FLOPs) during streaming and leads to higher wall clock time. Another limitation of these encoders is poorer generalization to partial inputs Madureira and Schlangen (2020), stemming from these models being trained only on complete (and offline) inputs. Despite these disadvantages, such bidirectional models offer better offline performance than unidirectional models across several NLP tasks like sequence tagging Kahardipraja et al. (2021).

We address these issues by proposing HEAR – Hybrid Encoder with Adaptive Restarts, where a separate Adaptive Restart Module (ARM) predicts when to restart the encoder. The encoder in HEAR is a Hybrid encoder that reduces the computational cost of running the models in streaming settings. In a Hybrid encoder, the earlier layers are unidirectional and the deeper layers are bidirectional. This design allows early contextualization, and limits the need for restart to the later layers. While Hybrid encoder reduce streaming computational overhead, restarting them at every new token may not be required. Thus, we propose a lightweight Adaptive Restart Module (ARM) to guide restarts, by predicting whether restarting the bidirectional layers of the Hybrid encoder will benefit the model performance. This module is also compatible with fully bidirectional encoders.

Table 1 showcases the strength of HEAR on one of the sequence tagging dataset we consider. In terms of streaming computation, measured in terms of FLOPs (lower is better), Hybrid encoder offers significant saving from purely bidirectional (BiDi) encoders and FLOPs savings improve upon incorporating ARM in HEAR. In terms of offline performance, Hybrid encoder and HEAR achieves parity with BiDi encoders and in terms of streaming performance, HEAR significantly outperforms BiDi encoders.

Following are our key contributions:

• •

  We introduce Hybrid encoder for computationally cheaper streaming processing (§3.1), that maintains the offline F1 score of bidirectional encoders, while reducing FLOPs by an average of 40.2% across four tasks.

• •

  We propose the ARM module (§3.2) to decide when to restart. The ARM reduces FLOP of Hybrid encoder by 32.3% and improves streaming predictions by +4.23 Exact Match.

• •

  Our best model HEAR combines Hybrid encoder with ARM (§5) to achieve strong streaming performance while saving FLOPs and offering competitive offline performance across four sequence tagging tasks.

In the streaming sequence tagging task, we assume that at time $1\leq t\leq n$, we have received the first $t$ tokens $\mathbf{x}_{t}=(x_{1},\cdots x_{t})$ as input from a stream of transcribed speech or user-typed input.¹¹1We assume that new tokens get added without changing previous tokens (contrary to some ASR systems), even though, our method can be used in such settings. The model then predicts the tags for all $t$ tokens $\hat{\mathbf{y}}_{t}=(\hat{y}_{1,t},\hat{y}_{2,t}\cdots\hat{y}_{t,t})$. The models predictions $\hat{\mathbf{y}}_{n}$ over the offline input $\mathbf{x}=\mathbf{x}_{n}$ is the offline tag sequence prediction which is evaluated against the ground truth $\mathbf{y}^{*}_{n}=(y^{*}_{1},y^{*}_{2}\cdots y^{*}_{n})$. However, in the streaming settings, we are concerned with predicted label sequences over all of the prefixes $\hat{\mathbf{Y}}=(\hat{\mathbf{y}}_{1},\hat{\mathbf{y}}_{2}\cdots\hat{\mathbf{y}}_{n})$.

During training we only have access to the offline ground-truth label sequence over the offline input sequence $\mathbf{y}^{*}_{n}$, even though ground truth labels for tokens may change as additional context is received in the future timesteps.

In order to motivate HEAR, consider how existing BiDi encoders would be used in streaming sequence tagging. At time $t$ and the input sequence $\mathbf{x}_{t}$, the model restarts to predicts the label sequence $\hat{\mathbf{y}}_{t}$ i.e. it re-computes all its layers for all tokens, without leveraging any of its previous computation or any previously predicted labels. Consequently, running a typical $O(n^{2})$ encoder would have $O(n^{3})$ computations in streaming settings, over $n$ input tokens. Therefore, naively using existing BiDi encoders for streaming settings is highly inefficient. Previous works tackle inefficiency by modifying the streaming model to infer only once for each word, after $k$ future tokens have been received Oda et al. (2015); Grissom II et al. (2014). This leads to a $k$-delayed output with possibility for revisions after additional future words have been received. However, this may lead to poor performance for tasks involving long-range dependency (e.g., SRL) and a higher output lag.

HEAR is a system consisting of a trained Hybrid encoder model and an ARM that is trained over the Hybrid encoder. The early unidirectional layers in Hybrid encoder, reduces the computational cost of restart of the encoder. The ARM guides when to restart the Hybrid encoder. It saves computational cost by keeping the model from restarting at each timestep and also improves streaming performance by avoiding unnecessary label flips stemming from avoidable restarts. Figure 2 shows running of HEAR in streaming settings over an example. For each new token in the streaming input, all the unidirectional layers and a part of first bidirectional layer is ran for the token, to obtain its unidirectional encodings and updated cache. These are then used by ARM to predict whether to restart the bidirectional layer or not. If the bidirectional layer is to be restarted, then we obtain updated labels for all the tokens received in the stream. Otherwise, the auxiliary predictor is ran over the unidirectional encoder for the current token to obtain its label and the labels from previous timestep is copied for all other tokens.

We now formally introduce Hybrid encoder (§3.1), followed by ARM (§3.2).

In this section, we provide empirical results to answer the following questions:

1. (a)

   In offline setting, does Hybrid encoder achieve parity with BiDi encoders?

2. (b)

   In streaming setting, by how much does Hybrid encoder reduce GFLOPs count?

3. (c)

   Does HEAR improve streaming performance and save on GFLOP count?

In this section, we perform various analysis pertaining to the Hybrid encoder and ARM.

Streaming (or incremental) setting has been widely studied in machine translation and parsing, dating as far back as two decades Larchevêque (1995); Lane and Henderson (2001). Specifically, for incremental parsing, there are two broad approaches that have been studied: transition-based and graph-based. Transition-based incremental parsers allow for limited backtracking and correcting parsing over prefixes by keeping track of multiple parse segments Buckman et al. (2016) or via beams search Bhargava and Penn (2020). Such methods can fail on garden-path sentences and long-range dependencies and only work proficiently with a large beam. Graph-based parsers incrementally assign scores to edges of the graph, discarding those edges that cause conflicts to tree-structure of the graph. Yang and Deng (2020) proposed an attach-juxtapose system to grow the tree, requiring restart at each new token over streaming input. Kitaev et al. (2022) improved on its efficiency by proposing a information bottleneck. However, these methods rely on the structured nature of parsing output can not be adapted to incremental sequence tagging tasks without restarting at each token.

Recently,  Madureira and Schlangen (2020) benchmarked the modern encoders on streaming sequence labelling and observed poor streaming performance of pretrained transformer models. They explored improvement strategies by adopting techniques from other streaming streaming like chunked training  Dalvi et al. (2018), truncated training  Köhn (2018) (training model on heuristically-aligned partial input-output pairs), and prophecy Alinejad et al. (2018) (autocompleting input using a separate language model). They observed performance degradation from truncated training unless used with prophecies. However, running a language model at each timestep for prophecy is computationally infeasible. Therefore, unlike our approach, neither methods can improve performance feasibly.

Previous works have attempted to improve computation efficiency in BiDi encoders. Monotonic attention moves away from $O(n^{2})$ soft attention overhead by restricting attention to monotonically increase across timesteps Raffel et al. (2017); Chiu and Raffel (2018); Ma et al. (2020). However, unlike HEAR, such attempts can’t maintain offline performance. Recently, Kahardipraja et al. (2021) used linear transformer  Katharopoulos et al. (2020) as unidirectional model using masking for streaming sequence tagging and classification. This approach performs well under the assumption of delayed output  Grissom II et al. (2014); Oda et al. (2015), a relaxation, where the model waits for additional tokens before predicting. Furthermore, the unidirectional model could not revise its output, rendering the model incapable of handling long-range dependency or tasks that go backward like SRL — an ability common to any model with some bidirectionality, such as HEAR. Similar drawbacks were in the partial bidirectional encoder, a bidirectional attention with restricted window, proposed by Iranzo Sanchez et al. (2022).

Improving efficiency through adaptive computing has been independently studied for reasoning-based tasks  Eyzaguirre and Soto (2020), text generation models  Arumae and Bhatia (2020); Eyzaguirre et al. (2022) and diffusion models  Ye et al. (2022). These works are restricted to offline settings and can be readily incorporated within the proposed overall approach of HEAR.

We propose HEAR for sequence tagging in streaming setting where the input is received one token at a time to the model. The encoder in our model is Hybrid encoder where early layers are unidirectional and later are bidirectional. It reduces the computational cost in streaming settings, by reducing the need of restart only to the later bidirectional layers while preserving the offline performance of the model. HEAR additionally consists of an ARM to predict when to restart the model. Using ARM leads to reduced number of restart of the encoder, leading to better streaming performance and further savings in computation. Compared to BiDi encoders, our model, HEAR (Hybrid encoder + ARM) reduces the computation by upto 71% in streaming settings while maintaining the performance of the BiDi encoders across various sequence tagging tasks. HEAR improved streaming EM by upto +10.0% and reduced unnecessary edits by upto -12.0%.

An additional but small training cycle is required to train the lightweight ARM module of HEAR in order to reap the benefits of extra savings in efficiency and streaming performance. Also, even though we do not assume any language specific-design choices, we benchmarked on the standard streaming sequence labelling benchmark datasets, all of which were in English.