Refining Covariance Matrix Estimation in Stochastic Gradient Descent Through Bias Reduction

1. Problématique

L'algorithme de Descente de Gradient Stochastique (SGD) est un pilier de l'apprentissage automatique à grande échelle. Bien que la convergence des itérés de SGD vers le paramètre optimal $x^*$ soit bien établie, l'inférence statistique (comme la construction d'intervalles de confiance) nécessite une estimation précise de la matrice de covariance asymptotique $\Sigma$ de la moyenne des itérés $\bar{x}_n$.

Selon le théorème de Polyak et Juditsky (1992), sous des conditions de régularité et avec un taux d'apprentissage polynomial $\eta_i = \eta i^{-\alpha}$ ($\alpha \in (0.5, 1)$), on a :
$$\sqrt{n}(\bar{x}_n - x^*) \Rightarrow \mathcal{N}(0, \Sigma)$$
où $\Sigma = A^{-1}SA^{-1}$ (forme "sandwich"), avec $A = \nabla^2 F(x^*)$ et $S = \mathbb{E}[\nabla f \nabla f^T]$.

Les défis actuels :

• Méthodes "Plug-in" : Elles nécessitent le calcul de la Hessienne stochastique ($A$), ce qui est souvent inaccessible ou coûteux en calcul ($O(d^3)$).
• Méthodes Bootstrap : Elles sont très coûteuses en temps de calcul car elles nécessitent de multiples ré-exécutions de l'algorithme.
• Méthodes "Batch-Means" (Moyennes par lots) : Ce sont des méthodes sans Hessienne (Hessian-free) qui fonctionnent en ligne. Cependant, elles souffrent d'un biais significatif et d'un taux de convergence lent, généralement de l'ordre de $O(n^{(\alpha-1)/4})$.

L'objectif de cet article est de développer un estimateur de covariance entièrement en ligne, sans Hessienne, qui corrige ce biais et améliore considérablement le taux de convergence.

---

2. Méthodologie

Les auteurs proposent un estimateur de covariance débiaisé (de-biased estimator) basé sur une technique de réduction de biais appliquée à la méthode des moyennes par lots.

A. Motivation et Formulation

Au lieu d'estimer directement la variance de la moyenne via une covariance d'échantillon classique (qui sous-estime la variance en raison de la forte corrélation des itérés SGD), les auteurs exploitent une identité théorique pour la variance d'une moyenne :
$$\sigma_n = \frac{1}{n} \sum_{i=1}^n \mathbb{E} \left[ 2(X_i - \mathbb{E}X_i) \sum_{k=1}^i (X_k - \mathbb{E}X_k) - (X_i - \mathbb{E}X_i)^2 \right]$$
Cette identité permet de construire un estimateur empirique qui, contrairement aux méthodes précédentes, ne repose pas sur une structure de variance biaisée, mais sur une expansion directe de la définition de la variance.

L'estimateur proposé $\hat{\Sigma}_n$ est défini comme :
$$\hat{\Sigma}_n = \frac{1}{n} \sum_{i=1}^n \left[ (x_i - \bar{x}_n) \left( \sum_{k=i-\ell_i+1}^i x_k - \ell_i \bar{x}_n \right)^T + \dots - (x_i - \bar{x}_n)(x_i - \bar{x}_n)^T \right]$$
où $\ell_i$ est la taille du "lot" (batch) à l'itération $i$.

B. Schéma de Lotting par Blocs (Block-based Batching)

Pour rendre l'estimateur compatible avec une mise à jour récursive et en ligne (sans stocker tout l'historique), les auteurs introduisent un schéma de lotting spécifique :

• Les lots ne sont pas de taille fixe, mais croissent selon $O(i^\alpha \log i)$.
• Une séquence d'indices $\{a_m\}$ définit le début de nouveaux blocs.
• À chaque itération $i$, le lot $B_i$ est construit pour inclure les deux blocs les plus récents, assurant ainsi que la taille du lot $\ell_i$ satisfait $\lfloor i^\alpha \log i \rfloor \le |\ell_i| \le 2\lfloor i^\alpha \log i \rfloor$.
• Cela permet de mettre à jour la somme des lots et la taille du lot de manière récursive avec un coût de $O(d^2)$ par itération.

C. Algorithme

L'algorithme (Algorithme 1) maintient en mémoire :

1. L'itéré SGD moyen $\bar{x}_n$.
2. La somme des lots actuels et précédents ($S_0, S_1$).
3. Des accumulateurs pour les termes croisés et quadratiques nécessaires au calcul de $\hat{\Sigma}_n$.
   Le coût de mise à jour est minimal ($O(d^2)$), ce qui est optimal pour estimer une matrice $d \times d$.

---

3. Contributions Clés

1. Nouvel Estimateur Débiaisé : Introduction d'un estimateur de covariance qui élimine le biais inhérent aux méthodes "Batch-Means" classiques en utilisant une formulation basée sur l'identité de la variance de la moyenne.
2. Amélioration Théorique du Taux de Convergence :
   • Les méthodes existantes (Hessian-free) convergent à un taux de $O(n^{(\alpha-1)/4})$.
   • L'estimateur proposé atteint un taux de convergence de $O(n^{(\alpha-1)/2} \sqrt{\log n})$.
   • Ce taux est prouvé être optimal (tight) pour cette classe de problèmes sans Hessienne.
3. Efficacité Computationnelle : L'algorithme est entièrement en ligne (single-pass), ne nécessite pas le calcul de la Hessienne, et a une complexité par itération de $O(d^2)$.
4. Validité Théorique Générale : La consistance de l'estimateur est prouvée sous des hypothèses standards de convexité forte et de régularité du gradient, couvrant des cas non-gaussiens et multidimensionnels.

---

4. Résultats Expérimentaux

Les auteurs ont évalué leur méthode sur trois modèles de régression (Linéaire, Logistique, et Expectile) avec différentes dimensions ($d=5, 20, 50$) et comparé leurs résultats à la méthode "Online Batch-Means" (non-overlap) de Zhu et al. (2023).

• Précision de l'estimation : Les tables 1 à 3 montrent que l'erreur d'estimation (norme de Frobenius $\|\hat{\Sigma}_n - \Sigma\|_F$) de la méthode débiaisée est systématiquement inférieure à celle de la méthode Batch-Means, et ce, pour toutes les dimensions et tailles d'échantillon.
• Taux de convergence : Les graphiques log-log (Figure 3) confirment visuellement une pente plus raide (convergence plus rapide) pour la méthode proposée.
• Inférence Statistique : La construction d'intervalles de confiance utilisant $\hat{\Sigma}_n$ débiaisé donne des taux de couverture empirique beaucoup plus proches du niveau nominal (95%) que la méthode de référence (Tableau 4). Par exemple, pour la régression logistique avec $d=50$, la couverture passe de ~83% (Batch-Means) à ~85% (Débiaisé), se rapprochant de l'objectif théorique.

---

5. Signification et Impact

Ce travail résout un problème fondamental en apprentissage statistique en ligne : comment quantifier l'incertitude des estimateurs SGD sans recourir à des calculs prohibitifs (Hessienne) ou à des méthodes lourdes (Bootstrap).

• Avance Théorique : En améliorant le taux de convergence de $n^{(\alpha-1)/4}$ à $n^{(\alpha-1)/2}$, l'article repousse les limites de ce qui est possible avec des méthodes sans Hessienne, s'approchant de la performance des méthodes plug-in tout en restant computationnellement efficace.
• Applicabilité Pratique : La méthode est directement applicable dans des scénarios réels où la mémoire est limitée et où les données arrivent en flux continu. Elle permet une inférence statistique rigoureuse (intervalles de confiance, tests d'hypothèses) sur les modèles entraînés par SGD.
• Généralité : L'approche de réduction de biais et le schéma de lotting adaptatif offrent des perspectives pour d'autres problèmes d'estimation de covariance dans des processus non-stationnaires et non-homogènes, au-delà du cadre strict du SGD.

En résumé, cet article propose une solution élégante et efficace pour rendre l'inférence statistique sur les modèles d'apprentissage profond à grande échelle à la fois précise et économiquement viable en termes de calcul.