Technische Samenvatting: Budgeted LoRA

Probleemstelling

Grote Taalmodellen (LLMs) bereiken sterke prestaties door schaling, maar dit gaat gepaard met aanzienlijke implementatiekosten. Hoewel Kennisdistillatie (KD) een oplossing biedt door kennis over te dragen van grote docenten naar kleinere studenten, falen bestaande pipelines vaak om modellen te produceren die structureel efficiënt zijn tijdens inferentie.

Een primaire beperking ligt in de wijdverbreide toepassing van parameter-efficiënte fijnafstelmethode zoals LoRA (Low-Rank Adaptation). Hoewel LoRA het aantal trainbare parameters vermindert door de dichte ruggengraat te bevriezen en laag-rang updates toe te voegen, laat het de dichte ruggengraat onveranderd. Bijgevolg levert LoRA-gebaseerde distillatie studenten op die goedkoper zijn om te trainen, maar de dominante inferentiekost niet verminderen, aangezien de dichte matrixvermenigvuldigingen intact blijven. Huidige benaderingen optimaliseren vaak voor het aantal parameters of trainings-efficiëntie zonder expliciet te controleren hoe computatie binnen de studentenarchitectuur tijdens inferentie wordt toegewezen.

Methodologie: Budgeted LoRA

De auteurs stellen Budgeted LoRA voor, een distillatiekader dat modelcompressie herformuleert als een gestructureerd computetoewijzingsprobleem. In plaats van de studentenarchitectuur vast te leggen, introduceert de methode een globaal computebudget dat het uiteindelijke fractie van dichte computatie die behouden moet blijven, dicteert.

De aanpak werkt in drie fasen:

1. Begrotingsbewuste Gated LoRA-modules:
   Het studentenmodel vervangt geselecteerde lineaire projecties (bijv. attention $q/k/v/o$ en MLP $up/gate/down$) door gated modules. Deze modules berekenen:
   $$y = d \cdot Wx + \frac{\alpha}{r_{max}} B((Ax) \odot g)$$
   Waarbij:

   • $W$ de bevroren dichte gewicht is.
   • $d \in [0, 1]$ een dicht-pad behoudcoëfficiënt is, gecontroleerd door een budgetplanner.
   • $A, B$ LoRA-factoren zijn, en $g$ geleerde per-rang poorten voorstelt.
   • Deze formulering stelt het model in staat om de actieve capaciteit van zowel dichte als laag-rang paden dynamisch te variëren over verschillende modules.

2. Begrotingsplanning en Controller:
   Een globaal budgetparameter $F \in [0, 1]$ specificeert de doel-fractie van dichte computatie die tegen het einde van de training behouden moet blijven. Een controller handhaaft dit budget door:

   • Een doel-dicht-kosttraject $b(t)$ te definiëren met behulp van een cosinus-schema.
   • De totale behouden dichte kost $C(t)$ te berekenen op basis van het MAC-aantal van actieve dichte projecties.
   • Gierig module-niveau behoudcoëfficiënten ($d$) te verlagen om het doelbudget te halen, met prioriteit voor het verwijderen van modules met lagere kosten eerst.
   • Dit creëert een gecontroleerde overdracht van computatie van dichte paden naar laag-rang paden tijdens de training.

3. Post-training Compressie:
   Na de training ondergaat het model een consolidatiestap om daadwerkelijke inferentiewinsten te realiseren. Op basis van de uiteindelijke behoudcoëfficiënt $d$, wordt elke module toegewezen aan een van drie implementatievormen:

   • Geval 1 (Dicht Verwijderen): Als $d < \epsilon_{zero}$, wordt het dichte pad volledig verwijderd, waardoor alleen het ingekorte laag-rang component overblijft.
   • Geval 2 (SVD-benadering): Als $\epsilon_{zero} \le d < \epsilon_{lr}$, wordt de resterende dichte bijdrage benaderd via een laag-rang SVD, waarbij deze wordt samengevoegd met de LoRA-factoren.
   • Geval 3 (Dicht Samenvoegen): Als $d \ge \epsilon_{lr}$, wordt het dichte pad behouden en samengevoegd met de LoRA-update tot een standaard dichte lineaire laag.

Deze pipeline produceert een familie van studentenmodellen geparametriseerd door de enkele budgetknop $F$, waardoor een directe afweging tussen efficiëntie en prestaties mogelijk wordt.

Belangrijkste Bijdragen

• Nieuwe Formulering: Het artikel stelt voor om distillatie te behandelen als een gestructureerd computetoewijzingsprobleem onder een globaal budgetbeperking, in plaats van een optimalisatie met een vaste architectuur.
• Budgeted LoRA Kader: Introduceert een methode die controller-afgedwongen dichte behoud combineert met geleerde laag-rang toewijzing, gevolgd door post-training compressie om implementatie-efficiënte studenten te creëren.
• Controleerbare Afweging: Demonstreert dat een enkele budgetparameter een controleerbare kwaliteits-efficiëntie afweging oplevert, vormend een Pareto-grens over trainings- en inferentiekosten.
• ICL Behoud: Toont aan dat Budgeted LoRA aanzienlijk hogere prestaties behoudt dan standaard LoRA-gebaseerde distillatie op functionele In-Context Learning (ICL) probes.

Resultaten

De auteurs evalueerden Budgeted LoRA met behulp van een 6-laags student geïnitieerd vanuit een 0,13B model, gedistilleerd van docenten van variërende schalen (0,13B, 0,5B, 1B).

• Kwaliteit-Efficiëntie Afweging:

  • Bij een gematigd budget ($F=0,4$) komt Budgeted LoRA overeen met de perplexiteit van standaard KD-LoRA, terwijl het een 1,74× snelheidswinst bereikt in inferentie met samengevoegde modules en de trainingscomputatie met 0,59× vermindert.
  • Bij een agressief budget ($F=0,0$) bereikt de methode een 4,05× snelheidswinst in inferentie met samengevoegde modules en 0,38× trainingscomputatie, met slechts een gematigde verslechtering van de perplexiteit.
  • De methode traceert een duidelijke Pareto-grens, waardoor lagere kosten-operatiepunten worden geboden met gecontroleerde perplexiteitsverslechtering in vergelijking met standaard LoRA.

• In-Context Learning (ICL) Behoud:

  • Op functionele ICL-probes (geadaptieerd van de Function Vectors suite) presteert Budgeted LoRA aanzienlijk beter dan standaard KD-LoRA, vooral bij distillatie van grotere docenten.
  • Bijvoorbeeld, met een 1B docent behaalde KD-LoRA een macro-gemiddelde ICL-score van 2,7, terwijl Budgeted LoRA ($F=0,0$) 5,0 bereikte en ($F=0,4$) 5,6 bereikte.
  • Dit suggereert dat het behoud van hoger niveau gedrag niet alleen afhangt van het matchen van perplexiteit, maar van het handhaven van de onderliggende computatiestructuur door middel van budgetgedrege dichte behoud.

Betekenis en Claims

Het artikel beweert dat bij computegelimiteerde distillatie het behouden van gedrag minder gaat over het matchen van perplexiteit of simpelweg het verwijderen van parameters; het gaat over het controleren hoe dichte computatie wordt overgedragen naar laag-rang paden.

De auteurs betogen dat standaard LoRA faalt om inferentiewinsten te leveren omdat het de dichte ruggengraat intact laat. Daarentegen behandelt Budgeted LoRA dichte en laag-rang paden als concurrerende bronnen. De resultaten suggereren dat:

1. Gestructureerde Toewijzing Uitmaakt: Een geleidelijke, budgetgedrege vermindering van dichte capaciteit zorgt voor een soepelere overdracht van capaciteit naar laag-rang paden in vergelijking met abrupte verwijdering.
2. Voorbij Perplexiteit: Hogeniveau gedrag (zoals ICL) wordt beter behouden wanneer het trainingsproces expliciet de structuur van computatie beheert, zelfs als de uiteindelijke perplexiteit vergelijkbaar is met standaard LoRA.
3. Implementatie-efficiëntie: De methode slaagt erin om budgetsignalen tijdens de training om te zetten in daadwerkelijke inferentiewinsten via post-training compressie, waardoor een spectrum van modellen wordt geboden van dichte tot volledig laag-rang.

Het werk positioneert zichzelf als een stap in de richting van het zien van efficiëntie niet louter als een probleem van parameteraantal, maar als een probleem van computetoewijzing tijdens het trainings- en distillatieproces.