Calculer les dérivées sans remplir les matrices

OOndine

arXiv présente asdex, une méthode de différentiation automatique creuse pour JAX. Elle vise à réduire le coût du calcul de Jacobiennes et de Hessiennes en exploitant les dépendances limitées entre les entrées et les sorties. La méthode peut ainsi calculer ces dérivées sans construire une matrice dense, c’est-à-dire sans calculer toutes les cases possibles. L’idée s’appuie sur la structure du calcul. Imaginez un grand tableau où seules quelques cases peuvent contenir une valeur utile. Au lieu de traiter chaque case séparément, asdex repère lesquelles peuvent être regroupées. Cette comparaison aide à comprendre le principe : la méthode utilise les dépendances entre entrées et sorties pour éviter de calculer les éléments absents du motif de parcimonie. Le calcul suit 4 étapes. D’abord, asdex détecte ce motif : il repère quels éléments de la matrice de dérivées peuvent être non nuls. Ensuite, la méthode colore un graphe pour regrouper des colonnes ou des lignes compatibles. Cette coloration sert à identifier les groupes qui peuvent partager un même passage de différentiation. Puis, le calcul produit une dérivée compressée, avec un passage par couleur. Enfin, asdex décompresse le résultat et le replace dans le motif initial. Le nombre de couleurs détermine donc le nombre de passages nécessaires. Pour une Jacobienne en bandes, l’exemple donné dans la publication indique que ce nombre peut rester fixe, même si la taille du problème augmente. La méthode exploite ainsi la forme particulière de la matrice, plutôt que de traiter celle-ci comme si toutes ses cases étaient utiles. Il faut toutefois rester prudent sur la portée du résultat. Le principe est posé ; son intérêt concret dépendra des résultats mesurés sur des calculs réels.