onóna: embedding gemelli con unembedding condivisi

Abstract degli autori

La condivisione dei pesi riduce il costo in parametri dei modelli linguistici, ma la consueta testa di output condivisa è limitata a una singola softmax. Introduciamo onóna, una parametrizzazione a due tabelle che aumenta l’espressività dell’output mantenendo la condivisione dei pesi tra input e output. Il Transformer riceve in input la media di due tabelle di embedding e la testa di output combina le distribuzioni softmax definite da quelle stesse tabelle. Un gate condizionato esclusivamente al token di input corrente determina il peso della miscela. I suoi valori possono essere precalcolati per l’inferenza e il corpo del Transformer resta invariato. Valutiamo onóna con architetture Dense, ibride Gated DeltaNet, mixture-of-experts e Attention Residual, alle scale di 142M, 330M e 1B. A parità di budget di token, onóna riduce la loss finale di validazione e la perplexity esterna in 11 confronti su 12 e aumenta l’accuratezza media su otto task downstream in 10 confronti su 12. L’unica eccezione per la loss è il modello MoE da 1B. Il costo è una tabella di embedding aggiuntiva e una proiezione aggiuntiva sul vocabolario.