Résumé technique : Mise à l'échelle de la vérification de réseaux de neurones avec le parallélisme de tenseurs et le parallélisme de données entièrement fragmenté

1. Énoncé du problème

La vérification formelle des réseaux de neurones vise à prouver mathématiquement qu'un réseau respecte des propriétés de sécurité pour toutes les entrées au sein d'un domaine spécifié. Bien que les algorithmes de propagation de bornes comme IBP, CROWN et $\alpha$-CROWN soient l'état de l'art pour cette tâche, leur application pratique est sévèrement contrainte par la mémoire GPU. Les implémentations standards exigent que les matrices de poids et les matrices de relaxation linéaire $A$ résident entièrement sur un seul accélérateur. Pour les réseaux de grande taille (par exemple, une seule couche cachée avec une largeur $h=262\,144$), les matrices de poids seules peuvent dépasser 4 Go en float32, ne laissant pas suffisamment de mémoire pour les matrices $A$, le batching ou les tenseurs intermédiaires requis pour la vérification. Ce document traite de la nécessité de lever ce goulot d'étranglement de la mémoire d'un seul GPU en adaptant les stratégies de fragmentation de paramètres issues de l'entraînement de modèles à grande échelle au cadre de vérification.

2. Méthodologie

Les auteurs adaptent deux techniques de parallélisme — le parallélisme de tenseurs (TP) et le parallélisme de données entièrement fragmenté (FSDP) — au cadre de vérification auto LiRPA / $\alpha, \beta$-CROWN.

2.1 Parallélisme de tenseurs (TP)

Le TP divise les multiplications matricielles au sein d'une seule couche entre $P$ GPU.

• Implémentation : Les auteurs ont implémenté des opérateurs personnalisés BoundLinearTP Col et BoundLinearTP Row. Les poids sont fragmentés selon les dimensions de sortie (colonne) ou d'entrée (ligne).
• Mécanisme : Dans une paire Colonne–Ligne, le calcul est distribué, nécessitant une opération AllReduce pour récupérer le résultat.
• Limitation : Lorsque des activations ReLU se produisent à l'intérieur d'une zone fragmentée (entre les couches Colonne et Ligne), le passage arrière (backward pass) de CROWN nécessite des bornes intermédiaires pour déterminer les pentes de relaxation ($\lambda_i$). Calculer ces bornes exactement nécessiterait une communication inter-GPU à chaque étape. Pour préserver la localité, les auteurs substituent les bornes intermédiaires par la méthode plus simple de propagation de bornes par intervalle (IBP). Cela introduit un « effet d'enveloppement » (wrapping effect), provoquant une dégradation de la précision des bornes à mesure que la profondeur de la zone fragmentée augmente.

2.2 Parallélisme de données entièrement fragmenté (FSDP)

Le FSDP fragmente les tenseurs de paramètres de manière égale sur $P$ GPU mais les reconstruit pour le calcul.

• Implémentation : Implémenté avec environ 100 lignes de Python, le FSDP remplace les tenseurs BoundParams par des fragments. Avant le calcul d'une couche, un AllGather assemble la matrice de poids complète ; immédiatement après, le tenseur complet est libéré, ne conservant que le fragment.
• Mécanisme : Cette approche garantit qu'au plus une matrice de poids complète réside en mémoire GPU à un instant donné. Crucialement, le FSDP ne modifie ni l'ordre ni le type des opérations arithmétiques ; il change simplement l'emplacement de stockage des poids.
• Correction : Parce que l'arithmétique reste identique à la base de référence sur un seul GPU, le FSDP produit des bornes bit à bit identiques ($\Delta = 0$).
• Intégration : Le FSDP s'intègre de manière transparente avec la vérification complète ($\beta$-CROWN + Branch-and-Bound) et les couches convolutives (BoundConv).

2.3 Configuration expérimentale

Les expériences ont été menées sur des GPU NVIDIA A40 utilisant PyTorch 2.x et la bibliothèque auto LiRPA. Les auteurs ont introduit un drapeau bab.force_synchronous pour assurer des comparaisons équitables entre les modes mono-GPU et distribués en désactivant le redimensionnement automatique des lots et l'élagage (pruning) dans le chemin mono-GPU, garantissant ainsi des charges de travail identiques.

3. Contributions clés

1. Implémentation de TP pour la vérification : Les auteurs ont implémenté le TP pour auto LiRPA en utilisant des opérateurs personnalisés, atteignant une réduction de la mémoire de pointe d'environ $2\times$ pour $P=2$. Ils ont formellement analysé le compromis, confirmant que si la validité (soundness) est préservée, la précision des bornes se dégrade avec le nombre de zones fragmentées en raison de la substitution par l'IBP.
2. Implémentation de FSDP : Ils ont implémenté le FSDP pour la propagation des bornes, obtenant une réduction de la mémoire de base de 80 à 90 % et une réduction de la mémoire de pointe de 34 à 39 % sur des MLP larges. Crucialement, cette méthode produit des bornes bit à bit identiques à la base de référence mono-GPU.
3. Intégration avec la vérification complète : Le FSDD a été intégré avec succès à $\beta$-CROWN et au Branch-and-Bound (BaB), permettant la première vérification complète (résultat UNSAT) d'un modèle ResNet-large de CIFAR-100.
4. Identification du goulot d'étranglement : À travers des expérimentations approfondies, les auteurs ont identifié que dans le mode $\alpha$-CROWN + BaB, le principal goulot d'étranglement de la mémoire n'est pas les matrices de poids, mais les tenseurs alpha par neurone (pentes ReLU).

4. Résultats

4.1 Réduction de la mémoire

• TP : A atteint un facteur de réduction de la mémoire de pointe d'environ $1,985$ (proche du théorique $P=2$) pour un MLP large ($h=262\,144$).
• FSDP :
  • Mémoire de base : Réduite de 80 à 90 % sur diverses tailles de modèles, car seuls les fragments sont alloués avant le calcul des bornes.
  • Mémoire de pointe : Réduite de 34 à 39 % sur des MLP larges ($h \ge 4096$). Les économies sont inférieures aux économies de base car les matrices $A$ non fragmentées (qui passent à l'échelle avec le carré de la largeur) restent en mémoire.

4.2 Précision et validité des bornes

• TP :
  • Les modèles avec une seule zone fragmentée (sans ReLU intermédiaire) ont atteint un accord de précision machine ($\sim 10^{-8}$) avec la référence mono-GPU.
  • Les modèles avec plusieurs zones fragmentées ont montré une déviation significative (ex: $\sim 750$ unités pour un modèle à 3 zones) due à la substitution par l'IBP, bien que la validité (sur-approximation unilatérale) soit préservée.
• FSDP :
  • A obtenu une déviation exactement nulle ($\Delta lb = \Delta ub = 0,00$) sur tous les modèles testés (MLP et benchmarks ONNX) pour l'IBP et CROWN.
  • A vérifié avec succès un modèle Transformer (ViT) et des ResNets convolutifs avec des résultats bit à bit identiques à la base de référence.

4.3 Vérification complète (BaB)

En mode $\beta$-CROWN + BaB, la consommation de mémoire est dominée par les tenseurs alpha (ex: 1,5 Go pour les tenseurs alpha contre 1,5 Mo pour les poids dans un réseau MNIST spécifique).

• La fragmentation des poids via FSDP a fourni des économies de mémoire de pointe négligeables dans ce mode car les tenseurs alpha n'étaient pas fragmentés.
• Cependant, le FSDP a permis avec succès la vérification d'un modèle ResNet-large (VNN-COMP 2024) qui serait autrement infaisable ou difficile à gérer, confirmant une propriété UNSAT sans branchement BaB.

5. Signification et limites

L'article démontre que les stratégies de fragmentation de paramètres issues de l'entraînement peuvent être transférées à la vérification formelle avec un effort d'ingénierie modeste.

• Signification : Le FSDP offre une voie pour vérifier des réseaux plus grands avec des bornes exactes, tandis que le TP offre la plus grande réduction de mémoire au prix de la précision des bornes pour les réseaux profonds. Ce travail établit que le FSDP est compatible avec les exigences complexes de la vérification complète (BaB) et de diverses architectures (CNN, Transformers).
• Limites :
  • Scalabilité : Les expériences ont été limitées à $P=2$. Les auteurs notent qu'à mesure que $P$ augmente, les matrices $A$ non fragmentées deviendront probablement le coût dominant, entraînant des rendements décroissants.
  • Précision du TP : La dégradation de la précision des bornes dans le TP pour les réseaux profonds limite son utilité pour la vérification précise d'architectures profondes.
  • Déplacement du goulot d'étranglement : Les expériences révèlent que la recherche future sur la vérification multi-GPU doit se concentrer sur la fragmentation des tenseurs alpha par neurone, et non seulement sur les poids, pour obtenir une mise à l'échelle significative dans les modes de vérification complète.

Les auteurs concluent que si le FSDP résout le goulot d'étranglement du stockage des poids, le goulot d'étranglement des tenseurs alpha reste le défi critique pour la mise à l'échelle de la vérification complète vers des modèles plus grands.