FlexAct: Why Learn when you can Pick?

Résumé Technique : Flex-Act

Énoncé du Problème
Malgré le rôle critique que jouent les fonctions d'activation dans la propagation des signaux, la capacité d'apprentissage et la généralisation au sein des réseaux de neurones profonds, celles-ci sont généralement traitées comme des hyperparamètres fixes (par exemple, ReLU, GELU, Swish) appliqués uniformément à toutes les couches et à toutes les tâches. Ce paradigme statique suppose qu'une seule non-linéarité est optimale, quel que soit la profondeur de la couche ou le régime de données, ce qui peut conduire à un biais inductif sous-optimal, à un éclatement des gradients (shattered gradients) ou à des problèmes de saturation. Bien qu'il existe des activations paramétrées (comme PReLU, SPLASH) et des approximateurs de fonctions apprenables, elles introduisent souvent une complexité d'optimisation élevée, une interprétabilité réduite ou nécessitent de s'engager dans une famille fonctionnelle spécifique. Les auteurs soutiennent qu'aucune non-linéarité n'est universellement optimale et que les réseaux profonds bénéficieraient de la capacité de sélectionner dynamiquement la fonction d'activation la plus appropriée pour chaque couche sans réglage manuel des hyperparamètres.

Méthodologie
Le papier introduit Flex-Act, un cadre qui permet la sélection discrète, par couche, de fonctions d'activation à partir d'un ensemble prédéfini pendant l'entraînement. Le mécanisme central repose sur l'astuce Gumbel-Softmax pour créer une relaxation différentiable des choix discrets.

1. Routage Discret via Gumbel-Softmax :
   Pour chaque couche $i$, le modèle calcule une combinaison convexe de fonctions d'activation candidates $\{\sigma^{(1)}, \dots, \sigma^{(p)}\}$ appliquées à la transformation affine $h_i$. Les probabilités de sélection $p_i$ sont dérivées de logits entraînables $\pi_i$ en utilisant la distribution Gumbel-Softmax :
   $$X_i = \sum_{j=1}^{p} p_i^{(j)} \cdot \sigma^{(j)}(h_i)$$
   où $p_i \sim \text{GumbelSoftmax}(\log \pi_i, \tau)$. À mesure que la température $\tau$ est refroidie (annealed), le modèle passe de l'exploration d'activations mélangées à l'engagement vers une fonction discrète spécifique par couche. Cela permet au processus de sélection d'être optimisé de bout en bout via la descente de gradient.

2. Correction de Biais via la Normalisation du Gradient :
   Les auteurs identient une faille critique dans le routage naïf par Gumbel-Softmax : un biais d'échelle favorisant les fonctions d'activation non bornées (comme ReLU). Parce que la rétropropagation met à l'échelle les gradients par la magnitude de l'activation, les fonctions ayant des sorties plus importantes dominent le processus de sélection, indépendamment de leur pertinence pour les données.
   Pour contrer cela, Flex-Act introduit un régularisateur basé sur la norme du gradient. Il calcule les normes de gradient de chaque activation candidate par rapport à l'entrée et les convertit en pseudo-probabilités. La fonction de perte totale combine la perte de la tâche principale ($L_{task}$) avec un terme de divergence de Kullback-Leibler ($L_{KL}$) qui aligne les probabilités de routage apprises avec ces pseudo-étiquettes basées sur le gradient :
   $$L = L_{task} + \alpha L_{KL}$$
   Cette régularisation empêche le modèle de se rabattre sur des fonctions non bornées simplement en raison de différences de magnitude, garantissant que la sélection est dictée par la pertinence fonctionnelle.

Contributions Clés

• Routage d'Activation Discret : Un mécanisme novateur utilisant la reparamétrisation Gumbel-Softmax pour permettre l'optimisation de bout en bout des choix de fonctions d'activation discrètes par couche.
• Dérivations Analytiques de Gradient : Dérivation formelle des équations de rétropropagation pour les réseaux employant un routage d'activation discret, résolvant la non-différentiabilité du processus de sélection.
• Régularisation par Norme de Gradient : Une technique spécifique pour corriger le biais inhérent envers les activations non bornées dans le routage discret, améliorant considérablement la précision de la sélection.
• Interprétabilité et Modularité : En traitant la sélection de l'activation comme une variable discrète et apprenable, le cadre fournit des informations sur les non-linéarités utilisées à différentes profondeurs, offrant une alternative modulaire aux familles paramétrées.

Résultats Expérimentaux
Le papier évalue Flex-Act sur des tâches de régression synthétiques et sur des benchmarks du monde réel :

• Régression Synthétique : Dans des tâches où la fonction d'activation de vérité terrain est connue (ReLU, Sigmoid, Tanh, LeakyReLU, Identity), Flex-Act converge systématiquement vers la fonction correcte.
  • Sans régularisation ($\alpha=0$), le modèle échoue à sélectionner les fonctions bornées (Sigmoid/Tanh), par défaut vers ReLU en raison du biais de magnitude de gradient.
  • Avec régularisation ($\alpha=0.3$), Flex-Act atteint une erreur quadratique moyenne (MSE) proche de zéro dans tous les scénarios de vérité terrain, égalant ou surpassant les bases à activation fixe sans nécessiter de multiples cycles d'entraînement ou de balayages d'activation.
• Benchmarks du Monde Réel :
  • Classification d'Images (CIFAR-10/100) : Intégré dans ResNet-18 et ResNet-34, Flex-Act (appliqué à la couche pénultième et à toutes les couches) obtient des améliorations statistiquement significatives par rapport aux bases ReLU fixes et aux autres méthodes paramétrées (PReLU, Swish, AReLU). Par exemple, ResNet-34 avec Flex-Act (toutes les couches) atteint une précision de 95,75 % sur CIFAR-10 et 79,42 % sur CIFAR-100, surpassant la base ReLU.
  • Classification de Texte : Le cadre a également été testé sur les benchmarks GLUE avec des modèles BERT-base, démontrant sa compatibilité avec les architectures de type Transformer.

Signification et Revendications
Le papier affirme que Flex-Act offre une solution « plug-and-play » pour l'apprentissage profond adaptatif. Il soutient qu'en permettant aux réseaux de récupérer dynamiquement la non-linéarité appropriée pour une tâche ou une couche spécifique, les praticiens peuvent éviter la surcharge de calcul liée au balayage manuel des hyperparamètres ou à l'entraînement de multiples modèles avec différentes activations fixes. Le cadre comble le fossé entre la flexibilité théorique et l'utilité pratique, fournissant une approche robuste, interprétable et modulaire de la conception de la non-linéarité qui passe de tâches synthétiques peu profondes à des architectures profondes du monde réel. Les auteurs soulignent que leur méthode atteint ces gains sans augmenter significativement le nombre de paramètres, en s'appuyant plutôt sur la sélection intelligente de formes fonctionnelles existantes et bien caractérisées.