Technische Samenvatting: Leren Selecteren, Niet Opnieuw Leren: Hard-Routed Mixtures van Reasoning LoRA's

Probleemstelling

Het paper behandelt de uitdaging van het samenstellen van onafhankelijk getrainde Low-Rank Adaptation (LoRA) adapters tot één Large Language Model (LLM) voor multi-domein adaptatie, met name wanneer de oorspronkelijke trainingsdata niet gedeeld kan worden vanwege privacy- of regelgevingsbeperkingen.

Een kritisch probleem doet zich voor bij het toepassen van standaard Mixture-of-Experts (MoE) soft routing op bevroren, onafhankelijk getrainde LoRA adapters. LoRA-modules worden getraind onder de assumptie van een unit-scale additive update ($W_{updated} = W_{pretrained} + AB$). Echter, standaard soft MoE routing combineert experts met behulp van gewogen sommen ($W_{updated} = W_{pretrained} + \sum \omega_i A_i B_i$), waarbij de routinggewichten $\omega_i < 1$ zijn. Dit creëert een schaalverschil (scale mismatch): de effectieve update die wordt toegepast tijdens integratie verschilt van de unit-scale update waaronder de expert oorspronkelijk is getraind.

Bestaande oplossingen, zoals LoRA-Mixer, proberen dit op te lossen door de adapters opnieuw te trainen met hulp van auxiliary preservation losses. Deze aanpak verhoogt het aantal trainbare parameters, vereist meer data en brengt het risico met zich mee dat de kennis die in de originele experts is gecodeerd, wordt aangetast. Bovendien is deze uitdaging bijzonder prominent bij instruction-tuned modellen met Chain-of-Thought (CoT) redeneren, waarbij conventionele fine-tuning de interne redeneergedragingen kan verslechteren, en bestaande compositiemethoden vaak falen in het behouden van deze gespecialiseerde redeneersporen.

Methodologie

De auteurs stellen Hard-Routed MoR-LoRA voor, een tweestaps framework ontworpen om bevroren reasoning LoRA-experts te componeren via unit-scale hard selection, waardoor hertraining van de experts wordt vermeden.

Fase I: Training van Reasoning Experts

Domeinspecifieke LoRA adapters worden onafhankelijk getraind met behulp van Reinforcement Learning from Verifiable Feedback (RLVF). In tegen tegenstelling tot supervised fine-tuning (SFT), dat imitatie van referentie-outputs aanmoedigt, optimaliseert RLVF voor correctheid met behulp van automatisch verifieerbare beloningen (bijv. het controleren van het uiteindelijke antwoord). Dit paradigma stimuleert de opkomst van gestructureerde CoT-redenering zonder de gebalanceerde redeneerpatronen van het basismodel te overschrijven. De resulterende experts worden vervolgens bevroren.

Fase II: Hard-Routed Integratie

De bevroren experts worden geïntegreerd in een gedeeld model met behulp van een Hard-Routed Mixture.

• Hard Top-1 Routing: Voor elk token selecteert de router precies één expert. Dit zorgt ervoor dat de geselecteerde expert wordt toegepast met unit scale, waardoor de oorspronkelijke LoRA-assumptie wordt behouden ($W_{updated} = W_{pretrained} + 1 \cdot A_j B_j$).
• Straight-Through Estimator (STE): Omdat discrete selectie (argmax) niet differentieerbaar is, gebruiken de auteurs STE. In de forward pass wordt de harde beslissing gebruikt; in de backward pass worden gradiënten gepropageerd door de soft waarschijnlijkheden. Dit maakt end-to-end training van de router mogelijk terwijl het discrete, unit-scale gedrag behouden blijft tijdens inferentie.
• Parameter Efficiëntie: Alleen een lichtgewicht gedeelde lineaire router en een kleine attention LoRA-module worden getraind. De voorgetrainde experts blijven bevroren.
• Distillatie: De router wordt getraind via Supervised Fine-Tuning (SFT) op een kleine dataset van reasoning traces die zijn gedestilleerd uit de bevroren experts. Het doel is om te leren welke expert te selecteren, niet om de redeneercapaciteit zelf opnieuw te leren.

Belangrijkste Bijdragen

1. Identificatie van Schaalverschil: Het paper identificeert dat het toepassen van soft MoE routing op onafhankelijk getrainde bevroren LoRA adapters een schaalverschil introduceert dat de effectieve functie van de adapters verandert.
2. Hard Unit-Scale Routing: De auteurs stellen het gebruik van hard top-1 routing met een Straight-Through Estimator voor om de unit-scale toepassing van experts te behouden, wat modulaire compositie zonder hertraining mogelijk maakt.
3. Tweestaps Framework: Een nieuw pipeline die de acquisitie van redeneren (via RLVF) scheidt van expertselectie (via hard routing), waardoor experts bevroren blijven tijdens integratie.
4. Empirische Validatie: De studie demonstreert dat deze aanpak het expertgedrag behoudt terwijl er aanzienlijk minder trainbare parameters nodig zijn dan bij soft-routing baselines.

Experimentele Resultaten

Experimenten werden uitgevoerd over vijf benchmarks (GSM8K, ARC-Challenge, MedQA, BoolQ, CoLA) en meerdere model-schalen (1B, 3B, 8B) met behulp van de LLaMA-3, Gemma-3 en Qwen-2.5 families.

• RLVF Superioriteit: RLVF-getrainde experts presteerden aanzienlijk beter dan SFT-getrainde experts op taken met veel redenering voor capabele modellen (3B en 8B), terwijl SFT bij grote instruction-tuned modellen soms de redeneerprestaties verslechterde.
• Prestaties vs. Parameters: Hard-Routed MoR-LoRA behaalde vergelijkbare of betere prestaties dan soft-routing baselines (LoRAMixer, MoLE) terwijl het aanzienlijk minder parameters trainde. Bijvoorbeeld, op het 3B model met RLVF-experts trainde de voorgestelde methode ~73M parameters om een gemiddelde nauwkeurigheid van 72,07 te bereiken, terwijl LoRAMixer ~606M parameters nodig had om 65,67 te bereiken.
• Analyse van Routinggedrag: Analyse van soft top-2 routing toonde aan dat dit vaak convergeert naar een bijna top-1 selectie (gemiddeld dominant gewicht ~0,71), wat suggereert dat hard routing het beoogde gedrag directer en efficiënter vastlegt.
• Data Efficiëntie: De integratiefase (Fase II) vereiste slechts ~1.000 samples per dataset om een prestatie te bereiken die dicht bij het trainen op de volledige dataset ligt, wat aangeeft dat de router selectiepatronen leert in plaats van domeinkennis.
• Mixed-Domain Capaciteit: In mixed-domain evaluaties (bijv. een prompt die zowel wiskunde als leesvaardigheid bevat), presteerde token-level hard routing beter dan prompt-level routing, wat flexibiliteit toont in het selecteren van verschillende experts binnen een enkele input.

Betekenis en Claims

Het paper claimt dat modulaire adaptatie kan worden bereikt door te leren wanneer elke expert te gebruiken, in plaats van te modificeren wat elke expert weet.

De auteurs beargumenteren dat voor bevroren voorgetrainde LoRA experts, hard unit-scale routing een eenvoudige en efficiënte abstractie biedt. Door integratie primair als een routingprobleem te behanden, vermijdt de methode de computationele kosten en de potentiële degradatie van kennis die gepaard gaat met het hertrainen van experts. De resultaten suggereren dat het primaire voordeel van compositie voortkomt uit het selecteren van de juiste bevroren expert voor de context, in plaats van het uitgebreid modificeren van de parameters van de experts.

Het werk benadrukt dat hoewel prompt-level routing voldoende is voor zuivere single-domain inputs, token-level hard routing de noodzakelijke flexibiliteit biedt voor complexe, mixed-domain inputs. De auteurs merken bescheiden op dat de methode niet volledig training-vrij is, aangezien Fase II enige gelabelde data vereist om het routingbeleid te leren, en het momenteel geen echte zero-shot integratie van nieuwe domeinen ondersteunt zonder aanvullende routingdata.