Technische Samenvatting: Het leren van een Zeroth-Order Optimizer voor het Fine-tunen van LLM's

Probleemstelling

Het fine-tunen van grote taalmodellen (LLM's) met standaard first-order optimizers (bijv. Adam) brengt een prohibitieve geheugenoverhead met zich mee, waarbij vaak tot wel 12 keer zoveel geheugen nodig is als voor inferentie door de opslag van optimizer-toestanden en backpropagation-activaties. Hoewel zeroth-order (ZO) optimizers zoals MeZO zijn opgekomen als een geheugenefficiënt alternatief door uitsluitend te vertrouwen op forward passes, maken bestaande methoden doorgaans gebruik van handmatige, statische bemonsteringsstrategieën (bijv. standaard normale verdelingen). Deze statische strategieën zijn niet aanpasbaar aan de specifie specifieke structurele eigenschappen van verschillende modellen of taken, vereisen vaak uitgebreide hyperparameter-afstemming en maken geen volledig gebruik van de lokale geometrie van het loss-landschap.

Methodologie: ZO Fine-tuner

De auteurs stellen de ZO Fine-tuner voor, een op leren gebaseerde zeroth-order optimizer die is ontworpen om automatisch efficiënte perturbatiestrategieën te leren via een "Learning to Learn" (L2L) framework. De kerninnovatie ligt in de overstap van statische, uniforme perturbatie naar adaptieve, blok-gewijze perturbatie.

1. Blok-gewijze Perturbatie Architectuur

Geïnspireerd door recente analyses die suggereren dat Transformer Hessians een benaderend blok-diagonaal structuur vertonen, vermijdt deze methode coördinaat-gewijze parametrisatie (wat op schaal geheugenprohibitief is) ten gunste van blok-gewijze adaptatie.

• PertNN: Voor elk parameterblok $i$ van de LLM voorspelt een lichtgewicht hulp-neuraal netwerk (PertNN) een gedeelde perturbatie-variantie $\sigma_t^{(i)}$.
• Inputs: PertNN neemt als input compacte samenvattende statistieken van de huidige parameters (gemiddelde en variantie van het blok), de vorige perturbatie-variantie en de loss-waarden van de vorige stap.
• Output: Het netwerk geeft een bijgewerkte variantie $\sigma_t^{(i)}$ uit, die wordt gebruikt om een diagonale variantie-matrix $\Sigma_t = \text{diag}(\sigma_t^{(1)}I, \dots, \sigma_t^{(n)}I)$ te construeren.
• Normalisatie: Om de geleerde variantie te ontkoppelen van de effectieve leersnelheid, normaliseert de methode $\Sigma_t$ zodanig dat de Frobenius-norm constant blijft ($\|\Sigma_t\|_F = \sqrt{d}$). Dit zorgt ervoor dat de perturbatie-amplitude stabiel blijft terwijl de optimizer in staat is om relatieve blok-gewijze varianties te leren.

2. Training via Learning to Learn (L2L)

De optimizer wordt getraind met een meta-doelstelling waarbij het fine-tuning traject van de LLM zelf dient als supervisie.

• Meta-doelstelling: Het doel is om de post-update loss $L(\theta_1)$ na een enkele ZO-stap te minimaliseren. De gradiënt van deze loss met betrekking tot de PertNN-parameters $\omega$ wordt berekend met behulp van de reparameterisatietruc (waarbij $z \sim N(0, I)$ wordt gesampled en $u = \Sigma z$ wordt ingesteld), waardoor het perturbatieproces differentieerbaar wordt.
• Trainingsstrategie: De PertNN wordt getraind op een enkele dataset (bijv. COPA) met behulp van een first-order optimizer om modeltoestanden te genereren. Om overfitting naar gebieden met een lage loss te voorkomen, introduceren de auteurs een periodieke reset mechanisme, waarbij de LLM-parameters na bepaalde intervallen worden teruggezet naar hun staat vóór de fine-tuning.
• Generalisatie: Zodra de ZO Fine-tuner is getraind voor een specifieke basismodel, kan deze zonder verdere training worden hergebruikt voor diverse downstream-taken en afgeleide checkpoints.

Belangrijkste Bijdragen

1. Extensie van L2L naar LLM's: Het artikel demonstreert dat een enkele geleerde optimizer, getraind op een basismodel, effectief kan generaliseren naar diverse downstream-taken en afgeleide checkpoints, wat een "train once, reuse widely" workflow valideert.
2. Blok-gewijze Variantie Leerproces: De auteurs stellen een geheugenefficiënte blok-gewijze parametrisatie voor met een gedeelde variantie per blok. Dit ontwerp is theoretisch onderbouwd door de blok-diagonale structuur van Transformer Hessians, wat L2L mogelijk maakt op de schaal van foundation models zonder prohibitieve kosten voor het geheugen.
3. Empirische Prestaties: Over vier LLM's (LLaMA-3.2-1B, LLaMA-3.1-8B, Qwen2.5-14B, OPT-30B) en zeven datasets presteert ZO Fine-tuner beter dan de sterkste zeroth-order baselines (MeZO, MeZO-AdamU, HIZOO, LOZO) in 82,1% van de taak-model combinaties wat betreft geconvergeerde loss, met een gemiddelde nauwkeurigheidsverbetering van 2,5%.

Experimentele Resultaten

• Convergentie en Nauwkeurigheid: ZO Fine-tuner bereikt consequent sneller een lagere loss dan de baselines. Het vertoont superieure robuustheid tegen de keuze van de leersnelheid, waarbij het vaak vergelijkbare of betere prestaties behaalt bij lagere leersnelheden (bijv. $10^{-8}$) waar baselines niet convergeren of instabiel worden.
• Transfereerbaarheid: Een optimizer getraind op de COPA-dataset slaagt erin te transfereren naar het fine-tunen van modellen op SST-2, CB, SQuAD, WSC, BoolQ en DROP, evenals naar wiskundige redeneringstaken (MetaMathQA).
• Afgeleide Checkpoints: De methode generaliseert van een basismodel (LLaMA-3.1-8B) naar de instructie-getunede variant (LLaMA-3.1-8B-Instruct), waarbij het zowel in loss als nauwkeurigheid MeZO overtreft.
• Efficiëntie:
  • Geheugen: De geheugenoverhead is verwaarloosbaar vergeleken met de modelgrootte (bijv. <2MB voor hulp-netwerken op OPT-30B). De totale geheugenvoetafdruk blijft vergelijkbaar met MeZO en is aanzienlijk lager dan first-order methoden (bijv. 62GB vs. 312GB voor OPT-30B).
  • Tijd: De tijdsoverhead geïntroduceerd door het bevragen van de PertNN is minimaal (<3,4% voor kleinere modellen, afnemend voor grotere modellen), aangezien de dominante kosten blijven liggen bij de forward passes die vereist zijn voor de zeroth-order schatting.

Betekenis en Claims

Het artikel claimt dat ZO Fine-tuner de kritieke flessenhals van geheugenefficiëntie bij het fine-tunen van LLM's aanpakt, terwijl het de starheid van statische zeroth-order strategieën overwint. Door gebruik te maken van het "train once, reuse widely" paradigma, stellen de auteurs dat modelmakers een getrainde pretrained learned finetuner samen met de basismodellen kunnen leveren, waardoor downstream-gebruikers efficiënte, hoogwaardige fine-tuning kunnen uitvoeren met bijna de geheugenkosten van inferentie. Het werk vestigt dat het leren van adaptieve perturbatieverdelingen haalbaar is op de schaal van foundation models en biedt een praktisch pad naar schaalbare, geheugenefficiënte aanpassing van LLM's.