Résumé Technique : SAFE-KD

Énoncé du Problème

Les architectures de vision modernes atteignent une précision élevée mais engendrent des coûts d'inférence importants, ceما limite leur déploiement sur les appareils de bord (edge devices) et les systèmes sensibles à la latence. Les réseaux à sortie précoce (early-exit networks) répondent à cela en attachant des classificateurs intermédiaires, permettant aux entrées « faciles » d'interrompre le calcul plus tôt. Cependant, le déploiement actuel repose sur des règles d'arrêt heuristiques (ex. : seuils de confiance ou d'entropie) qui manquent de garanties formelles. Ces heuristiques sont souvent peu fiables en raison d'un mauvais étalonnage du modèle et se dégradent rapidement sous l'effet de changements de distribution, rendant difficile l'échange de la précision contre la latence de manière prévisible. Il existe un besoin critique pour une politique d'inférence adaptative capable de fournir des garanties sur échantillons finis concernant le risque des sorties précoces.

Méthodologie

SAFE-KD (Risk-Controlled Early-Exit Distillation) est un wrapper multi-sortie universel conçu pour permettre des sorties étalonnées avec des garanties statistiques. La méthode intègre trois composants fondamentaux :

1. Architecture Multi-Sortie Universelle

SAFE-KD attache des têtes de sortie légères à des profondeurs intermédiaires de backbones modernes (CNN et Vision Transformers) avec des hypothèses architecturales minimales :

• Placement : Pour les CNN (ex. : ResNet, ConvNeXt), les sorties sont placées après des groupes de blocs résiduels sélectionnés. Pour les ViT (ex. : ViT-S, Swin-T), les sorties s'attachent après des blocs de transformateurs sélectionnés, en effectuant un pooling des embeddings CLS ou des tokens moyens.
• Conception de la Tête : Une petite tête de prédiction (Global Average Pooling ou Token Pooling suivi d'un éventuel MLP/LayerNorm et d'un classificateur linéaire) assure un surcoût minimal tout en produisant des logits bien formés.
• Signal : Le principal signal de confiance est la Probabilité Maximale de Softmax (MSP).

2. Distillation Hiérarchique

Pour remédier à la faiblesse et au bruit des sorties intermédiaires, SAFE-KD emploie une stratégie de distillation hiérarchique utilisant un enseignant par Moyenne Mobile Exponentielle (EMA) :

• Distillation de Connaissance Découplée (DKD) : Chaque sortie est supervisée par l'enseignant via la DKD, qui sépare les composantes de la classe cible et des classes non-cibles dans la perte de distillation afin de mieux exploiter la « connaissance sombre » (dark knowledge).
• Cohérence Profonde-Vers-Superficielle : Une perte de cohérence aligne les sorties précoces avec la sortie finale, assurant la stabilité des prédictions à travers la profondeur du réseau.
• Objectif : La perte totale combine l'entropie croisée standard, la perte DKD par rapport à l'enseignant, et la perte de cohérence par rapport à la sortie finale, pondérées par les hyperparamètres $\alpha$ et $\beta$.

3. Sortie Contrôlée par le Risque via le Conformal Risk Control (CRC)

Au lieu de seuils heuristiques, SAFE-KD utilise le Conformal Risk Control pour étalonner les seuils d'arrêt par sortie sur un ensemble de calibration tenu à l'écart ($D_{cal}$) :

• Définition du Risque Sélectif : La sortie précoce est formulée comme une classification sélective. Le score d'incertitude est défini par $r_j(x) = 1 - \max_c p_j(c|x)$.
• Calibration : Pour un niveau de risque $\delta$ spécifié par l'utilisateur, le CRC calcule un seuil $\hat{\tau}_j(\delta)$ pour chaque sortie $j$. Ce seuil garantit que le risque de classification sélective erronée (la probabilité d'erreur parmi les échantillons acceptés à la sortie $j$) est borné par $\delta + O(1/|D_{cal}|)$ sous l'hypothèse d'échangeabilité.
• Politique d'Inférence : Au moment du test, le modèle sort à la tête $j$ la plus précoce où le score d'incertitude $r_j(x) \leq \hat{\tau}_j(\delta)$. Si aucune sortie précoce ne satisfait la condition, le modèle progresse vers la sortie finale.

Principales Contributions

1. Wrapper Multi-Sortie Universel : Une conception modulaire qui ajoute des sorties intermédiaires aux CNN et ViT avec un surcoût minimal, préservant les caractéristiques du backbone.
2. Distillation Hiérarchique : Une approche d'entraînement utilisant la DKD et la cohérence profonde-vers-superficielle pour améliorer la précision et la stabilité des prédictions intermédiaires.
3. Sortie Contrôlée par le Risque : Une procédure de calibration qui convertit les scores de confiance en seuils de sortie bornés par le risque et ajustables par l'utilisateur, remplaçant les heuristiques fragiles par un test statistiquement valide.
4. Validation Empirique : Démontrée sur plusieurs jeux de données et architectures, montrant des compromis précision-calcul améliorés, un meilleur étalonnage et une plus grande robustesse face aux corruptions.

Résultats Expérimentaux

Les auteurs ont évalué SAFE-KD sur des jeux de données incluant CIFAR-10/100, STL-10, Oxford-IIIT Pets, Flowers102 et FGVC Aircraft, en utilisant des backbones tels que ResNet-50, MobileNetV3-S, EfficientNet-B0, ConvNeXt-T, ViT-S et Swin-T.

• Précision et Efficacité : SAFE-KD égale ou dépasse la précision de la sortie finale de la DKD standard et surpasse significativement les lignes de base MultiExit. Sur ViT-S avec CIFAR-100, il a obtenu une amélioration de précision de +0,80 % par rapport à l'ERM.
• Dominance de Pareto : SAFE-KD atteint une frontière de Pareto strictement dominante par rapport au partitionnement basé sur la confiance fixe et l'entropie. Par exemple, sur ResNet-50 avec un risque cible de $\delta=0,05$, SAFE-KD a réduit la profondeur de calcul de 41 % tout en maintenant 94,1 % de précision relative.
• Validité du Risque : Les courbes de risque sélectif empiriques suivent systématiquement la diagonale $y=x$ ou restent légèrement en dessous, confirmant la validité sur échantillon fini. En revanche, les lignes de base heuristiques (ex. : confiance fixe de 0,95) ont souvent violé les contraintes de risque, entraînant des risques observés (ex. : 12 % sur Flowers102) dépassant largement la cible.
• Robustesse : Sous la corruption CIFAR-10-C, la dégradation de SAFE-KD est plus graduelle que celle des modèles standards. Les seuils calibrés rejettent naturellement les échantillons corrompus lors des sorties précoces, forçant le modèle à utiliser des couches plus profondes et plus robustes.
• Ablation : Supprimer la DKD a provoqué une chute significative de la précision (-2,1 %) et a imposé des seuils plus conservateurs (augmentant la profondeur à 0,85). Supprimer la cohérence a augmenté la variance du risque. La méthode complète a fourni le meilleur équilibre entre faible profondeur et haute précision.

Signification et Revendications

L'article pose que SAFE-KD transforme l'inférence adaptative d'un problème d'optimisation heuristique en un système fiable et contrôlé par le risque. En bornant mathématiquement le taux d'erreur des sorties précoces, il fournit une interface fiable entre les sorties du modèle et les exigences du système. Cela permet aux concepteurs de systèmes de traiter le paramètre de risque $\delta$ comme une contrainte stricte plutôt que comme un hyperparamètre ajustable, rendant les réseaux à sortie précoce viables pour des déploiements critiques où la latence prévisible et les limites d'erreur sont essentielles.

Les auteurs notent des limites, spécifiquement la dépendance à un ensemble de calibration (les petits ensembles mènent à des seuils conservateurs) et la dépendance à l'hypothèse d'échangeabilité, qui peut être rompue lors de changements de distribution sévères. Cependant, dans ces limites, SAFE-KD offre un cadre rigoureux pour une inférence efficace et sûre.