Technische Samenvatting: Een Perspectief van Beperkte Optimalisatie op Unrolled Transformers

Probleemstelling

Hoewel recente literatuur de interpretatie geeft dat Transformer-architecturen iteratieve descent-algoritmen zijn die optimalisatieproblemen oplossen (bijv. Yang et al., 2022; De Weerdt et al., 2023), faalt standaardtraining via Empirical Risk Minimization (ERM) er vaak in om het verwachte monotone descent-gedrag af te dwingen. Wanneer getraind met ERM, vertonen de intermediaire representaties van unrolled Transformers niet noodzakelijkerwijs een monotone afname van de verliesfunctie over de lagen heen. Deze inconsistentie ontstaat omdat standaardtraining niet garandeert dat geleerde parameters voldoen aan de voorwaarden die vereist zijn voor energie-minimalisatie, noch houdt het rekening met de bilevel-aard van unrolled architectuurtraining. Bijgevolg kunnen deze modellen niet-monotone verliespatronen en een verminderde robuustheid tegen inputperturbaties en distributieverschuivingen vertonen.

Methodologie

De auteurs stellen een framework voor dat Transformer-training behandelt als een beperkt leerprobleem (constrained learning problem). In plaats van alleen de gemiddelde loss over de volledige netwerkoutput te minimaliseren, dwingt de methode af dat de verwachte loss bij elke laag monotoon afneemt.

1. Beperkte Formulering (Constrained Formulation)

De kernbijdrage is de formulering van een beperkt optimalisatieprobleem waarbij de output van elke laag $l$ de verwachte loss $f$ met een factor $(1 - \alpha_l)$ moet verminderen ten opzichte van de vorige laag $l-1$:
$$\begin{aligned} T^* = \arg\min_{T} \quad & \mathbb{E}[f(X, \Phi(X; T))] \\ \text{onder de voorwaarde dat} \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}$$
Hierbij stelt $\Phi_l$ de output van laag $l$ voor, en $\alpha_l$ is een stepsize-schema. Deze formulering dwingt de Transformer om zich te gedragen als een descent-algoritme, waardoor wordt gegarandeerd dat intermediaire representaties de objectieven progressief verbeteren.

2. Primal-Dual Trainingsalgoritme

Om dit niet-convexe beperkte probleem op te lossen, maken de auteurs gebruik van een primal-dual trainingschema in het duale domein. Zij construeren een Lagrange-functie:
$$\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))]$$
De trainingsprocedure (Algoritme 1) wisselt af tussen:

• Primal Update: Het minimaliseren van de empirische Lagrange-functie met betrekking tot de Transformer-parameters $T$ middels stochastic gradient descent.
• Dual Update: Het maximaliseren van de duale functie met betrekking tot de Lagrange-multiplicatoren $\lambda$ middels projected gradient ascent.

De auteurs maken gebruik van recente resultaten die aantonen dat diepe neurale netwerken kleine dualiteitsgaten vertonen (Chamon et al., 2023), waardoor het duale probleem bijna optimale en bijna haalbare oplossingen oplevert. Om strikte haalbaarheid te garanderen, is een veerkrachtige relaxatie met een slack-variabele en een kwadratische penalty geïntroduceerd.

Belangrijkste Bijdragen

Het artikel schetst vier primaire bijdragen:

1. Beperkte Leerformulering: De auteurs formuleren een leerprobleem waarbij de Transformer-lagen expliciet worden beperkt om de verwachte loss monotoon te reduceren, wat afwijkt van standaard ERM.
2. Dual Trainingsalgoritme & Convergentiegaranties: Zij ontwikkelen een primal-dual trainingsalgoritme en bieden theoretische garanties (Theorem 2) dat de beperkte unrolled Transformer asymptotisch convergeert naar de optimale waarde van de statistische loss. Zij stellen verder Out-of-Distribution (OOD) generalisatiegaranties vast (Theorem 4), waarbij bewezen wordt dat de descent-eigenschappen behouden blijven onder distributieverschuivingen.
3. Robuustheid tegen Geperturbeerde Embeddings: Bij tekstclassificatietaken (IMDb, MNLI) demonstreren de auteurs dat de beperkte modellen een hogere nauwkeurigheid behouden onder Gaussische, uniforme en discrete perturbaties van token-embeddings vergeleken met onbeperkte baselines.
4. Robuustheid in Video Denoising: Toegepast op video denoising (CUHK Avenue, UCSD, ShanghaiTech), levert de methode modellen op die robuuster zijn tegen OOD-ruinniveaus, gemeten via Root Mean Squared Error (RMSE), terwijl de in-distribution prestaties behouden blijven.

Experimentele Resultaten

De auteurs evalueerden hun methode op zowel conventionele pre-trained Transformers (DistilBERT, RoBERTa, Llama 3.1) als unrolled architecturen (Deep Unfolded Sequential Transformer - DUST, Unrolled Transformer - UT).

• Monotone Descent: De beperkte modellen vertoonden een vloeiende, monotone afname van de laag-specifieke loss, terwijl onbeperkte baselines niet-monotone patronen met pieken vertoonden (bijv. een piek bij laag 19 in RoBERTa).
• OOD Robuustheid:
  • Tekstclassificatie: Naarmate de perturbatieniveaus ($\gamma$) toenamen, degradeerden de beperkte modellen meer geleidelijk. Zo behield het beperkte model op RoBERTa met IMDb een significant hogere nauwkeurigheid onder hoge ruisniveaus vergeleken met de onbeperkte versie.
  • Video Denoising: De beperkte modellen behaalden consistent een lagere RMSE onder hoge test-perturbatieniveaus ($\gamma \geq 1.0$) over diverse datasets en architecturen.
  • LLM Fine-tuning: Op Llama 3.1 8B behield het beperkte model de in-distribution prestaties (50% win rate op AlpacaEval zonder ruis), terwijl het de onbeperkte versie aanzienlijk overtrof onder ruis (69.93% win rate met $\gamma=1.5$).
• Trade-offs: De methode behield over het algemeen de in-distribution prestaties met verwaarloosbare degradatie (vaak < 0,3%), wat suggereert dat het afdwingen van descent-beperkingen de standaard nauwkeurigheid niet opoffert voor robuustheid.

Betekenis en Claims

Het artikel claimt dat het afdwingen van descent-beperkingen de interne dynamiek van Transformers afstemt op de stabiliteitseigenschappen van klassieke optimalisatiealgoritmen. Deze afstemming biedt een theoretische en empirische basis voor verbeterde robuustheid tegen inputperturbaties en distributieverschuivingen.

De auteurs positioneren dit werk als een brug tussen algoritmische unrolling (waarbij netwerken optimalisatieproblemen oplossen) en constrained learning. Zij argumenteren dat men door Transformers expliciet te trainen om te dalen (descend), modellen kan bereiken die inherent robuuster zijn tegen ruis, een eigenschap die bijzonder waardevol is voor toepassingen die betrokken zijn bij local differential privacy, waar inputperturbaties noodzakelijk zijn voor gebruikersprivacy. Het werk suggereert dat het "unrolled" perspectief niet slechts een theoretische interpretatie is, maar een praktische trainingsdoelstelling die tastbare verbeteringen in generalisatie en robuustheid oplevert zonder architecturale wijzigingen te vereisen.