A Framework for Variational Inference of Lightweight Bayesian Neural Networks with Heteroscedastic Uncertainties

1. Énoncé du problème

Les réseaux de neurones bayésiens (BNN) sont essentiels pour quantifier l'incertitude prédictive, en distinguant l'incertitude épistémique (incertitude du modèle due à des données limitées) de l'incertitude aléatoire (bruit des données).

• Le défi : Dans de nombreuses applications, en particulier celles nécessitant des incertitudes hétéroscédastiques (où les niveaux de bruit varient selon l'entrée), les approches standard exigent que le BNN produise la variance aléatoire comme une tête de réseau supplémentaire en plus de la moyenne prédictive.
• La limitation : L'ajout d'un nœud de sortie pour la variance augmente le nombre de paramètres apprenables et la complexité architecturale. Cela est préjudiciable pour les BNN légers destinés au déploiement sur des systèmes aux ressources matérielles contraintes ou où le temps de calcul est limité.
• Le vide : Les méthodes existantes imposent souvent un compromis entre la taille du modèle et la capacité à capturer des structures de bruit complexes dépendantes de l'entrée.

2. Méthodologie

Les auteurs proposent un cadre d'inférence novateur qui élimine le besoin d'un nœud de sortie séparé pour l'incertitude aléatoire en intégrant directement les incertitudes aléatoire et épistémique dans les variances des paramètres apprenables du BNN.

A. Réinterprétation des distributions de paramètres

Au lieu de traiter la variance des paramètres des neurones bayésiens comme représentant uniquement l'incertitude épistémique, les auteurs réinterprètent ces distributions pour encoder l'incertitude prédictive totale (aléatoire + épistémique).

• Approche standard : $Var(\hat{y}) = \sigma^2_{épistémique}$. La variance aléatoire ($\sigma^2_{aléatoire}$) est prédite comme une sortie séparée.
• Approche proposée : La variance totale $\sigma^2_k$ pour une entrée donnée $k$ est définie comme la somme des variances aléatoire et épistémique :
  $$\sigma^2_k = \sigma^2_{a,k} + \sigma^2_{e,k}$$
  Pendant l'entraînement, le réseau minimise la borne inférieure de l'évidence (ELBO) où la perte de reconstruction suppose que la distribution de sortie est gaussienne avec une variance $\sigma^2_k$. La variance $\sigma^2_k$ est calculée comme la variance de la sortie du BNN dérivée des distributions de paramètres, absorbant efficacement le bruit aléatoire dans les variances des paramètres.

B. Inférence variationnelle sans échantillonnage par propagation des moments

Pour maintenir l'efficacité computationnelle et éviter le coût de l'échantillonnage de Monte Carlo, les auteurs emploient un schéma de propagation des moments.

• Mécanisme : La moyenne et la variance de l'entrée de chaque couche sont propagées analytiquement à travers le réseau pour calculer la moyenne et la variance de la sortie.
• Gestion des couches :
  • Pleinement connectées et convolutives : Propagation analytique exacte des moyennes et des variances en supposant l'indépendance entre les neurones.
  • Leaky-ReLU : Propagation approximative en supposant que l'entrée de la fonction d'activation suit une distribution normale. Les auteurs dérivent des expressions en forme fermée pour la moyenne et la variance de la sortie en utilisant la loi de l'espérance totale et de la variance, impliquant des distributions normales tronquées.
• Avantage : Cela permet une inférence déterministe sans besoin d'échantillonnage répété, rendant le cadre adapté aux applications légères et en temps réel.

3. Contributions clés

1. Intégration de la variance des paramètres : Un changement conceptuel où les variances des poids et des biais du BNN sont entraînées pour représenter la somme des incertitudes aléatoire et épistémique, éliminant le besoin d'une tête de sortie dédiée à la variance.
2. Architecture légère : Le cadre permet aux BNN de capturer des incertitudes hétéroscédastiques avec moins de paramètres apprenables par rapport aux approches standard de "variance apprise".
3. Inférence sans échantillonnage : Intégration d'une technique simplifiée de propagation des moments fournissant des estimations d'incertitude analytiques, évitant la surcharge computationnelle du dropout de Monte Carlo ou de l'échantillonnage.
4. Cadre unifié : Un pipeline complet combinant la réinterprétation des sources d'incertitude avec une propagation efficace des moments pour l'entraînement et l'inférence.

4. Résultats expérimentaux

Les auteurs ont évalué leur approche "Variance intégrée" par rapport à l'approche standard "Variance apprise" (où la variance est une sortie séparée) en utilisant une tâche de régression polynomiale avec un bruit gaussien hétéroscédastique.

• Performance dans les régimes légers :
  • Avec un petit nombre de neurones cachés (par exemple 4), le modèle à Variance intégrée a prédit avec précision à la fois la tendance et l'incertitude pour les données de la distribution d'entraînement et a correctement exprimé une forte incertitude pour les données hors distribution.
  • Le modèle à Variance apprise n'a pas réussi à retrouver l'incertitude réelle pour les données de la distribution d'entraînement et a considérablement sous-estimé l'incertitude pour les données hors distribution avec la même petite architecture.
• Efficacité des paramètres :
  • Pour un nombre fixe de neurones, le modèle à "Variance apprise" nécessitait environ 34 % de paramètres apprenables en plus que le modèle à "Variance intégrée" pour atteindre des performances comparables.
  • L'approche à "Variance intégrée" a obtenu une perte de reconstruction plus faible (négatif de la vraisemblance logarithmique) dans le régime à faible nombre de paramètres.
• Généralisation :
  • Les modèles à "Variance intégrée" ont démontré une meilleure généralisation aux données hors distribution, montrant une augmentation plus progressive de la perte de reconstruction à mesure que les données s'éloignaient de la distribution d'entraînement, par rapport aux modèles à "Variance apprise".
• Évolutivité : À mesure que la taille du réseau augmentait (par exemple 1024 neurones), l'écart de performance s'est réduit et les deux méthodes ont performé de manière similaire, suggérant que la méthode proposée est plus critique pour les modèles légers aux ressources contraintes.

5. Importance

Ce travail aborde un goulot d'étranglement critique dans le déploiement des réseaux de neurones bayésiens pour l'informatique de pointe et les systèmes critiques pour la sécurité.

• Efficacité : En supprimant le besoin de nœuds de sortie supplémentaires pour prédire le bruit, la méthode réduit considérablement la taille du modèle et le coût computationnel, rendant les BNN viables pour le matériel aux contraintes strictes.
• Robustesse : La capacité à capturer des incertitudes hétéroscédastiques sans augmenter la complexité du modèle garantit que les modèles légers ne sacrifient pas la fiabilité ou les évaluations de sécurité.
• Praticité : La nature sans échantillonnage de l'inférence permet des temps de prédiction plus rapides, ce qui est essentiel pour les applications en temps réel où la latence est un problème.

En résumé, l'article démontre que l'incertitude totale peut être efficacement "intégrée" dans les distributions de paramètres d'un BNN, offrant une alternative supérieure aux têtes de sortie multiples traditionnelles pour la quantification légère et hétéroscédastique de l'incertitude.