Résumé Technique : Cubit – Mélangeur de Tokens avec Régression Ridge à Noyau

Énoncé du Problème

Depuis l'introduction de l'architecture Transformer en 2017, elle est devenue le paradigme dominant en apprentissage profond pour le traitement du langage naturel, la vision par ordinateur et les tâches multimodales. Cependant, le mécanisme central de mélange de tokens dans les Transformers repose sur l'auto-attention, qui peut être théoriquement interprété comme une régression de Nadaraya–Watson (NW) avec un noyau exponentiel (via softmax) et une normalisation L1.

Bien que efficace, cette approche basée sur NW présente des limitations inhérentes :

1. Complexité Computatoire : L'attention standard engendre des coûts computatoires quadratiques ($O(N^2)$) par rapport à la longueur de la séquence, limitant l'évolutivité pour les longues séquences.
2. Contraintes Théoriques : La régression NW manque de régularisation explicite pour le compromis biais-variance et peut être sensible au bruit et aux effets de bordure.
3. Évolutivité des Améliorations : Bien que de nombreuses modifications (attention parcimonieuse, attention linéaire, mélange d'experts) aient été proposées, le paradigme computatoire fondamental reste ancré dans la régression NW.

Les auteurs postulent que remplacer le fondement de la régression NW par un cadre statistique plus robuste pourrait produire des architectures aux fondements mathématiques plus solides et aux capacités de modélisation de longues séquences améliorées.

Méthodologie : L'Architecture Cubit

L'article propose Cubit, une architecture novatrice qui remplace la régression NW du module d'attention par la Régression Ridge à Noyau (KRR).

1. Fondement Théorique : De NW à KRR

• Nadaraya–Watson (Transformer) : Dans l'attention standard, le token de sortie $z_i$ est une somme pondérée des vecteurs de valeur $v_j$, où les poids sont déterminés par des similarités normalisées par softmax entre les requêtes et les clés. Cela correspond à la régression NW, qui estime une valeur de fonction basée sur une moyenne pondérée locale sans résoudre un système global.
• Régression Ridge à Noyau (Cubit) : La KRR résout un problème de minimisation du risque empirique régularisé. La prédiction pour une nouvelle entrée implique l'inverse de la matrice de noyau plus un terme de régularisation : $(K + \lambda I)^{-1}$.
  • Différence Clé : Contrairement à NW, la KRR introduit un paramètre de régularisation explicite $\lambda$ pour contrôler le compromis biais-variance et utilise l'inverse de la matrice de Gram pour coupler l'information à travers tous les échantillons de données, offrant une plus grande robustesse au bruit et aux effets de bordure.

2. Conception Architecturale

Cubit modifie le calcul d'attention classique en incorporant la solution fermée de la KRR :

• Construction du Noyau : La matrice de noyau est construite à l'aide de cartes de caractéristiques dérivées des embeddings de clés. Les auteurs paramètrent la carte de caractéristiques $\psi(x)$ en utilisant une matrice de projection $W_K$ entraînables et une fonction d'activation prescrite $\phi$.
• Normalisation via Matrice Inverse : Au lieu d'une normalisation softmax ligne par ligne, Cubit calcule une matrice de normalisation $\Sigma^{(h)} = (K^{(h)}K^{(h)\top} + \lambda I)^{-1}$. Cette matrice agit comme un opérateur de couplage global.
• Formulation de la Sortie : La sortie est calculée comme $Z^{(h)} = A^{(h)} \Sigma^{(h)} V^{(h)}$, où $A^{(h)}$ représente la matrice de similarité du noyau (interactions requête-clé) et $\Sigma^{(h)}$ est la matrice de normalisation inspirée de la KRR.
• Contrainte Causale : Pour gérer la génération auto-régressive, la construction du noyau est adaptée de sorte que le $i$-ème token ne s'attende qu'à lui-même et aux tokens précédents, garantissant que la régression est effectuée sur le contexte disponible.

3. Redimensionnement à Plage Limitée (LRR)

Pour résoudre les problèmes de stabilité d'entraînement inhérents à l'inversion de grandes matrices, les auteurs introduisent le Redimensionnement à Plage Limitée (LRR) :

• Un vecteur d'échelle $s^{(h)}$ apprenable est appliqué pour ajuster adaptativement la précision du noyau.
• Les facteurs d'échelle sont contraints dans une plage bornée $(\beta, \alpha + \beta)$ en utilisant une fonction sigmoïde.
• Cela assure l'inversibilité de la matrice d'échelle et borne les éléments inverses, empêchant l'instabilité numérique lors de la résolution du système linéaire.

4. Distinction par rapport aux Travaux Connexes

L'article distingue explicitement Cubit de DeltaFormer (qui utilise également des matrices inverses) :

• Motivation : DeltaFormer combine DeltaNet et les Transformers ; Cubit est motivé par le changement théorique de NW vers KRR.
• Implémentation : DeltaFormer utilise deux embeddings différents pour construire l'inverse et force les valeurs diagonales à 1. Cubit utilise un embedding spécifique unique (par exemple, uniquement $K$) pour assurer l'inversibilité naturellement, sans forcer les valeurs diagonales à 1.
• Stabilité : Cubit propose LRR pour améliorer la stabilité, un composant absent dans DeltaFormer.

Contributions Clés

1. Cadre Théorique Unifié : Les auteurs établissent un lien entre les mécanismes de mélange de tokens et les méthodes de régression classiques, analysant systématiquement la régression NW et la KRR pour offrir une nouvelle perspective dans le développement d'architectures de modèles.
2. Architecture Cubit : La proposition d'un mélangeur de tokens basé sur KRR qui remplace l'estimateur NW dans les Transformers, intégrant le mécanisme de Redimensionnement à Plage Limitée (LRR) pour améliorer la stabilité de l'entraînement.
3. Validation Empirique : Des expériences complètes démontrant que Cubit présente des capacités de modélisation de longues séquences plus fortes que les Transformers standards et DeltaFormer, avec des gains de performance augmentant à mesure que la longueur de la séquence croît.

Résultats Expérimentaux

Les auteurs ont évalué Cubit par rapport aux Transformers et DeltaFormer sur divers jeux de données (Arxiv, Books3, FinWeb-Edu), longueurs de séquences (512 à 8192) et tailles de modèles (125M à 1,3B de paramètres).

• Performance sur la Modélisation du Langage : Cubit a constamment obtenu une perplexité plus faible (perte de validation) que les bases de référence. Par exemple, sur le jeu de données Arxiv (modèle 125M, 512 tokens), Cubit a atteint une perte de 1,8802 contre 1,9003 pour le Transformer.
• Modélisation de Longues Séquences : L'écart de performance entre Cubit et le Transformer s'est élargi à mesure que la longueur de la séquence d'entraînement augmentait. À 8192 tokens, l'écart de perte s'est stabilisé entre 0,05 et 0,06, suggérant que le mécanisme basé sur le noyau de Cubit capture plus efficacement les dépendances à long terme.
• Propriétés d'Évolutivité : Cubit a démontré un comportement d'évolutivité favorable. Alors que l'avantage de DeltaFormer sur le Transformer diminuait à mesure que la taille du modèle augmentait (de 350M à 1,3B), les gains de performance de Cubit persistaient et s'amélioraient.
• Tâches en Aval : Sur les évaluations zero-shot (ARC, HellaSwag, PIQA, etc.), Cubit a constamment surpassé à la fois le Transformer et DeltaFormer sur différentes tailles de données (10B à 50B de tokens) et échelles de modèles.
• Études d'Ablation :
  • Conception Architecturale : Les gains de performance ont été attribués à la conception structurelle (formulation KRR) plutôt qu'à l'expansion des paramètres.
  • Efficacité du LRR : L'inclusion du LRR a fourni des améliorations mesurables, "Cubit-NoLRR" surpassant toujours le Transformer standard, mais "Cubit" (avec LRR) obtenant les meilleurs résultats.

Signification et Revendications

L'article revendique que Cubit offre une fondation mathématique plus solide que le Transformer vanilla en tirant parti des avantages théoriques de la Régression Ridge à Noyau, notamment la régularisation explicite, des taux de convergence plus rapides dans les Espaces de Hilbert à Noyau Reproductible (RKHS) et une robustesse au bruit.

Les auteurs suggèrent que Cubit pourrait représenter une architecture de modèle fondation de nouvelle génération, en particulier pour les tâches nécessitant une compréhension de contexte long. Ils notent que les gains de performance ne sont pas simplement dus à une augmentation des tokens d'entraînement mais découlent de l'efficacité architecturale sur les longues séquences. Bien que l'article reconnaisse que le potentiel d'abus de modèles plus puissants existe, il met l'accent sur les bénéfices sociétaux potentiels d'une amélioration des capacités de modélisation.

Le travail indique également des directions futures, telles que l'exploration de la Régression Linéaire Locale (qui traite le biais de bordure mais est computatoirement plus lente) et d'autres cadres basés sur la régression pour le mélange de tokens, indiquant que la perspective basée sur la régression de l'attention est un terrain fertile pour l'innovation architecturale future.