ReMix: Fixing Routing Weight Collapse in Mixtures of LoRAs with Constant Weights and RLOO Reinforcement Learning

Forum topic · 小凯 · 2026-05-08

Summary

ReMix, a collaboration between UIUC and Meta researchers, identifies a critical flaw in Mixture-of-LoRAs methods: although k LoRA experts are activated, learnable softmax routing weights collapse during training so that effectively only one expert does the work—an effect the authors prove theoretically via an effective support size (ESS) bound. ReMix removes learnable routing weights entirely, forcing all activated LoRAs to share equal constant weights, and reformulates the now non-differentiable routing decision as a reinforcement learning problem trained with the RLOO (REINFORCE Leave-One-Out) unbiased gradient estimator. Fine-tuning Llama 3 8B, ReMix achieves 65.66% on GSM8K (+3.19), 32.93% on HumanEval (+1.83), and an average accuracy 2.82% above prior SOTA while activating only 0.07B parameters. Inference uses top-k selection, which a theorem shows is optimal once sampled routing accuracy exceeds 50%. Increasing the number of RL samples M further improves accuracy, revealing train-time scalability unique to the RL formulation.

ReMix: Reinforcement Routing for Mixtures of LoRAs in LLM Finetuning

Paper: ReMix: Reinforcement Routing for Mixtures of LoRAs in LLM Finetuning (arXiv:2603.10160), UIUC + Meta AI + Washington University, Lifelong Agents Workshop @ ICLR 2026.

Key points

Routing weight collapse bug: In Mixture-of-LoRAs, activating k LoRA experts does not mean k experts work. Learnable softmax routing weights collapse so one expert takes nearly all weight (e.g., weight 0.933 vs 0.003–0.022 for the rest; ESS ≈ 1.15 with n=8), and the collapse is irreversible after ~1000 steps.
Theory (Theorem 1): With standard Gaussian initialization of router parameters, the effective support size $ESS(\pi) = 1/\sum_i \pi_i^2$ is provably small with high probability—softmax routing is a winner-take-all operation, so collapse is a mathematical inevitability, not an engineering bug.
Fix 1 — Constant weights: ReMix replaces learnable weights with a constant ω for all activated LoRAs ( $\omega = 2/(kr)$ for LoRA-type, $2/\sqrt{kr}$ for rsLoRA-type), forcing ESS = k. Only *which* experts are selected remains learned.
Fix 2 — RL reformulation: Since selecting k experts is a discrete, non-differentiable decision, routing is treated as a policy trained with reward = negative SFT loss, using sampling without replacement.
Fix 3 — RLOO estimator: The REINFORCE Leave-One-Out gradient

\hat{G}_{P^{(l)}} = \frac{1}{M-1} \sum_{m=1}^M (L(J_m) - \bar{L}) \nabla_{P^{(l)}} \log Q(J_m)

is unbiased and reduces variance versus vanilla REINFORCE, which is essential for stable training.

Inference (Theorem 2): If sampled routing hits the optimal expert subset with probability > 1/2, deterministic top-k selection is guaranteed to hit it.

Results (Llama 3 8B, matched parameter budgets)

| Method | GSM8K | HumanEval | ARC-c | Avg | Active params | |---|---|---|---|---|---| | LoRA | 59.21 | 26.83 | 83.05 | 56.36 | 0.112B | | rsLoRA | 62.47 | 28.66 | 82.71 | 57.95 | 0.028B | | DoRA | 55.34 | 31.10 | 83.39 | 56.61 | 0.127B | | MixLoRA | 61.87 | 28.05 | 82.37 | 57.43 | 0.101B | | HydraLoRA | 62.47 | 20.12 | 82.71 | 55.10 | 0.084B | | ReMix | 65.66 | 32.93 | 83.73 | 60.77 | 0.070B |

Average accuracy is 2.82% above the strongest baseline with 31% fewer active parameters than MixLoRA.
Ablations confirm both RLOO (vs. high-variance REINFORCE) and top-k inference matter; LoRA-subset diversity at k=4 reaches 64.22% vs 59.21% for an equal-parameter single rank-kr LoRA, proving subsets genuinely vary.
Train-time scaling: increasing RL samples M from 2 to 32 raises GSM8K accuracy from 56.03% to 58.83%—an RL-native scalability knob traditional MoE-LoRA methods lack.
Efficiency: only ~10% more training time than MixLoRA for a ~16% relative accuracy gain.

Interpretation and caveats

The root cause is softmax being applied across *experts* (competition between different inputs), unlike attention where sharpness across *tokens* within one input is desirable.
Common misreadings to avoid: constant weights do not make all LoRAs equally important overall—different inputs still select different subsets; and this differs from MoE load-balancing auxiliary losses.
Limitations: gains from larger M diminish around M=32; results shown only on Llama 3 8B; the router is a simple linear projection; optimal (n, k) settings remain unexplored.

References

Qiu et al., *ReMix*, arXiv:2603.10160 (Lifelong Agents @ ICLR 2026)
Hu et al., *LoRA*, ICLR 2022
Li et al., *MixLoRA*, arXiv:2404.15159
Tian et al., *HydraLoRA*, NeurIPS 2024
Kool, van Hoof & Welling, *Buy 4 REINFORCE samples, get a baseline for free!* (RLOO), ICLR 2019 Workshop
Kalajdzievski, *rsLoRA*, arXiv:2312.03732
Liu et al., *DoRA*, ICML 2024
Zeng et al., *S'MoRE*, arXiv:2504.06426