Dynamic Encoder Size Based on Data-Driven Layer-wise Pruning for Speech Recognition

TL;DR

This work tackles the need for ASR models that fit varying hardware constraints by enabling multiple encoder sizes to be trained in one go. It introduces a dynamic encoder size framework that jointly trains a full supernet and several subnets using score-based layer pruning to select the per-subnet depth from a total of $L$ layers, with masks $\mathbf{z}_m$ and layer scores $\mathbf{s}$. Two data-driven pruning methods, Simple-Top-k and Iterative-Zero-Out, learn $\mathbf{s}$ and generate subnet masks, followed by a Step 2 joint training using the sandwich rule; experiments on Librispeech and TED-LIUM-v2 with a Conformer-CTC model show that subnets achieve on-par WER with independently trained models, and the full supernet also gains a modest overall improvement. The results indicate that convolutional layers are often retained in subnets and that this approach can significantly reduce training redundancy while preserving accuracy across multiple encoder sizes, enabling practical deployment under diverse constraints.

Abstract

Varying-size models are often required to deploy ASR systems under different hardware and/or application constraints such as memory and latency. To avoid redundant training and optimization efforts for individual models of different sizes, we present the dynamic encoder size approach, which jointly trains multiple performant models within one supernet from scratch. These subnets of various sizes are layer-wise pruned from the supernet, and thus, enjoy full parameter sharing. By combining score-based pruning with supernet training, we propose two novel methods, Simple-Top-k and Iterative-Zero-Out, to automatically select the best-performing subnets in a data-driven manner, avoiding resource-intensive search efforts. Our experiments using CTC on both Librispeech and TED-LIUM-v2 corpora show that our methods can achieve on-par performance as individually trained models of each size category. Also, our approach consistently brings small performance improvements for the full-size supernet.

Figures (2)

• Figure 1: Illustration of Simple-Top-k, STE uses a relaxed k-hot vector to estimate the gradients of the binary mask.
• Figure 2: The distribution of selected layers for the models with 12, 24 and 36 layers shown in Table \ref{['tab:tedlium_four_model']}.