BicKD: Bilateral Contrastive Knowledge Distillation

1. Énoncé du problème

La distillation de connaissances (KD) vise à transférer les connaissances d'un modèle « enseignant » volumineux et performant vers un modèle « étudiant » plus petit et efficace. Bien que la méthode KD classique (Hinton et al.) domine le domaine, elle souffre de deux limitations critiques :

• Absence de comparaison inter-classe : La KD classique effectue uniquement un alignement des probabilités au niveau de l'échantillon (minimisation de la divergence KL entre les prédictions de l'enseignant et de l'étudiant pour le même échantillon). Elle échoue à modéliser explicitement les relations ou les différences structurelles entre les différentes classes.
• Absence de contraintes structurelles : La KD classique n'impose aucune contrainte sur la structure géométrique de l'espace des probabilités. Idéalement, les distributions de probabilité des différentes classes devraient être orthogonales (c'est-à-dire que les classes distinctes occupent des directions bien séparées dans l'espace vectoriel). Les méthodes existantes n'imposent pas explicitement cette orthogonalité pour imiter la structure géométrique apprise par l'enseignant.

2. Méthodologie : Distillation de connaissances contrastive bilatérale (BicKD)

Les auteurs proposent BicKD, un cadre basé sur les logits qui introduit une fonction de perte contrastive bilatérale pour imposer l'orthogonalité et la cohérence structurelle. La méthode opère selon deux dimensions :

A. Concepts fondamentaux

• Orthogonalité : La méthode suppose que, dans un modèle bien entraîné, les vecteurs de prédiction pour différentes classes devraient être orthogonaux (produit scalaire $\approx$ 0), tandis que les vecteurs pour la même classe devraient être cohérents.
• Contraste bilatéral : La fonction de perte contraste les prédictions de deux manières :
  1. Au niveau de l'échantillon : Comparaison des prédictions de différents échantillons au sein d'un lot.
  2. Au niveau de la classe : Comparaison des prédictions entre différentes classes (colonnes dans la matrice de probabilité).

B. Fonctions de perte

La perte d'entraînement totale est une somme pondérée de trois composantes :
$$L_{train} = \alpha L_{CE} + \beta L_{SC} + \gamma L_{CC}$$

1. Perte contrastive au niveau de l'échantillon ($L_{SC}$) :

   • Amplification de l'orthogonalité au niveau de l'échantillon (SOA) : Pour les échantillons ayant des étiquettes de vérité terrain différentes ($i, j$ où $y_i \neq y_j$), la méthode maximise la distance cosinus entre la prédiction de l'étudiant sur l'échantillon $i$ et la prédiction de l'enseignant sur l'échantillon $j$. Cela repousse les représentations de classes différentes.
   • Alignement au niveau de l'échantillon : Utilise la divergence KL standard pour garantir que l'étudiant imite l'enseignant sur le même échantillon.

2. Perte contrastive au niveau de la classe ($L_{CC}$) :

   • Amplification de l'orthogonalité au niveau de la classe (COA) : Au lieu d'examiner les échantillons individuels, cela examine les colonnes de classes. Elle maximise la distance cosinus entre le vecteur de prédiction de l'enseignant pour la classe $k$ et le vecteur de prédiction de l'étudiant pour une classe différente $j$. Cela impose que les directions des différentes classes restent orthogonales.
   • Alignement au niveau de la classe : Utilise la distance L1 pour minimiser la divergence entre l'enseignant et l'étudiant pour la même colonne de classe.

C. Avantages par rapport aux méthodes existantes

• Légèreté : Contrairement aux méthodes basées sur les caractéristiques (par exemple, CRD) qui nécessitent le stockage de représentations de caractéristiques et l'échantillonnage de paires (mémoire/calcul élevés), BicKD opère uniquement sur les logits (faible dimension), ne nécessitant aucun tampon de mémoire supplémentaire.
• Structure explicite : Elle modélise explicitement la structure géométrique du simplexe de probabilité, guidant l'étudiant à apprendre la logique de séparation inter-classe de l'enseignant.

3. Contributions clés

1. Cadre novateur : Proposition de BicKD, la première méthode KD introduisant un contraste bilatéral (au niveau de l'échantillon et de la classe) pour imposer l'orthogonalité dans l'espace des probabilités.
2. Régularisation géométrique : Introduction d'un mécanisme pour régulariser la structure géométrique de la distribution prédictive, garantissant que les classes distinctes sont bien séparées (orthogonales) tout en maintenant une cohérence intra-classe.
3. Évaluation complète : Démonstration que BicKD surpasse les méthodes de l'état de l'art (SOTA) sur diverses architectures (ResNet, VGG, WideResNet, ShuffleNet, MobileNet) et jeux de données (CIFAR-100, Tiny-ImageNet).
4. Robustesse : Démonstration de performances supérieures dans des scénarios difficiles, notamment l'apprentissage Few-Shot (CIFAR-100-FS) et les distributions à longue traîne (CIFAR-100-LT), où la rareté des données et le déséquilibre dégradent généralement les performances de la KD.

4. Résultats expérimentaux

• Benchmarks standards :
  • CIFAR-100 : BicKD a réalisé une amélioration moyenne de 1,28 % par rapport aux méthodes de référence. Notamment, lors de la distillation de ResNet-110 vers WRN-16-2, BicKD (75,15 %) a surpassé le modèle enseignant (74,37 %).
  • Tiny-ImageNet : BicKD a constamment atteint la plus haute précision Top-1 pour toutes les paires enseignant-étudiant testées, surpassant la deuxième meilleure méthode (CRD) de 0,48 % dans le cadre ResNet-50 $\to$ WRN-40-1.
• Significativité statistique : Des tests t appariés ont confirmé que les améliorations de BicKD par rapport à la KD classique sont statistiquement significatives ($p < 0,05$).
• Efficacité : BicKD est computationnellement efficace. Les temps d'entraînement par lot étaient comparables à la KD classique (par exemple, 0,0450 s sur CIFAR-100) et nettement plus rapides que les méthodes basées sur les caractéristiques comme CRD (0,1752 s).
• Few-Shot et Longue traîne :
  • Dans les scénarios few-shot (35 à 100 échantillons par classe), BicKD a maintenu un écart de précision significatif par rapport à la KD classique alors que le volume de données diminuait.
  • Dans les configurations à longue traîne ($\rho=100$), BicKD a considérablement amélioré la précision sur les classes minoritaires (par exemple, doublant la précision sur les classes n'ayant que 5 échantillons) et réduit l'écart-type entre les classes, indiquant une meilleure cohérence.

5. Importance

Le papier comble une lacune fondamentale dans la distillation de connaissances : l'absence de contraintes structurelles explicites sur l'espace des probabilités. En déplaçant l'accent de la simple correspondance des probabilités de sortie vers l'imitation de l'orthogonalité géométrique des représentations de classes de l'enseignant, BicKD permet aux étudiants d'apprendre des caractéristiques plus riches et plus discriminatives.

Son importance réside dans :

• Simplicité : Elle ne nécessite que des logits, ce qui la rend facile à implémenter comme un remplacement « plug-and-play » pour les pipelines KD existants.
• Généralisation : Elle prouve que l'imposition de contraintes géométriques (orthogonalité) est cruciale pour transférer efficacement les connaissances, en particulier lorsque la capacité du modèle étudiant est limitée ou que les données sont rares.
• Orientation future : Elle ouvre une nouvelle voie de recherche sur la « Neural Collapse » et la relation entre l'orthogonalité et la structure du Cadre Équiangulaire Tendu Simplexe (ETF) dans l'apprentissage profond.