Immagine AI

MiniMax ha pubblicato MiniMax Sparse Attention, o MSA, una nuova architettura di attenzione sparsa progettata per ridurre il costo computazionale dei modelli linguistici che lavorano su contesti molto estesi. La tecnologia è stata integrata in MiniMax-M3, modello multimodale nativo con finestra di contesto fino a un milione di token, circa 428 miliardi di parametri complessivi e 23 miliardi di parametri attivati per token grazie a una struttura Mixture of Experts.

Il problema affrontato da MSA riguarda il comportamento dell’attenzione standard quando aumenta la lunghezza della sequenza. Nei transformer convenzionali, ogni token deve confrontarsi con un numero crescente di token precedenti. Con contesti da centinaia di migliaia o milioni di token, il costo della lettura dei key-value cache, della selezione delle informazioni rilevanti e del calcolo Softmax diventa una delle principali limitazioni per inferenza, agenti software, analisi di repository e gestione di documenti molto lunghi.

MSA parte dalla struttura Grouped Query Attention, già diffusa nei modelli di grandi dimensioni per ridurre la memoria associata alle chiavi e ai valori. L’architettura introduce due percorsi distinti. Il primo è l’Index Branch, una componente leggera che esamina l’intero contesto e attribuisce un punteggio ai blocchi di chiavi. Il secondo è il Main Branch, che esegue l’attenzione completa soltanto sui blocchi selezionati dall’indice.

Il contesto viene suddiviso in blocchi da 128 token. Per ciascuna query e per ciascun gruppo GQA, l’indexer seleziona i 16 blocchi con punteggio più alto attraverso un’operazione Top-k. Il Main Branch effettua quindi la normale attenzione Softmax solo su questi segmenti. In questa configurazione, ogni query può consultare fino a 2.048 token key-value, invece di analizzare l’intera sequenza disponibile nel contesto.

La selezione non avviene eliminando casualmente parti del testo. L’Index Branch viene addestrato per approssimare la distribuzione di attenzione prodotta dal Main Branch. Per superare il fatto che il Top-k non è direttamente differenziabile, MiniMax utilizza una perdita di allineamento basata sulla divergenza di Kullback-Leibler, con cui l’indexer impara a scegliere blocchi coerenti con quelli più rilevanti per l’attenzione completa. Durante l’addestramento è previsto anche un passaggio graduale dalla full attention alla sparse attention, così da evitare instabilità nella fase iniziale.

L’architettura mantiene inoltre i blocchi vicini alla posizione della query, anche quando non rientrano nella selezione globale più alta. Questa scelta serve a non perdere le dipendenze locali, come istruzioni recenti, porzioni immediatamente precedenti di codice, riferimenti grammaticali o dati appena introdotti in una conversazione. Il sistema combina quindi recupero selettivo di informazioni lontane e continuità del contesto locale.

Nei test riportati da MiniMax, MSA ha mantenuto risultati vicini ai modelli con full attention su comprensione linguistica, ragionamento, matematica, generazione di codice, comprensione di immagini e video e attività agentiche. Le valutazioni su benchmark di long-context retrieval come RULER e HELMET sono state utilizzate per misurare la capacità del modello di individuare informazioni in sequenze molto lunghe, senza un peggioramento marcato rispetto al riferimento denso.

Il vantaggio principale emerge sul piano dell’inferenza. A un milione di token di contesto, MiniMax indica una riduzione di circa 28,4 volte dei FLOPs di attenzione per token rispetto a una configurazione GQA convenzionale. Nei test riportati, la fase di prefill, cioè l’elaborazione iniziale dell’intero prompt, ha raggiunto un’accelerazione fino a 14,2 volte. La fase di decoding, nella quale il modello genera nuovi token usando il contesto già elaborato, ha registrato un incremento fino a 7,6 volte.

Per rendere utilizzabile l’algoritmo su hardware reale, MiniMax ha pubblicato anche il repository MSA con licenza MIT. Il progetto include kernel per NVIDIA SM100, con implementazioni sia di FlashAttention densa sia di sparse top-k attention. Il codice supporta BF16, FP8, NVFP4 e FP4, formati a precisione ridotta rilevanti per ridurre memoria e tempi di calcolo nei sistemi di inferenza ad alta densità.

Il repository utilizza una dimensione di pagina di 128 token e una selezione Top-k pari a 16 come configurazione di riferimento. Include inoltre un kernel sparse_topk_select per trasformare i punteggi dei blocchi negli indici delle pagine key-value da processare, oltre a un adattatore che collega i percorsi di sparse prefill alle API del kernel FMHA. La compilazione avviene in parte tramite JIT e richiede Linux x86_64, Python 3.10 o superiore, CUDA e una GPU NVIDIA compatibile con architettura SM100.

MiniMax-M3 utilizza MSA per estendere il contesto fino a un milione di token senza sostenere il costo della full attention su ogni token del prompt. Il modello supporta input testuali, immagini e video ed espone tre modalità di ragionamento: enabled, con reasoning sempre attivo; adaptive, in cui il modello decide quando approfondire il ragionamento; e disabled, pensata per ridurre latenza e aumentare il throughput. La pubblicazione dell’architettura e dei kernel consente di separare il modello dal componente infrastrutturale, rendendo MSA potenzialmente adottabile anche in altri sistemi GQA orientati a contesti lunghi.

Di Fantasy