Transformer ARC‑AGI‑1 atteint 44 % en 1,5 h sur RTX 5090

Contexte et objectifs

L’auteur entraîne un petit transformeur de 8 couches sur le benchmark ARC‑AGI‑1, composé de 1 000 puzzles de méta‑apprentissage. L’objectif déclaré est de mesurer la limite d’efficacité d’échantillonnage des transformeurs classiques tout en réduisant le coût d’expérimentation. Le modèle est entraîné depuis zéro pendant 1,5 heure sur une carte RTX 5090, pour un coût estimé à 0,67 $, et obtient 44 % de réussite sur la partie ARC‑1 et 7 % sur ARC‑2.

Architecture et améliorations

Le pipeline convertit chaque paire entrée‑sortie en une séquence de tokens, puis applique un entraînement autorégressif. Deux innovations majeures sont introduites : un embedding additif propre à chaque puzzle et des embeddings positionnels 3D RoPE, qui combinent coordonnées spatiales et profondeur du token. Les séquences sont également augmentées par des permutations de couleur et diédriques, puis les deux sorties les plus fréquentes sont soumises (AAIVR).

Parmi les changements d’architecture, on trouve le remplacement de GELU par SwiGlu, l’usage de RMSnorm à la place de LayerNorm, et le passage de 4 à 8 couches. Au niveau de l’optimisation, AdamW est remplacé par l’optimiseur NorMuon, et le modèle utilise flash‑attention avec entraînement à longueur variable et kernels flex‑attention pour l’inférence. Enfin, la perte ne porte plus sur les tokens d’entrée ; seul le token de sortie est considéré, ce qui rend l’entraînement supervisé.

optimizer = NorMuon(model.parameters(), lr=1e-3, weight_decay=0.01)

Analyse des performances et ablations

Le passage à 44 % provient principalement de la combinaison d’une meilleure représentation (3D RoPE + embedding par tâche) et d’une architecture modernisée. Les ablations montrent que remplacer 3D RoPE par une version 1D fait chuter le score à ~24 %, tout comme la suppression des embeddings par tâche. Entraîner sur les tokens d’entrée réduit le score à ~39 %, confirmant que la supervision partielle améliore la généralisation malgré une perte de test plus élevée.

Lorsque le jeu de données est limité à ARC‑1 + ConceptARC, le score reste autour de 40 %, indiquant que l’ajout des tâches non chevauchantes d’ARC‑2 n’est pas indispensable mais double le temps de calcul si l’on veut atteindre 44 %. Les variantes CompressARC, qui entraînent le modèle séparément sur chaque puzzle, obtiennent seulement 15‑18 % de réussite, soulignant l’importance du partage de représentation entre tâches.

Implications et perspectives

Ces résultats suggèrent que, même avec un transformeur de taille modeste, il est possible d’atteindre des performances compétitives sur un benchmark de méta‑apprentissage sans recourir à des données synthétiques massives. La réduction du nombre d’augmentations et l’utilisation d’un optimiseur dédié permettent de diminuer le coût énergétique, ouvrant la voie à des itérations rapides pour la communauté. L’auteur propose d’explorer de nouveaux schémas d’embeddings (PoPE) et d’optimiser davantage le code GPU afin de réduire le facteur de coût d’environ 10×. Atteindre 65 % reste présenté comme un objectif réaliste, à condition de conserver la structure d’embedding et les améliorations d’architecture identifiées.