Sintesi Tecnica: Una Prospettiva di Ottimizzazione Vincolata per i Transformer Unrolled

Definizione del Problema

Sebbene la letteratura recente abbia interpretato le architetture Transformer come algoritmi di discesa iterativa che risolvono problemi di ottimizzazione (ad es., Yang et al., 2022; De Weerdt et al., 2023), l'addestramento standard tramite la Minimizzazione del Rischio Empirico (ERM) spesso non riesce a imporre il comportamento di discesa monotona atteso da tali ottimizzatori. Quando addestrati con ERM, le rappresentazioni intermedie dei Transformer unrolled non necessariamente riducono la funzione di perdita in modo monotono attraverso i layer. Questa inconsistenza deriva dal fatto che l'addestramento standard non garantisce che i parametri appresi soddisfino le condizioni necessarie per la minimizzazione dell'energia, né tiene conto della natura bilevel dell'addestramento di un'architettura unrolled. Di conseguenza, questi modelli possono esibire pattern di perdita non monotoni e una ridotta robustezza rispetto alle perturbazioni degli input e ai cambiamenti di distribuzione.

Metodologia

Gli autori propongono un framework che tratta l'addestramento dei Transformer come un problema di apprendimento vincolato. Invece di minimizzare la perdita media sull'output dell'intera rete, il metodo impone che la perdita attesa diminuisca monotonicamente ad ogni layer.

1. Formulazione Vincolata

Il contributo principale è la formulazione di un problema di ottimizzazione vincolata in cui l'output di ogni layer $l$ è richiesto di ridurre la perdita attesa $f$ di un fattore $(1 - \alpha_l)$ rispetto al layer precedente $l-1$:
$$\begin{aligned} T^* = \arg\min_{T} \quad & \mathbb{E}[f(X, \Phi(X; T))] \\ \text{soggetto a} \quad & \mathbb{E}[f(X, \Phi_l(X; T))] \leq (1 - \alpha_l) \mathbb{E}[f(X, \Phi_{l-1}(X; T))], \quad \forall l \end{aligned}$$
Qui, $\Phi_l$ rappresenta l'output del layer $l$, e $\alpha_l$ è uno schema di stepsize. Questa formulazione forza il Transformer a comportarsi come un algoritmo di discesa, garantendo che le rappresentazioni intermedie migliorino progressivamente l'obiettivo.

2. Algoritmo di Addestramento Primal-Dual

Per risolvere questo problema di ottimizzazione non convesso, gli autori impiegano uno schema di addestramento primal-dual nel dominio duale. Costruiscono una funzione Lagrangiana:
$$\mathcal{L}(T, \lambda) = \mathbb{E}[f(X, \Phi(X; T))] + \sum_{l=1}^L \lambda_l \mathbb{E}[f(X, \Phi_l(X; T)) - (1 - \alpha_l)f(X, \Phi_{l-1}(X; T))]$$
La procedura di addestramento (Algoritmo 1) alterna tra:

• Aggiornamento Primal: Minimizzazione della Lagrangiana empirica rispetto ai parametri del Transformer $T$ tramite discesa del gradiente stocastica.
• Aggiornamento Duale: Massimizzazione della funzione duale rispetto ai moltiplicatori di Lagrange $\lambda$ tramite ascesa del gradiente proiettata.

Gli autori sfruttano risultati recenti che dimostrano come le reti neurali profonde esibiscano piccoli gap di dualità (Chamon et al., 2023), permettendo al problema duale di fornire soluzioni quasi ottimali e quasi ammissibili. Per garantire la stretta ammissibilità, viene introdotta una rilassazione resiliente con una variabile di slack e una penalità quadratica.

Contributi Chiave

Il documento delinea quattro contributi primari:

1. Formulazione di Apprendimento Vincolato: Gli autori formulano un problema di apprendimento in cui i layer dei Transformer sono esplicitamente vincolati a ridurre la perdita attesa monotonicamente, divergendo dall'ERM standard.
2. Algoritmo di Addestramento Duale e Garanzie di Convergenza: Sviluppano un algoritmo di addestramento primal-dual e forniscono garanzie teoriche (Teorema 2) secondo cui il Transformer unrolled vincolato converge asintoticamente al valore ottimale della perdita statistica. Stabiliscono inoltre garanzie di generalizzazione Out-of-Distribution (OOD) (Teorema 4), provando che le proprietà di discesa sono mantenute sotto cambiamenti di distribuzione.
3. Robustezza alle Perturbazioni degli Embedding: Nei compiti di classificazione del testo (IMDb, MNLI), gli autori dimostrano che i modelli vincolati mantengono un'accuratezza superiore sotto perturbazioni gaussiane, uniformi e discrete agli embedding dei token rispetto ai baseline non vincolati.
4. Robustezza nella Denoisazione Video: Applicato alla denoisazione video (CUHK Avenue, UCSD, ShanghaiTech), il metodo produce modelli che sono più robusti ai livelli di rumore OOD, misurati tramite Root Mean Squared Error (RMSE), pur preservando le prestazioni in-distribution.

Risultati Sperimentali

Gli autori hanno valutato il loro metodo sia su Transformer pre-addestrati convenzionali (DistilBERT, RoBERTa, Llama 3.1) che su architetture unrolled (Deep Unrolled Sequential Transformer - DUST, Unrolled Transformer - UT).

• Discesa Monotona: I modelli vincolati hanno esibito diminuzioni della perdita per layer fluide e monotone, mentre i baseline non vincolati hanno mostrato pattern non monotoni con picchi (ad es., un picco al layer 19 in RoBERTa).
• Robustezza OOD:
  • Classificazione del Testo: All'aumentare dei livelli di perturbazione ($\gamma$), i modelli vincolati sono degradati in modo più graduale. Ad esempio, su RoBERTa con IMDb, il modello vincolato ha mantenuto un'accuratezza significativamente più alta sotto alti livelli di rumore rispetto alla versione non vincolata.
  • Denoisazione Video: I modelli vincolati hanno costantemente ottenuto un RMSE inferiore sotto alti livelli di perturbazione del test ($\gamma \geq 1.0$) attraverso vari dataset e architetture.
  • Fine-tuning di LLM: Su Llama 3.1 8B, il modello vincolato ha mantenuto le prestazioni in-distribution (50% di win rate su AlpacaEval senza rumore) pur superando significativamente il modello non vincolato sotto rumore (69.93% di win rate con $\gamma=1.5$).
• Trade-off: Il metodo generalmente preserva le prestazioni in-distribution con una degradazione trascurabile (spesso < 0.3%), suggerendo che l'imposizione di vincoli di discesa non sacrifica l'accuratezza standard in favore della robustezza.

Significato e Rivendicazioni

Il documento sostiene che l'imposizione di vincoli di discesa allinea la dinamica interna dei Transformer con le proprietà di stabilità degli algoritmi di ottimizzazione classici. Questo allineamento fornisce una base teorica ed empirica per una migliore robustezza contro le perturbazioni degli input e i cambiamenti di distribuzione.

Gli autori posizionano questo lavoro come un ponte tra l'unrolling algoritmico (dove le reti risolvono problemi di ottimizzazione) e l'apprendimento vincolato. Argomentano che, addestrando esplicitamente i Transformer per la discesa, si possono ottenere modelli intrinsecamente più robusti al rumore, una proprietà particolarmente preziosa per applicazioni che coinvolgono la privacy differenziale locale, dove le perturbazioni dell'input sono necessarie per la privacy dell'utente. Il lavoro suggerisce che la prospettiva "unrolled" non sia solo un'interpretazione teorica, ma un obiettivo di addestramento pratico che produce miglioramenti tangibili in generalizzazione e robustezza senza richiedere modifiche architettoniche.