SMELT: leggi di scala per transformer MoE con strati ripetuti a parità di calcolo

Abstract degli autori

I transformer con strati ripetuti aumentano la profondità effettiva eseguendo più passaggi attraverso un blocco di strati condiviso. La maggior parte delle valutazioni, però, confronta modelli della stessa dimensione, confondendo il vantaggio dell’architettura con l’impiego di FLOPs aggiuntivi. Studiamo la ripetizione degli strati nei transformer Mixture-of-Experts, mantenendo strettamente comparabili i FLOPs per token, il numero totale di parametri esclusi gli embedding e la cache KV. Attraverso una serie di studi di ablazione, arriviamo a una configurazione che chiamiamo SMELT (Sparse MoE Transformer, middle layers Loop Twice): esegue due passaggi attraverso la metà centrale degli strati, rispettando gli stessi tre vincoli del modello Baseline senza ripetizioni. Applichiamo SMELT a quattro dimensioni, fino a 54 miliardi di parametri esclusi gli embedding, e stimiamo per ciascuna architettura una distinta legge di scala sul modello di Chinchilla. All’aumentare del calcolo, la loss di SMELT diminuisce più rapidamente, con un risparmio del 6,8–18,0\% dei FLOPs di addestramento sulla frontiera ottimale rispetto al calcolo. Il vantaggio si ritrova nei benchmark downstream in misura superiore a quanto previsto dalla loss di validazione, è massimo su Code e cresce con la lunghezza dei campioni e il numero di esempi nel contesto. Un’analisi dei meccanismi mostra che il secondo passaggio riduce l’attention sink e reindirizza la massa dell’attenzione verso i token pertinenti al contenuto: un bias induttivo che potrebbe spiegare i miglioramenti osservati nelle prestazioni. Questi risultati mostrano che ripetere gli strati può migliorare i transformer anche a parità di risorse, offrendo una configurazione pratica che trasforma il riutilizzo della profondità in vantaggi misurabili.