Moonshot AI (Kimi) ha pubblicato un articolo in cui presenta una nuova architettura per i modelli Transformer, basata su 'Attention Residuals'. Questa architettura sostituisce le tradizionali connessioni residuali, utilizzate fin dal 2015 (ResNet).

Il problema della diluizione

Le connessioni residuali standard tendono ad accumulare gli output di tutti i livelli precedenti. Nei livelli più profondi, ad esempio il quarantesimo, si sommano gli output dei livelli da 1 a 39. Questo, secondo Kimi, porta a una diluizione delle informazioni provenienti dai primi livelli.

La soluzione: Attention Residuals

La soluzione proposta consiste nel permettere a ciascun livello di 'attenzionare' selettivamente gli output dei livelli precedenti, anziché sommarli semplicemente. In pratica, ogni livello può scegliere quali livelli precedenti sono più importanti per l'input corrente, utilizzando pesi di attenzione appresi.

Risultati

I benchmark di Moonshot AI mostrano:

  • Miglioramenti da 3 a 7.5 punti in test di ragionamento matematico, generazione di codice e task che richiedono un contesto esteso.
  • Risparmio di calcolo di circa 1.25x con la variante 'block attention residual'.
  • Overhead di training inferiore al 4% e aumento della latenza di inference inferiore al 2%.
  • Scalabilità: i modelli più grandi beneficiano maggiormente di questa architettura.

È stata anche sviluppata una variante 'block attention residual', in cui i livelli sono raggruppati in blocchi. All'interno di un blocco si utilizzano le normali connessioni residuali, mentre tra i blocchi si utilizza l'attenzione. Questo approccio mantiene gran parte dei benefici, riducendo i costi di esecuzione.

Confronto con DeepSeek

DeepSeek aveva recentemente proposto un'altra soluzione per migliorare le connessioni residuali (mHC), ma con un approccio completamente diverso. Mentre DeepSeek aggiunge flussi paralleli, Kimi introduce l'attenzione selettiva. Secondo alcuni confronti, l'approccio di Kimi richiede circa 1/6 della bandwidth di memoria rispetto a DeepSeek mHC, ottenendo risultati simili o migliori.

Implicazioni pratiche

La versione di Kimi è potenzialmente 'drop-in replaceable': si sostituisce il modulo residuale, si mantiene tutto il resto invariato, si riesegue il training e si ottengono miglioramenti. DeepSeek mHC, invece, richiede una ristrutturazione completa dell'architettura del modello.

Considerazioni finali

Karpathy ha commentato che l'attenzione potrebbe essere applicata in più punti del Transformer di quanto si pensasse. Per chi sviluppa modelli in locale, questa innovazione potrebbe portare a miglioramenti significativi nella qualità senza la necessità di modelli più grandi: stesso numero di parametri, migliore flusso di informazioni, risultati migliori. Resta da valutare l'interazione con la quantization, dato che i pesi di attenzione tra i livelli potrebbero essere sensibili alla precisione.