Technische Samenvatting: Wanneer Verliezen Aligneren: Op Gradiënten Gebaseerde Gewichtsafstelling voor Samengestelde Verliezen voor Efficiënte Pretraining

Probleemformulering

Moderne deep learning-systemen maken vaak gebruik van een tweestapsparadigma: pretraining op grote schaal, vaak ongelabelde data, gevolgd door adaptatie aan downstream-taken. In veel pretraining-pijplijnen is de objectief Functie niet een enkel verlies, maar een gewogen combinatie van meerdere verlies termen (bijvoorbeeld next-token prediction, contrastieve verliezen of regressiedoelen). De relatieve gewichten van deze termen zijn kritieke hyperparameters die de kwaliteit van de geleerde representaties en de downstream-prestaties aanzienlijk beïnvloeden.

Huidig is het selecteren van deze gewichten afhankelijk van handmatige afstelling, willekeurige zoekopdrachten of Bayesiaanse optimalisatie. Deze methoden zijn computatief onhaalbaar voor pretraining op grote schaal omdat ze talloze onafhankelijke trainingen vereisen. Bovendien falen naïeve standaardwaarden, zoals uniforme weging, vaak om rekening te houden met de variërende bijdragen en conflicten tussen verschillende verliestermen. Het artikel adresseert de behoefte aan een efficiënte methode om lineaire verliesgewichten in samengestelde pretraining-doelen af te stellen om downstream-taakprestaties te maximaliseren zonder de kosten van herhaalde volledige trainingen te incasseren.

Methodologie: Gradiënt-gealigneerde Pretraining (GraP)

De auteurs stellen Gradiënt-gealigneerde Pretraining (GraP) voor, een op gradiënten gebaseerde bilevel-optimalisatiemethode die verliesgewichten online leert. Het kernidee is om de samengestelde pretraining-gradiënt te aligneren met de gradiënt van een downstream-doel.

1. Bilevel-optimalisatieformulering

Het probleem wordt geformuleerd als een bilevel-optimalisatietak waar de buitenste lus de verliesgewichten $w$ optimaliseert om het downstream-verlies $L_{down}$ te minimaliseren, en de binnenste lus modelparameters $\theta$ optimaliseert om het samengestelde pretraining-verlies $L_{pre}(w)$ te minimaliseren.
$$\min_{w} L_{down}(\theta^*(w)) \quad \text{onder voorwaarde} \quad \theta^*(w) = \arg\min_{\theta} \sum_{k=1}^K w_k L_k(\theta)$$

2. Op Gradiënten Gebaseerde Update-regel

Met behulp van een afgeknotte backpropagation-aanpak (single-step SGD) leiden de auteurs een update-regel voor de gewichten af. De gradiënt van het downstream-verlies met betrekking tot de gewichten is evenredig met het inproduct tussen de downstream-gradiënt en de individuele pretraining-gradiënten:
$$\frac{\partial L_{down}}{\partial w_k} \propto -\eta \, g_{down}^\top g_k$$
waarbij $g_{down}$ de downstream-gradiënt is en $g_k$ de gradiënt van het $k$-de pretraining-verlies. Het doel is om de alignatie tussen de samengestelde gradiënt $\sum w_k g_k$ en de downstream-gradiënt te maximaliseren.

3. Benadering in Inbeddingsruimte

Een naïeve implementatie zou vereisen dat $K$ volledige backpropagaties door het model worden uitgevoerd om $g_k$ te verkrijgen. Om de rekentijd te verminderen, berekent GraP gradiënten met betrekking tot de gedeelde backbone-inbedding $z$ in plaats van de volledige parameter ruimte.
$$\tilde{g}_k = \nabla_z L_k(h_k(z))$$
Aangezien de verliezen worden toegepast via aparte heads bovenop dezelfde inbedding, is het berekenen van deze gradiënten aanzienlijk goedkoper dan volledige backpropagaties. Het optimalisatiedoel wordt het maximaliseren van de alignatie in deze gedeelde representatieruimte:
$$\arg\max_{w} \frac{w^\top \tilde{G} \tilde{g}_{down}}{\|w^\top \tilde{G}\|}$$

4. Normalisatiestrategie

Het artikel pleit tegen het normaliseren van de gewichten zelf (bijvoorbeeld $\sum w_k = 1$) of hun $L_2$-norm, aangezien deze de optimale richting van de samengestelde gradiënt kunnen vervormen. In plaats daarvan legt GraP een beperking op aan de $L_2$-norm van de samengestelde gradiënt zelf ($\|\sum w_k g_k\| = 1$). Dit zorgt ervoor dat de optimalisatie zich richt op de richting van het update-vector in plaats van zijn grootte, waardoor de optimale alignatierichting onder het lokale doel wordt hersteld.

5. Algorithmische Efficiëntie

Het resulterende algoritme vereist slechts één forward pass en één backward pass door de backbone per iteratie. De gradiënten die voor de heads worden berekend, worden hergebruikt voor zowel het updaten van de verliesgewichten als het updaten van de backbone-parameters. Dit vermindert de overhead van hyperparameter-afstelling tot ongeveer 30% boven een enkele training, vergeleken met de kosten in de orde van grootte van Bayesiaanse optimalisatie.

Belangrijkste Bijdragen

1. Probleemformulering: De auteurs formuleren de selectie van pretraining-verliesgewichten als een bilevel-optimalisatieprobleem dat expliciet gericht is op downstream-prestaties.
2. Inzicht in Normalisatie: Ze analyseren normalisatiestrategieën en tonen aan dat het beperken van de norm van het update-vector (samengestelde gradiënt) superieur is aan het beperken van de gewichten zelf, waardoor de optimale alignatierichting behouden blijft.
3. Efficiënt Algoritme: Ze stellen GraP voor, een algoritme met $O(1)$ complexiteit wat betreft het aantal forward-backward-stappen ten opzichte van het aantal hyperparameters. Het maakt gebruik van gradiëntalignatie in de gedeelde representatieruimte om meerdere backward passes te vermijden.
4. Empirische Efficiëntie: De methode bereikt prestaties die vergelijkbaar zijn met of beter zijn dan standaard hyperparameter-optimalisatie (bijvoorbeeld Bayesiaanse zoekopdracht), terwijl het bijna 50$\times$ efficiënter is, waardoor afstelling mogelijk is in situaties waar traditionele methoden computatief onhaalbaar zijn.
5. Identificatie van Redundantie: De aanpak demonstreert het vermogen om redundante verliestermen te identificeren en minder zwaar te wegen (bijvoorbeeld het toekennen van een bijna nul-gewicht aan specifieke contrastieve verliezen in visietaken), wat effectief werkt als een automatisch verlies-selectiemechanisme.

Experimentele Resultaten

De methode werd geëvalueerd op twee domeinen: modelle ring van gebeurtenissequenties en zelftoezichtende computer vision.

Modellering van Gebeurtenissequenties

• Datasets: Churn, AgePred, AlfaBattle, MIMIC-III en Taobao.
• Opzet: Pretraining omvatte een samengesteld doel van next-token prediction (NTP) over meerdere velden en een contrastief verlies.
• Prestaties: GraP kwam overeen met of verbeterde zorgvuldig afgestelde baselines (inclusief GradNorm, DWA, MGDA en PCGrad) en Optuna-afgestelde gewichten. Het behaalde consistent de beste resultaten op de Churn- en MIMIC-III-datasets voor RNN-, Transformer- en HT-Transformer-backbones.
• Efficiëntie: In vergelijking met Bayesiaanse optimalisatie (Optuna), die tientallen onafhankelijke runs vereist, bereikte GraP vergelijkbare prestaties met één enkele training en een gematigde overhead (gemiddelde 44% vertraging in doorvoer).

Zelftoezichtende Computer Vision

• Opzet: Geëvalueerd op CIFAR-10, CIFAR-100 en ImageNet-100 met behulp van het All4One-framework (combinatie van Barlow Twins-, BYOL- en NNCLR-doelen).
• Prestaties: In deze setting, waar doelen vaak goed in evenwicht zijn, bereikte GraP prestaties die vergelijkbaar waren met de standaard All4One-parameters en andere multi-task learning-methoden.
• Verliesselectie: Opmerkelijk is dat GraP een bijna nul-gewicht toekende aan het Attn-NNCLR-verlies. Het handmatig verwijderen van dit verlies leverde vergelijkbare prestaties op, wat de capaciteit van de methode bevestigt om laag-impact of redundante doelen te identificeren.

Ablatiestudies

• Labellekage: Experimenten waarbij het aandeel gelabelde data voor gewichtsafstelling werd gevarieerd, toonden robuustheid, wat suggereert dat labels een zwakke leidingsignaal bieden in plaats van directe supervisie voor de backbone.
• Normalisatie: De voorgestelde normalisatie van de samengestelde gradiënt presteerde beter dan de standaard normalisatie van de gewichtsom.
• Auxiliaire Training: Pretraining met GraP gevolgd door supervised fine-tuning presteerde beter dan gezamenlijke auxiliaire training, waarschijnlijk vanwege verminderde overfitting op beperkte gelabelde data.

Betekenis en Claims

Het artikel claimt dat op gradiënten gebaseerde verliesgewicht-optimalisatie een praktische en efficiënt alternatief is voor klassieke multi-task learning en zoekgebaseerde hyperparameter-optimalisatie. De primaire betekenis ligt in het haalbaar maken van hyperparameter-afstelling voor pretraining-scenario's op grote schaal met meerdere verliestermen, waar traditionele methoden computatief onhaalbaar zijn.

De auteurs benadrukken dat GraP geen meerdere trainingen vereist, waardoor het geschikt is voor situaties met beperkte rekenbudgetten. Hoewel de methode uitgaat van aparte heads voor elk verlies, behandelt het effectief de afwegingen tussen conflicterende doelen door de pretraining-update richting te aligneren met het downstream-doel. Het werk suggereert dat in complexe pretraining-pijplijnen automatische, online alignatie van verliesgewichten representaties kan opleveren die net zo goed, of beter, overdragen als die verkregen via dure handmatige afstelling of black-box zoekopdrachten.