Robust and Fast Training via Per-Sample Clipping

Résumé Technique : Entraînement Robuste et Rapide via le Écrêtage par Échantillon (Per-Sample Clipping)

Énoncé du Problème
Le papier traite le problème d'optimisation stochastique $\min_{x \in \mathbb{R}^d} F(x)$, où $F(x) = \mathbb{E}_{\xi \sim \mathcal{D}}[f(x, \xi)]$. Bien que la descente de gradient stochastique (SGD) soit le solveur standard, ses garanties de convergence reposent généralement sur l'hypothèse que le bruit du gradient possède une variance bornée (distributions à queue légère). Cependant, des travaux empiriques et théoriques récents indiquent que dans de nombreuses applications d'apprentissage automatique, le bruit du gradient peut être à queue lourde, possédant des moments d'ordre $p$ finis pour $p \in (1, 2]$ mais manquant d'une variance bornée. Sous de telles conditions, la SGD standard peut mal performer. Les alternatives robustes existantes, telles que l'écrêtage de gradient standard (Clip-SGD) et la SGD Normalisée, échouent soit à atteindre les bornes inférieures optimales de complexité d'échantillonnage, soit nécessitent un réglage de paramètres spécifique et souvent peu pratique (par exemple, la dépendance envers l'écart d'initialisation $\Delta_1$ ou des facteurs logarithmiques spécifiques de la probabilité d'échec $\delta$).

Méthodologie
Les auteurs proposent la Per-Sample Clipped SGD (PS-Clip-SGD), un estimateur de gradient robuste qui applique un écrêtage aux échantillons individuels au sein d'un mini-lot avant l'agrégation, plutôt que d'écrêter le gradient agrégé du lot.

Pour un lot de taille $n$ avec des échantillons $\xi^{(1)}_t, \dots, \xi^{(n)}_t$, l'estimateur de gradient $G(x_t, \alpha, \beta)$ est défini comme :
$$G(x_t, \alpha, \beta) = \frac{1}{n} \sum_{k=1}^n \gamma^{(k)}_t \nabla f(x_t, \xi^{(k)}_t)$$
où le facteur d'écrêtage par échantillon est $\gamma^{(k)}_t := \min\left\{1, \frac{\alpha k^{1/\beta}}{|\nabla f(x_t, \xi^{(k)}_t)|}\right\}$. Les paramètres $\alpha$ et $\beta$ contrôlent le seuil d'écrêtage, $\alpha$ étant lié à la borne du moment $\sigma$ et $\beta$ à l'indice de queue $p$. La règle de mise à jour suit la SGD standard : $x_{t+1} = x_t - \eta_t G(x_t, \alpha, \beta)$.

Contributions Clés

1. Convergence Optimale en Espérance : Les auteurs prouvent que la PS-Clip-SGD atteint une complexité d'échantillonnage optimale pour l'optimisation non convexe sous un bruit à queue lourde. Plus précisément, la méthode atteint un taux de convergence de $O(1/T)$ pour la norme au carré du gradient moyen, correspondant aux bornes inférieures établies par Zhang et al. (2020). Cela améliore les résultats existants pour Clip-SGD et la SGD Normalisée, particulièrement parce que la borne porte sur la somme des normes au carré des gradients (un critère plus fort que la somme des normes utilisée dans certains travaux antérieurs) sans nécessiter la connaissance de l'écart d'initialisation $\Delta_1$ pour le pas de l'apprentissage.
2. Convergence à Haute Probabilité : Le papier établit des garanties de convergence à haute probabilité qui correspondent aux taux en espérance à des facteurs polylogarithmiques près de la probabilité d'échec $\delta$ et du nombre d'étapes $T$. Notamment, la dépendance en $\log(1/\delta)$ est améliorée par rapport aux bornes de l'état de l'art de la Clip-SGD, et la méthode ne nécessite pas que les pas d'apprentissage ou les seuils d'écrêtage soient explicitement réglés en fonction de $\delta$.
3. Aperçu de l'Accumulation de Gradient : Les auteurs contestent la pratique courante consistant à n'appliquer l'écrêtage du gradient qu'après toutes les étapes d'accumulation dans l'accumulation de gradient. Ils proposent et valident empiriquement que l'application de l'écrêtage au niveau du mini-lot (Mini-Batch Clipping) pendant l'accumulation peut améliorer les performances d'entraînement avec un coût computationnel supplémentaire virtuellement nul, contrastant avec le consensus standard selon lequel un écrêtage intermédiaire introduit un biais.

Résultats

• Théoriques : L'analyse théorique démontre que la PS-Clip-SGD fournit de solides garanties de convergence sous les hypothèses de $L$-lissé et de moments $p$-ièmes finis. Les bornes dérivées sont montrées comme étant plus serrées que celles de la SGD Normalisée et de la Clip-SGD en termes de dépendance à $\delta$ et de la force de la mesure de convergence (normes au carré vs normes).
• Optimisation Quadratique : Dans les expériences minimisant une fonction quadratique avec un bruit de type Pareto, la PS-Clip-SGD a nettement surpassé la Clip-SGD et la SGD Normalisée à travers divers régimes de bruit ($p \in \{1.2, 1.5, 1.8\}$), même avec des hyperparamètres non réglés. Des expériences à haute probabilité ont confirmé la dépendance sous-linéaire théorique en $\log(1/\delta)$.
• Classification d'Images (AlexNet sur CIFAR-100) : Lors de l'entraînement d'AlexNet avec momentum, la PS-Clip-SGD a atteint une précision de validation plus élevée (64,84 %) par rapport à la SGD vanilla avec momentum (58,90 %) et à la Clip-SGD avec momentum (58,94 %). Cette amélioration persiste même après avoir pris en compte le surcoût computationnel d'environ $1,3\times$ introduit par les calculs de la norme du gradient par échantillon. Les auteurs attribuent cela au contrôle efficace d'un petit nombre de gradients aberrants possédant des normes très élevées, que l'écrêtage standard au niveau du lot ne parvient pas à détecter.
• Modélisation de Langage (GPT-2) : Dans un contexte d'accumulation de gradient pour l'entraînement d'un modèle GPT-2 de 124 millions de paramètres, le Mini-Batch Clipping (écrêtage par étape d'accumulation) a surpassé la Clip-SGD standard (écrêtage uniquement après l'accumulation complète) sur plusieurs taux d'apprentissage, malgré la vue prédominante selon laquelle l'écrêtage intermédiaire est préjudiciable.

Signification et Revendications
Le papier affirme être le premier travail proposant l'écrêtage par échantillon spécifiquement comme une méthode pour améliorer les performances d'entraînement en présence de bruit à queue lourde, se distinguant de son utilisation établie en confidentialité différentielle. La signification réside dans le fait de fournir une méthode qui :

1. Atteint des taux de convergence théoriques optimaux pour les problèmes non convexes à queue lourde.
2. Offre une robustesse aux choix d'hyperparamètres, performant même avec des réglages non raffinés.
3. Démontre que l'écrêtage par échantillon peut engendrer des gains de performance pratiques dans des tâches d'apprentissage profond (comme AlexNet) qui l'emportent sur ses coûts computationnels.
4. Suggère un changement de paradigme dans les stratégies d'accumulation de gradient, montrant qu'un écrêtage intermédiaire peut être bénéfique plutôt que nuisible.

Les auteurs reconnaissent des limites, notant que leur analyse théorique nécessite la connaissance de la borne du moment $\sigma$ et de l'indice de queue $p$ pour définir les paramètres optimaux, et que l'hypothèse de moments $p$-ièmes centrés bornés est plus forte que celle des moments $p$-ièmes bornés. Ils notent également que bien que l'écrêtage par échantillon soit efficace pour des modèles plus petits comme AlexNet, le surcoût computationnel peut le rendre peu pratique pour des architectures très larges sans optimisation supplémentaire, bien que le mini-batch clipping offre une alternative rentable dans les scénarios d'accumulation.