Qu’est-ce que la distillation de connaissances ?

La distillation de connaissances tire parti du fait que les réseaux neuronaux artificiels sont des « approximateurs universels » : si les données d’entraînement et la taille de la couche cachée sont suffisantes, le réseau neuronal peut approximer n’importe quelle fonction avec une précision arbitraire.⁵

Dans le machine learning traditionnel, les « connaissances» du modèle entraîné sont déterminées par les paramètres appris : les poids variables (et les biais), appliqués aux différentes opérations mathématiques intervenant dans le réseau neuronal, qui amplifient ou diminuent l’influence d’une certaine partie des sorties du réseau sur une autre partie. Cette vision des connaissances fait qu’il est difficile de déterminer la manière dont un modèle peut absorber les connaissances d’un autre modèle de taille et de structure différentes.

Hinton et al ont choisi d’adopter une vision plus abstraite et plus souple des connaissances, qu'ils considèrent comme étant « un mappage appris entre vecteurs d’entrée et vecteurs de sortie ». En d’autres termes, les connaissances du modèle sont interprétées non pas en fonction des paramètres strictement mathématiques qu’il apprend lors de l’entraînement, mais en fonction de la manière dont elles se généralisent à de nouvelles données après cet entraînement.

Avec cette nouvelle approche, les méthodes de distillation des connaissances visent à entraîner les modèles élèves à imiter non seulement la sortie du modèle enseignant pour une entrée donnée, mais aussi les étapes de raisonnement que ce dernier suit pour y parvenir. Conceptuellement, cela fonctionne de la même manière que le réglage d’instructions à l’aide de prompts CoT (« chain-of-thought », ou «  chaîne de pensée », en français), qui permet d’améliorer la qualité des réponses en apprenant aux LLM à articuler leur raisonnement « étape par étape ».

Dans l’apprentissage conventionnel supervisé ou auto-supervisé, une fonction de perte produit un vecteur représentant la divergence (ou la perte) entre les sorties du modèle et les sorties « correctes » (ou la vérité terrain) sur différentes entrées. En ajustant les paramètres du modèle pour minimiser la pente (ou le gradient) de ce vecteur via un algorithme d’optimisation comme la descente de gradient, les sorties du modèle sont plus proches de ces sorties correctes. Bien que les étapes de raisonnement du modèle soient « importantes » du fait de leur influence sur la sortie finale, elles ne sont généralement pas mesurées par une fonction de perte conventionnelle.

Inversement, la distillation des connaissances entraîne le modèle élève à imiter le processus de raisonnement du modèle enseignant. Pour ce faire, on ajoute un type spécialisé de fonction de perte, la perte de distillation, qui consiste à utiliser les étapes de raisonnement discrètes comme cibles souples à des fins d’optimisation.

Les sorties d’un modèle d’IA, quel qu’il soit, peuvent être interprétées comme des prédictions : un LLM autorégressif prédit le ou les mots suivants dans une séquence donnée ; un modèle de vision par ordinateur utilisé pour la classification des images prédit la catégorie d'une certaine image. Pour arriver à ces prédictions, appelées « cibles dures » dans ce contexte, les modèles d’apprentissage en profondeur font généralement plusieurs prédictions préliminaires et emploient une fonction softmax pour générer la prédiction avec la probabilité la plus élevée. Pendant l’entraînement, une fonction de perte d’entropie croisée est utilisée pour maximiser la probabilité attribuée à la sortie correcte et minimiser la probabilité attribuée aux sorties incorrectes.

Par exemple, un modèle de classification des images prédit la probabilité qu’une image en entrée appartienne à chaque classe connue que le modèle est entraîné à reconnaître, puis génère la classe avec la valeur de probabilité la plus élevée. Dans le jargon mathématique du machine learning, ces prédictions individuelles par classe sont appelées logits. De la même manière, un LLM autorégressif prédit plusieurs possibilités pour chaque mot suivant et (selon son réglage de température ) échantillonne l’une de ces possibilités pour sa sortie.

Avec la distillation de connaissances, ces prédictions intermédiaires, ou « cibles douces générées par le modèle enseignant fournissent souvent les données d’entraînement principales au modèle élève. Les probabilités relatives attribuées à ces prédictions préliminaires fournissent des indications précieuses sur la manière dont le modèle enseignant tend à se généraliser. Par exemple, un modèle de classification des images est beaucoup plus susceptible de (mal) classer l’image d’un renard comme représentant un « chien » qu’un « sandwich ». Les cibles douces fournissent donc beaucoup plus d’informations pour chaque cas d’entraînement que les cibles dures.

Les cibles souples offrent également plus de cohérence que les cibles strictes : la prédiction finale d’un modèle peut finalement dépendre d’une différence minime entre deux valeurs logit, mais les valeurs logit elles-mêmes présentent une variance beaucoup moins importante dans le gradient entre chaque exemple d’entraînement.

En raison de la richesse et de la stabilité des informations fournies par les cibles souples, le modèle élève peut être entraîné sur moins d’exemples, avec un taux d’apprentissage plus élevé que celui utilisé pour entraîner le modèle enseignant.

Pour rapprocher les tendances de généralisation du réseau élève de celles du réseau enseignant, la distillation des connaissances utilise généralement deux fonctions de perte. La première est une fonction de perte standard qui opère sur la « perte dure », mesurant les sorties finales du modèle élève par rapport aux étiquettes de vérité terrain (pour l’apprentissage supervisé) ou à l’échantillon de données d’origine (pour l’apprentissage auto-supervisé). La seconde est la perte par distillation, une « perte souple » qui mesure les objectifs faibles du modèle élève par rapport à ceux de du modèle enseignant.

Comme il peut y avoir plusieurs cibles souples pour chaque exemple d’entraînement, la perte de conversion mesure la différence entre la distribution des probabilités des cibles souples du réseau enseignant et la distribution des probabilités du réseau élève. La divergence de Kullback-Leibler (ou « divergence K-L ») est couramment utilisée à cette fin.

Ces différentes formes de connaissances relèvent généralement de l’une de ces trois catégories : connaissances basées sur les réponses, connaissances basées sur les caractéristiques et connaissances basées sur les relations.

Les connaissances basées sur les réponses, le genre le plus courant de distillation des connaissances, se concentre sur le transfert d’informations depuis la couche de sortie finale du modèle enseignant. Dans une méthode KD typique basée sur les réponses, le modèle élève est entraîné à produire des logits qui correspondent aux prédictions du modèle enseignant.

Lorsque les cibles souples du modèle enseignant ont une faible entropie, c’est-à-dire lorsque les prédictions sont extrêmement « confiantes », par exemple, si un modèle de classification produit un logit très proche de 1 (représentant la certitude) pour une classe et des logits proches de 0 pour toutes les autres, elles ne fournissent pas autant d’informations. Les méthodes basées sur les réponses utilisent donc souvent une température élevée pour les sorties des modèles, ce qui augmente l’entropie des prédictions des modèles. Cela garantit une distribution de probabilités plus variable et donc une plus grande quantité d’informations provenant de chaque exemple d’entraînement.

Les connaissances basées sur les caractéristiques se concentrent sur les informations véhiculées dans les couches intermédiaires d’un réseau neuronal, ou « couches cachées ». C’est à ce stade que les réseaux neuronaux effectuent généralement l’extraction des caractéristiques, c’est-à-dire l’identification des caractéristiques et des schémas distincts des données d’entrée qui sont pertinents pour la tâche à accomplir.

Par exemple, dans les réseaux neuronaux convolutifs utilisés principalement pour des tâches de vision par ordinateur comme la segmentation d’image, chaque couche cachée successive capture des détails de plus en plus riches à mesure que les données sont transmises sur le réseau. Par exemple, dans un modèle utilisé pour classer les images d’animaux par espèce, les premières couches cachées ne discernent que la présence d’une forme animale dans une partie de la photo, les couches médianes cachées discernent que l’animal est un oiseau, les dernières couches cachées, juste avant la couche de sortie, discernent les détails qui différencient une espèce d’oiseau d’une autre espèce étroitement apparentée.

L’objectif des méthodes de distillation des connaissances basées sur les caractéristiques est donc d’entraîner le modèle élève à apprendre les mêmes caractéristiques que le réseau enseignant. Des fonctions de perte par distillation basées sur les caractéristiques permettent de mesurer, puis de minimiser la différence entre les activations de caractéristiques des deux réseaux.

Alors que les connaissances basées sur les réponses et sur les caractéristiques se concentrent sur les sorties de couches de modèles spécifiques, la distillation des connaissances basée sur les relations porte sur les relations entre différentes couches ou entre les cartes de caractéristiques représentant les activations à différents niveaux ou emplacements.

En substance, les connaissances basées sur les relations représentent peut-être l’approche globale d’entraînement du réseau élève à imiter le « processus de pensée » du modèle enseignant. Ces relations et corrélations peuvent être modélisées de différentes manières, notamment les corrélations entre les cartes de caractéristiques, les matrices représentant la similarité entre les différentes couches, les embeddings de caractéristiques ou les distributions probabilistes des représentations de caractéristiques.

Dans le cas d’une distillation hors ligne, le réseau enseignant est déjà pré-entraîné et les poids de son modèle sont gelés pour éviter tout changement ultérieur. La distillation hors ligne est typique de nombreuses approches KD pour les LLM, dans lesquelles l’enseignant est souvent un modèle propriétaire plus grand dont les poids ne peuvent pas être modifiés.

Dans certaines circonstances, un modèle enseignant correctement pré-entraîné et performant peut ne pas être disponible, ou un data scientist souhaitera peut-être adapter le réseau enseignant à son cas d’utilisation spécifique. Les programmes de distillation en ligne visent à entraîner simultanément les réseaux enseignant et élève.

Par exemple, Cioppa et al a proposé un programme de distillation en ligne pour les modèles de segmentation sémantique utilisés dans les événements sportifs en direct, où les circonstances visuelles peuvent changer tout au long d’un match. Son objectif était de contourner le compromis entre la vitesse d’un petit réseau et la précision d’un plus grand réseau en entraînant en permanence un modèle lent et performant sur des données de match en direct, tout en distillant simultanément les connaissances de ce plus grand modèle dans un modèle plus petit et plus rapide déployé pour générer des résultats en temps réel.⁶

Dans l’autodistillation, un réseau agit à la fois comme enseignant et comme élève. Alors que la distillation conventionnelle des connaissances implique le transfert de connaissances d’un modèle à un autre, l’autodistillation peut être comprise comme le transfert de connaissances des couches profondes d’un réseau vers les couches superficielles de ce même réseau.⁷

Dans l’autodistillation, plusieurs « classificateurs peu profonds » basés sur l’attention sont ajoutés aux couches intermédiaires du modèle à différentes profondeurs. Pendant l’entraînement, les classificateurs plus profonds servent de modèles enseignant et orientent l’entraînement des autres modules basés sur l’attention en fonction de deux types de pertes par distillation : une perte métrique de divergence KL sur les sorties et une perte de régularisation L2 sur les cartes de caractéristiques.

Une fois le modèle entraîné et prêt pour l’inférence, tous ces classificateurs superficiels sont supprimés du modèle. Cela permet d’obtenir un modèle plus grand dont la capacité de reconnaissance des schémas pendant l’entraînement est supérieure, mais aussi d’un modèle plus petit, donc plus rapide et plus efficace lorsqu’il est déployé.

Les LLM les plus performants sont, dans la plupart des cas, trop coûteux et trop gourmands en ressources de calcul pour être accessibles à de nombreux utilisateurs potentiels, comme les amateurs, les start-ups ou les instituts de recherche. De plus, malgré leurs performances avancées et leurs capacités uniques, les LLM propriétaires, de par leur nature, ne peuvent pas être adaptés à des applications de niche et à des cas d’utilisation spécifiques.

Par ailleurs, la plupart des LLM commercialement viables sont trop volumineux et exigeants en calcul pour être utilisés localement sur les téléphones mobiles ou d’autres appareils de périphérie. Un modèle plus petit exécuté directement sur les appareils mobiles permet de contourner ces problèmes logistique, d’informatique et de confidentialité. La compression des modèles de KD présente donc un moyen prometteur de transférer les qualités émergentes de grands modèles à des modèles suffisamment petits pour être exécutés sur les appareils.