[Feature Request]: Streaming ASR (WIP) #1970

🚀 The feature

Streaming ASR being a goal for SpeechBrain has already been discussed internally among the core team members, but it seems worthwhile to share the info/personal notes/WIP I have and to keep this publicly updated. I'll try to semi-regularly update the issue according to how things evolve, and it will be closed once a working implementation is merged.

Solution outline

The current plan is to base off work on #1782, first meaning changes to certain modules have to be done to enable streaming and having a working recipe. Implementing the actual interfaces for streaming ASR can come after once there is a proven to be working model. Certain features (like beam search decoding or LM rescoring) are non-essential for a proof-of-concept.

Some notes:

• k2/icefall appears to use https://arxiv.org/pdf/2012.05481.pdf aka a form of "chunked attention"
  • during training, at a given batch, a random chunk size is emulated by aligning visible future context for an output token to the chunk size (the paper explains it better)
  • because the model is trained with multiple chunk sizes, it can be adjusted for a latency-performance tradeoff at runtime
  • DCConv is an interesting variant (implemented and detailed further down): https://arxiv.org/pdf/2304.09325.pdf

Implementation status:

tl;dr on the streaming architecture changes:

• Feature extraction sees little changes

  • Padding is adjusted/performed manually in order to make it easier to align frames exactly as we want, because center==True in STFT could cause problems otherwise with certain sequence lengths. No other change.

• Conformer changes

  • Attention is masked in "chunks": All outputs within a chunk (of e.g. 16) can "see" the keys of its chunk + eventually some left context. For example, assuming infinite left context:
    • ALL outputs 0..15 will depend on inputs 0..15, but cannot depend on 16..
    • ALL outputs 16..31 will depend on inputs 0..31, but cannot depend on 32..
    • ALL outputs 32..47 wil ldepend on inputs 0..47, but cannot depend on 48..
    • etc.
    • Thus, if we were to have no left context, each output would reliably only depend on 16 inputs.
  • The convolution requires masking, because with no special care its perceptive field would peek into future frames. I use the DCConv approach (AWS paper) which effectively "masks" frames future to each chunk.
    • The way it is done is that the convolution operator is replaced to split the input into chunks (the same way as attention effectively does). Then, we add the left context (of size (kernel_size-1)/2) to serve as left padding to the convolution. The right padding will be applied as normally, and will not be future chunk data (but zero-padding or reflected padding or whatever).
    • This means that we need to perform the padding manually for best efficiency. Otherwise, the convolution would add left padding to our left padding, and output a tensor of length (chunk_size + left_context) even though we only need the last (chunk_size) elements.
  • Dynamic Chunk Training is implemented: For each back, 50% chance of full context (as usual) or 50% chance of picking a random chunk size between e.g. 8 and 32 chunks. Papers mention this as a requirement for proper convergence.
    • Randomly switching between chunk sizes means the model should be able to adapt to them, meaning the chunk size can be selected at runtime for a tradeoff between accuracy and latency.
  • Left context is to be specified in number of chunks as it is the most natural fit (consider inter-chunk dependencies across layers) and we probably do not need finer granularity. Left context size effectively only affects memory usage due to what needs to be cached (at MHA and Conv levels), and linear cost because of the MHA visible context.
  • Depending on certain batching and masking circumstances, it is possible for some sequences to be entirely masked, i.e. their keys are set to -inf in attentional score calculation. This is a problem, because it means that all inputs to the softmax layer in attention can become -inf, which causes all outputs to become NaN (because the divisor in the softmax formula sums to 0 as e^-inf = 0). This regularly broke training with NaN losses that never recovered. Adding an epsilon value to the divisor mitigates the issue as it will become > 0.

Computing chunk sizes, required left contexts, etc. and effective frame counts is entirely predictable here, but it does depend on deep understanding of the models inputs. Here, 8 chunks corresponds roughly to a latency of 365ms (to double check).

tl;dr on what will be required for the model when it comes to streaming interfaces:

• For feature extraction, chunking is relatively easy, but carries left dependencies due to the striding nature of the FBanks and CNN.
• At the output of each conformer layer, the input-output dependency properties mentioned before for attention are respected. Thus, left context of chunks can be cached at each layer of the conformer, that is, the values of our current change will never affect prior chunks at any conformer layer. This left context encompasses:
  • Left context keys and values for the MHA: as many chunks as desired, for EACH layer. (Doesn't have to be the same amount for each conformer layer!)
  • Left context for the DCConv: only dependent on kernel size; might however span multiple chunks ultimately.
• For everything involved, some adjustments may prove necessary for chunk-by-chunk inference.

tl;dr on what work would be needed to port this approach to the Branchformer:

• At a glance, it does not seem like many conceptual changes are necessary compared to the Conformer model here. DCConv-like masking should be applied in the Conv1D in the CSGU (aside of the attn changes) and that should be about it.
• Investigate if any other "bricks" were modified that could cause issues; adjust contexts and chunk size effective frame count calculations. Use the dependency detection functions I made.

Additional context

No response