HDMoLE: Mixture of LoRA Experts with Hierarchical Routing and Dynamic Thresholds for Fine-Tuning LLM-based ASR Models

Low-Rank Adaptation. LoRA [17] is an exceptional parameter-efficient fine-tuning method for adapting pre-trained large models to specific domains. It reduces the number of trainable parameters by updating low-rank decomposition matrix pairs while maintaining the original weights unchanged. Specifically, for a given linear layer with the weight matrix $\mathbf{W}_{\mathrm{0}}$, LoRA employs two low-rank matrices $\mathbf{A}$ and $\mathbf{B}$ with rank $r$, where $\mathbf{W}_{\mathrm{0}}\in\mathbb{R}^{d_{out}\times d_{in}}$, $\mathbf{A}\in\mathbb{R}^{r\times d_{in}}$, $\mathbf{B}\in\mathbb{R}^{d_{out}\times r}$, and $r\ll\min(d_{in},d_{out})$. With the application of LoRA, the forward process for the given linear layer $\mathbf{l}=\mathbf{W}_{\mathrm{0}}\mathbf{x}+\mathbf{b}$ can be impressed as follows:

where low-rank matrices $\mathbf{A}$ and $\mathbf{B}$ are trainable while the original weight $\mathbf{W}_{\mathrm{0}}$ and bias $\mathbf{b}$ remain unchanged during training. The matrix $\mathbf{A}$ is initialized with a random Gaussian distribution, and matrix $\mathbf{B}$ starts from zero. Scaling $\Delta\mathbf{Wx}$ by $\alpha/r$ controls the extent of adjustments to the original weights imposed by LoRA, with $\alpha$ and $r$ representing constants.

Mixture of Experts. MoE [18, 19, 20] framework scales model capacity and complexity by incorporating multiple sub-network experts, each potentially addressing specific domains or tasks. Within an MoE layer, $N$ independent experts $\left\{\mathbf{E}_{i}\right\}_{i=1}^{N}$ are coordinated by a gating network router, which applies a trainable matrix to generate a probability distribution for weighting the outputs of these experts, employing a softmax function for normalization. For the given input vector $\mathbf{x}$, the output probability distribution $\mathbf{p}$ of the router can be impressed as:

where $\mathbf{W}_{g}$ represents the trainable weights of the gating network router. The final output $\mathbf{y}$ from the MoE layer is a weighted sum of the outputs from the top $K$ experts:

where $\mathbf{p}_{i}$ and $\mathbf{p}_{j}$ are the weights of $i\text{-th}$ and $j\text{-th}$ expert in MoE layer, respectively. The TopK function identifies and retains the highest $K$ weights, setting the rest to zero. The weights retained by the TopK function are normalized to ensure their sum equals one.

The original intention of MoE is that the routers determine weights for each expert according to the input sample, with higher weights given to experts focused on the input domain, thereby establishing a clear correspondence between experts and domains. However, routers in MoE layers tend to converge to a state where it always produces large weights for early-stage well-performing experts, leading to only a handful of experts having a significant impact [10]. In other words, the experts assigned larger weights by the MoE router essentially remain constant for inputs across different domains. Consequently, this imbalance of the experts’ utilization problem causes ambiguity in the correspondence between experts and domains. To address this, we propose a hierarchical routing strategy to clarify the correspondence between experts and domains. Specifically, hierarchical routing comprises global and local routing with a pre-trained accent recognition (AR) model as the global router and individual MoE layer routers as local routers. The global routing explicitly guides each expert to focus on a specific accent domain, clarifying the correspondence between experts and accent domains and establishing them as domain-specific experts. Meanwhile, the local routing implicitly guides the experts within the MoE layer to collaborate across multiple accent domains through a learnable gating network router. The input speech features $\mathbf{X}_{\mathrm{in}}$ generate global weights $\mathbf{P}_{\mathrm{g}}$ through the global router, which can be impressed as follows:

where the global router is frozen during training and inferring. Subsequently, the global weights are input into each MoE layer. In each MoE layer, the hidden speech features $\mathbf{H}_{\mathrm{in}}$ produce local weights $\mathbf{P}_{\mathrm{l}}$ via the local router, which can be impressed as follows:

where the local router is a trainable linear layer.

The standard MoE layer employs the static Top-K expert selection strategy, choosing $K$ experts with the highest weights determined by the router. Given that each MoE layer concentrates on distinct aspects of domains, varying numbers of experts are required to participate in different MoE layers. Therefore, we propose a dynamic threshold expert selection strategy, replacing the static Top-K approach by implementing dynamic thresholds. The dynamic thresholds strategy allows each MoE layer to select the experts that need to be activated flexibly. This strategy requires defining a threshold, where experts with weights exceeding the threshold are selected. It is essential to carefully initialize the threshold, as a highly initialized threshold may cause all expert weights to fall below it, resulting in no experts being selected. Therefore, appropriate initialization of the threshold is indispensable. Here, the threshold is initialized at $1/N$ to ensure that at least one expert is selected. In each HDMoLE, we assign two independent dynamic thresholds to the global and local weights, respectively. The quantities of global and local thresholds are identical. Through the global threshold $\tau_{g}$, we can obtain the global adapted weights $\mathbf{P}_{\mathrm{ga}}$, which can be impressed as follows:

where $\mathbb{E}$(condition) equals one if the condition is true and zero otherwise, $\mathbf{P}_{\mathrm{g}}^{i}$ represents the weight of the $i\text{-th}$ expert in the global weights. Additionally, scaling the adapted weights by $\tau_{g}$ guarantees that the $\tau_{g}$ remains learnable during backpropagation. Similarly, the local threshold $\tau_{l}$ allows us to obtain the local adapted weights $\mathbf{P}_{\mathrm{la}}$, which can be impressed as follows:

where $\mathbf{P}_{\mathrm{l}}^{i}$ represents the weight of the $i\text{-th}$ expert in the local weights. The final adapted weights $\mathbf{P}_{\mathrm{a}}$ are the sum of the global and local adapted weights, which can be defined as follows:

MoLE substitutes conventional dense layer experts with LoRA experts, rendering it a parameter-efficient fine-tuning approach. The final output of MoLE is a combination of the weighted outputs from the LoRA experts and the original model. After applying hierarchical routing and dynamic thresholds, the final adapted weights $\mathbf{P}_{\mathrm{a}}$ are obtained. Therefore, the output of the MoLE layer $\mathbf{H}_{\mathrm{out}}$ can be expressed as follows:

Notably, when the weight of a certain LoRA expert is simultaneously lower than the global and local thresholds, the LoRA expert’s weight in the final adapted weights $\mathbf{P}_{\mathrm{a}}$ becomes zero.

Datasets. We conduct experiments using the multi-accent Mandarin KeSpeech [30] and the standard Mandarin AISHELL-2 [31] datasets to evaluate the effectiveness of HDMoLE. KeSpeech involves 1,542 hours of speech recorded by 27,237 speakers from 34 cities in China, and the pronunciation includes standard Mandarin and eight major accented Mandarin. All of our HDMoLE experiments are trained solely on the KeSpeech dataset. We use the KeSpeech test dataset to evaluate the performance of HDMoLE in the target multi-accent domains. In contrast, the AISHELL-2 test dataset measures performance degradation in the source general domain.

Settings. HDMoLE is implemented in the projector module of a pre-trained LLM-based ASR model with the Hubert+Baichuan2 structure [7]. This model uses over 11,000 hours of standard Mandarin Chinese data from four general domain corpora as training datasets, including WenetSpeech [32], AISHELL-1 [33], AISHELL-2 [31], and AISHELL-4 [34], excluding multi-accent Mandarin corpora. The projector module is a 4-layer Transformer [35] with 51M parameters, where feed-forward networks (FFN) have 2560 dimensions, multi-head self-attentions (MHSA) have 256 dimensions, and the attention head is 4. Each HDMoLE has 8 LoRA experts, each with a rank and an alpha value of 8. We apply HDMoLE to the FFN layers and the four weight matrices $(\mathbf{W}_{\mathrm{q}},\mathbf{W}_{\mathrm{k}},\mathbf{W}_{\mathrm{v}},\mathbf{W}_{\mathrm{o}})$ of the MHSA in the projector module. The global router pre-trained frozen AR model is a 12-layer Conformer [36] with 31M parameters trained solely on KeSpeech, where the attention head is 4, and the dimensions of FFN and MHSA are respectively set to 2048 and 256, achieving an AR accuracy of 81.44% on the KeSpeech test dataset.

Main Results. Table I presents the CER results on the KeSpeech and AISHELL-2 test datasets across various methods. The first line shows the inference results on the KeSpeech and AISHELL-2 test datasets using the pre-trained Hubert+Baichuan2 model, with poor performance on KeSpeech attributed to the lack of accented speech corpora during pre-training. The second line displays the results of fully fine-tuning the Hubert+Baichuan2 model projector module. The significant improvement on the KeSpeech test dataset is the target domain topline for HDMoLE. Meanwhile, the regression in AISHELL-2 performance indicates that full fine-tuning causes degradation in the model’s source domain performance. Lines 3 to 7 present the results of the original LoRA and various MoLE methods on KeSpeech and AISHELL-2, none of which match the performance of HDMoLE in both the target and source domains. Lines 8 to 11 present the ablation study of HDMoLE, demonstrating the necessity of the two strategies employed. In summary, HDMoLE fine-tunes the Hubert+Baichuan2 model projector module using fewer trainable parameters, achieving performance close to full fine-tuning in the target domain while minimizing regression in the source domain.

Various Global Routers. Table II illustrates the influence of global routers with different performances on the efficacy of HDMoLE. Our experiments involved three kinds of global routers: non-convergent AR model, convergent AR model, and ground-truth accent labels. We observe that improvements in the global router’s performance lead to enhanced performance for HDMoLE. This phenomenon indicates that improved performance of the global router clarifies the correspondence between LoRA experts and accent domains, enabling each expert to concentrate more effectively on its specific domain.

Rank of LoRA Experts. Table III examines the impact of LoRA rank variations on the performance of HDMoLE. With higher LoRA ranks, the parameter for each LoRA expert grows, consequently elevating the number of trainable parameters in HDMoLE, leading to noticeable improvements in performance. However, when the LoRA rank increases beyond 16, the performance improvement of HDMoLE diminishes, and even regression may occur.

Figure 2 shows the average number of activated LoRA experts for every projector layer. HDMoLE in each projector layer concentrates on distinct aspects, requiring a dynamic activation of LoRA experts based on the layer’s focus. HDMoLE in lower projector layers tends to activate more experts, signifying that these layers are essential for processing general speech features, which are more complex and diverse, thus requiring comprehensive domain knowledge. Conversely, HDMoLE in higher projector layers tends to activate fewer experts, indicating that these layers focus on processing domain-specific speech features, demanding more specialized domain knowledge.