ReMix: Reinforcement routing for mixtures of LoRAs in LLM finetuning

Dit paper introduceert ReMix, een nieuwe routeringsmethode voor Mixture-of-LoRAs die het probleem van onbalans in leerbare routingsgewichten oplost door niet-leerbare gewichten te combineren met een onbevooroordeelde gradiëntschatter op basis van reinforcement learning, wat leidt tot een aanzienlijk betere prestatie dan bestaande parameter-efficiënte finetuning-methoden.

De kern

Stel je voor dat je een gigantische, superintelligente robot (een Large Language Model of LLM) hebt die al alles weet over de wereld. Maar nu wil je deze robot leren om specifieke taken te doen, zoals wiskundige raadsels oplossen of computercode schrijven.

Je kunt de robot niet volledig herschrijven; dat kost te veel tijd en energie. In plaats daarvan plak je kleine, slimme "stikselstukjes" (de LoRAs) op de robot. Deze stukjes leren de specifieke vaardigheid.

Het probleem met de huidige methode:
Stel je voor dat je een team van 8 van deze stikselstukjes hebt. De robot heeft een "manager" (de router) die moet beslissen welk stukje hij gebruikt voor een bepaalde vraag.
In de huidige systemen probeert de manager te leren welk stukje het beste is. Maar er gebeurt iets raars: de manager wordt te zeker van zichzelf. Hij begint te denken: "Oh, stukje nummer 4 is geweldig! Ik zal 99% van mijn aandacht op hem richten en de andere 7 stukjes negeren."

Dit noemen de auteurs Router Collapse. Het is alsof je een orkest hebt met 8 muzikanten, maar de dirigent laat alleen de trompettist spelen. De andere 7 muzikanten (die misschien ook prachtige viool- of fluitstukjes kunnen spelen) worden nutteloos. Je betaalt voor 8 muzikanten, maar hoort maar één geluid.

De oplossing: ReMix (Reinforcement Routing for Mixture-of-LoRAs)
De auteurs van dit paper, ReMix, zeggen: "Stop met die manager die probeert te kiezen. Laat hem gewoon iedereen gelijk behandelen!"

Hier is hoe hun nieuwe systeem werkt, vertaald naar een eenvoudig verhaal:

1. De "Gelijke Kansen" Regeling

In plaats van dat de manager beslist hoeveel gewicht elk stukje krijgt (en zo één stukje laat domineren), zegt ReMix: "Oké, we kiezen 4 stukjes. Maar we geven ze allemaal precies hetzelfde stemgeluid."
Dit zorgt ervoor dat alle actieve stukjes echt worden gebruikt. Het is alsof je in een vergadering zegt: "We luisteren naar 4 mensen, en we geven ze allemaal evenveel tijd om te spreken." Niemand wordt overstemd.

2. Het Gokspel (Reinforcement Learning)

Maar wacht, als de manager geen gewichten kan aanpassen, hoe leert hij dan welk stukje hij moet kiezen?
Stel je voor dat de manager een gokker is in een casino.

• De oude manier: Hij probeerde te raden welke machine het meeste geld uitkeerde door de knoppen zachtjes te duwen (gradiëntafdaal). Maar omdat hij soms één knop te hard duwde, verloor hij de andere opties.
• De ReMix-methode: De manager doet een gok. Hij kiest willekeurig 4 stukjes. Als het resultaat goed is, krijgt hij een beloning. Als het slecht is, krijgt hij een boete.
  • Ze gebruiken een slimme truc genaamd RLOO (Reinforce Leave-One-Out). Dit is alsof de manager zegt: "Ik heb 4 keer geprobeerd. Laten we kijken wat er gebeurt als ik één van die pogingen niet meetel, om te zien of die ene poging echt het verschil maakte."
  • Dit helpt de manager om te leren welke combinatie van stukjes het beste werkt, zonder dat hij de "gelijkheid" tussen de stukjes verstoort.

3. De Finale Keuze (Top-k)

Tijdens het trainen (het leren) laat de manager soms wat willekeur toe, zodat hij kan ontdekken wat er werkt. Maar als de robot klaar is om de echte wereld in te gaan (tijdens het gebruik), doet hij iets slim:
Hij kijkt naar alle 8 de stukjes en kiest gewoon de 4 beste op basis van wat hij heeft geleerd. Omdat hij tijdens het trainen eerlijk heeft getraind, weet hij nu precies welke 4 het beste bij de vraag passen.

Waarom is dit geweldig?

• Geen verspilling: Je gebruikt echt alle stukjes waarvoor je betaalt.
• Beter resultaat: Omdat ze samenwerken in plaats van dat één de baas is, wordt de robot slimmer. In tests bleek ReMix beter te zijn dan alle andere methoden, zelfs met minder rekenkracht.
• Schaalbaar: Je kunt de robot nog slimmer maken door gewoon meer "gokpogingen" te doen tijdens het trainen.

Kortom:
De huidige systemen laten één expert de leiding nemen en vergeten de anderen. ReMix zorgt voor een eerlijk team, waar iedereen even hard mag werken, en gebruikt slimme goktechnieken om het beste team samen te stellen. Het resultaat is een slimmere, efficiëntere robot die beter presteert zonder dat je meer hoeft te betalen.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "ReMix: Reinforcement Routing for Mixtures of LoRAs in LLM Finetuning", geschreven in het Nederlands.

1. Het Probleem: Routing Weight Collapse

De auteurs identificeren een fundamenteel tekortkoming in bestaande Mixture-of-LoRAs (MoLoRA) modellen. Hoewel deze modellen bedoeld zijn om de expressiviteit te vergroten door meerdere Low-Rank Adapters (LoRAs) dynamisch per laag te activeren, vertonen ze een fenomeen dat ze "routing weight collapse" noemen.

• Observatie: Bestaande routers gebruiken leerbare gewichten (via een softmax-functie) om LoRAs te selecteren. Theoretisch en empirisch blijkt dat deze gewichten tijdens het finetunen extreem onbalans raken.
• Het Effect: Zelfs als $k > 1$ LoRAs geactiveerd zijn, krijgt vaak slechts één LoRA een dominant gewicht (dicht bij 1), terwijl de gewichten van de andere $k-1$ LoRAs verwaarloosbaar klein worden.
• Gevolg: De berekening van de andere LoRAs wordt effectief verspild. Het model gedraagt zich alsof er maar één LoRA actief is ($k=1$), wat de beoogde expressiviteit en prestaties van het mengsel ondermijnt.
• Oorzaak: De auteurs tonen aan dat dit een inherente eigenschap is van het gebruik van leerbare, continue routing-gewichten die via gradient descent worden geoptimaliseerd.

2. Methodologie: ReMix (Reinforcement Routing for Mixtures)

Om dit probleem op te lossen, stellen de auteurs ReMix voor, een nieuwe routeringsarchitectuur die twee kerncomponenten combineert:

A. Niet-leerbare Routing Gewichten (Constante Weegfactoren)

In plaats van te leren welke LoRA het belangrijkst is voor een specifieke input, wijst ReMix constante, niet-leerbare gewichten toe aan alle geactiveerde LoRAs.

• Als $k$ LoRAs worden geselecteerd, krijgt elk een gewicht $\omega$ (bijvoorbeeld $\omega = 2/kr$ of $\omega = 2/\sqrt{kr}$).
• Dit garandeert dat het Effective Support Size (ESS) altijd gelijk is aan $k$, waardoor alle actieve LoRAs evenveel bijdragen en geen enkele LoRA de andere domineert.
• Nadeel: Omdat de gewichten constant zijn, is de router niet differentieerbaar via backpropagation, wat standaard training onmogelijk maakt.

B. Reinforcement Learning met RLOO Gradient Estimator

Om de router (die nu LoRAs selecteert op basis van een kansverdeling) toch te kunnen trainen, formuleren de auteurs het probleem als Reinforcement Learning (RL):

• Policy: De router is het beleid dat een subset van $k$ LoRAs selecteert.
• Reward: De supervised finetuning loss (SFT loss) wordt gezien als de negatieve beloning (dus een lagere loss is een betere beloning).
• Gradient Estimator: Omdat de selectie discreet is, kunnen ze geen directe gradienten berekenen. Ze gebruiken een onbevooroordeelde gradient estimator gebaseerd op de Reinforce Leave-One-Out (RLOO) techniek.
  • Ze sample meerdere selecties ($M$) van LoRAs.
  • Ze berekenen de loss voor elke selectie en gebruiken het gemiddelde als baseline om de variance van de gradient schatting te verlagen.
  • Dit maakt het mogelijk om de router te optimaliseren zonder dat de gewichten zelf worden aangepast, maar wel de selectiekansen.

C. Inference: Top-k Selectie

Tijdens het trainen worden LoRAs gesampleerd volgens de router-verdeling. Tijdens de inferentie (testfase) gebruiken ze echter Top-k selectie.

• De auteurs bewijzen (Theorema 2) dat als de router goed getraind is (d.w.z. de kans op de optimale subset > 50%), het nemen van de $k$ LoRAs met de hoogste waarschijnlijkheid de optimale subset garandeert. Dit elimineert de noodzaak voor sampling tijdens inferentie en verbetert de deterministische prestaties.

3. Belangrijkste Bijdragen

1. Theoretische Inzicht: Het aantonen dat leerbare routing-gewichten in MoLoRA-modellen bijna altijd leiden tot "collapse" naar één dominante LoRA, wat de effectiviteit van het mengsel beperkt.
2. Nieuwe Architectuur (ReMix): Een simpele maar effectieve router die constante gewichten gebruikt voor geactiveerde LoRAs, waardoor een evenwichtige bijdrage wordt gegarandeerd zonder extra inferentie-kosten.
3. RL-gebaseerde Training: Een oplossing voor het trainen van niet-differentieerbare routers door het probleem te herschrijven als RL met een RLOO-based gradient estimator.
4. Schaalbaarheid: Het vermogen om de trainingscompute te schalen (door meer samples $M$ te nemen) om de prestaties verder te verbeteren, iets wat bij deterministische methoden niet mogelijk is.

4. Resultaten

De auteurs hebben ReMix getest op diverse benchmarks (GSM8K voor wiskundig redeneren, HumanEval voor codegeneratie, en ARC-c voor kennisherinnering) met een Llama 3 8B model.

• Prestaties: ReMix overtreft consistent state-of-the-art PEFT-methoden (zoals LoRA, DoRA, rsLoRA, MixLoRA, en HydraLoRA).
  • Gemiddelde verbetering van 2.82% ten opzichte van de sterkste concurrent.
  • Op GSM8K: 65.66% accuraatheid (vs. 62.47% voor de beste concurrent).
  • Op HumanEval: 32.93% Pass@1 (vs. 31.10% voor de beste concurrent).
• Parameter-efficiëntie: ReMix bereikt deze resultaten met slechts 0.070B trainbare parameters. Dit is een reductie van 90% ten opzichte van VB-LoRA en 31% ten opzichte van MixLoRA, terwijl het toch beter presteert.
• Ablatie Studies:
  • Het verwijderen van RLOO of Top-k selectie leidt tot een aanzienlijke daling in prestaties, wat aantoont dat beide componenten cruciaal zijn.
  • De methode is niet gevoelig voor de specifieke keuze van de constante gewicht $\omega$.
• Schaalbaarheid: Het verhogen van het aantal samples ($M$) tijdens training leidt tot lineaire verbeteringen in de accuraatheid, wat de efficiëntie van de RL-benadering bevestigt.

5. Betekenis en Impact

Dit paper biedt een paradigmaverschuiving in het ontwerpen van routers voor Mixture-of-Experts (MoE) en Mixture-of-LoRAs.

• Fundamenteel Inzicht: Het laat zien dat "leren" welke expert het beste is, vaak leidt tot het negeren van andere experts, en dat een eenvoudige, niet-lerende gelijkverdeling soms superieur is.
• Efficiëntie: Het biedt een manier om de expressiviteit van grote modellen te vergroten zonder de parameter-efficiëntie van LoRA te verliezen.
• Toekomst: De combinatie van RL-technieken (zoals RLOO) met PEFT-methoden opent nieuwe wegen voor het trainen van complexe, niet-differentieerbare componenten in grote taalmodellen, met name in scenario's waar rekenkracht schaars is maar prestaties cruciaal zijn.

Kortom, ReMix lost het probleem van "routing collapse" op door de router te simplifiseren en de trainingsdynamiek over te nemen via Reinforcement Learning, wat resulteert in een robuuster en efficiënter finetuning-framework.