Preaddestrare i transformer con softmax quantizzata nell’attenzione

Abstract degli autori

I sistemi transformer a bassa precisione quantizzano sempre più spesso le moltiplicazioni tra matrici nell’attenzione, mentre la softmax rimane spesso a una precisione maggiore. Durante il preaddestramento, una softmax approssimata modifica sia i gradienti usati per addestrare il modello sia il calcolo nella passata in avanti. Studiamo questa interazione con l’attenzione a K intervalli, che approssima l’esponenziale usando K+1 valori su una griglia. Variamo la calibrazione della griglia per ciascuna riga, la scelta tra interpolazione e arrotondamento netto e la posizione di un surrogato straight-through rispetto alla normalizzazione. Deriviamo le corrispondenti regole per la passata all’indietro, incluse le derivate della calibrazione, e confrontiamo queste scelte in esperimenti di preaddestramento a parità di modello, dati e ottimizzatore. Escludere gli estremi di ciascuna riga dal calcolo del gradiente lascia invariato il calcolo nella passata in avanti, ma provoca un aumento ritardato della perdita di validazione. Con arrotondamento netto e K=4, la calibrazione min-max e un surrogato prima della normalizzazione producono un ampio divario nella perdita; modificare una delle due scelte lo riduce sensibilmente. Con 124 milioni di parametri e 2,5 miliardi di token di addestramento, la calibrazione a finestra fissa con un surrogato dopo la normalizzazione produce, a K=4, una differenza nella perdita di validazione di +0,019 nat rispetto alla softmax; con un surrogato prima della normalizzazione, la differenza è di +0,004 nat a K=16.