Robust and Fast Training via Per-Sample Clipping

Technische Samenvatting: Robuuste en Snelle Training via Per-Sample Clipping

Probleemstelling
Het artikel behandelt het stochastische optimalisatieprobleem $\min_{x \in \mathbb{R}^d} F(x)$, waarbij $F(x) = \mathbb{E}_{\xi \sim \mathcal{D}}[f(x, \xi)]$. Hoewel Stochastic Gradient Descent (SGD) de standaard oplosser is, rusten de convergentiegaranties er meestal op dat de gradiëntruis een begrensde variantie heeft (lichte staartverdelingen). Echter, recent empirisch en theoretisch werk wijst uit dat in veel toepassingen voor machine learning de gradiëntruis zware staarten kan hebben, met eindige $p$-de momenten voor $p \in (1, 2]$, maar zonder begrensde variantie. Onder dergelijke omstandigheden kan standaard SGD slecht presteren. Bestaande robuuste alternatieven, zoals standaard gradiënt clipping (Clip-SGD) en genormaliseerde SGD, falen of vereisen specifieke, vaak onpraktische parameterinstellingen (bijv. afhankelijkheid van de initialisatiekloof $\Delta_1$ of specifieke logaritmische factoren van de faalkans $\delta$).

Methodologie
De auteurs stellen Per-Sample Clipped SGD (PS-Clip-SGD) voor, een robuuste gradiëntschatter die clipping toepast op individuele monsters binnen een mini-batch vóór aggregatie, in plaats van de geaggregeerde batchgradiënt te clippen.

Voor een batch van grootte $n$ met monsters $\xi^{(1)}_t, \dots, \xi^{(n)}_t$ wordt de gradiëntschatter $G(x_t, \alpha, \beta)$ gedefinieerd als:
$$G(x_t, \alpha, \beta) = \frac{1}{n} \sum_{k=1}^n \gamma^{(k)}_t \nabla f(x_t, \xi^{(k)}_t)$$
waarbij de per-sample clippingfactor $\gamma^{(k)}_t := \min\left\{1, \frac{\alpha k^{1/\beta}}{|\nabla f(x_t, \xi^{(k)}_t)|}\right\}$ is. De parameters $\alpha$ en $\beta$ controleren de clipping-drempel, waarbij $\alpha$ gerelateerd is aan de momentenlimiet $\sigma$ en $\beta$ aan de staartindex $p$. De update-regel volgt standaard SGD: $x_{t+1} = x_t - \eta_t G(x_t, \alpha, \beta)$.

Belangrijkste Bijdragen

1. Optimale Convergentie in Verwachting: De auteurs bewijzen dat PS-Clip-SGD een optimale steekproefcomplexiteit bereikt voor niet-convexe optimalisatie onder heavy-tailed ruis. Specifiek bereikt de methode een convergentiesnelheid van $O(1/T)$ voor de gemiddelde gekwadrateerde gradiëntnorm, wat overeenkomt met de ondergrenzen vastgesteld door Zhang et al. (2020). Dit verbetert de resultaten voor Clip-SGD en genormaliseerde SGD, met name omdat de bovengrens geldt voor de som van de gekwadrateerde gradiëntnormen (een sterker criterium dan de som van de normen die in sommige eerdere werken wordt gebruikt) zonder dat de kennis van de initialisatiekloof $\Delta_1$ nodig is voor de stapgrootte.
2. Convergentie met Hoge Waarschijnlijkheid: Het artikel stelt convergentiegaranties met hoge waarschijnlijkheid vast die de convergentie in verwachting matchen tot polylogaritmische factoren in de faalkans $\delta$ en het aantal stappen $T$. Merkwaardig genoeg is de afhankelijkheid van $\log(1/\delta)$ verbeterd ten opzichte van de state-of-the-art Clip-SGD bounds, en vereist de methode niet dat stapgroottes of clipping-drempels expliciet worden afgestemd als functies van $\delta$.
3. Inzicht in Gradiëntaccumulatie: De auteurs dagen de gangbare praktijk uit om gradiënt clipping alleen toe te passen na alle accumulatiestappen in gradiëntaccumulatie. Ze stellen voor en valideren empirisch dat het toepassen van clipping op het niveau van de mini-batch (Mini-Batch Clipping) tijdens de accumulatie de trainingsprestaties kan verbeteren met vrijwel geen extra computationele kosten, in tegenstelling tot de standaard consensus dat dergelijke tussenliggende clipping bias introduceert.

Resultaten

• Theoretisch: Theoretische analyse demonstreert dat PS-Clip-SGD sterke convergentiegaranties biedt onder aannames van $L$-gladheid en eindige $p$-de momenten. De afgeleide bounds blijken nauwer te zijn dan die voor genormaliseerde SGD en Clip-SGD in termen van de afhankelijkheid van $\delta$ en de kracht van de convergentiemaat (gekwadrateerde normen versus normen).
• Kwadratische Optimalisatie: In experimenten waarbij een kwadratische functie met Pareto-verdeelde ruis werd geminimaliseerd, presteerde PS-Clip-s SGD aanzienlijk beter dan zowel Clip-SGD als genormaliseerde SGD over diverse ruisregimes ($p \in \{1.2, 1.5, 1.8\}$), zelfs met ongetunede hyperparameters. Experimenten met hoge waarschijnlijkheid bevestigden de theoretische sublineaire afhankelijkheid van $\log(1/\delta)$.
• Beeldclassificatie (AlexNet op CIFAR-100): Bij het trainen van AlexNet met momentum bereikte PS-Clip-SGD een hogere validatienauwkeurigheid (64,84%) vergeleken met vanilla SGD met momentum (58,90%) en Clip-SGD met momentum (58,94%). Deze verbetering bleef bestaan zelfs na rekening te houden met de $\approx 1,3\times$ computationele overhead geïntroduceerd door per-sample gradiëntnormberekeningen. De auteurs schrijven dit toe aan de effectieve controle van een klein aantal uitschieter-gradiënten met zeer grote normen, die standaard batch-level clipping niet kan detecteren.
• Taalmodellering (GPT-2): In een gradiëntaccumulatie-setting bij het trainen van een 124M parameter GPT-2 model, presteerde Mini-Batch Clipping (clipping per accumulatiestap) beter dan standaard Clip-SGD (clipping pas na volledige accumulatie) over meerdere leersnelheden, ondanks de heersende opvatting dat tussenliggende clipping nadelig is.

Betekenis en Claims
Het artikel claimt het eerste werk te zijn dat per-sample clipping specif으로 voorstellen als een methode om de trainingsprestaties te verbeteren in aanwezigheid van heavy-tailed ruis, verschillend van het gevestigde gebruik in differential privacy. De significantie ligt in het bieden van een methode die:

1. Optimale theoretische convergentiesnelheden bereikt voor niet-convexe heavy-tailed problemen.
2. Robuustheid biedt tegen hyperparameterkeuzes, en goed presteert zelfs met onverfijnde instellingen.
3. Aantoont dat per-sample clipping praktische prestatiewinsten kan opleveren voor deep learning taken (zoals AlexNet) die de computationele kosten overtreffen.
4. Een paradigmaverschuiving suggereert in gradiëntaccumulatiestrategieën, door te laten zien dat tussenliggende clipping gunstig kan zijn in plaats van schadelijk.

De auteurs erkennen beperkingen, waarbij zij opmerken dat hun theoretische analyse kennis vereist van de momentenlimiet $\sigma$ en de staartindex $p$ om optimale parameters te definiëren, en dat de aanname van eindige gecentreerde $p$-de momenten sterker is dan die van eindige $p$-de momenten. Ze merken ook op dat hoewel per-sample clipping effectief is voor kleinere modellen zoals AlexNet, de computationele overhead het onpraktisch kan maken voor zeer grote architecturen zonder verdere optimalisatie, hoewel mini-batch clipping een kosteneffectief alternatief biedt in accumulatiescenario's.