El investigador Vakde ha presentado una nueva versión de su modelo para el benchmark ARC-AGI-1, un pequeño transformer entrenado desde cero en solo 1,5 horas con una GPU Nvidia RTX 5090. El sistema alcanza un 44% de precisión en ARC-AGI-1 y un 7% en ARC-AGI-2, con un coste total de aproximadamente 67 centavos por ejecución. El rendimiento iguala al de modelos como TRM y HRM y supera al de varios modelos de lenguaje de gran tamaño, según explica el autor en su blog técnico.
El enfoque se basa en entrenar de forma autónoma un transformer reducido sobre cada puzle, convirtiendo los pares de entrada y salida en secuencias de tokens. Cada puzle recibe su propio embedding aditivo y se usan codificaciones posicionales 3D RoPE para representar las dos cuadrículas bidimensionales. Durante la inferencia se aplican aumentaciones de color y permutaciones diédricas, y se envían las dos salidas más frecuentes. Entre las mejoras respecto a la versión anterior destacan la adopción de SwiGLU, RMSNorm, 8 capas en lugar de 4, el optimizador NorMuon y kernels de Flash Attention y Flex Attention para acelerar el entrenamiento y la inferencia.
El autor también añade al entrenamiento los puzles no solapados de ARC-2, filtrando cuidadosamente las 773 tareas repetidas para evitar fuga de datos. Las ablaciones muestran que las representaciones 3D RoPE y los embeddings por tarea son responsables de la mayor parte del rendimiento. El código está disponible en GitHub como proyecto de código abierto, con el objetivo de impulsar la investigación en eficiencia de muestras y reducir los costes de experimentación en ARC-AGI.