J'ai passé quelques heures à parcourir le dépôt /karpathy/autoresearch ligne par ligne. L'angle des "agents IA faisant de la recherche" est celui qui attire toute l'attention, mais je pense que ce qui est vraiment intéressant, c'est ce qu'il y a à l'intérieur du script d'entraînement et les décisions d'ingénierie qui rendent la boucle de recherche efficace. C'est l'un des setups d'entraînement en fichier unique les plus denses que j'ai lus. Commençons par ce qui rend tout le projet possible : le budget de temps est fixé à 300 secondes en temps réel. Pas de pas fixes, pas de tokens fixes, pas de flops fixes. Des secondes en temps réel. Cela peut sembler un détail mineur, mais c'est la raison principale pour laquelle la boucle autonome fonctionne. L'agent peut rendre le modèle 3 fois plus grand, réduire la taille du lot de moitié, échanger une architecture complètement différente, et le résultat reste directement comparable à chaque autre expérience car elles ont toutes eu exactement 5 minutes d'entraînement sur le même GPU. Si vous fixez les pas à la place, un modèle plus grand obtiendrait moins de mises à jour de gradient par seconde et vous le pénaliseriez de manière injuste. Si vous fixez les tokens, vous auriez le même problème. Fixer le temps réel signifie que vous posez la bonne question : étant donné ce matériel et ce temps, quel est le meilleur modèle que vous pouvez produire ? Tout le reste est une variable libre. L'agent peut explorer toute la surface de Pareto de la taille du modèle par rapport au débit par rapport à la vitesse de convergence sans que ces compromis soient confondus par le protocole d'évaluation. La métrique est également soigneusement choisie. C'est des bits par octet, pas de perte d'entropie croisée. La perte d'entropie dépend de la taille de votre vocabulaire. Un modèle avec 32k tokens et un modèle avec 8k tokens auront des valeurs de perte très différentes même s'ils compressent les données de manière équivalente. Le bpb normalise cela en additionnant l'entropie croisée par token en nats, en additionnant les longueurs en octets UTF-8 des tokens cibles, et en convertissant les nats par octet en bits par octet. Donc même si l'agent change quelque chose qui affecte la distribution effective des tokens, la comparaison reste équitable. Ces deux choix, le temps réel fixe et une métrique invariante par vocabulaire, transforment ce qui serait une recherche difficilement comparable en un problème d'optimisation clair. Maintenant, le modèle lui-même. C'est un GPT mais avec un tas de techniques modernes qui valent la peine d'être comprises. D'abord, RMSnorm partout. Sur les entrées de bloc (pré-norm), et aussi sur les requêtes et les clés juste avant le produit scalaire d'attention. Ce truc de QK-norm est important car sans lui, les normes de q et k peuvent croître de manière illimitée pendant l'entraînement, provoquant un affûtage des logits d'attention et une saturation du softmax. Normaliser q et k maintient les produits scalaires dans une plage stable, peu importe la profondeur du réseau ou comment les dynamiques d'entraînement évoluent. L'attention elle-même est FA 3, chargée via la bibliothèque kernels. Elle utilise l'implémentation de varunneal sur hopper (sm_90) et revient à une version communautaire sur des GPU plus anciens. Le modèle d'attention est "SSSL", ce qui signifie trois couches d'attention à fenêtre glissante (fenêtre = la moitié de la longueur de la séquence) suivies d'une couche d'attention causale complète, répétée. C'est le modèle sparse-to-dense que vous voyez dans mistral et gemma2. Les couches d'attention locale sont peu coûteuses sur le plan computationnel car la matrice d'attention est en bande, et la couche globale périodique permet à l'information de circuler à travers le contexte complet. Avec 8 couches et un motif de 4 caractères, vous obtenez les couches 0,1,2 locales, la couche 3 globale, les couches 4,5,6 locales, la couche 7 globale. La dernière couche est forcée globale indépendamment du motif. La chose d'embedding de valeur est subtile et je pense sous-estimée. Chaque autre couche obtient sa propre table d'embedding, complètement séparée de l'embedding principal des tokens, qui mappe les ids de tokens directement à des vecteurs de dimension de valeur. Ceux-ci sont mélangés dans les valeurs d'attention à travers une porte apprise : v = v + 2 * sigmoid(W_gate @ x:32) * ve. Le poids de la porte est initialisé à zéro, donc sigmoid(0) = 0.5, multiplié par 2 donne 1.0, ce qui est un point de départ neutre. Au cours de l'entraînement, le modèle peut apprendre à amplifier ou à supprimer l'embedding de valeur par tête en fonction des 32 premières dimensions de l'état caché. Cela vient de la ligne de travail ResFormer et l'intuition est que cela donne à l'attention un raccourci direct vers l'identité du token. Les vecteurs de valeur peuvent porter des informations sur "quel token est à cette position" sans que cette information ait à survivre aux transformations du flux résiduel des couches précédentes. C'est essentiellement une connexion de saut de l'entrée directement dans les valeurs d'attention, contrôlée pour que le modèle puisse décider quand c'est utile. Il y a aussi des scalaires apprenables par couche sur le flux résiduel : x = lambda_residi * x + lambda_x0i * x0, où x0 est l'embedding normalisé de la couche 0. Chaque couche peut contrôler indépendamment combien elle écoute le résiduel en cours par rapport à l'entrée originale. Les lambdas résiduels commencent à 1.0, les lambdas x0 commencent à 0.1. C'est une version douce de l'idée de "résiduel désentrelacé". Dans un transformateur standard, le flux résiduel est une somme de toutes les sorties des couches précédentes et il devient de plus en plus pollué à mesure que vous allez plus profondément. Donner à chaque couche accès à l'embedding original propre signifie qu'elle n'a pas à apprendre à "défaire" les couches précédentes pour récupérer des informations de bas niveau. Les logits sont softcaps à 15 via tanh(logits/15)*15, ce qui empêche le modèle d'être trop confiant au début de l'entraînement lorsque les représentations sont encore bruyantes. Mais honnêtement, la partie la plus intéressante de tout le fichier est l'optimiseur. MuonAdamW est un optimiseur combiné qui dispatches différentes règles de mise à jour en fonction du groupe de paramètres. Les embeddings (embedding de token, embeddings de valeur, tête de désembedding) et les scalaires par couche obtiennent un AdamW standard avec des taux d'apprentissage différents pour chaque groupe. L'écart est sauvage. Le taux d'apprentissage d'embedding est de 0.6, le taux d'apprentissage de désembedding est de 0.004, c'est une différence de 150x, et c'est intentionnel. La matrice d'embedding voit chaque token et doit se mettre à jour de manière agressive. La matrice de désembedding est une sonde linéaire sur la représentation finale et bénéficie de la stabilité. Les taux d'apprentissage d'embedding, d'embedding de valeur et de désembedding sont tous ajustés par (d_model / 768)^(-0.5), ce qui est une correction inspirée par muP. À mesure que la largeur du modèle change, ces taux d'apprentissage s'ajustent pour maintenir les dynamiques d'apprentissage des caractéristiques invariantes à l'échelle. Les taux d'apprentissage scalaires pour les lambdas par couche sont gérés séparément et ne reçoivent pas cette mise à l'échelle. Les matrices de poids 2D dans le transformateur, les projections d'attention et les poids MLP, obtiennent Muon, et c'est là que cela devient vraiment intéressant. Muon prend le gradient, applique un moment de Nesterov, puis exécute une itération de Newton-Schulz pour approximer la décomposition polaire de la matrice de gradient. La décomposition polaire factorise une matrice G en G = U * S où U est orthogonal et S est semi-défini positif symétrique. Muon calcule U, la matrice orthogonale la plus proche du gradient, et l'utilise comme direction de mise à jour. L'itération de Newton-Schulz est de 5 étapes. Pour les matrices hautes (plus de lignes que de colonnes), A = X^T @ X puis X -> aX + X @ (bA + cA^2). Pour les matrices larges, A = X @ X^T puis X -> aX + (bA + cA^2) @ X. Les coefficients sont codés en dur à partir d'une pré-computation. Ils l'appellent "polar express". Le tout se compile en un seul noyau fusionné via torch.compile. Pourquoi cela importe-t-il ? Parce que pour les matrices de poids, le gradient de norme de Frobenius (ce que Adam et SGD utilisent) est géométriquement incorrect. La direction de descente la plus raide "correcte" pour une matrice de poids est celle qui minimise la perte sous la contrainte que la mise à jour a une norme spectrale unitaire, pas une norme de Frobenius unitaire. Le facteur polaire orthogonal vous donne exactement cela. En pratique, cela signifie que Muon effectue des mises à jour effectives beaucoup plus grandes car il ne gaspille pas la taille de pas à échelle des valeurs singulières. Il ne fait que les faire tourner. C'est pourquoi Muon converge significativement plus vite qu'Adam sur les matrices de poids des transformateurs. Muon maintient des tampons de moment par élément (même forme que les paramètres, empilés à travers chaque groupe de forme), mais contrairement à Adam, il ne suit pas les secondes moments par élément. Les estimations de seconde moment sont par ligne ou par colonne après orthogonalisation, pas par élément. C'est là que NorMuon entre en jeu. Au-dessus de la base Muon, il y a NorMuon, un schéma de réduction de variance. Après orthogonalisation, il calcule des estimations de seconde moment par ligne (ou par colonne selon le rapport d'aspect), maintient une moyenne mobile exponentielle de celles-ci, et redimensionne la mise à jour afin que chaque dimension de sortie obtienne sa propre taille de pas adaptative. C'est essentiellement l'idée d'adaptivité d'Adam mais appliquée dans le système de coordonnées orthogonalisé plutôt que dans l'espace de paramètres brut. La décadence de poids est également non standard. Elle est "cautieuse", ce qui signifie qu'elle ne décaye que les paramètres où la direction de mise à jour de Muon est d'accord avec le signe du paramètre : mask = (g * params) >= 0. Cela évite le mode de défaillance connu où la décadence de poids pousse les paramètres vers zéro contre les souhaits de la mise à jour, ce qui peut déstabiliser l'entraînement. Un petit détail que j'ai apprécié : après la toute première étape d'entraînement, le code appelle gc.collect(), gc.freeze(), gc.disable() pour complètement désactiver le ramasse-miettes de Python. Le GC de Python s'exécute périodiquement et provoque des pauses d'environ 500 ms. Lorsque votre budget total est de 300 secondes et que chaque étape prend peut-être 300 ms, une pause aléatoire du GC vous coûte presque 2 étapes d'entraînement. Ils déclenchent manuellement gc.collect() tous les 5000 étapes comme compromis. C'est le genre de chose que vous apprenez seulement en profilant de véritables exécutions d'entraînement et en remarquant des baisses mystérieuses de débit. Les 11 premières étapes (0 à 10) ne sont pas comptées dans le budget de temps non plus. C'est le réchauffement où torch.compile fait son travail et les noyaux CUDA sont JITés. Sans cette exclusion, différentes expériences obtiendraient des montants différents de "vrai" entraînement en fonction de la durée de la compilation pour cette configuration de modèle particulière. Encore une fois, un choix de conception qui semble petit mais qui est critique pour rendre les expériences comparables. Maintenant, zoomons. La véritable boucle d'autorecherche est : l'agent lit program.md (un fichier markdown qui décrit son travail), modifie train.py, s'engage, exécute pendant 5 minutes, vérifie si val_bpb s'est amélioré, garde ou revient en arrière, répète. program.md dit explicitement "NE JAMAIS S'ARRÊTER". L'agent fonctionne indéfiniment jusqu'à ce qu'un humain l'arrête. ~12 expériences par heure, ~100 pendant que vous dormez. ...