Résumé Technique : CoFrGeNet – Architectures de Fractions Continues pour la Génération de Langage

Énoncé du Problème

Les Transformers sont devenus l'architecture dominante pour la génération de langage, reposant sur l'Attention Multi-têtes et les Réseaux de Neurones à Propagation Avant (FFN) avec de grandes couches cachées. Cependant, ces composants introduisent une surcharge computationnelle significative, notamment en termes de nombre de paramètres et de temps d'entraînement. Bien que diverses approximations (par exemple, attention parcimonieuse, attention linéaire) aient été proposées pour améliorer l'efficacité, il subsiste un besoin de classes de fonctions fondamentalement différentes capables de remplacer ces composants de base tout en maintenant ou en améliorant les performances.

Des travaux antérieurs ont introduit les CoFrNets (Réseaux de Fractions Continues) pour l'apprentissage supervisé, utilisant une structure en échelle inspirée des fractions continues. Cependant, l'adaptation de cette architecture à la modélisation générative présente trois défis spécifiques :

1. Sorties multidimensionnelles : Les modèles génératifs nécessitent la production de sorties de haute dimension, contrairement aux sorties scalaires souvent observées dans les contextes supervisés.
2. Apprentissage de fonctions plus riches : L'architecture doit apprendre des représentations complexes adaptées à la modélisation de séquences.
3. Contraintes causales : Le modèle doit apprendre des paramètres dépendant uniquement des tokens précédents pour assurer une génération cohérente.
4. Inefficacité computationnelle : La non-linéarité standard $1/x$ dans les fractions continues nécessite $d$ divisions pour une échelle de profondeur $d$. Étant donné que la division est nettement plus lente que la multiplication sur le matériel moderne, cela crée un goulot d'étranglement pour l'entraînement et l'inférence, en particulier lorsque la profondeur ($d$) et le nombre d'échelles ($L$) augmentent.

Méthodologie

1. Réseaux Génératifs de Fractions Continues (CoFrGeNets)

Les auteurs proposent les CoFrGeNets, une famille d'architectures qui remplacent les couches d'Attention et/ou de FFN dans les blocs Transformers par des composants basés sur les fractions continues. La classe de fonction de base est définie comme suit :
$$w_0x + \frac{1}{w_1x + \frac{1}{w_2x + \dots}}$$
où $x$ est l'entrée et $w_k$ sont des vecteurs de paramètres apprenables.

2. Composants Architecturaux

L'article introduit des conceptions spécifiques pour remplacer les composants Transformers tout en maintenant les contraintes causales :

• Remplacement de l'Attention Causale (CAttnU et CAttnM) :

  • CAttnU : Utilise des échelles CoFrNet univariées où les entrées sont transposées pour mélanger les dimensions à travers les tokens. Pour imposer la causalité (empêchant la fuite d'informations vers les tokens futurs), l'architecture emploie des couches linéaires triangulaires supérieures au sein des ensembles. La multiplication élémentaire des sorties de deux ensembles génère des termes croisés, améliorant la puissance de représentation sans violer le flux causal.
  • CAttnM : Utilise $L$ échelles CoFrNet pour mapper les entrées vers des embeddings de longueur de séquence, qui servent de poids d'attention. La causalité est maintenue en appliquant un softmax causal (combinaison convexe des seuls tokens précédents) avant de pondérer la matrice de valeurs.

• Remplacement du FFN (Cffn) :

  • Puisque les FFN nécessitent un mélange de caractéristiques plutôt qu'un mélange de tokens, l'architecture utilise des ensembles d'échelles p-variées (où $p$ est la dimension d'embedding) sans transposer l'entrée.
  • Une couche linéaire à la fin projette la sortie vers la dimension cible.
  • Crucialement, l'entrée de ces échelles est une représentation gated, non étendue ($\alpha=1$), évitant l'explosion des paramètres typique des FFN standards des Transformers qui utilisent souvent un facteur d'expansion $\alpha \gg 1$.

3. Implémentation Efficace via les Continus

Pour remédier à l'inefficacité computationnelle des divisions répétées, les auteurs dérivent une implémentation personnalisée basée sur les continus (polynômes $K_k$).

• Passage Avant : Au lieu de calculer la fraction couche par couche avec $d$ divisions, la fraction continue est calculée comme le rapport de deux polynômes : $\tilde{f}(a) = K_{d-1} / K_d$. Cela nécessite de calculer les polynômes récursivement en utilisant uniquement des multiplications et des additions, suivis d'une seule division ($1/K_d$).
• Passage Arrière : Les auteurs dérivent des expressions en forme close pour les gradients en utilisant les propriétés des continus. Le gradient $\frac{\partial \tilde{f}}{\partial a_k}$ est exprimé comme un rapport de continus au carré.
• Optimisation : En calculant $K_d$ une fois et en réutilisant $1/K_d$ pour tous les calculs de gradient, le nombre de divisions est réduit de $d$ à 1. Cela accélère considérablement l'entraînement et l'inférence.

4. Plan d'Entraînement Personnalisé

Les auteurs proposent un plan de mise à jour des paramètres dyadique pour stabiliser l'entraînement. Au lieu de mettre à jour tous les paramètres simultanément :

• Les composants linéaires sont mis à jour dès le début.
• Les paramètres aux profondeurs supérieures sont gelés initialement.
• Les mises à jour pour les paramètres de profondeur $i$ ne commencent qu'après $t/2^i$ itérations (où $t$ est le nombre total d'itérations).
  Cette approche incrémentale prévient l'instabilité lors des phases précoces de l'entraînement.

Contributions Clés

1. Architectures Nouvelles : Introduction des composants CoFrGeNet en tant que remplacements plug-in pour les couches d'Attention et de FFN dans les Transformers, offrant une nouvelle classe de fonctions pour la modélisation générative.
2. Formulation Efficace des Gradients : Dérivation de formules de gradient personnalisées utilisant les continus, réduisant le nombre d'opérations de division de $d$ à 1 par échelle, améliorant drastiquement l'efficacité matérielle.
3. Stabilité de l'Entraînement : Développement d'un plan d'entraînement dyadique personnalisé qui stabilise l'entraînement d'architectures profondes de fractions continues.
4. Validation Empirique : Pré-entraînement et évaluation sur des modèles à grande échelle (GPT2-xl et Llama3) démontrant des performances compétitives avec une réduction du nombre de paramètres.

Résultats Expérimentaux

Configuration

• Modèles : Les expériences ont été menées sur GPT2-xl (1,5B) pré-entraîné sur OpenWebText et GneissWeb, et Llama3 (3,2B) pré-entraîné sur le mélange de données docling (2T de tokens).
• Lignes de Base : Les comparaisons incluaient les Transformers originaux, le Dense Synthesizer et l'Attention Parcimonieuse.
• Variantes : CoFrGeNet-F (FFN remplacé), CoFrGeNet-A (Attention remplacée) et CoFrGeNet (les deux remplacés).

Performance

• Tâches en Aval (GLUE) : Les modèles CoFrGeNet, bien qu'ayant 2/3 à 1/2 des paramètres des modèles originaux, ont atteint une précision compétitive ou supérieure sur les tâches de classification, de questions-réponses et de raisonnement. Par exemple, CoFrGeNet-F (985M) a surpassé le GPT2-xl de 1,5B sur la plupart des benchmarks GLUE.
• Perplexité : Sur les benchmarks de modélisation du langage (PTB, Wikitext, Lambada), les variantes CoFrGeNet ont généralement atteint une perplexité inférieure à celle des Transformers de base et d'autres lignes de base d'attention efficaces.
• Efficacité :
  • Temps d'Entraînement : L'implémentation basée sur les continus a réduit considérablement le temps d'entraînement par rapport à une implémentation "littérale" (CoFrGeNetB) n'utilisant pas de continus. Pour GPT2-xl, le temps d'entraînement est passé d'environ 203 heures (B) à environ 178 heures (CoFrGeNet).
  • Temps d'Inférence : La vitesse d'inférence s'est améliorée d'un ordre de grandeur presque complet pour la version basée sur les continus par rapport à l'implémentation littérale.
  • Débit : Sur Llama3, les variantes CoFrGeNet ont montré un débit plus élevé en tokens/jour et un temps d'entraînement total réduit (par exemple, CoFrGeNet-F entraîné en 6,6 jours contre 8,5 jours pour la ligne de base).

Importance et Revendications

L'article affirme que les CoFrGeNets représentent une alternative viable et économe en paramètres aux composants Transformers standards. En exploitant les propriétés mathématiques des fractions continues et en optimisant leur implémentation via les continus, les auteurs démontrent que :

• Des modèles de langage haute performance peuvent être construits avec nettement moins de paramètres.
• Le goulot d'étranglement computationnel des opérations de division dans les fractions continues peut être efficacement atténué, rendant l'architecture pratique pour les flux de travail industriels à grande échelle.
• L'approche est "plug-and-play", nécessitant des modifications minimales aux pipelines d'entraînement et d'inférence existants.

Les auteurs concluent que ce travail jette les bases des architectures génératives inspirées des fractions continues, potentiellement en démocratisant l'IA en permettant à des entités disposant de moins de ressources de pré-entraîner des modèles de haute qualité. Ils notent cependant que ces modèles ne possèdent pas intrinsèquement de garde-fous de sécurité et restent vulnérables aux hallucinations et aux attaques adverses, suggérant que les recherches futures devraient se concentrer sur la résolution de ces défis au sein de la forme fonctionnelle spécifique.