Mise en œuvre de la sparsification avec la fonction Top-K
L'opération la plus fondamentale effectuée par Q-Sparse est l'application de la fonction de sparsification Top-K au tenseur d'entrée.
Plus précisément, l'architecture Transformer utilise des couches linéaires nn.Linear (multiplication matricielle) dans les couches d'attention et de feed-forward pour effectuer des projections, qui peuvent être représentées par Y=X·W^T. (Où X est le tenseur d'entrée, W représente ses poids, et Y est le tenseur de sortie)
Dans Q-Sparse, pour un tenseur d'activation d'entrée X, on calcule d'abord sa valeur absolue |X| et on la trie pour trouver les K éléments ayant les plus grandes valeurs absolues.
Ici, K est un hyperparamètre prédéfini qui détermine le degré de sparsification.
Ensuite, Q-Sparse crée un tenseur de masque binaire M de même forme que X, où les positions correspondant aux K éléments ayant les plus grandes valeurs absolues dans |X| sont mises à 1, et les autres positions sont mises à 0.
Puis, le tenseur d'entrée X est multiplié élément par élément (produit de Hadamard) avec le tenseur de masque M pour obtenir le tenseur sparsifié X_sparse.
Lors de la propagation avant, le tenseur sparsifié X_sparse remplace le tenseur d'entrée original X dans les calculs suivants (comme la multiplication matricielle).
Comme la plupart des éléments de X_sparse sont maintenant mis à zéro, cela réduit significativement la quantité de calculs et les besoins en bande passante mémoire.
Lors de la rétropropagation, Q-Sparse utilise l'estimateur direct (Straight-Through Estimator, STE) pour calculer le gradient de la fonction Top-K.
Dans les méthodes d'entraînement traditionnelles, il faut généralement calculer le gradient de la fonction de perte par rapport aux paramètres du réseau et utiliser la descente de gradient pour mettre à jour les paramètres afin de minimiser la perte.
Cependant, lorsque le réseau contient des opérations non différentiables comme la quantification ou le Top-K, le calcul du gradient pose problème car la sortie de ces opérations a un gradient nul par rapport à l'entrée en la plupart des points, empêchant une propagation efficace du gradient.
Le STE évite le problème de disparition du gradient en transmettant directement le gradient au tenseur avant sparsification.
Dans une rétropropagation normale, le gradient de la fonction de perte L par rapport à x est ∂L/∂x=∂L/∂y⋅∂y/∂x, mais il ne peut pas être calculé directement en raison de la non-différentiabilité.
La solution du STE est de calculer uniquement le gradient de la fonction de perte par rapport au tenseur sparsifié y, puis de le copier directement sur le tenseur original x, c'est-à-dire d'utiliser directement ∂L/∂y comme estimation de ∂L/∂x.
Pour les couches feed-forward, Q-Sparse utilise la fonction ReLU au carré au lieu de la fonction d'activation ReLU classique, l'opération de mise au carré augmentant davantage la sparsité de l'activation (⊙ représente le produit de Hadamard).
De plus, pour s'adapter aux modèles quantifiés, Q-Sparse quantifie d'abord le tenseur d'entrée avant d'appliquer la sparsification Top-K, afin d'assurer la compatibilité de l'opération de sparsification avec la représentation quantifiée. Sa fonction est représentée comme suit :
Où ε est une petite constante utilisée pour éviter une division par zéro.
En particulier, pour les poids quantifiés sur 1 bit, Q-Sparse utilise la fonction de quantification suivante, où α est la valeur absolue moyenne du tenseur de poids W.
60% des paramètres d'activation atteignent le même effet
Les expériences comparatives montrent que Q-Sparse surpasse significativement les méthodes ReLU précédentes, tant en termes de taux de sparsité que de performance du modèle.
Concernant les effets spécifiques de Q-Sparse, les auteurs ont évalué ses performances sur trois tâches : l'entraînement à partir de zéro, la poursuite de l'entraînement et le fine-tuning.
L'expérience d'entraînement à partir de zéro a utilisé le modèle Llama, et les résultats montrent que sur les modèles de 700M et 7B, Q-Sparse avec 70% top-K (soit un taux de sparsité global de 40%) peut atteindre une perte d'entraînement comparable à celle de la baseline dense.
L'objectif de la poursuite de l'entraînement est de sparsifier un modèle dense, et l'objet de l'expérience ici était Mistral-7B.
Les résultats montrent qu'avec 2,9B et 3,8B de paramètres d'activation, les scores du modèle sur des ensembles de données tels que ARC et MMLU n'ont pas connu de baisse significative.
Dans l'expérience de fine-tuning, pour les modèles Qwen-7B et Mistral-7B, Q-Sparse a montré des résultats similaires à ceux de la poursuite de l'entraînement, réalisant des performances très proches de celles des modèles denses avec environ 60% des paramètres d'activation.
Ces résultats signifient qu'à performance égale, les modèles à activation sparse peuvent réduire significativement les paramètres d'activation lors de l'inférence par rapport aux modèles denses, réduisant ainsi le nombre de FLOPS consommés.
Pour les modèles quantifiés, l'équipe a appliqué Q-Sparse à leur modèle BitNet b1.58 et l'a entraîné et évalué sur plusieurs ensembles de données.
On peut voir que pour les deux échelles de 700M et 7B, la vitesse de convergence et la valeur finale de la fonction de perte des modèles quantifiés utilisant Q-Sparse sont comparables à celles des modèles quantifiés sans Q-Sparse (BitNet b1.58).
Cela indique que Q-Sparse peut être intégré de manière transparente dans les modèles quantifiés sans affecter significativement l'entraînement et la convergence du modèle.
Sur cette base, les auteurs estiment que la combinaison de Q-Sparse avec les techniques de quantification peut améliorer davantage l'efficacité des grands modèles de langage lors de la phase d'inférence.
Découverte d'une nouvelle "Loi d'échelle" pour l'optimisation de l'inférence
En plus d'évaluer les performances de ces modèles lors de l'adoption d'une activation sparse, les auteurs ont également exploré la relation entre les performances du modèle, sa taille et le taux de sparsité, et ont fait de nouvelles découvertes.
Loi d'échelle des performances pour les modèles à activation sparse : Les auteurs ont découvert que, comme pour les modèles denses, les performances des modèles à activation sparse suivent également une relation d'échelle en loi de puissance.
Plus précisément, étant donné un taux de sparsité S, la valeur de la fonction de perte L(N,S) du modèle à la convergence peut être approximée par la formule suivante :