Résumé Technique : Optimisation en Forme Close de la Dernière Couche

Énoncé du Problème

Les réseaux de neurones profonds sont généralement optimisés à l'aide de la Descente de Gradient Stochastique (SGD) ou de ses variantes, traitant toutes les couches de manière uniforme. Cependant, sous une fonction de perte quadratique, les poids optimaux de la dernière couche linéaire d'un réseau de neurones admettent une solution en forme close (spécifiquement, une solution de régression ridge) étant donné les caractéristiques produites par le backbone.

La SGD standard traite la dernière couche comme un ensemble de paramètres mis à jour via des étapes de gradient, ignorant cette optimalité analytique. Cela crée un décalage : tandis que les paramètres du backbone ($\theta$) nécessitent une optimisation itérative, les poids de la dernière couche ($W$) pourraient théoriquement être calculés directement. De plus, dans des contextes stochastiques avec des mini-lots, le calcul naïf de la solution en forme close pour chaque lot entraîne un surajustement sévère à ce lot spécifique, provoquant une instabilité. Les approches existantes tentant de tirer parti de cela (par exemple, en inférence causale) nécessitent souvent une rétropropagation à travers l'inverse en forme close, ce qui est coûteux en calcul, ou découplent la dernière couche du backbone, conduisant à des instabilités d'optimisation.

Méthodologie

Les auteurs proposent une procédure d'entraînement qui traite la dernière couche comme une fonction implicite et déterministe des paramètres du backbone, optimisant uniquement le backbone tout en garantissant que la dernière couche reste optimale à chaque étape.

1. Fondement Théorique (Cas Déterministe)

Pour un paramètre de backbone fixe $\theta$, la dernière couche optimale $W^\star(\theta)$ minimise la perte quadratique régularisée :
$$W^\star(\theta) = \arg\min_W \sum \|y_i - W\phi_\theta(x_i)\|^2 + \beta\|W\|_F^2$$
Les auteurs exploitent le Théorème de l'Enveloppe pour montrer que le gradient de la perte composite $L^\star(\theta) = L(W^\star(\theta), \theta)$ par rapport à $\theta$ est identique au gradient de la perte originale où $W$ est fixé à $W^\star(\theta)$ :
$$\nabla_\theta L^\star(\theta) = \nabla_\theta L(W^\star(\theta), \theta)$$
Ce résultat critique implique que l'on n'a pas besoin de rétropropager à travers l'inversion de matrice requise pour calculer $W^\star(\theta)$. Au lieu de cela, l'algorithme alterne entre :

1. Le calcul de $W_{t+1} = W^\star(\theta_t)$ (mise à jour en forme close).
2. La mise à jour de $\theta_{t+1} = \theta_t - \alpha \nabla_\theta L(W_{t+1}, \theta_t)$ (rétropropagation standard à travers le backbone).

2. Adaptation Pratique (Mini-Lot Stochastique)

Dans le contexte stochastique, le calcul de $W^\star$ sur un seul mini-lot provoque un surajustement. Pour y remédier, les auteurs introduisent un régularisateur proximal qui pénalise la distance entre la solution en forme close actuelle et l'estimation précédente $W_t$. L'objectif devient :
$$L_{B_t, W_t}^{\text{prox}}(W, \theta) = \sum_{(x,y) \in B_t} \|y - W\phi_\theta(x)\|^2 + \lambda \|W - W_t\|_F^2$$
La solution en forme close pour cette perte proximale est :
$$W^\star_{B_t, W_t}(\theta) = (Y_t \phi_\theta(X_t)^\top + \lambda W_t)(\phi_\theta(X_t)\phi_\theta(X_t)^\top + \lambda I)^{-1}$$
Cette formulation implémente efficacement un filtre de Kalman approximatif sur la dernière couche, accumulant les informations des lots précédents tout en s'adaptant au lot courant. L'algorithme met à jour le backbone via SGD, puis réajuste la dernière couche en utilisant cette solution proximale en forme close.

3. Extension à la Classification

Bien que principalement conçue pour la régression, la méthode est étendue à la classification en traitant les étiquettes one-hot comme des cibles pour une perte quadratique. La sortie est interprétée via une opération $\arg\max$. Les auteurs notent que, bien que cette heuristique fonctionne pour des nombres de classes faibles à modérés, elle ne garantit pas la normalisation des probabilités.

Contributions Clés

1. Cadre d'Optimisation : Un algorithme pratique qui intègre des mises à jour en forme close de la dernière couche dans les pipelines SGD standards sans nécessiter de rétropropagation à travers l'opération d'inversion de matrice, en exploitant le théorème de l'enveloppe.
2. Régularisation Stochastique : Une formulation proximale qui prévient le surajustement dans les contextes de mini-lots en régularisant la dernière couche vers son estimation précédente, interprétée comme un filtrage de Kalman approximatif.
3. Analyse Théorique :
   • Régime NTK : Preuve de convergence vers un minimum global dans la limite de largeur infinie du Neural Tangent Kernel (NTK), malgré la non-convexité de la perte dans l'espace des fonctions.
   • Analyse en Une Étape : Une borne théorique quantifiant la supériorité de l'étape en forme close par rapport à une étape de descente de gradient, montrant que le gain est lié au nombre de conditionnement (variance des valeurs singulières) de la matrice de caractéristiques.
4. Validation Empirique : Des améliorations démontrées par rapport à la SGD standard et à la régression ridge en forme close naïve sur diverses tâches.

Résultats Expérimentaux

La méthode ("$\ell_2$ c.f. proximal ($\lambda$)") a été évaluée par rapport à la SGD/Adam standard avec perte quadratique et à une base de régression ridge en forme close naïve ("$\ell_2$ c.f. ridge ($\beta$)").

• Chimie Quantique (QM9) : La méthode proximale a constamment surpassé la SGD standard, en particulier à de petites tailles de lot où la base ridge naïve échouait à converger. Elle a montré une robustesse au choix de l'optimiseur de backbone (SGD vs Adam).
• Opérateurs de Neurones de Fourier (FNO) : Dans l'apprentissage d'opérateurs pour l'équation de Burgers visqueuse, la méthode proximale est restée efficace sur toutes les tailles de lot, tandis que la base ridge naïve échouait sur les petits lots. La méthode a géré efficacement les degrés de liberté spatiaux de haute dimension.
• Variable Instrumentale à Caractéristiques Profondes (DFIV) : En inférence causale, la méthode a éliminé le besoin de réestimation coûteuse des dernières couches sur l'ensemble des données après l'entraînement, atteignant des performances comparables aux stratégies de réestimation tout en étant efficace en calcul.
• Classification (CIFAR-100, ImageNet) : Sur CIFAR-100, la méthode a atteint une précision comparable ou supérieure à l'Entropie Croisée pour des tailles de lot faibles à modérées. Cependant, sur ImageNet (grand nombre de classes), elle a sous-performé l'Entropie Croisée, mettant en évidence une limitation de la perte quadratique pour la classification à haute cardinalité.

Importance et Revendications

L'article revendique que l'exploitation de la solution en forme close pour la dernière couche est une stratégie viable et efficace pour accélérer l'optimisation et améliorer la stabilité, en particulier dans les tâches de régression et les régimes de petits lots.

• Importance Théorique : Ce travail fournit la première analyse des points critiques de la perte modifiée dans l'espace des fonctions et prouve la convergence vers un minimum global dans le régime NTK, abordant la non-convexité inhérente au problème.
• Importance Pratique : La formulation proximale résout le problème de surajustement inhérent aux mises à jour en forme close par mini-lot, rendant la méthode robuste et compatible avec les pipelines d'entraînement standards.
• Limites : Les auteurs déclarent explicitement que la méthode est actuellement limitée à la perte quadratique. Bien qu'elle fonctionne bien sur la classification avec peu de classes, elle sous-performe l'Entropie Croisée sur la classification à grande échelle (ImageNet), suggérant que l'adaptation de l'approche à la perte d'entropie croisée est une direction nécessaire pour les travaux futurs. La méthode engendre également une surcharge de calcul (inversion de matrice) qui évolue avec la dimension de la dernière couche, la rendant plus lente que la SGD standard pour des très grandes dernières couches dans des régimes de petits lots, bien que ce coût diminue à mesure que la taille du backbone augmente.