Kalman Delta Networks: memoria associativa consapevole dell’incertezza

Abstract degli autori

L’attenzione lineare è sempre più utilizzata nei modelli linguistici di frontiera per rendere efficiente l’inferenza su contesti lunghi e consentire la decodifica con memoria costante. La sua memoria ricorrente di dimensione fissa richiede però una decisione online a ogni token: che cosa scrivere e con quale intensità sovrascrivere le associazioni esistenti, prima di sapere quali informazioni saranno necessarie alle query future. I modelli basati sulla regola Delta apprendono questa intensità dall’embedding del token corrente, ma non tengono traccia della fiducia nella stima della memoria, impedendo così di adattare ogni scrittura alle evidenze accumulate. Per rappresentare esplicitamente questa incertezza, riformuliamo la memoria associativa ricorrente come un modello a spazio di stati lineare-gaussiano, per il quale il filtro di Kalman è lo stimatore ricorsivo ottimale, e introduciamo una nuova famiglia di modelli, Kalman Delta Networks (KDN). Nelle KDN, la transizione propaga sia lo stato della memoria sia la sua incertezza, consentendo al guadagno di Kalman di pesare ogni scrittura del residuo in base alle evidenze accumulate e all’affidabilità dell’osservazione. In questa formulazione, gli aggiornamenti in stile Delta emergono come caso particolare che sostituisce la covarianza predittiva con un’approssimazione isotropa distinta per token e omette il tracciamento della covarianza. Il tracciamento esatto, tuttavia, comporta una ricorsione di Riccati densa e dipendente dallo stato, poco adatta alle scansioni di attenzione lineare parallele su GPU. Per affrontare il problema, introduciamo due approssimazioni KDN compatibili con le scansioni. Diagonal KDN proietta ogni distribuzione a posteriori a un passo sulla famiglia delle distribuzioni gaussiane diagonali mediante inferenza variazionale mean-field online, mentre Isotropic KDN usa un’approssimazione isotropa con un solo scalare di incertezza per testa. Le loro ricorrenze dell’incertezza sono mappe di Mobius, che consentono scansioni associative con profondità parallela logaritmica. In esperimenti controllati di preaddestramento con 750 milioni e 1,3 miliardi di parametri, le varianti KDN migliorano sistematicamente la perplexity e l’accuratezza media nei compiti downstream rispetto ai modelli di attenzione lineare allo stato dell’arte.