Pruned RNN-T for fast, memory-efficient ASR training

TL;DR

This paper proposes a pruning method for RNN-T loss computation that significantly improves speed and reduces memory usage, enabling more practical training for large-vocabulary speech recognition systems.

Contribution

It introduces a pruning approach that efficiently bounds RNN-T recursion, facilitating faster and more memory-efficient training, especially for large vocabularies.

Findings

Achieves faster RNN-T loss computation

Reduces GPU memory usage during training

Enables training with larger vocabularies

Abstract

The RNN-Transducer (RNN-T) framework for speech recognition has been growing in popularity, particularly for deployed real-time ASR systems, because it combines high accuracy with naturally streaming recognition. One of the drawbacks of RNN-T is that its loss function is relatively slow to compute, and can use a lot of memory. Excessive GPU memory usage can make it impractical to use RNN-T loss in cases where the vocabulary size is large: for example, for Chinese character-based ASR. We introduce a method for faster and more memory-efficient RNN-T loss computation. We first obtain pruning bounds for the RNN-T recursion using a simple joiner network that is linear in the encoder and decoder embeddings; we can evaluate this without using much memory. We then use those pruning bounds to evaluate the full, non-linear joiner network.