Riepilogo Tecnico: Quando le Perdite si Allineano: Ponderazione Composita delle Perdite Basata sul Gradiente per il Pre-addestramento Efficiente

Formulazione del Problema

I moderni sistemi di deep learning adottano frequentemente un paradigma a due stadi: pre-addestramento su dati su larga scala, spesso non etichettati, seguito da adattamento a compiti downstream. In molte pipeline di pre-addestramento, la funzione obiettivo non è una singola perdita, ma una combinazione ponderata di più termini di perdita (ad esempio, previsione del token successivo, perdite contrastive o obiettivi di regressione). I pesi relativi di questi termini sono iperparametri critici che influenzano significativamente la qualità delle rappresentazioni apprese e le prestazioni downstream.

Attualmente, la selezione di questi pesi si basa su regolazione manuale, ricerca casuale o ottimizzazione bayesiana. Questi metodi sono computazionalmente proibitivi per il pre-addestramento su larga scala perché richiedono numerosi training indipendenti. Inoltre, valori predefiniti ingenui, come la ponderazione uniforme, spesso non tengono conto delle diverse contribuzioni e dei conflitti tra i vari termini di perdita. Il lavoro affronta la necessità di un metodo efficiente per regolare i pesi lineari delle perdite negli obiettivi di pre-addestramento compositi, al fine di massimizzare le prestazioni sui compiti downstream senza sostenere il costo di ripetuti training completi.

Metodologia: Pre-addestramento Allineato al Gradiente (GraP)

Gli autori propongono Gradient-aligned Pretraining (GraP), un metodo di ottimizzazione bilevel basato sul gradiente che apprende i pesi delle perdite online. L'idea centrale è allineare il gradiente di pre-addestramento composito con il gradiente di un obiettivo downstream.

1. Formulazione dell'Ottimizzazione Bilevel

Il problema è formulato come un compito di ottimizzazione bilevel in cui il ciclo esterno ottimizza i pesi delle perdite $w$ per minimizzare la perdita downstream $L_{down}$, e il ciclo interno ottimizza i parametri del modello $\theta$ per minimizzare la perdita di pre-addestramento composita $L_{pre}(w)$.
$$\min_{w} L_{down}(\theta^*(w)) \quad \text{soggetto a} \quad \theta^*(w) = \arg\min_{\theta} \sum_{k=1}^K w_k L_k(\theta)$$

2. Regola di Aggiornamento Basata sul Gradiente

Utilizzando un approccio di backpropagation troncato (SGD a singolo passo), gli autori derivano una regola di aggiornamento per i pesi. Il gradiente della perdita downstream rispetto ai pesi è proporzionale al prodotto scalare tra il gradiente downstream e i singoli gradienti di pre-addestramento:
$$\frac{\partial L_{down}}{\partial w_k} \propto -\eta \, g_{down}^\top g_k$$
dove $g_{down}$ è il gradiente downstream e $g_k$ è il gradiente della $k$-esima perdita di pre-addestramento. L'obiettivo è massimizzare l'allineamento tra il gradiente composito $\sum w_k g_k$ e il gradiente downstream.

3. Approssimazione nello Spazio degli Embedding

Un'implementazione ingenua richiederebbe il calcolo di $K$ passaggi di backpropagation completi attraverso il modello per ottenere $g_k$. Per ridurre il costo computazionale, GraP calcola i gradienti rispetto all'embedding condiviso del backbone $z$ piuttosto che rispetto all'intero spazio dei parametri.
$$\tilde{g}_k = \nabla_z L_k(h_k(z))$$
Poiché le perdite sono applicate tramite testine separate sopra lo stesso embedding, il calcolo di questi gradienti è significativamente più economico rispetto ai passaggi di backpropagation completi. L'obiettivo di ottimizzazione diventa massimizzare l'allineamento in questo spazio di rappresentazione condiviso:
$$\arg\max_{w} \frac{w^\top \tilde{G} \tilde{g}_{down}}{\|w^\top \tilde{G}\|}$$

4. Strategia di Normalizzazione

Il lavoro si oppone alla normalizzazione dei pesi stessi (ad esempio, $\sum w_k = 1$) o della loro norma $L_2$, poiché queste possono distorcere la direzione ottimale del gradiente composito. Invece, GraP impone un vincolo sulla norma $L_2$ del gradiente composito stesso ($\|\sum w_k g_k\| = 1$). Ciò garantisce che l'ottimizzazione si concentri sulla direzione del vettore di aggiornamento piuttosto che sulla sua magnitudine, recuperando la direzione di allineamento ottimale sotto l'obiettivo locale.

5. Efficienza dell'Algoritmo

L'algoritmo risultante richiede solo un passaggio in avanti e un passaggio all'indietro attraverso il backbone per iterazione. I gradienti calcolati per le testine vengono riutilizzati sia per aggiornare i pesi delle perdite sia per aggiornare i parametri del backbone. Ciò riduce l'overhead della regolazione degli iperparametri a circa il 30% sopra un singolo training, rispetto ai costi di ordini di grandezza dell'ottimizzazione bayesiana.

Contributi Chiave

1. Formulazione del Problema: Gli autori formulano la selezione dei pesi delle perdite di pre-addestramento come un problema di ottimizzazione bilevel che mira esplicitamente alle prestazioni downstream.
2. Insight sulla Normalizzazione: Analizzano le strategie di normalizzazione e dimostrano che vincolare la norma del vettore di aggiornamento (gradiente composito) è superiore al vincolare i pesi stessi, preservando la direzione di allineamento ottimale.
3. Algoritmo Efficiente: Propongono GraP, un algoritmo con complessità $O(1)$ riguardo al numero di passaggi avanti-indietro rispetto al numero di iperparametri. Sfrutta l'allineamento dei gradienti nello spazio di rappresentazione condiviso per evitare molteplici passaggi di backpropagation.
4. Efficienza Empirica: Il metodo raggiunge prestazioni pari o superiori rispetto all'ottimizzazione standard degli iperparametri (ad esempio, ricerca bayesiana) essendo quasi 50$\times$ più efficiente, consentendo la regolazione in contesti dove i metodi tradizionali sono computazionalmente non fattibili.
5. Identificazione della Ridondanza: L'approccio dimostra la capacità di identificare e ridurre il peso dei termini di perdita ridondanti (ad esempio, assegnando un peso vicino allo zero a specifiche perdite contrastive nei compiti di visione), agendo efficacemente come un meccanismo automatico di selezione delle perdite.

Risultati Sperimentali

Il metodo è stato valutato su due domini: modellazione di sequenze di eventi e visione artificiale auto-supervisionata.

Modellazione di Sequenze di Eventi

• Dataset: Churn, AgePred, AlfaBattle, MIMIC-III e Taobao.
• Configurazione: Il pre-addestramento ha coinvolto un obiettivo composito di previsione del token successivo (NTP) su più campi e una perdita contrastiva.
• Prestazioni: GraP ha eguagliato o migliorato le baseline attentamente regolate (inclusi GradNorm, DWA, MGDA e PCGrad) e i pesi ottimizzati con Optuna. Ha costantemente ottenuto i migliori risultati sui dataset Churn e MIMIC-III attraverso backbone RNN, Transformer e HT-Transformer.
• Efficienza: Rispetto all'ottimizzazione bayesiana (Optuna), che richiede dozzine di esecuzioni indipendenti, GraP ha raggiunto prestazioni simili con un singolo training e un overhead moderato (rallentamento medio del 44% nel throughput).

Visione Artificiale Auto-Supervisionata

• Configurazione: Valutato su CIFAR-10, CIFAR-100 e ImageNet-100 utilizzando il framework All4One (che combina gli obiettivi Barlow Twins, BYOL e NNCLR).
• Prestazioni: In questo contesto, dove gli obiettivi sono spesso ben bilanciati, GraP ha raggiunto prestazioni comparabili ai parametri predefiniti di All4One e ad altri metodi di apprendimento multi-task.
• Selezione delle Perdite: In modo notevole, GraP ha assegnato un peso vicino allo zero alla perdita Attn-NNCLR. La rimozione manuale di questa perdita ha prodotto prestazioni simili, validando la capacità del metodo di identificare obiettivi a basso impatto o ridondanti.

Studi di Ablazione

• Fuga di Etichette (Label Leakage): Esperimenti che variavano la frazione di dati etichettati utilizzati per la regolazione dei pesi hanno mostrato robustezza, suggerendo che le etichette forniscono un segnale di guida debole piuttosto che una supervisione diretta per il backbone.
• Normalizzazione: La normalizzazione del gradiente composito proposta ha superato la normalizzazione standard della somma dei pesi.
• Training Ausiliario: Il pre-addestramento con GraP seguito da fine-tuning supervisionato ha superato il training ausiliario congiunto, probabilmente a causa di un ridotto overfitting su dati etichettati limitati.

Significato e Affermazioni

Il lavoro afferma che l'ottimizzazione dei pesi delle perdite basata sul gradiente è un'alternativa pratica ed efficiente all'apprendimento multi-task classico e all'ottimizzazione degli iperparametri basata sulla ricerca. Il suo significato principale risiede nel rendere fattibile la regolazione degli iperparametri per scenari di pre-addestramento su larga scala con più termini di perdita, dove i metodi tradizionali sono computazionalmente proibitivi.

Gli autori sottolineano che GraP non richiede molteplici training, rendendolo adatto a contesti con budget computazionali limitati. Sebbene il metodo presupponga testine separate per ogni perdita, gestisce efficacemente i compromessi tra obiettivi conflittuali allineando la direzione di aggiornamento del pre-addestramento con l'obiettivo downstream. Il lavoro suggerisce che in pipeline di pre-addestramento complesse, l'allineamento automatico e online dei pesi delle perdite può produrre rappresentazioni che trasferiscono tanto bene, o meglio, di quelle ottenute tramite regolazione manuale costosa o ricerca black-box.