Dentro il loop: uno studio meccanicistico sui transformer con pesi condivisi per compiti algoritmici vincolati dalla profondità

Abstract degli autori

Risolvere compiti algoritmici con molti passaggi sequenziali richiede in genere di aumentare la profondità dei transformer o la lunghezza della catena di pensiero, con costi rispettivamente sfavorevoli in termini di parametri e di calcolo durante l’inferenza. I transformer con ciclo, che riutilizzano un blocco di layer per molte iterazioni, sono in grado di affrontare questo problema, poiché aumentano la profondità effettiva senza aggiungere parametri o token. Tuttavia, non è ancora chiaro che cosa calcoli effettivamente ogni iterazione, né se il ciclo si limiti a simulare una pila più profonda di layer senza condivisione dei pesi. Qui studiamo l’attraversamento di permutazioni dirette, in cui un modello segue $H$ passaggi a partire da un nodo iniziale e indica il nodo raggiunto, così che il risultato corretto di ogni passaggio intermedio sia noto. Scopriamo che per risolvere il compito è necessaria una profondità effettiva minima, oltre la quale il fattore determinante diventa il grado di condivisione dei pesi. Per esempio, un transformer con ciclo e un blocco di quattro layer, iterato otto volte, risolve il compito con ogni seed, mentre un transformer senza condivisione dei pesi, addestrato in modo identico e con la stessa profondità effettiva, non ci riesce mai: ha successo solo con un learning rate più basso o supervisionando i suoi layer intermedi. Esaminando e modificando il flusso residuo, scopriamo che il modello con ciclo procede gradualmente attraverso i nodi intermedi, formando una scala che si adatta al suo numero di iterazioni, mentre il modello senza condivisione dei pesi, con quel learning rate più basso, li rende visibili solo nei layer finali. Nel complesso, i nostri risultati descrivono i transformer con ciclo come modelli che apprendono programmi iterativi: allocano le proprie iterazioni ai passaggi successivi dell’algoritmo e risolvono compiti che i modelli senza condivisione dei pesi, addestrati in modo identico e con la stessa profondità effettiva, non riescono a risolvere. Il codice è disponibile all’indirizzo https://github.com/magnus-sesodia/insidetheloop.