PaperlayerView exact revision on arXiv

Source-linked AI summary

DR-LoRA: Dynamic Rank LoRA for Fine-Tuning Mixture-of-Experts Models

Guanzhi Deng, Bo Li, Ronghao Chen, Xiujin Liu, Zhuo Han, Huacan Wang, Lijie Wen, Linqi Song

arXiv:2601.04823v6cs.AIcs.CL

TL;DR

Uniform LoRA ranks can mismatch the heterogeneous demands of MoE experts, under-provisioning task-relevant experts while wasting capacity elsewhere. DR-LoRA grows expert ranks using routing frequency and gradient-based importance, and it consistently outperforms LoRA and other strong baselines across three models and six tasks.

  • Problem

    Uniform LoRA rank allocation overlooks heterogeneous expert usage in MoE models, under-provisioning frequently used experts and wasting capacity on low-frequency experts.

  • Method

    DR-LoRA initializes small active ranks and periodically grows high-saliency expert ranks using routing frequency and gradient-based rank importance.

  • Results

    DR-LoRA consistently outperforms LoRA and other strong PEFT baselines across three MoE models and six tasks.

  • Takeaways & Limitations

    Task-adaptive heterogeneous rank allocation improves active capacity utilization in MoE fine-tuning.

  • Takeaways & Limitations

    The saliency score is a heuristic based on a first-order approximation and mean-field factorization, and transfer to other routing paradigms or multimodal MoE architectures remains open.

Abstract

from arXiv · show

Mixture-of-Experts (MoE) has become a prominent paradigm for scaling Large Language Models (LLMs). Parameter-efficient fine-tuning methods, such as LoRA, are widely adopted to adapt pretrained MoE LLMs to downstream tasks. However, existing approaches typically assign identical LoRA ranks to all expert modules, ignoring the heterogeneous specialization of pretrained experts. This uniform allocation leads to a resource mismatch: task-relevant experts are under-provisioned, while less relevant ones receive redundant parameters. To address this, we propose DR-LoRA, a Dynamic Rank LoRA framework for fine-tuning pretrained MoE models. Specifically, DR-LoRA initializes all expert LoRA modules with a small active rank and uses an expert saliency score, which combines routing frequency and gradient-based rank importance, to identify which experts would benefit most from additional capacity. It then periodically expands the active ranks of the task-critical expert LoRA, progressively constructing a heterogeneous rank distribution tailored to the target task. Experiments on three MoE models across six tasks show that DR-LoRA consistently outperforms LoRA and other strong baselines, demonstrating that task-adaptive heterogeneous rank allocation is an effective strategy to improve active capacity utilization in MoE fine-tuning.

1 Introduction

MoE fine-tuning with uniform LoRA ranks overlooks experts’ heterogeneous activation and specialization, creating capacity mismatch. DR-LoRA addresses this with saliency-guided rank growth and consistently outperforms strong PEFT baselines across three models and six tasks.

  • Uniform LoRA ranks under-provision frequently used experts while wasting capacity on rarely used experts, creating expert-level capacity mismatch.
  • DR-LoRA initializes experts with small active ranks and periodically expands high-saliency experts, forming a task-adaptive heterogeneous rank structure.
  • DR-LoRA combines routing frequency with gradient-based rank importance to identify experts that should receive additional rank capacity.
  • DR-LoRA consistently outperforms LoRA and other strong PEFT baselines across three MoE models and six tasks.
  • The work identifies heterogeneous expert usage as the basis for concentrating trainable capacity on task-critical experts.

2 Related Work

Existing LoRA-based MoE adaptation methods commonly use uniform expert ranks, despite task-dependent activation patterns. DR-LoRA instead combines routing frequency and gradient-based importance with growth-based allocation.

  • LoRA methods for MoE models typically inject adapters into experts with uniform ranks, assuming similar fine-tuning capacity is needed everywhere.
  • Uniform allocation overlooks task-dependent imbalance in expert activation induced by sparse routing.
  • DR-LoRA combines routing frequency and gradient-based rank importance to guide expert-level rank allocation.
  • Its growth-based mechanism matches the gradual emergence of expert-specific learning signals under sparse routing.

3 Methodology

DR-LoRA dynamically allocates expert LoRA ranks using routing frequency and gradient-based learning intensity. It grows high-saliency experts progressively under per-layer quotas while controlling monopolization and transient fluctuations.

  • 3.1 Preliminaries and Problem Formulation: Each expert’s LoRA module uses low-rank injection into its adapted feed-forward projections, with an expert-level active rank.
  • 3.3 Dynamic Rank Allocation: Experts pre-allocate rank capacity but begin with only a small active subspace used in the forward pass.
  • 3.2 Expert Saliency Scoring: DR-LoRA scores experts using routing frequency and gradient-based rank importance, targeting experts that are frequently routed and still learning.
  • 3.2 Expert Saliency Scoring: The saliency numerator estimates expansion benefit, while a rank-dependent denominator reduces priority for experts that already have higher ranks.
  • 3.3 Dynamic Rank Allocation: Rank growth occurs periodically with a fixed per-layer quota, reaching the target active-rank budget over the growth window.
  • 3.3 Dynamic Rank Allocation: Per-layer greedy allocation assigns new ranks to descending-saliency experts, caps per-event growth, and resets rank importance after each event.

4 Experiments

Experiments evaluate DR-LoRA across three MoE models and six domains against parameter-efficient baselines. Under matched budgets, DR-LoRA achieves the strongest overall results, while ablations and masking analyses support heterogeneous expert capacity allocation.

  • 4.2 Main Results: DR-LoRA achieves the best average performance on all three MoE models, improving over the strongest baseline by 2.2, 1.6, and 1.8 points, respectively.Its mean scores are 44.1 on OLMoE, 62.4 on Qwen1.5-MoE, and 38.7 on LLaMA-MoE; it also wins all 18 model–task settings with a significant aggregate improvement (p < 0.001).
  • 4.2 Main Results: Under matched trainable-parameter counts, DR-LoRA consistently outperforms fixed-rank LoRA, showing the advantage of fine-grained expert-level capacity allocation.LoRA uses r = 32 to match DR-LoRA’s trainable parameter count and training-time memory footprint.
  • 4.2 Main Results: Masking large-rank DR-LoRA experts causes greater degradation than masking small-rank experts, unlike the less differentiated pattern of uniform-rank LoRA.The masking comparison uses budgets of 5%, 10%, and 20% of the total rank budget per layer.
  • 4.3 Ablation Study: DR-LoRA maintains lower loss and higher average accuracy throughout GSM8K training, outperforming fixed-rank LoRA and AdaLoRA at all stages.The training-dynamics comparison indicates faster convergence for DR-LoRA; AdaLoRA remains ahead of fixed-rank LoRA but behind DR-LoRA.
  • 4.3 Ablation Study: Full DR-LoRA reaches an average ablation score of 23.9, exceeding versions using only gradient-based importance or only routing frequency by 1.7 points.The result supports combining the two saliency components rather than relying on either signal alone.
  • 4.3 Ablation Study: Unfreezing the router yields further gains for DR-LoRA and produces the best overall performance.The router is frozen during warmup and then trained jointly with the LoRA modules.

5 Analysis

The analyses show that DR-LoRA learns task-relevant heterogeneous ranks and benefits from progressive growth, while rank penalties balance concentration against differentiation and performance remains robust across hyperparameters.

  • Masking analysis: Masking the top 25% large-rank experts causes substantially greater degradation than masking the remaining 75% in DR-LoRA across GSM8K and MedMCQA.Uniform-rank LoRA shows less differentiated degradation under the same masking budgets.
  • Growth versus pruning: DR-LoRA converges faster and maintains lower training loss and higher evaluation accuracy than fixed-rank LoRA and pruning-based AdaLoRA throughout GSM8K training.The analysis attributes this advantage to growth after expert-specific signals accumulate, unlike early pruning under sparse routing.
  • Progressive growth versus fixed heterogeneous allocation: 46.3 average score for full DR-LoRA exceeds 45.1 for Static-Final-Rank and 42.8 for LoRA (r = 32), separating benefits from heterogeneous allocation and progressive growth.Progressive growth contributes an additional 1.2 points beyond the final rank structure.
  • Rank penalty analysis: At γ=1.2, DR-LoRA reaches a 23.9 average score with Gini = 0.31, whereas γ=0 yields 20.5 average and Gini = 0.62.Increasing γ from 0 to 1.2 improves performance and reduces concentration; excessive γ makes allocation too uniform and lowers performance.
  • Hyperparameter robustness: All tested growth intervals outperform LoRA, while the default Tgrow=200 achieves the highest average score of 23.9 and β–γ configurations vary within a narrow performance range.The main experiments use β=0.9, γ=1.2, and Tgrow=200 without per-task tuning.

6 Conclusion

DR-LoRA dynamically allocates LoRA capacity across experts using routing frequency and gradient-based importance rather than uniform ranks. Across three open MoE LLMs and six benchmarks, it outperforms uniform-rank and adaptive baselines in accuracy, efficiency, and robustness.

  • Conclusion: DR-LoRA combines routing frequency and gradient-based rank importance with growth-based allocation to adapt expert capacity during MoE fine-tuning.The framework is designed to match capacity to each expert’s downstream task relevance.
  • Conclusion: Across three open MoE LLMs and six benchmarks, DR-LoRA outperforms uniform-rank baselines and state-of-the-art adaptive methods in accuracy, efficiency, and robustness.Ablations and diagnostic visualizations support the method’s claims about heterogeneous rank specialization.

Limitation

The paper identifies limitations involving the saliency heuristic, routing and modality scope, and the overhead of pre-allocating rank space.

  • Limitation: The saliency score is a practical heuristic based on a first-order approximation and mean-field factorization, not a strict optimal estimator of expert utility.Its behavior under strongly coupled routing–gradient dynamics remains to be studied.
  • Limitation: Experiments primarily use text MoE LLMs with top-k routing, leaving transfer to expert-choice routing and multimodal MoE architectures open.The paper does not establish equal effectiveness for those alternative settings.
  • Limitation: Dynamic rank growth requires pre-allocated rank space during training, introducing extra training-time overhead whose larger-scale trade-off remains to be characterized.

A.1 From Capacity Expansion to the f ·g Factorization

The analysis models the benefit of adding expert rank through routing frequency and conditional gradient intensity, then uses a concave rank-growth objective to allocate capacity progressively.

  • From Capacity Expansion to the f ·g Factorization: Expert gradients vanish when routing excludes an expert, so capacity expansion must account for both expert usage and activated gradient intensity.The routing weight modulates the gradient, while q measures local gradient intensity when the expert is activated.
  • From Capacity Expansion to the f ·g Factorization: The expected expansion benefit is approximated as the product of routing frequency f and conditional learning intensity g after neglecting their covariance.The factorization is presented as a theoretically motivated heuristic rather than an exact independence claim.
  • From Capacity Expansion to the f ·g Factorization: The multiplicative score requires both LoRA factors to carry nontrivial gradient-weight products, avoiding importance estimates dominated by a converged factor.An additive alternative can overestimate a rank dimension when one side contributes little.
  • From Capacity Expansion to the f ·g Factorization: Mean aggregation estimates per-dimension learning intensity, whereas summation favors already-wide experts and maximization is sensitive to outliers and sparse noise.The mean supports rank-normalized comparisons across experts with different current ranks.
  • From Capacity Expansion to the f ·g Factorization: A concave benefit function makes each additional rank less valuable, motivating greedy allocation to the expert with the largest incremental gain.The default penalty exponent is γ = 1.2, which controls how strongly growth priority declines with existing rank.
  • From Capacity Expansion to the f ·g Factorization: DR-LoRA training progressively concentrates higher ranks in task-relevant experts as training proceeds.Figure 6 depicts rank evolution with darker colors indicating higher average expert ranks.

B.1 Training Datasets

The experiments fine-tune MoE models on six domain-specific datasets spanning reasoning, code, instruction following, medicine, translation, and legal understanding, using a scheduled growth procedure.

  • Training Datasets: The training suite covers mathematical reasoning, code instruction following, general instruction following, medical QA, machine translation, and legal provisions.MetaMathQA and CodeAlpaca-20k are used alongside OLMoE SFT Mix, MedMCQA, WMT-DA-Human-Evaluation, and LEDGAR.
  • Training Datasets: WMT-DA-Human-Evaluation spans English–Czech, English–German, and English–Chinese pairs with 20k examples per pair.The dataset uses direct-assessment human annotations from WMT News Translation shared tasks.
  • Growth Schedule and Router Training: Rank growth begins after warmup and ends 200 steps before training completion, giving newly activated ranks time to optimize.The growth window is defined between the warmup endpoint and the final 200-step buffer.
  • Growth Schedule and Router Training: All layers grow simultaneously, with the per-layer quota determined by the target-to-initial rank gap and the number of scheduled growth events.For OLMoE, N=128 denotes 64 experts multiplied by two adapted projections.
  • Growth Schedule and Router Training: The router remains frozen during warmup, then is trained jointly with LoRA modules through completion.This schedule is intended to stabilize initial LoRA training before router adaptation begins.
  • Training Datasets: Evaluations use three random seeds and temperature τ = 0.2, with vLLM used except for LLaMA-MoE.LLaMA-MoE is evaluated with HuggingFace Transformers because vLLM does not support it.

C.1 Statistical Significance of the Main Results

Across 18 model–task settings, DR-LoRA consistently exceeds the strongest competing baseline, with statistically significant aggregate improvement.

  • Statistical Significance of the Main Results: DR-LoRA outperforms the best competing baseline in all 18 model–task settings, with gains ranging from 1.0 to 2.5 points.The mean improvement is 1.8 points, and the aggregate paired test reports p < 0.001.
  • Statistical Significance of the Main Results: Twelve of the 18 individual comparisons are significant at p < 0.05, while the remaining differences are mainly associated with WMT23 and MedMCQA.The passage attributes the smaller differences in those settings to relatively close performance among the strongest methods.

C.2 Initial Rank Ratio Analysis

DR-LoRA remains robust across initial-rank choices, with the default half-target initialization providing the best trade-off while progressively forming heterogeneous expert ranks.

  • Initial Rank Ratio Analysis: All evaluated initial-rank settings outperform LoRA, and rinit = rtarget/2 provides the best trade-off.The study fixes rtarget = 16 and rmax = 32 for the OLMoE analysis.
  • Initial Rank Ratio Analysis: A smaller initial rank increases room for saliency-guided differentiation, whereas a larger initial rank stabilizes early training but leaves less growth budget.The comparison evaluates three initial-rank settings under the same target and maximum ranks.
  • Initial Rank Ratio Analysis: During training, most experts initially remain at rinit = 8, then task-relevant experts receive additional ranks and a heterogeneous pattern emerges.Figure 6 shows this progression from early to mid-training stages.

D.1 Memory Analysis

DR-LoRA matches LoRA (r=32) in memory usage and LoRA FLOPs while adding modest training-time overhead. Because base expert computation dominates total FLOPs, its overall computational increase remains limited.

  • DR-LoRA matches LoRA (r=32) exactly in per-GPU memory usage because both allocate identical parameter space.Compared with LoRA (r=16), DR-LoRA adds approximately 2 GB per GPU due to 100.6M reserved rank dimensions.
  • DR-LoRA incurs 5% and 4% training-time overhead over LoRA (r = 32) on MetaMathQA and CodeAlpaca-20K, respectively.The overhead comes from importance scoring, expert usage tracking, and periodic rank growth.
  • Despite near-equivalent training cost, DR-LoRA yields larger accuracy gains than LoRA (r = 32): +15.5 versus +12.3 on GSM8K and +3.1 versus +1.2 on HumanEval.
  • Base expert computation accounts for 98.5% of total FLOPs, limiting DR-LoRA’s overall increase over the no-LoRA baseline to 3.1%.DR-LoRA’s LoRA FLOPs are estimated using the maximum rank r=32 as a conservative upper bound.
Loading 2601.04823v6…