Spartan : Économie différentiable via le transport régularisé

Nous présentons Spartan, une méthode d'entraînement de modèles de réseaux de neurones creux (sparse) avec un niveau de sparsité prédéterminé. Spartan repose sur une combinaison de deux techniques : (1) un masquage soft top-k des paramètres à faible magnitude via un problème d’optimisation de transport optimal régularisé, et (2) des mises à jour de paramètres basées sur une moyenne duale, accompagnées d’une sparsification rigide (hard) lors du passage avant. Ce schéma réalise un compromis entre exploration et exploitation : au début de l’entraînement, l’apprenant peut explorer diverses configurations de sparsité, et au fur et à mesure que l’approximation soft top-k se précise progressivement, l’équilibre se déplace vers une optimisation fine des paramètres sous un masque de sparsité fixe. Spartan est suffisamment flexible pour intégrer une variété de politiques d’allocation de sparsité, y compris à la fois la sparsité non structurée et la sparsité structurée par blocs, ainsi que des allocations de sparsité générales sensibles aux coûts, modélisées par des fonctions linéaires des coûts par paramètre. Sur la classification ImageNet-1K, Spartan permet d’obtenir des modèles ResNet-50 à 95 % de sparsité et des modèles ViT-B/16 à 90 % de sparsité par blocs, tout en subissant une perte absolue en précision top-1 inférieure à 1 % par rapport à un entraînement entièrement dense.