Migliorare lo scaling in fase di test con transformer a cicli adattivi

Abstract degli autori

I transformer a cicli hanno mostrato un’efficienza promettente nell’uso dei parametri riutilizzando i livelli per il calcolo latente. Gli studi precedenti confrontano modelli con e senza cicli a parità di parametri o di FLOP per token. Tuttavia, per quanto ne sappiamo, resta poco studiato se i cicli migliorino lo scaling in fase di test quando gli output diventano più lunghi. Sottoponendo i transformer a cicli a post-addestramento, studiamo la pendenza della curva accuratezza-calcolo, misurata come il guadagno di accuratezza a ogni raddoppio dei FLOP di decodifica in fase di test. Scopriamo che i transformer a cicli esistenti presentano spesso pendenze maggiori rispetto ai corrispondenti modelli senza cicli, ma ottengono risultati peggiori a parità di calcolo. Mentre i cicli a profondità fissa impiegano iterazioni aggiuntive per ogni token, la nostra analisi mostra che molti token non ne traggono beneficio. Proponiamo quindi TaH2, che consente al modello di concentrare le iterazioni aggiuntive sui token che beneficiano dei cicli. TaH2 sottopone congiuntamente a post-addestramento il backbone e un modulo che decide il numero di iterazioni, usando una supervisione della profondità basata su una valutazione anticipata e su etichette generate online che indicano se un’ulteriore iterazione migliora la previsione. TaH2 migliora sia l’efficienza sia l’accuratezza raggiungibile dello scaling in fase di test. Nei difficili benchmark AIME, TaH2 migliora la pendenza della curva accuratezza-calcolo del 53% (2,74 contro 1,79) rispetto al modello senza cicli, superando di circa 3,4 punti la massima accuratezza di quest’ultimo a parità di calcolo in fase di test. All’aumentare della profondità massima delle iterazioni, le prestazioni dei modelli a cicli esistenti si stabilizzano in larga misura, mentre il vantaggio di TaH2 rispetto al modello senza cicli continua a crescere, da +2,8 punti alla profondità 2 a +3,9 punti alla profondità 8. Il nostro codice è disponibile all’indirizzo https://github.com/thu-nics/TaH.