Kronecker Delta Rule per l’attenzione lineare su contesti lunghi

Abstract degli autori

Proponiamo KronLSTM, un’architettura di attenzione lineare basata sulla Kronecker Delta Rule, una nuova regola di aggiornamento per le memorie matriciali associative. KronLSTM estende la famiglia xLSTM con un meccanismo di revisione della memoria che corregge associazioni mirate modificando la memoria solo lungo un’unica direzione nello spazio delle matrici. Rispetto agli aggiornamenti basati sulla regola del delta, questa revisione minimamente invasiva riduce le interferenze con le informazioni già memorizzate, rendendo KronLSTM particolarmente adatto alla modellazione del linguaggio su contesti lunghi. La dinamica della memoria che ne risulta può essere descritta matematicamente come una catena di Markov che converge verso una distribuzione stazionaria, mentre il valore atteso della norma della matrice di memoria rimane uniformemente limitato per tutta la sequenza di input. KronLSTM garantisce quindi una dinamica stabile della memoria senza richiedere ulteriori meccanismi di normalizzazione e stabilizzazione dello stato della memoria. Gli esperimenti mostrano che KronLSTM è competitivo rispetto alle architetture di attenzione lineare esistenti, e spesso le supera, compresi i modelli della famiglia (Gated) DeltaNet, con prestazioni particolarmente elevate nei benchmark su contesti lunghi. Questi risultati indicano che la Kronecker Delta Rule costituisce una base efficace e trattabile sul piano analitico per una memoria associativa stabile nelle architetture ricorrenti di attenzione lineare. Il nostro codice (https://github.com/NX-AI/KronLSTM) e i nostri modelli (https://huggingface.co/collections/NX-AI/kronlstm) sono open source.