TODM: Train Once Deploy Many Efficient Supernet-Based RNN-T Compression For On-device ASR Models

TL;DR

TODM introduces an efficient Supernet-based training method for RNN-T models, enabling the deployment of multiple hardware-optimized on-device ASR models with minimal additional training cost.

Contribution

It presents a novel Supernet training approach with techniques like adaptive dropouts and knowledge distillation, achieving comparable or better performance than individually tuned models.

Findings

Supernet matches or surpasses manually tuned models in WER by up to 3%.

Training many models costs roughly the same as a single training job.

The approach is validated on LibriSpeech with promising results.

Abstract

Automatic Speech Recognition (ASR) models need to be optimized for specific hardware before they can be deployed on devices. This can be done by tuning the model's hyperparameters or exploring variations in its architecture. Re-training and re-validating models after making these changes can be a resource-intensive task. This paper presents TODM (Train Once Deploy Many), a new approach to efficiently train many sizes of hardware-friendly on-device ASR models with comparable GPU-hours to that of a single training job. TODM leverages insights from prior work on Supernet, where Recurrent Neural Network Transducer (RNN-T) models share weights within a Supernet. It reduces layer sizes and widths of the Supernet to obtain subnetworks, making them smaller models suitable for all hardware types. We introduce a novel combination of three techniques to improve the outcomes of the TODM Supernet: adaptive dropouts, an in-place Alpha-divergence knowledge distillation, and the use of ScaledAdam optimizer. We validate our approach by comparing Supernet-trained versus individually tuned Multi-Head State Space Model (MH-SSM) RNN-T using LibriSpeech. Results demonstrate that our TODM Supernet either matches or surpasses the performance of manually tuned models by up to a relative of 3% better in word error rate (WER), while efficiently keeping the cost of training many models at a small constant.