Transformer ricorrente a cicli

Abstract degli autori

Tenere traccia dello stato richiede un aggiornamento a ogni input, ma il numero di strati che un Transformer applica a ciascun token è fisso, indipendentemente dalla lunghezza della sequenza. Presentiamo il Recurrent Looped Transformer (RLT), che ripartisce i suoi strati tra un encoder causale parallelo e un decoder ricorrente. A ogni token, il decoder combina l’output dell’encoder con lo stato finale del decoder relativo al token precedente: così il percorso di calcolo cresce con la lunghezza della sequenza, mentre il costo per token resta fisso. Su sei compiti algoritmici, confrontiamo cinque ripartizioni di otto strati con un Transformer a otto strati, usando tre seed. Dopo un addestramento su sequenze di al massimo 40 bit, due ripartizioni di RLT generalizzano il calcolo della parità a 256 bit con un’accuratezza del 100% per ciascun seed, mentre il Transformer resta al livello del caso. Nel tracciamento delle permutazioni di $S_5$ basato su scambi, su sequenze otto volte più lunghe di quelle di addestramento, RLT raggiunge un’accuratezza dello stato finale del 97%, contro meno dell’1% del Transformer; l’accuratezza aumenta con il numero di strati del decoder. Nell’aritmetica modulare su sequenze più lunghe di quelle di addestramento, RLT arriva fino al 93%, contro il 33% del Transformer. Gli studi di ablazione mostrano che questi miglioramenti dipendono dalla retroazione: eliminarla porta la parità e il tracciamento delle permutazioni di $S_5$ basato su scambi al livello del caso per ogni ripartizione. Aggiornare la retroazione una volta per ogni blocco di quattro token consente di elaborare in parallelo i token noti di un blocco e mantiene al 99% l’accuratezza della parità su 64 bit. Il tracciamento delle permutazioni, invece, dipende dalla retroazione a ogni token: la suddivisione in blocchi riduce l’accuratezza nel tracciamento delle permutazioni di $S_5$ basato su scambi, con sequenze di lunghezza 64, dal 100% al 20%.