TODM: Train Once Deploy Many
Efficient Supernet-Based RNN-T Compression for on-device ASR models

End-to-end (E2E) all-neural automatic speech recognition (ASR) has gained attention for its compatibility with edge devices [1, 2, 3]. Recurrent neural network transducer (RNN-T) is one of the most popular architectures for on-device E2E ASR [4, 5, 6, 7]. Modern applications of RNN-T ASR run on a variety of hardware, including central processing units (CPUs) on phones, tensor processing units (TPUs) [8], and other neural accelerators. The process of optimizing on-device ASR models efficiently for each distinct hardware configuration necessitates fine-tuning model training hyperparameters and streamlining the model’s architecture [9]. This demands a substantial amount of computational resources, primarily in the form of GPU-hours or GPU cluster energy expenses in MegaWatt-hour (MWh). Researchers must balance accuracy and size while efficiently managing training resources.

A potential approach is to leverage the Supernet [10, 11, 12, 13, 14]. A Supernet is a weight-sharing neural network graph that can generate smaller subnetworks tailored for specific applications. Prior work has attempted to facilitate the simultaneous training of a fixed number of models by implementing weight sharing capabilities. Examples include work on RNN-T cascaded encoders [8, 15]. These efforts, however, typically involve a limited number of networks, often around three, that only share portions of the RNN-T encoders. Their primary objective is to streamline management for a few models, while utilizing distinct model decoders to improve model performances. Some also combine non-causal and causal encoders to enhance model accuracy. The training of such cascaded encoder systems consumes much more resources compared to a single model-training task, and it can not scale to a large number of different sizes of models. Besides cascaded encoders, our previous work explored the concept of Supernet in the Omni-Sparsity DNN framework [16], in which a large number of ASR subnetworks could be derived from one Supernet training job. Furthermore, there have been efforts to expand the use of Omni-Sparsity DNN to train ASR models while accommodating various latency constraints for the subnetworks. This was demonstrated by training a non-streaming dense Supernet with sparse streaming subnetworks in [17]. Omni-Sparsity DNNs’ reliance on structured sparsity, however, unnecessarily limits their hardware compatibility.

This paper presents Train Once Deploy Many (TODM), an approach to efficiently train many hardware-friendly on-device ASR models with comparable GPU-hours to a single training job. TODM leverages insights from prior work on Omni-sparsity DNN Supernet, but does not rely on sparsity. Instead, it drops layers and reduces layer widths in the Supernet to obtain optimized subnetworks for all hardware types. We explain TODM Supernet in Section 2. We improve its training outcome by introducing adaptive dropout, an in-place logit-sampled Alpha-divergence knowledge distillation mechanism, and ScaledAdam optimizer (Section 3). We validate our approach by generating Supernet-trained sub-models using Evolutionary Search (described in Section 2.2) on the validation set and comparing their performances to individually optimized ASR models using LibriSpeech [18] in Section 4. We discuss the results and why Supernet can function efficiently in Section 5. This work is the first to use Supernet training to efficiently produce a variety of on-device ASR models. With similar training resources as one model, we can deploy many of different sizes, each with the optimal quality at its desired size.

In this section, we describe how a Supernet RNN-T generates subnetwork RNN-Ts of dynamic sizes by varying the width and height of the encoders. It then uses evolutionary search [16] to discover the best encoder architecture of RNN-T at different size constraints.

In this section, we outline three key strategies for improving the quality of Supernet-trained models. While each technique has been applied in non-Supernet contexts, their combination is novel.

1. 1.

   Adaptive Dropout: different subnetworks contribute different magnitudes of gradient L2-norm during Supernet training. Inspired by the adaptive dropout used in Omni-sparsity DNN training [16], we adjust the magnitude of dropouts in the FFN layers, which immediately follows the modules with reduced output channel during training. Intuitively, a module with reduced dimensionality requires less regularization in order to produce a similar output and gradient norm. Therefore, $dropout_{c_{i}}=dropout_{c_{m_{i}}}\times\frac{c_{i}}{m_{i}}$, immediately follows the reduced output channel size of module $i$, where $m_{i}$ is the maximum output channel size, and $c_{i}$ is the current channel size.

2. 2.

   Sampled In-place Knowledge Distillation (KD): during each step of the Supernet training, we use the entire Supernet as the “teacher” model, and force in-place distillation of the max-network’s output probability distributions onto the sampled subnetwork, i.e. the “student”, with data from each mini-batch. We experiment with two types of output probability divergence functions: (1) the Kullback–Leibler divergence (KLD), which is known to improve the accuracy of a compressed ASR during knowledge distillation [20, 21]; and (2) the Alpha-divergence [14] (AlphaD), which has been shown to better capture the teacher network’s uncertainties in the output probability distributions than KL in a Supernet training setting. Distilling the teacher’s output probability distribution over the entire RNN-T lattice is memory-intensive and slows down training by requiring a smaller batch size. To overcome that, we improve the training memory by subsampling top $j$ output probabilities from the teacher. In a previous study on efficient RNN-T distillation [22], $j$ was set to 2, which meant that the distillation focused on only 3 logit dimensions: the target token, the blank token, and the cumulative sum of the remaining probabilities. We hypothesize that this insufficiently captures the uncertainty of the teacher output probabilities. Therefore, in this work, we compare $j$ values of $[2,10,100]$ (distillation probability dimensions = j+1). The loss function is now:

   	$\displaystyle\mathcal{L}$	$\displaystyle=\lambda\mathcal{L}_{KD}+\mathcal{L}_{rnnt}$		(4)
   	$\displaystyle\underset{s\in\bar{\mathbb{Z}}}{\mathcal{L}_{KD}(\Theta_{s};\Theta)}$	$\displaystyle=\mathbb{E}_{x\in D_{mini-batch}}\left[f(p_{j}(x;\Theta)||q_{j}(x;\Theta_{s})\right]$		(5)

   where we use default $\lambda=1.0$, $f$ is KLD or AlphaD, computed over the sampled probability of top $j$ tokens, and the sum of the rest ($p_{j}$ being the teacher’s distributions, $q_{j}$ the student’s). For Alpha-divergence, we use the default setting of $\alpha_{-}$=-1, $\alpha_{+}$=1 and $\beta$=5.0.

3. 3.

   ScaledAdam Optimizer: we explore training the Supernet with the ScaledAdam optimizer [23]. As a variant of Adam optimizer, the ScaledAdam scales each parameter’s update based on its norm. Intuitively, the ScaledAdam optimizer improves gradient stability of the Supernet training, where each subnetwork contributes drastically different gradient norms to model parameters during mini-batch forward-backward step.

In this paper, we validate the results of TODM Supernet on a non-streaming Multi-Head State Space Model (MH-SSM) RNN-T. The encoder of our RNN-T is built with 16 MH-SSM layers; these layers are attention free and have been demonstrated to achieve on-par performances compared to transformer-based ASR [24]. Each MH-SSM layer consists of 2-stacked, 1-headed MH-SSM, and 4096 feed-forward net (FFN) dimension. Our RNN-T architecture follows the example set in prior work [24]: three 512-dimensional LSTM layers plus a linear layer in the predictor; a linear project layer and a ReLU gate in the joiner. In total, our non-streaming MH-SSM RNN-T has 99.9 million parameters. During Supernet training, we selectively reduce entire layers of MH-SSM, or alter the size of the FFN module to $C_{i}\in[512,1024,2048,4096]$. The FFN consists of 77.8% of the total number of parameters in a MH-SSM layer, and has high potential of parameter redundancy. Layer reduction in MH-SSM supernet is performed as reducing top 0, 3, or 7 MH-SSM layers. See Fig 1(a) for illustration. All our baseline MH-SSM RNN-T ($A1,A2,A3$ in Table 1) are trained with the same hyper-parameter set up: a total of 180 epochs of training, under a fixed 0.006 learning rate, force-anneal at 60 epochs with shrinking factor 0.96; 0.1 weight decay; Adam optimizer with $\beta_{1}=0.9,\beta_{2}=0.999$.

We train the RNN-Ts with the LibriSpeech 960 hr training data [25]. We obtain 80-dimensional log Mel-filterbank features from each 25 ms audio window, sliding the window ahead every 10 ms. We use a pre-trained 4096-dimensional sentence piece vocab [26], plus a ‘blank’ symbol, as the RNN-T target. After Supernet training, we use the 10.7h LibriSpeech dev-clean and dev-other data to identify best subnetwork architectures during evolutionary search.

All model training is done with 32 NVIDIA-A100 GPUs; to prevent GPU queue preemption, CPU scheduling wait-times, and other data center down-times from polluting the GPU-hour computation, we report the total energy expenses of each model training in terms of MegaWatt-hour (MWh), measured by the energy consumed and recorded by our GPU cluster.

We compare our TODM Supernet (Model F) with 4 baselines:

1. 1.

   A1 represents a collection of individually trained and tuned models, the sizes of which are determined by reducing layer sizes, channel sizes, or both, from the largest MH-SSM model. The energy used in training scales proportionally to the number and sizes of the models, we are thus capable to explore only limited architectures.

2. 2.

   A2 represents the results of the largest model, trained with auxiliary cross-entropy loss [27] at layers 8, 12, and 16; auxiliary training introduced 1.15 million more parameters to the model, but did not end up improving the model performances in Table 1.

3. 3.

   A3 is the results of the largest model, trained with auxiliary RNN-T loss at layers 8, 12, and 16 (i.e. layer-drop) – each of these layers is accompanied with a unique combination of linear layer, gates and layer-norm, which is then fed into a shared predictor and joiner [27]. For $k$ subnetworks in the auxiliary RNN-T loss training, there is $k\times$ resources increase during training, and an additional model size of $k\times$0.53 million parameters. Note that A2 and A3 exhibit comparable training energy consumption. We use them to investigate the extra regularization effects when loss functions are applied to the intermediate layers of the RNN-T model.

4. 4.

   B0 represents a channel-reduction-only Supernet. It is a generalized version of the slimmable network [13], in which subnetworks of different layer-wise widths co-exist in a Supernet.

This paper presents the TODM framework, which can train and discover optimized, dense RNN-Ts of various sizes from a single Supernet, with comparable training resources to a single-model training job. We introduce three strategies to TODM to improve Supernet outcomes: adaptive dropout, in-place sampled knowledge distillation, and ScaledAdam optimizer fine-tuning. TODM discovers many models along the Pareto front of accuracy vs. size, which would have be resource-intensive to find and train manually via trial-and-error.