Résumé Technique : PowerSoftmax : Vers une Inférence Sécurisée de LLM sur des Données Chiffrées

Énoncé du Problème

Les modèles de langage de grande taille (LLM) modernes préservant la vie privée reposent souvent sur le chiffrement homomorphe (HE), et plus spécifiquement sur le schéma CKKS, qui permet d'effectuer des calculs sur des données chiffrées sans déchiffrement. Cependant, CKKS ne prend en charge que les opérations polynomiales. Les architectures Transformer standard dépendent fortement de composants non polynomiaux, notamment la fonction Softmax dans l'attention auto, qui implique des exponentielles, des divisions et des opérations logarithmiques (par exemple, pour la stabilité numérique via l'astuce log-sum-exp).

Les approches existantes visant à combler ce fossé font face à des limitations significatives :

1. Approximation Post-Entraînement (PTA) : Approximation directe des fonctions non polynomiales dans les modèles pré-entraînés à l'aide de polynômes de haut degré. Cela entraîne souvent une inférence inefficace en raison d'une profondeur multiplicative élevée et d'un bruit accumulé.
2. Entraînement Conscient de l'Approximation (AAT) : Remplacement des primitives non polynomiales par des alternatives plus faciles à approximer (par exemple, l'attention pointuelle ou le ReLU mis à l'échelle) avant l'entraînement. Bien que efficaces pour les petits modèles, ces méthodes ont eu du mal avec la stabilité et l'évolutivité, échouant à prendre en charge les LLM comportant des milliards de paramètres ou des architectures profondes (par exemple, 32 couches).

Le défi central consiste à concevoir une architecture Transformer intrinsèquement compatible avec les contraintes du HE (polynômes de bas degré, plages numériques stables) tout en maintenant le pouvoir expressif, les capacités de raisonnement et l'apprentissage en contexte (ICL) des Transformers standard basés sur Softmax à l'échelle du milliard de paramètres.

Méthodologie

Les auteurs proposent PowerSoftmax, une variante novatrice et compatible avec le HE de l'attention auto, conçue pour reproduire les principes fondamentaux de Softmax tout en restant adaptable à l'approximation polynomiale. La méthodologie comprend trois étapes distinctes :

1. Conception d'une Attention Compatible HE

Les auteurs distillent les propriétés essentielles de Softmax : la normalisation dans l'intervalle $[0, 1]$, l'amplification des différences de scores et la préservation de l'ordre. Ils remplacent la fonction exponentielle $\exp(x)$ dans Softmax par une fonction puissance $x^p$ (où $p$ est un entier pair positif).
Le mécanisme d'attention proposé est défini comme suit :
$$\text{PowerSoftmax}(x)_j = \frac{x_j^p}{\sum_i x_i^p}$$
Cette formulation évite les exponentielles et les problèmes de division par zéro inhérents aux approximations polynomiales directes de Softmax.

2. Stabilité de l'Entraînement et Robustesse Numérique

Pour remédier à l'instabilité de l'entraînement causée par la fonction puissance (où les entrées $>1$ explosent et les entrées $<1$ s'annulent), les auteurs introduisent deux modifications clés :

• Variante Stable : Inspirée de l'astuce log-sum-exp, ils normalisent les entrées par une constante $c = \|x\|_\infty + \delta$ avant d'appliquer la fonction puissance. Cela garantit que les entrées restent dans une plage gérable $(0, 1)$, empêchant les débordements/sous-débordements et préservant la précision en virgule flottante.
• Division Bornée par $\epsilon$ : Pour gérer l'opération de division, qui est non polynomiale et sensible près de zéro, ils modifient le dénominateur en $\epsilon + \sum x_i^p$. Cela crée une fonction continue de Lipschitz en $1/\epsilon^2$, rendant l'approximation polynomiale subséquente de la division nettement plus traitable et réduisant la profondeur de multiplication requise.

3. Approximation Indépendante de la Longueur

L'approximation standard de la division (par exemple, l'algorithme de Goldschmidt) peine avec la somme non bornée des scores d'attention, qui évolue linéairement avec la longueur de séquence $L$. Les auteurs proposent un PowerSoftmax Indépendant de la Longueur :
$$\text{Indépendant de la Longueur}(x)_j = \frac{x_j^p}{L \cdot \text{Moyenne}(x^p)}$$
Puisque $L$ (longueur de séquence) est public, le terme $1/L$ peut être calculé directement sans approximation. Le modèle approxime la division de la moyenne plutôt que celle de la somme, garantissant que le domaine d'approximation reste borné quelle que soit la longueur de la séquence.

4. Pipeline de Construction (Algorithme 1)

Le pipeline complet pour créer un LLM polynomial implique :

1. Modification Architecturale : Remplacement de Softmax par la variante Stable PowerSoftmax et entraînement à partir de zéro (ou via un pré-entraînement continu) avec des objectifs de minimisation de plage.
2. Minimisation de Plage : Application d'une fonction de perte spécialisée pour contraindre les plages d'entrée des couches non polynomiales (PowerSoftmax, GELU, LayerNorm) afin de garantir que les approximations polynomiales de bas degré soient précises.
3. Remplacement Polynomiale : Après l'entraînement, remplacement des opérations non polynomiales restantes (division, racine carrée inverse) par des approximations polynomiales de bas degré (par exemple, méthode de Goldschmidt) pour obtenir un modèle entièrement polynomial.

Contributions Clés

1. Mécanisme PowerSoftmax : Une nouvelle variante d'attention auto qui remplace l'échelle exponentielle par une échelle polynomiale, spécifiquement conçue pour les environnements HE tout en conservant les principes de conception fondamentaux de l'attention.
2. LLM Polynomiaux Évolutifs : Le développement des premiers LLM polynomiaux dépassant un milliard de paramètres (spécifiquement 1,4 milliard de paramètres avec 32 couches). Ces modèles démontrent des capacités de raisonnement et d'apprentissage en contexte (ICL) comparables à celles des Transformers standard de même taille.
3. Stratégies d'Entraînement et d'Inférence Stables : Introduction de variantes d'entraînement numériquement stables et de stratégies d'inférence indépendantes de la longueur qui surmontent les problèmes d'instabilité et d'évolutivité des approches précédentes de Transformers polynomiaux.
4. Évaluation Complète : Une analyse détaillée de la latence pour les calculs chiffrés et des études d'ablation validant les choix de conception, y compris l'impact du paramètre de puissance $p$ et de la division bornée par $\epsilon$.

Résultats Expérimentaux

Les auteurs ont évalué leur méthode sur des architectures causales (GPT/Pythia) et bidirectionnelles (RoBERTa) :

• Performance : Le modèle polynomial de 1,4 milliard de paramètres a obtenu des performances comparables aux modèles Pythia originaux basés sur Softmax sur des benchmarks zero-shot et few-shot (par exemple, Lambada, PIQA, WinoGrande, ARC). Notamment, le modèle a démontré des performances compétitives sur des benchmarks de raisonnement comme ARC-Challenge, une capacité précédemment non observée dans les LLM polynomiaux.
• Modèles Bidirectionnels : Une variante polynomial de RoBERTa-Base n'a montré qu'une dégradation d'environ 1 % par rapport au RoBERTa original basé sur Softmax sur les tâches GLUE (SST-2, QNLI, MNLI), surpassant les approches polynomiales précédentes comme Nexus.
• Inférence HE : Des expériences sur le modèle Pythia 70M utilisant le framework HElayers ont confirmé que l'inférence chiffrée produit des résultats presque identiques à l'inférence en clair (perte MSE maximale de 0,005).
• Efficacité : La méthode proposée réduit considérablement la profondeur de multiplication par rapport aux alternatives de l'état de l'art (Nexus, Cho et al.). Pour une séquence de 128 tokens, PowerSoftmax a réduit la latence d'un facteur d'environ 9,7 par rapport aux approches Taylor/Goldschmidt et d'environ 20 par rapport aux approches de stratégie carrée.

Importance et Revendications

L'article revendique une percée dans le domaine de l'apprentissage automatique préservant la vie privée en démontrant que les LLM entièrement polynomiaux peuvent passer à l'échelle du régime du milliard de paramètres sans sacrifier les capacités de raisonnement ou d'apprentissage en contexte.

• Évolutivité : Contrairement aux travaux précédents limités à de petits modèles ou à des tâches de classification simples, cette approche entraîne avec succès des Transformers profonds (32 couches) avec 1,4 milliard de paramètres.
• Sécurité Non Interactive : En utilisant le chiffrement homomorphe totalement (FHE) avec un modèle entièrement polynomial, la solution permet une inférence non interactive, évitant la surcharge de communication et les vulnérabilités potentielles associées au calcul multipartite (MPC) ou aux protocoles interactifs.
• Praticité : L'analyse détaillée de la latence et le déploiement réussi d'un modèle de 1,4 milliard de paramètres suggèrent que l'inférence sécurisée et privée de LLM devient pratiquement réalisable, ouvrant la voie à des applications dans des secteurs sensibles aux données comme la santé et la finance.

Les auteurs reconnaissent que, bien que les modèles fonctionnent bien, une évaluation complète des capacités génératives auto-régressives dans le décodage séquentiel sur des données chiffrées reste une direction pour un travail futur.