Efficient Supernet Training with Orthogonal Softmax for Scalable ASR Model Compression

TL;DR

This work tackles efficient ASR model compression under hardware constraints by proposing a two-step training pipeline that trains a single supernet with multiple subnets. The core contribution, OrthoSoftmax, applies $N$ softmax functions on a learnable score matrix to generate orthogonal, sparse subnet masks, enabling fine-grained, FLOPs-aware component-wise subnet selection without expensive search. Empirical results on Librispeech and TED-LIUM-v2 using a Conformer-CTC setup show that FLOPs-aware component-wise subnet selection achieves WERs comparable to or better than individually trained models, with substantial reductions in training time. The paper also provides detailed analyses of retained components, revealing that Convs are highly important and that the first Conformer block and MHSA heads have distinctive roles, offering practical guidance for deployment in resource-constrained settings.

Abstract

ASR systems are deployed across diverse environments, each with specific hardware constraints. We use supernet training to jointly train multiple encoders of varying sizes, enabling dynamic model size adjustment to fit hardware constraints without redundant training. Moreover, we introduce a novel method called OrthoSoftmax, which applies multiple orthogonal softmax functions to efficiently identify optimal subnets within the supernet, avoiding resource-intensive search. This approach also enables more flexible and precise subnet selection by allowing selection based on various criteria and levels of granularity. Our results with CTC on Librispeech and TED-LIUM-v2 show that FLOPs-aware component-wise selection achieves the best overall performance. With the same number of training updates from one single job, WERs for all model sizes are comparable to or slightly better than those of individually trained models. Furthermore, we analyze patterns in the selected components and reveal interesting insights.

Figures (2)

• Figure 1:
• Figure 2: The remaining ratio for each layer in encoders of varying sizes, trained using the FLOPs-aware component-wise OrthoSoftmax method, as shown in Fig.\ref{['fig:joint_train_five_models']}b.