MassAlloc Attention: lasciare che l’attenzione distribuisca le proprie risorse di calcolo

Abstract degli autori

FullAttn assegna spesso una massa normalizzata trascurabile a gran parte dello spazio dei punteggi causali, eppure i kernel densi eseguono l’intera sequenza di operazioni successive al calcolo dei punteggi dopo aver formato ogni blocco QK. Presentiamo MALA, una primitiva di attenzione integrata che preserva l’accesso ai punteggi di tutte le interazioni causali consentite e usa il contributo normalizzato per distribuire il calcolo successivo. Nel passaggio forward usa il fattore di normalizzazione della softmax online, che si aggiorna progressivamente; nel passaggio backward riutilizza il fattore di normalizzazione definitivo per ricavare supporti conservati annidati, usando solo lo stato standard dell’attenzione. Una stessa soglia di tolleranza regola l’addestramento e l’inferenza, consentendo di selezionare in modo adattivo le operazioni da eseguire. MALA riduce il calcolo successivo ai punteggi associato ai contributi più bassi. Uno studio a 8K, condotto a parità di lavoro, isola il vantaggio di distribuire il calcolo in modo adattivo in base alla distribuzione: a parità esatta di lavoro complessivo successivo al calcolo dei punteggi, MALA si avvicina a un oracolo basato sulla massa di riferimento per ciascuna istanza, con una massa omessa media dello 0,0188% contro lo 0,0182%. Con lunghezze di contesto da 1K a 32K token, la stessa soglia di tolleranza mantiene bassi gli errori negli output e nei gradienti rispetto al riferimento. In un confronto controllato più ampio su un compito di richiamo associativo, MALA segue da vicino FullAttn all’aumentare del contesto, raggiungendo un’accuratezza dell’89,67% a 8K, contro l’89,97% di FullAttn. In un benchmark dell’operatore di attenzione a 128K token con parallelismo dei tensori, MALA riduce, rispetto a FullAttn, la latenza dei passaggi forward e backward durante l’addestramento di un fattore rispettivamente pari a 2,2 e 3,0 e quella della decodifica durante l’inferenza di un fattore pari a 1,6. Nell’addestramento secondo leggi di scala, da 0,6B a 14B parametri, MALA resta vicino a FullAttn in termini di perplexity, riducendo al contempo i FLOPs complessivi dell’addestramento. I modelli da 14B così ottenuti e i modelli da 32B frutto di un addestramento continuato separato raggiungono punteggi comparabili a quelli di FullAttn nella conoscenza, nel ragionamento e nel recupero di informazioni in contesti lunghi. Questi risultati indicano che distribuire il calcolo successivo ai punteggi in base ai contributi normalizzati dell’attenzione può preservare le capacità valutate di FullAttn, riducendo al contempo il calcolo richiesto dall’attenzione.