Prédire une séquence simple avec un transformer conçu à la main, sans apprentissage

• Il s’agit d’une expérience où l’on conçoit directement les poids d’un transformer décodeur-only proche de GPT-2, sans apprentissage, afin de lui faire prédire le motif (aab)*
• La prédiction de "aabaabaabaab..." nécessite de regarder les deux tokens précédents, ce qui met mieux en évidence le fonctionnement de l’attention qu’un simple motif alterné
• Le modèle est volontairement petit, avec N_CTX=5, N_VOCAB=2, N_EMBED=8, et utilise la tokenisation a=0, b=1 ainsi que des embeddings one-hot
• Une unique tête d’attention répartit l’attention à parts égales, 0.5, sur les deux tokens les plus récents, et calcule le token suivant en exploitant l’annulation d’un encodage a=1, b=-1
• Dans les contextes non ambigus, il atteint une précision de 100,0 % (27/27), mais une seule prédiction avec un contexte de 5 tokens demande environ 4 000 FLOPs, ce qui le rend bien moins efficace qu’une règle implémentée directement

---

Mini GPT-2 aux poids fixés à la main, sans apprentissage

• L’objectif est d’acquérir une compréhension intuitive de ce que font réellement les différents composants d’un transformer et de l’attention
• Le modèle n’est pas entraîné, n’utilise pas de poids préentraînés, et a été construit en assignant directement chaque poids à la main en une soirée
• L’architecture est un transformer décodeur-only proche de GPT-2, avec une implémentation simplifiée à partir de l’implémentation picoGPT de jaymody
  • suppression du layer norm
  • utilisation d’une seule tête au lieu de la multi-head attention
  • suppression de la couche feed-forward mlp du bloc transformer

Pourquoi choisir la séquence (aab)*

• Au départ, l’idée était de prédire une séquence comme "ababababab", mais comme le transformer prédit une shifted sequence, la tâche devenait trop simple
  • il suffit de prédire b si c’est a, sinon a, sans avoir besoin d’utiliser les embeddings de position
• La tâche finale a donc été fixée à la prédiction de "aabaabaabaab...", c’est-à-dire la séquence (aab)*
  • si les deux tokens précédents sont ab ou ba, le token suivant est a
  • si les deux tokens précédents sont aa, le token suivant est b
  • bb est traité comme un cas hors du périmètre de la tâche
• La tokenisation est volontairement simple et ne manipule que deux symboles
  • a vaut 0
  • b vaut 1

Dimensions du modèle et flux de calcul

• Trois paramètres de modèle ont été choisis
  • N_CTX = 5 : longueur maximale de contexte vue par le modèle en une fois
  • N_VOCAB = 2 : les deux tokens a et b
  • N_EMBED = 8 : taille d’embedding servant à stocker tokens, positions et espace de calcul
• En pratique, seuls les deux tokens précédents sont nécessaires, mais N_CTX=5 permet d’inclure le cas où il faut ignorer des tokens non pertinents
• La fonction gpt s’exécute dans l’ordre suivant
  • elle additionne embeddings de tokens et embeddings de position avec wte[inputs] + wpe[range(len(inputs))]
  • elle fait passer le résultat dans un unique bloc transformer
  • elle produit enfin les logits dans l’espace du vocabulaire via x @ wte.T

Des embeddings one-hot pour position et token

• wpe est l’embedding de position, et les 5 premières dimensions d’embedding servent de one-hot de position
  • la position 0 est [1, 0, 0, 0, 0]
  • la position 4 est [0, 0, 0, 0, 1]
• wte est l’embedding de token, et les 2 dimensions suivantes servent de one-hot de token
  • le token a est représenté par [1, 0] sur ces dimensions
  • le token b est représenté par [0, 1]
• La 8e position de l’embedding n’est pas utilisée au départ et sert ensuite de scratch space dans le bloc transformer
• Par exemple, "aabaa" est représenté par une matrice d’embedding 5 x 8, dont chaque ligne contient à la fois le one-hot de position et le one-hot de token

Comment la tête d’attention sélectionne les deux tokens les plus récents

• Le bloc transformer se compose d’une seule tête d’attention et d’une couche linéaire c_proj qui reprojette le résultat de l’attention dans l’espace d’embedding
• c_attn est une couche linéaire de taille embed_size x (embed_size * 3) qui transforme l’embedding d’entrée en matrices qkv, puis les sépare en q, k, v
• k isole les embeddings de position afin de représenter l’information de position portée par chaque token
• q représente la plage de positions que chaque position souhaite consulter, et q @ k.T produit ainsi la matrice des scores d’attention
• Après softmax et causal mask, la matrice d’attention possède les propriétés suivantes
  • la première ligne met 100 % de l’attention sur le premier token uniquement
  • les lignes suivantes répartissent chacune l’attention à 0.5 sur les deux tokens les plus récents accessibles
• Le causal mask empêche de voir les futurs tokens en ajoutant une très petite valeur aux positions futures, -1e10 dans le code réel
  • ce modèle conçu à la main n’est pas fait pour tricher en regardant dans le futur, mais le masque est conservé pour rester plus proche de l’architecture GPT-2
• La normalisation par np.sqrt(q.shape[-1]) aide à améliorer les gradients dans un entraînement réel, mais n’a pas d’effet ici sur ce modèle artisanal

Encodage de v et prédiction par annulation additive

• v transforme le one-hot du token en un encodage a=1, b=-1
• Comme le résultat de l’attention moyenne les deux tokens les plus récents avec un poids de 0.5 chacun, cet encodage calcule les règles suivantes
  • a, b → 0.5 * 1 + 0.5 * (-1) = 0
  • b, a → 0.5 * (-1) + 0.5 * 1 = 0
  • a, a → 0.5 * 1 + 0.5 * 1 = 1
• En conséquence, la 7e position de la ligne contient la valeur suivante
  • 0 lorsque le bon token à prédire est a
  • 1 lorsque le bon token à prédire est b
• Avec l’entrée "aabaa", la première ligne peut produire une prédiction b faute d’information suffisante, mais les prédictions suivantes respectent bien la règle (aab)*

Reprojeter la prédiction dans l’espace du vocabulaire

• c_proj reconvertit la valeur de la 7e position du résultat d’attention en format one-hot de token
• Plutôt que de produire simplement [..., 1, 0, ...] ou [..., 0, 1, ...], il produit un one-hot mis à l’échelle par 1024
  • embedding[row, 5] = 1024 + (-1024) * prediction
  • embedding[row, 6] = 0 + 1024 * prediction
• Le bloc transformer comporte une residual connection, donc x = x + causal_self_attention(...) ajoute le résultat au tenseur d’origine
• Comme ce signal résiduel resterait sinon inutilement présent, l’échelle 1024 sert à l’écraser
• Enfin, on calcule x @ wte.T pour obtenir les logits, puis on applique le softmax
  • dans le contexte "aabaa", la ligne de prédiction finale pointe vers b
  • pendant l’entraînement, les prédictions de toutes les lignes sont utiles, mais en inférence seule la dernière ligne est nécessaire

Résultats de génération et précision

• La fonction complete envoie au modèle les 5 derniers tokens au maximum, puis choisit le token suivant avec argmax sur la dernière ligne du résultat softmax
• Exemples de génération
  • complete("a") → a :: baabaabaab
  • complete("ba") → ba :: abaabaabaa
  • complete("abaab") → abaab :: aabaabaaba
• Même avec des entrées hors périmètre, le modèle peut parfois revenir à un motif répétitif
  • complete("ababa") → ababa :: abaabaabaa
  • complete("bbbbb") → bbbbb :: aabaabaaba
• Sur le test "aab" * 10, la précision atteint 100,0 % (27/27) si l’on n’évalue que les contextes non ambigus

4 000 FLOPs contre 8 instructions

• Lorsqu’il utilise l’intégralité d’un contexte de 5 tokens, ce modèle a besoin d’environ 4 000 floating point operations pour prédire un seul token
  • l’essentiel du coût provient du calcul de l’attention
  • il est possible de réduire ce coût avec une fenêtre de contexte plus petite, des fused multiply-add, du kv caching, etc.
  • malgré cela, une seule prédiction de token nécessite encore des centaines d’instructions machine
• Une implémentation directe de la même règle (aab)* en assembleur x64 calcule le token suivant en 8 instructions
• Cela laisse ouverte la question de savoir s’il est possible d’entraîner, pour la génération en langage naturel, un modèle de langage 1000 fois plus efficace que les modèles actuels

[GN⁺]

Avis de Hacker News

• Un travail connexe est « Thinking Like Transformers »
  Il présente un langage de programmation primitif appelé RASP, composé d’opérations modélisables avec des composants de Transformer, et montre qu’il permet d’écrire des programmes comme des histogrammes ou du tri.
  Il existe aussi d’excellents billets de blog de Sasha Rush et Gail Weiss, et des travaux ultérieurs ont montré que des programmes de type RASP pouvaient être compilés en véritables poids de modèle sans apprentissage.
  [1] https://arxiv.org/abs/2106.06981
  [2] https://srush.github.io/raspy/
  [3] https://arxiv.org/abs/2301.05062

  • J’aime vraiment beaucoup la famille RASP.
    Si ce domaine vous intéresse, mon travail HandCrafted Transformers vaut aussi le détour : j’y choisis directement les poids d’un modèle Transformer pour lui faire effectuer une addition posée d’une manière proche de celle qu’on apprend à l’école primaire.
    [1] https://colab.research.google.com/github/newhouseb/handcraft...
  • Ce genre de travail semble bien se prêter à un langage fonctionnel comme Haskell.
    Aller des réseaux de neurones vers le code serait aussi très intéressant du point de vue de l’explicabilité.

• Je pensais bien comprendre les Transformers, mais je n’en avais jamais implémenté un moi-même.
  Un jour, je m’y suis mis, et il ne fonctionnait ni ne s’entraînait aussi bien que le Transformer PyTorch standard ; j’ai fini par comprendre que la cause était que j’avais ignoré le dropout.
  Je l’entraînais à faire des additions de nombres et je pensais que le surapprentissage était impossible, puisque je ne lui avais jamais montré deux fois la même paire, mais le rôle du dropout s’est révélé bien plus important que prévu.
  En résumé, il vaut vraiment mieux implémenter soi-même un Transformer, et plus on part de zéro, mieux c’est.
  Toutes les personnes qui l’ont fait ont appris quelque chose d’inattendu, et ce qu’elles ont compris variait selon les cas, de la parallélisation de l’apprentissage au niveau des tokens jusqu’au fonctionnement réel de la rétropropagation.

  • Je me demande s’il existe des références utiles pour aborder ce travail.

• Les ressources de Karpathy sont très bonnes aussi, mais la vidéo qui m’a enfin fait comprendre les Transformers est celle-ci : https://youtu.be/kWLed8o5M2Y?si=SJT5_lCJ0hSR7Z_k

• J’ai eu des idées similaires pendant un moment.
  Ne pourrait-on pas créer une interface intuitive vers les poids d’un modèle, qui permette à un expert du domaine de les ajuster à la main pour accélérer l’apprentissage ?
  Par exemple, dans un modèle de vision qui détecte les cônes de signalisation, on augmenterait le groupe de poids correspondant à la « tendance à être orange ».
  Ainsi, au lieu de demander des milliers ou des millions d’exemples supplémentaires pour calibrer correctement cette « tendance à être orange », un humain pourrait accélérer le processus.
  Bien sûr, la difficulté est que cette interface devrait se mapper sur des groupes de poids ayant des significations différentes ; je me demande s’il y a une raison technique qui rende cela impossible.

  • « Un expert du domaine qui ajuste à la main les poids d’un modèle » ressemble à la manière dont fonctionnait la reconnaissance d’images avant le deep learning.
    [1] https://www.youtube.com/watch?v=8SF_h3xF3cE&t=1358s
  • La raison que vous cherchez s’appelle The Bitter Lesson.
    En bref, les méthodes où l’humain aide l’IA sont presque toujours moins rentables que celles qui mobilisent davantage de puissance de calcul.
    Pendant qu’un humain calibre une couche de poids pour détecter des cônes orange, un cluster de GPU entraîne l’IA à détecter des cônes de signalisation, des feux tricolores, des arbres, d’autres voitures et même des cônes d’un orange légèrement différent.
  • Le nombre de couches et de poids n’est pas à une échelle qu’un humain puisse mettre à jour manuellement, et même si c’était possible, gérer les effets en aval d’une modification des poids serait trop difficile.
    Même si l’on ajuste l’image pour mieux voir l’orange, si l’on ne peut pas surveiller en même temps la précision pour toutes les autres couleurs, on risque de créer sans s’en rendre compte des problèmes sur d’autres couleurs.
  • La raison pour laquelle c’est techniquement impossible, ou très difficile, est que les poids sont généralement très difficiles à interpréter.
    Ce n’est pas qu’un cluster précis de neurones corresponde à un concept précis ; dans l’ensemble, tout contribue un peu à tout.
  • Le mécanisme d’attention des Transformers ne semble pas se mapper facilement sur une sémantique compréhensible par les humains.
    Il y a trop de paramètres impliqués.

• L’article sur les Transformers est tellement technique que j’ai toujours voulu le comprendre, même superficiellement, mais c’était difficile.
  Ce billet m’a vraiment aidé à comprendre comment cela fonctionne, et au moins les exemples étaient très clairs.
  Il m’a aussi permis de me remémorer les matrices que j’avais apprises à l’université.

• N’est-ce pas une sorte de machine abstraite, comme une machine de Turing ou une machine qui parse des expressions régulières ?

  • En simplifiant un peu, c’est une « machine » qui mappe un ensemble d’entrées vers un ensemble de probabilités pour la sortie suivante.
    On commence par définir une liste de tokens ; pour faire simple, disons par exemple 24 caractères.
    Cette machine prend une séquence de tokens en entrée, effectue des opérations matricielles déterministes, puis produit une liste de probabilités pour tous les tokens.
    « L’apprentissage » n’est que le processus qui consiste à fixer une partie des nombres dans les matrices utilisées par ces opérations.
    Il est intéressant de noter que le code final ne contient qu’une seule instruction if, et encore, elle sert à évaluer l’exactitude du résultat.
    Toute la « logique » vient du résultat des opérations matricielles.
  • Il est assez difficile d’interpréter ce genre de choses comme des automates au sens où on l’entend généralement.
    Dans les réseaux de neurones, tout est en général un peu flou, et il n’y a presque pas de if/else, même s’il existe des cas comme dans l’exemple du Transformer où l’on « masque » des valeurs avec 0 ou -∞.
    La sortie est aussi presque toujours un ensemble de scores ou de probabilités : si un modèle qui distingue des photos de chats et de chiens produit quelque chose comme dog:0.95 cat:0.05, on dit qu’il a prédit un chien parce que le score du chien est plus élevé.
    Le mécanisme d’attention, qui est au cœur du Transformer, repose sur une sorte d’opération de recherche douce.
    Dans un système non flou, on parcourrait chaque token de la séquence pour vérifier s’il est pertinent par rapport au token courant, puis on ferait quelque chose s’il l’est ; dans un Transformer, la pertinence n’est pas une décision binaire.
    À la place, on calcule un score de pertinence continu entre toutes les paires de tokens de la séquence, puis on utilise ce score pour décider de la suite.
    Cela dit, certaines choses ne se généralisent pas facilement à partir de systèmes fondés directement sur des décisions binaires.
    Par exemple, ces scores de pertinence servent de poids pour calculer une moyenne pondérée sur les tokens du vocabulaire, ce qui donne un « token moyen » pour la position courante.
    Je ne vois pas vraiment comment interpréter facilement cela comme une extension d’un processus basé sur une logique de branchement.
  • Est-ce que ce ne serait pas juste l’AllSpark qui a touché un tas d’algèbre linéaire ?
  • Oui.
    Cet article, qui explique que les Linear Transformers sont en fait des Fast Weight Programmers, vaut le détour : https://arxiv.org/abs/2102.11174
  • Les réseaux de neurones sont des machines de Turing.
    En réglant soigneusement les poids, on peut leur faire effectuer n’importe quel calcul.
    Cela dit, ce serait bien d’avoir un compilateur qui ne repose pas sur l’approximation.

• Je me demande à quoi sert l’idée que « cela pourrait vous donner envie de fabriquer votre propre modèle », en dehors d’un exercice d’apprentissage destiné à satisfaire la curiosité.
  J’ai l’impression que les modèles de machine learning complexes commencent à devenir irréalistes à manipuler pour quelqu’un qui lit des blogs chez lui.

  • Dans nanoGPT, si l’on préentraîne un modèle sur Shakespeare, il atteint en 3 minutes un niveau de fidélité au matériau d’origine comparable au Jabberwocky de Lewis Carroll.
    Il génère beaucoup de mots pseudo-vieil anglais plausibles, et apprend les bases de la grammaire anglaise ainsi que la forme des pièces de théâtre, entre autres.
    J’ai trouvé assez étonnant qu’il aille aussi loin en si peu de temps.
    En entraînant localement plusieurs modèles jusqu’à un niveau de fidélité façon Shakespeare-from-Wish.com, on pourrait sans doute déterminer si l’on a trouvé une bonne architecture et s’il est temps d’essayer de la passer à l’échelle.
  • L’objectif est indiqué dans le premier paragraphe de l’article.
    L’auteur voulait mieux comprendre les Transformers et l’attention ; il avait lu The Illustrated Transformer, mais n’avait pas d’intuition claire sur ce que faisaient réellement les différents morceaux de l’attention.
    Il disait en substance ne pas comprendre la différence entre q et k, sans même parler de v.
  • C’est un excellent exercice d’apprentissage.
    Au-delà de la simple satisfaction de la curiosité, cela aide à construire et approfondir la compréhension.
  • Peut-être qu’on peut tout simplement vraiment aimer bidouiller ce genre de projet ? C’est bizarre, mais bon.

• Ce serait bien de pouvoir mettre une expression comme neural network dans le titre.
  Il ne s’agit pas d’un ensemble de bobines couplant électromagnétiquement deux circuits, mais de l’architecture « Transformer » du machine learning.