Addestramento di un trasformatore in 1,5 ore supera i LLM nel benchmark ARC-AGI

TL;DR

Il ricercatore indipendente Mithil Vakde ha addestrato un piccolo trasformatore da zero in appena 1,5 ore su una singola RTX 5090, raggiungendo il 44% nel benchmark ARC-AGI-1 per un costo computazionale totale di 67 centesimi. Il modello supera molti grandi modelli linguistici e eguaglia le prestazioni dei Tiny Recursive Models (TRM) e degli Hybrid Recursive Models (HRM), sfidando le ipotesi su scala ed efficienza dei campioni nell'AI.

La svolta: modello piccolo, grandi risultati

Vakde, laureato all'IIT Bombay, ha addestrato il suo trasformatore utilizzando un approccio innovativo che enfatizza l'efficienza dei campioni rispetto alla scala grezza. L'intero addestramento è costato solo 67 centesimi per il noleggio di GPU cloud, e il modello ottiene anche un 7% nel benchmark ARC-2, più difficile. Si tratta di un aggiornamento significativo rispetto al suo modello precedente, che era più lento, più costoso e meno accurato.

Il lavoro ha già attirato l'attenzione di ricercatori di spicco, tra cui Lucas Beyer, Jeremy Howard e Rohan Anil, con discussioni diventate virali su X (ex Twitter). Il codice è open source e disponibile su GitHub, rendendolo accessibile a ricercatori di tutto il mondo per sviluppi futuri.

Come funziona: dettagli tecnici

L'approccio converte ogni coppia input-output in una sequenza di token, che vengono poi addestrati in modo autoregressivo da un piccolo trasformatore. Questo viene fatto da zero al momento del test sia sul set di addestramento che su quello di valutazione, con le etichette di test nascoste. Ogni puzzle riceve un embedding additivo separato per consentire l'apprendimento tra compiti, e le informazioni posizionali sono codificate utilizzando embedding 3D RoPE.

I miglioramenti architetturali chiave rispetto al suo lavoro precedente includono:

• Architettura moderna: attivazione SwiGLU invece di GELU, RMSNorm invece di LayerNorm
• Scala aumentata: 8 strati invece di 4
• Migliore diversità dei dati: strategie migliorate di mescolamento e aumento
• Ottimizzatore NorMuon: sostituzione di AdamW, riducendo significativamente i costi di addestramento
• Attenzione flash: con addestramento varlen e kernel di attenzione flessibile per l'inferenza

Un risultato sorprendente è stato che l'addestramento solo sui token di output (stile supervisionato) funziona leggermente meglio rispetto all'inclusione dei token di input (non supervisionato), nonostante produca una perdita di test peggiore. Vakde ammette di non capire appieno il motivo, ma sospetta che sia legato alla capacità finita del modello.

Efficienza dei campioni: la motivazione principale

La motivazione principale di Vakde è affrontare quello che chiama "il problema più importante nell'AI oggi": l'efficienza dei campioni. ARC-AGI è un benchmark ideale per questo perché ha pochissimi campioni (solo 1.000 puzzle) in uno spazio ad alta dimensionalità, rendendolo un vero test di generalizzazione piuttosto che di memorizzazione.

Il benchmark è anche progettato per essere un test di meta-apprendimento, dove ogni puzzle utilizza una regola diversa ma condivide concetti comuni. Questo lo rende incredibilmente facile per gli umani ma impegnativo per i sistemi di AI che si basano sul riconoscimento di pattern da enormi dati di addestramento.

Comportamento interessante e ablazioni

Il modello mostra alcuni comportamenti controintuitivi. Nonostante una perdita di test peggiore, ottiene punteggi migliori nel benchmark, suggerendo che l'ottimizzazione per la perdita di validazione potrebbe non allinearsi sempre con le prestazioni effettive nei compiti di ragionamento. Questo indica un potenziale punto debole nelle attuali metodologie di ricerca sull'efficienza dei campioni.

Studi di ablazione hanno rivelato che i maggiori contributori alle prestazioni sono gli embedding posizionali 3D RoPE e gli embedding per compito. Rimuovere uno dei due fa scendere il punteggio dal 44% a circa il 24-25%. Altre ablazioni hanno mostrato che:

• L'addestramento sugli input funziona leggermente peggio (~39%)
• Limitare l'addestramento a ARC-1+ConceptARC funziona più o meno allo stesso modo (~40%)
• L'addestramento in stile CompressARC riduce le prestazioni a circa il 18%

Affrontare critiche e controversie

L'approccio ha affrontato critiche da parte di alcuni ricercatori, in particolare riguardo all'uso dell'addestramento al momento del test sugli input dei puzzle di valutazione. Vakde affronta direttamente queste preoccupazioni, sostenendo che si tratta di una forma di ragionamento transduttivo, un concetto ben consolidato nell'apprendimento automatico che risale a Vapnik.

Chiarisce inoltre che la politica di test ARC consente l'addestramento al momento del test e che la critica deriva da un fraintendimento della politica. L'intento della politica è prevenire che i progettisti umani utilizzino le informazioni del set di valutazione per costruire bias induttivi, non impedire ai sistemi di AI di apprendere durante l'inferenza.

Perché è importante

Questo risultato sfida la saggezza prevalente secondo cui risolvere ARC-AGI richiede modelli massicci, architetture ricorsive o dati sintetici estesi. Il trasformatore di Vakde, addestrato in meno di due ore per meno di un dollaro, eguaglia le prestazioni di modelli che utilizzano approcci significativamente più complessi.

Le implicazioni vanno oltre ARC-AGI. Suggerisce che l'efficienza dei campioni, non solo la scala, potrebbe essere la chiave per avanzare le capacità di ragionamento dell'AI. Per un'industria sempre più focalizzata su addestramenti massicci e budget computazionali enormi, questo lavoro dimostra che approcci più intelligenti possono ottenere risultati comparabili o migliori a una frazione del costo.

Prossimi passi

Vakde ritiene che il 65% su ARC-AGI-1 sia raggiungibile all'interno del framework del trasformatore e invita la comunità di ricerca a contribuire. Suggerisce diverse strade per il miglioramento, tra cui nuovi schemi di codifica posizionale come PoPE, un'ulteriore modernizzazione dell'architettura e l'eliminazione degli aumenti dei dati.

La natura open source del lavoro significa che ricercatori di tutto il mondo possono iterare sul suo approccio. Con costi computazionali così bassi, la barriera all'ingresso per la ricerca su ARC-AGI è drasticamente diminuita, democratizzando potenzialmente l'accesso alla ricerca all'avanguardia sul ragionamento nell'AI.