SAS: sparsificazione semplice dell’attenzione mediante ottimizzazione end-to-end dell’ordinamento del contesto

Abstract degli autori

La sparsificazione dell’attenzione dopo l’addestramento riduce il costo cumulativo quadratico dell’attenzione nei Transformer preaddestrati, selezionando per ogni query un piccolo insieme di unità di contesto (token o blocchi). I metodi addestrabili esistenti usano in genere un selettore leggero per assegnare un punteggio alle unità di contesto, seguito da una selezione rigida Top-K che impedisce ai gradienti della funzione di perdita della modellazione linguistica di raggiungere il selettore. Di conseguenza, questi metodi ricorrono comunemente alla distillazione delle distribuzioni di attenzione densa di ciascun livello. Questo incoraggia il selettore a ordinare le unità di contesto in base ai pesi dell’attenzione densa nel modello originale, ma tale ordinamento non è direttamente allineato al loro impatto sulle predizioni con un budget di attenzione fisso (ossia il numero di unità di contesto a cui ogni query presta attenzione). Il budget limitato può quindi essere speso per unità meno utili. Per risolvere questo disallineamento, proponiamo Simple Attention Sparsification (SAS), un meccanismo di attenzione sparsa con gate che ottimizza end-to-end l’ordinamento del contesto mediante la funzione di perdita della modellazione linguistica. L’idea centrale è inserire i punteggi continui del selettore nei logit dell’attenzione durante l’addestramento, consentendo alla funzione di perdita di aggiornare il selettore tramite la normale retropropagazione. Individuiamo alcune scelte cruciali perché questo semplice approccio funzioni bene nella pratica: collocare il gate all’interno della softmax dell’attenzione in forma logaritmica, usare gate softmax normalizzati per calibrare il contesto precedente rispetto al blocco corrente, che viene sempre mantenuto, e conservare i punteggi continui del selettore affinché il modello apprenda le priorità relative anziché soltanto le selezioni rigide. Per consentire l’addestramento su sequenze lunghe, implementiamo un kernel Triton efficiente nell’uso della memoria che integra SAS nel calcolo in stile FlashAttention. Nei compiti di ragionamento, comprensione di contesti lunghi e compiti agentici, SAS supera sistematicamente i metodi di riferimento addestrabili di attenzione sparsa con diversi budget di attenzione, con vantaggi particolarmente ampi quando il budget è ristretto. I risultati dimostrano un ordinamento del contesto più efficace per i compiti a valle.