Technische Zusammenfassung: Wenn Verluste übereinstimmen: Gradientenbasierte Gewichtung zusammengesetzter Verluste für effizientes Pretraining

Problemformulierung

Moderne Deep-Learning-Systeme verwenden häufig ein zweistufiges Paradigma: Pretraining auf großen, oft ungelabelten Daten, gefolgt von der Anpassung an nachgelagerte Aufgaben. In vielen Pretraining-Pipelines ist die Zielfunktion nicht ein einzelner Verlust, sondern eine gewichtete Kombination mehrerer Verlustterme (z. B. Next-Token-Vorhersage, kontrastive Verluste oder Regressionsziele). Die relativen Gewichte dieser Terme sind kritische Hyperparameter, die die Qualität der gelernten Repräsentationen und die Leistung in nachgelagerten Aufgaben erheblich beeinflussen.

Derzeit erfolgt die Auswahl dieser Gewichte durch manuelle Abstimmung, zufällige Suche oder bayessche Optimierung. Diese Methoden sind für groß angelegtes Pretraining rechnerisch prohibitiv, da sie zahlreiche unabhängige Trainingsläufe erfordern. Darüber hinaus versagen naive Standardwerte, wie eine gleichmäßige Gewichtung, oft darin, die unterschiedlichen Beiträge und Konflikte zwischen verschiedenen Verlusttermen zu berücksichtigen. Die Arbeit adressiert die Notwendigkeit einer effizienten Methode zur Abstimmung linearer Verlustgewichte in zusammengesetzten Pretraining-Zielen, um die Leistung nachgelagerter Aufgaben zu maximieren, ohne die Kosten wiederholter vollständiger Trainingsläufe in Kauf nehmen zu müssen.

Methodik: Gradientenabgestimmtes Pretraining (GraP)

Die Autoren schlagen Gradient-aligned Pretraining (GraP) vor, eine gradientenbasierte Bilevel-Optimierungsmethode, die Verlustgewichte online lernt. Die Kernidee besteht darin, den zusammengesetzten Pretraining-Gradienten mit dem Gradienten eines nachgelagerten Ziels auszurichten.

1. Formulierung der Bilevel-Optimierung

Das Problem wird als Bilevel-Optimierungsaufgabe formuliert, wobei die äußere Schleife die Verlustgewichte $w$ optimiert, um den nachgelagerten Verlust $L_{down}$ zu minimieren, und die innere Schleife die Modellparameter $\theta$ optimiert, um den zusammengesetzten Pretraining-Verlust $L_{pre}(w)$ zu minimieren.
$$\min_{w} L_{down}(\theta^*(w)) \quad \text{unter der Nebenbedingung} \quad \theta^*(w) = \arg\min_{\theta} \sum_{k=1}^K w_k L_k(\theta)$$

2. Gradientenbasierte Aktualisierungsregel

Mittels eines abgeschnittenen Backpropagation-Ansatzes (Single-Step-SGD) leiten die Autoren eine Aktualisierungsregel für die Gewichte her. Der Gradient des nachgelagerten Verlusts bezüglich der Gewichte ist proportional zum Skalarprodukt zwischen dem nachgelagerten Gradienten und den einzelnen Pretraining-Gradienten:
$$\frac{\partial L_{down}}{\partial w_k} \propto -\eta \, g_{down}^\top g_k$$
wobei $g_{down}$ der nachgelagerte Gradient und $g_k$ der Gradient des $k$-ten Pretraining-Verlusts ist. Das Ziel besteht darin, die Ausrichtung zwischen dem zusammengesetzten Gradienten $\sum w_k g_k$ und dem nachgelagerten Gradienten zu maximieren.

3. Approximation im Embedding-Raum

Eine naive Implementierung würde erfordern, $K$ vollständige Rückwärtsdurchläufe durch das Modell zu berechnen, um $g_k$ zu erhalten. Um die Rechenkosten zu senken, berechnet GraP Gradienten bezüglich des gemeinsamen Backbone-Embeddings $z$ anstatt des vollständigen Parameterraums.
$$\tilde{g}_k = \nabla_z L_k(h_k(z))$$
Da die Verluste über separate Köpfe auf demselben Embedding angewendet werden, ist die Berechnung dieser Gradienten deutlich günstiger als vollständige Rückwärtsdurchläufe. Das Optimierungsziel wird zur Maximierung der Ausrichtung in diesem gemeinsamen Repräsentationsraum:
$$\arg\max_{w} \frac{w^\top \tilde{G} \tilde{g}_{down}}{\|w^\top \tilde{G}\|}$$

4. Normalisierungsstrategie

Die Arbeit argumentiert gegen eine Normalisierung der Gewichte selbst (z. B. $\sum w_k = 1$) oder ihrer $L_2$-Norm, da diese die optimale Richtung des zusammengesetzten Gradienten verzerren können. Stattdessen setzt GraP eine Nebenbedingung für die $L_2$-Norm des zusammengesetzten Gradienten selbst ($\|\sum w_k g_k\| = 1$). Dies stellt sicher, dass sich die Optimierung auf die Richtung des Aktualisierungsvektors konzentriert und nicht auf seine Größe, wodurch die optimale Ausrichtung unter dem lokalen Ziel wiederhergestellt wird.

5. Algorithmische Effizienz

Der resultierende Algorithmus erfordert nur einen Vorwärtsdurchlauf und einen Rückwärtsdurchlauf durch den Backbone pro Iteration. Die für die Köpfe berechneten Gradienten werden sowohl für die Aktualisierung der Verlustgewichte als auch für die Aktualisierung der Backbone-Parameter wiederverwendet. Dies reduziert den Overhead der Hyperparameter-Abstimmung auf etwa 30 % über einem einzelnen Trainingslauf, im Vergleich zu den um Größenordnungen höheren Kosten der bayesschen Optimierung.

Hauptbeiträge

1. Problemformulierung: Die Autoren formulieren die Auswahl der Pretraining-Verlustgewichte als Bilevel-Optimierungsproblem, das explizit auf die Leistung nachgelagerter Aufgaben abzielt.
2. Erkenntnis zur Normalisierung: Sie analysieren Normalisierungsstrategien und zeigen, dass die Einschränkung der Norm des Aktualisierungsvektors (zusammengesetzter Gradient) der Einschränkung der Gewichte selbst überlegen ist, wodurch die optimale Ausrichtung erhalten bleibt.
3. Effizienter Algorithmus: Sie schlagen GraP vor, einen Algorithmus mit $O(1)$-Komplexität bezüglich der Anzahl der Vorwärts-Rückwärts-Schritte im Verhältnis zur Anzahl der Hyperparameter. Er nutzt Gradientenausrichtung im gemeinsamen Repräsentationsraum, um mehrere Rückwärtsdurchläufe zu vermeiden.
4. Empirische Effizienz: Die Methode erzielt eine Leistung, die der standardmäßigen Hyperparameter-Optimierung (z. B. bayessche Suche) entspricht oder diese übertrifft, ist dabei jedoch nahezu 50$\times$ effizienter, was eine Abstimmung in Szenarien ermöglicht, in denen traditionelle Methoden rechnerisch nicht machbar sind.
5. Identifikation von Redundanzen: Der Ansatz demonstriert die Fähigkeit, redundante Verlustterme zu identifizieren und herunterzuwiegen (z. B. Zuweisung eines nahezu null Gewichts für bestimmte kontrastive Verluste in Vision-Aufgaben) und wirkt effektiv als automatischer Verlustauswahlmechanismus.

Experimentelle Ergebnisse

Die Methode wurde in zwei Domänen evaluiert: Modellierung von Ereignissequenzen und selbstüberwachtes Computersehen.

Modellierung von Ereignissequenzen

• Datensätze: Churn, AgePred, AlfaBattle, MIMIC-III und Taobao.
• Aufbau: Das Pretraining umfasste ein zusammengesetztes Ziel aus Next-Token-Vorhersage (NTP) über mehrere Felder hinweg und einem kontrastiven Verlust.
• Leistung: GraP entsprach oder verbesserte sorgfältig abgestimmte Baselines (einschließlich GradNorm, DWA, MGDA und PCGrad) sowie Optuna-abgestimmte Gewichte. Es erzielte konsistent die besten Ergebnisse auf den Datensätzen Churn und MIMIC-III über RNN-, Transformer- und HT-Transformer-Backbones hinweg.
• Effizienz: Im Vergleich zur bayesschen Optimierung (Optuna), die Dutzende unabhängiger Läufe erfordert, erreichte GraP mit einem einzigen Trainingslauf und moderatem Overhead (durchschnittliche 44 % Verlangsamung des Durchsatzes) eine ähnliche Leistung.

Selbstüberwachtes Computersehen

• Aufbau: Evaluiert auf CIFAR-10, CIFAR-100 und ImageNet-100 unter Verwendung des All4One-Frameworks (Kombination von Barlow Twins-, BYOL- und NNCLR-Zielen).
• Leistung: In diesem Setting, in dem Ziele oft gut ausbalanciert sind, erreichte GraP eine Leistung, die den Standard-All4One-Parametern und anderen Multi-Task-Learning-Methoden entsprach.
• Verlustauswahl: Bemerkenswerterweise wies GraP dem Attn-NNCLR-Verlust ein nahezu null Gewicht zu. Das manuelle Entfernen dieses Verlusts ergab eine ähnliche Leistung, was die Fähigkeit der Methode validiert, geringe Auswirkungen oder redundante Ziele zu identifizieren.

Ablationsstudien

• Label-Leckage: Experimente, bei denen der Anteil der für die Gewichtsabstimmung verwendeten gelabelten Daten variiert wurde, zeigten Robustheit, was darauf hindeutet, dass Labels einen schwachen Guidance-Signal liefern und keine direkte Überwachung für den Backbone darstellen.
• Normalisierung: Die vorgeschlagene Normalisierung des zusammengesetzten Gradienten übertraf die Standardnormalisierung der Gewichtssumme.
• Auxiliäres Training: Pretraining mit GraP, gefolgt von überwachtem Fine-Tuning, übertraf gemeinsames auxiliäres Training, wahrscheinlich aufgrund von reduziertem Overfitting auf begrenzten gelabelten Daten.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass die gradientenbasierte Optimierung von Verlustgewichten eine praktische und effiziente Alternative zu klassischem Multi-Task-Learning und suchbasierten Hyperparameter-Optimierungen darstellt. Ihre primäre Bedeutung liegt darin, Hyperparameter-Abstimmung für groß angelegte Pretraining-Szenarien mit mehreren Verlusttermen machbar zu machen, wo traditionelle Methoden rechnerisch prohibitiv sind.

Die Autoren betonen, dass GraP keine mehreren Trainingsläufe erfordert, was es für Umgebungen mit begrenzten Rechenbudgets geeignet macht. Obwohl die Methode separate Köpfe für jeden Verlust voraussetzt, bewältigt sie die Kompromisse zwischen konfligierenden Zielen effektiv, indem sie die Pretraining-Aktualisierungsrichtung mit dem nachgelagerten Ziel ausrichtet. Die Arbeit legt nahe, dass in komplexen Pretraining-Pipelines eine automatische, online erfolgende Ausrichtung der Verlustgewichte Repräsentationen hervorbringen kann, die mindestens ebenso gut oder besser transferieren als solche, die durch teure manuelle Abstimmung oder Black-Box-Suche erzielt werden.