Complex KDA: comprendere e migliorare l’espressività di Kimi Delta Attention

Abstract degli autori

Le RNN lineari basate sulla delta-rule consentono di modellare sequenze in modo efficiente, ma i loro aggiornamenti lineari con una correzione di basso rango ne limitano l’espressività. Ricerche precedenti hanno mostrato che componendo due transizioni basate sulla delta-rule in un singolo aggiornamento ricorrente è possibile modellare una rotazione 2D, ma questo aumenta il rango e il costo degli aggiornamenti rispetto a una singola transizione. Mostriamo che Kimi Delta Attention (KDA) può realizzare rotazioni 2D combinando una singola trasformazione basata sulla delta-rule con una seconda riflessione fornita dal suo gate applicato canale per canale. Ciò richiede di estendere gli intervalli dei parametri di KDA combinando due estensioni già esistenti: consentire gate in $[-1,1]$ e un coefficiente della delta-rule $β$ in $[0,2]$. Chiamiamo il modello risultante Complex KDA (CKDA). Mantiene la stabilità e l’efficienza di KDA, con transizioni che restano diagonali più una matrice di rango uno e non-espansive, raggiungendo al contempo la capacità di tracciamento dello stato di DeltaProduct$_2$. Caratterizziamo l’espressività di CKDA e dimostriamo che ogni matrice ortogonale diagonale più una matrice di rango uno è esattamente una matrice di transizione CKDA. Un singolo strato CKDA può tracciare ogni gruppo finito isomorfo a un sottogruppo di $\mathrm{SO}(3)$, e molti risultati sul tracciamento dello stato richiedono uno strato in meno con CKDA rispetto ad altre RNN lineari diagonali più una matrice di rango uno. Sul piano empirico, la combinazione delle due estensioni offre la migliore estrapolazione della lunghezza tra le configurazioni degli intervalli di KDA testate su $S_3$, $S_4$ e nella continuazione di audio periodico. Nella modellazione del linguaggio, CKDA supera i Transformer e le altre RNN lineari, ottiene risultati simili a una baseline KDA e mostra un comportamento promettente in termini di scaling. Il nostro codice è open source all’indirizzo https://github.com/OpenEuroLLM/ComplexKDA e i nostri modelli sono disponibili all’indirizzo https://huggingface.co/collections/openeurollm/complexkda.