🤖 AI

Looped State-Space Language Models with Adaptive Exit-State Selection

Questo articolo dimostra che l'applicazione di architetture a ciclo (looped) ai modelli a spazio di stato basati su Mamba ne potenzia le capacità di ragionamento e l'efficienza dei parametri attraverso un aumento della profondità computazionale, introducendo al contempo una selezione adattiva dello stato di uscita per ottimizzare la profondità di predizione, sebbene rilevi che il reale risparmio in fase di inferenza richieda ulteriori meccanismi di gestione dello stato.

Autori originali: Zhenxuan Yu, Takeshi Kojima, Yutaka Matsuo, Yusuke Iwasawa

Pubblicato 2026-07-14

📖 6 min di lettura🧠 Approfondimento

CC BY 4.0

Autori originali: Zhenxuan Yu, Takeshi Kojima, Yutaka Matsuo, Yusuke Iwasawa

Articolo originale sotto licenza CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta né approvata dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

Immagina di avere un robot chef brillante e super intelligente chiamato Mamba. Questo chef è famoso per essere incredibilmente veloce ed efficiente nella memoria, capace di cucinare un banchetto senza bisogno di una cucina enorme piena di strumenti diversi. Ma c'è un problema: quando si trova di fronte a una ricetta davvero difficile — come un complesso puzzle matematico o un gioco del tipo "trova l'ingrediente nascosto" — Mamba a volte si blocca. Ha bisogno di pensare più a fondo, ma non ha abbastanza "strati" di pensiero per farlo senza costruire una nuova, gigantesca cucina (il che costa una fortuna).

Entrano in gioco i ricercatori della Università di Rikkyo e della Università di Tokyo. Hanno posto una domanda semplice e audace: E se non costruissimo una cucina più grande, ma facessimo cucinare a Mamba lo stesso piatto ripetutamente, affinando il sapore con ogni singolo passaggio?

La Cucina "a Cicli": Un Solo Chef, Molti Passaggi

Nel mondo dell'IA, la maggior parte dei modelli è come una catena di montaggio in una fabbrica. Un pezzo di dati (come una parola in una frase) si muove lungo la linea, passando attraverso 24 diverse stazioni (layer), e poi è finito. Se vuoi che il robot pensi più intensamente, di solito aggiungi semplicemente più stazioni.

Gli autori hanno provato qualcosa di diverso. Hanno costruito un Mamba a Cicli (Looped Mamba). Immagina che Mamba sia una singola stazione super talentuosa. Invece di inviare i dati lungo una lunga linea di 24 stazioni diverse, li inviano nuovamente alla stessa stazione, 24 volte di seguito.

La Magia: Il robot utilizza lo stesso cervello (parametri) per ogni singolo passaggio. È come uno studente che rilegge lo stesso paragrafo di un libro di testo quattro volte per capirlo, invece di leggere quattro paragrafi diversi e confusi.
Il Risultato: Su complessi enigmi logici chiamati Mano (aritmetica modulare) e p-hop induction (trovare schemi nascosti), questa strategia di "rilettura" ha dato risultati straordinari. Un modello che aveva solo 4 strati ma li ha ciclicamente ripetuti 6 volte (4 ⊗ 6) ha superato di gran lunga un modello standard con 24 strati unici. Questo suggerisce che per il ragionamento, la profondità del pensiero conta più del numero di strumenti unici.

Il "Cancello d'Uscita": Scegliere la Versione Migliore

Ma aspetta, e se il robot continuasse a rileggere il paragrafo all'infinito? Sarebbe uno spreco di tempo. I ricercatori hanno aggiunto una funzione intelligente chiamata Cancello d'Uscita (Exit Gate).

Pensa a questo cancello come a un selettore intelligente. Mentre Mamba elabora i dati, il cancello chiede: "La risposta sta diventando più chiara?"

Se la risposta è già chiara dopo 2 cicli, il cancello dice: "Fermati! Useremo il risultato di questo passaggio."
Se il problema è difficile, il cancello dice: "Continua!" e seleziona il risultato del 3° o 4° ciclo.

Questo è chiamato Selezione Adattiva dello Stato di Uscita (Adaptive Exit-State Selection). Il documento ha rilevato che per i modelli più piccoli (circa 140 milioni di parametri), questo cancello è stato un elemento decisivo. Ha permesso al modello di scegliere la quantità "giusta" di tempo di pensiero per ogni singola parola, spesso ottenendo prestazioni migliori rispetto al forzare il robot a usare sempre il risultato del numero massimo di cicli.

Tuttavia, c'è un colpo di scena. I ricercatori sono stati molto attenti a sottolineare che questo non fa risparmiare effettivamente elettricità o tempo sul computer per ora. Anche se il cancello seleziona un risultato precedente, il robot esegue comunque fisicamente tutti i cicli in background. Questo perché lo stato di memoria di Mamba è continuo: il risultato del ciclo 3 dipende dallo stato lasciato dal ciclo 2, che dipende dal ciclo 1. Non puoi semplicemente "far cadere" il robot dai cicli successivi senza rompere la catena per tutte le parole successive. Quindi, il cancello seleziona la profondità di predizione (quale versione della risposta usare), ma il calcolo effettivo nel tempo reale (wall-clock computation) rimane lo stesso. Per risparmiare davvero tempo, servirebbe un nuovo modo per gestire lo stato di memoria del robot, il che, dicono gli autori, è un compito per il lavoro futuro.

Il Controllo della Realtà: Cosa Non Ha Funzionato (e Cosa è Ancora un "Forse")

Il documento è rinfrescante nella sua onestà riguardo a ciò che non ha risolto.

Più grande non è sempre meglio (per questo trucco): Quando hanno confrontato il loro Mamba a Cicli con un modello standard non ciclico che era altrettanto "costoso" da eseguire (stesso numero di calcoli, o "iso-FLOPs"), il modello standard ha vinto ancora nel test base di "quanto è confuso il modello?" (perplessità). Il modello a cicli era ottimo nei compiti di ragionamento, ma il modello profondo standard era ancora leggermente migliore nel prevedere semplicemente la parola successiva in una frase.
Le dimensioni contano: Il "Cancello d'Uscita" ha funzionato magnificamente per i modelli da 140 milioni di parametri. Ma quando lo hanno provato sui modelli più grandi da 370 milioni di parametri, il cancello non ha migliorato molto le cose. I modelli più grandi sembravano preferire l'esecuzione di tutti i cicli comunque. Gli autori suggeriscono che i modelli più grandi potrebbero fare affidamento maggiormente su quel passaggio finale e profondo di pensiero.
Non è una bacchetta magica per tutto: I ricercatori hanno testato questo approccio su puzzle sintetici (come alberi matematici e ricerca di schemi). Sebbene questi dimostrino che il concetto funziona, ammettono che non sappiamo ancora se questo trucco del "looping" funzionerà altrettanto bene sul ragionamento aperto e del mondo reale, come scrivere un romanzo o risolvere un caso legale complesso.

Il Punto Fondamentale

Il documento suggerisce che riutilizzare lo stesso cervello intelligente ripetutamente è un modo potente per rendere l'IA più intelligente nel ragionamento senza costruire un cervello più grande ed costoso. È un asse di scalabilità "efficiente nei parametri".

Dimostrato: Su specifici enigmi logici, far girare in cicli un piccolo modello Mamba batte un modello standard delle stesse dimensioni e pareggia un modello molto più grande.
Suggerito: Questo approccio potrebbe essere la chiave per rendere l'IA futura più efficiente, specialmente per compiti che richiedono un pensiero profondo e passo dopo passo.
Non ancora risolto: Dobbiamo ancora capire come far sì che il "Cancello d'Uscita" faccia effettivamente risparmiare tempo di calcolo (invece di selezionare solo l'output) e se questo funzionerà per i massicci modelli da miliardi di parametri che alimentano le più grandi IA di oggi.

In breve, gli autori non hanno solo costruito un robot più grande; hanno insegnato a un robot più piccolo a pensare più profondamente permettendogli di dare un secondo (o terzo, o quarto) sguardo al problema. E per ora, questo secondo sguardo sembra piuttosto promettente.

Riepilogo Tecnico: Modelli di Linguaggio a Spazio-Stato con Cicli e Selezione Adattiva dello Stato di Uscita

Definizione del Problema

I recenti progressi nei Large Language Models (LLM) sono stati guidati principalmente dalla scalabilità della dimensione del modello, dei dati di addestramento e del calcolo intensivo durante l'inferenza. Tuttavia, questo paradigma di scalabilità introduce sfide significative per la distribuzione, tra cui l'aumento del consumo di memoria e i costi di inferenza in rapida crescita a causa del mantenimento di grandi cache key-value nei meccanismi di auto-attenzione basati su Transformer. Sebbene i modelli State Space (SSM), in particolare Mamba-2, offrano un'alternativa efficiente con una scalabilità favorevole per contesti lunghi, rimane poco chiaro se il principio di "profondità computazionale rispetto a parametri indipendenti" — recentemente dimostrato dai Transformer a ciclo (Looped Transformers) — si applichi anche agli SSM. Inoltre, la ricerca esistente sulle architetture a ciclo si è concentrata quasi esclusivamente sui backbone Transformer, lasciando un vuoto nella comprensione di come il riutilizzo ricorrente dei parametri funzioni all'interno di architetture Mamba e ibride Mamba-Transformer.

Metodologia

Gli autori propongono le architetture Looped Mamba e Looped Hybrid Mamba–Transformer. Il meccanismo centrale consiste nel condividere i parametri di un singolo blocco Mamba (o ibrido) attraverso i layer e applicarlo ripetutamente attraverso espliciti cicli a profondità finita.

Architettura

Notazione: Un modello è denotato come $N \otimes R$ , dove $N$ è il numero di layer con parametri indipendenti e $R$ è il numero di passi ricorrenti. La profondità effettiva è $NR$.
Looped Mamba: Applica ripetutamente un backbone Mamba $M_\theta$ $M_{θ}$ di $N$ $N$ layer con parametri condivisi.
- $h^{(r)} = M_\theta(h^{(r-1)})$ per $r = 1, \dots, R$ .
- L'output finale è derivato dallo stato nascosto al passo $R$ .
Looped Hybrid: Estende il ciclo a un backbone che contiene sia layer Mamba-2 sia un piccolo sottoinsieme di layer di auto-attenzione causale. Il ciclo si applica all'intero blocco ibrido, permettendo il raffinamento iterativo delle rappresentazioni nascoste attraverso sia aggiornamenti dello spazio-stato che interazioni basate sull'attenzione.
Protocollo di Addestramento:
1. Pre-training: Ottimizzato utilizzando un obiettivo di ciclo finale (loss calcolata solo al passo $R$ ).
2. Exit-Gate Stage 1: Introduce un obiettivo di computazione adattiva regolarizzato dall'entropia (adattato da Ouro) per apprendere una distribuzione a livello di token sulla profondità di uscita utilizzando gli output di tutti i passi ricorrenti.
3. Exit-Gate Stage 2: Congela il backbone e la testa, addestrando solo il gate di uscita per incoraggiare la continuazione solo quando il passo successivo migliora la perdita di predizione.

Selezione Adattiva dello Stato di Uscita

Il paper adatta il sistema a due stadi di exit gate di Ouro a Mamba. Durante l'inferenza, una soglia $q$ controlla la selezione del passo di uscita $z_i$ per ogni token in base alla probabilità cumulativa di uscita.

Nota Semantica: Gli autori chiariscono che negli SSM, la catena dello stato di tutti i $R$ passi deve essere preservata a causa dei vincoli di continuità dello stato. Pertore, l' "uscita" seleziona quale stato nascosto del passo ricorrente viene passato alla testa di linguaggio, ma non riduce intrinsecamente il tempo di calcolo (wall-clock) o il numero di passi eseguiti senza meccanismi aggiuntivi di gestione dello stato.

Contributi Chiave

Studio Empirico Controllato: Gli autori presentano un confronto rigoroso tra modelli Looped Mamba e modelli ibridi sotto protocolli di iso-parametro (stesso $N$ , variabile $R$ ) e iso-FLOPs (stessa profondità effettiva $NR$) rispetto a baseline Mamba e Transformer non cicliche.
Valutazione di Compiti di Ragionamento: Lo studio valida l'approccio su compiti di ragionamento sintetici (Mano e $p$ -hop induction) e sul pre-training di modelli linguistici con benchmark downstream.
Meccanismo di Uscita Adattiva: L'introduzione di un meccanismo di selezione dello stato di uscita adattivo controllato da una soglia per i modelli a spazio-stato ciclici, dimostrando la sua capacità di migliorare le prestazioni downstream a profondità intermedie.

Risultati Sperimentali

Compiti di Ragionamento Sintetico

Compito Mano (Aritmetica Modulare): Looped Mamba supera costantemente le baseline non cicliche con parametri corrispondenti. Nella configurazione più difficile ( $L=24$ ), un modello $4 \otimes 6$ ha raggiunto un'accuratezza del 99,35% rispetto al 7,50% della baseline $4 \otimes 1$ . In condizioni iso-FLOPs (profondità effettiva 24), il modello $4 \otimes 6$ ha superato significativamente la baseline non ciclica $24 \otimes 1$ (99,35% vs 37,45%), suggerendo che il riutilizzo ricorrente dei parametri fornisce un bias induttivo più forte per il calcolo composizionale rispetto all'aumento di layer non condivisi.
$p$ -hop Induction (Recupero Ricorsivo): I modelli Mamba non ciclici superficiali hanno faticato (es. $1 \otimes 1$ a circa 38-41%), mentre l'aggiunta di cicli ha migliorato drasticamente le prestazioni. Il modello $1 \otimes 12$ ha raggiunto un'accuratezza quasi perfetta (~99,89%), indicando che la profondità ricorsiva è più critica del numero di parametri distinti per il recupero ricorsivo.

Modellazione del Linguaggio e Benchmark Downstream

Iso-Parametro: La profondità ricorrente ha migliorato la perplessità di validazione (PPL) per i backbone Mamba-2 e LLaMA-2 sotto budget di parametri fissi.
Iso-FLOPs: I modelli non ciclici più profondi ( $24 \otimes 1$ ) mantengono un vantaggio nella perplessità di validazione rispetto ai modelli ciclici con profondità effettiva corrispondente. Tuttavia, i modelli ciclici sono rimasti competitivi nei benchmark downstream, talvolta superando le baseline non cicliche più profonde.
Uscita Adattiva:
- Scala 140M: L'uscita adattiva tramite gating ha migliorato le prestazioni downstream. Per il modello Mamba-2 da 140M, l'exit gate dello Stage 2 ha ottenuto un punteggio medio di 38,87 con un passo di uscita selezionato medio di 3,00, superando la baseline fixed-loop4 (37,79).
- Scala 370M: Le migliori prestazioni downstream sono rimaste concentrate vicino all'impostazione fixed-loop4. Le varianti con exit-gating non hanno superato i picchi di punteggio dei modelli a ciclo fisso, suggerendo che i modelli più grandi possano fare maggiore affidamento sull'ultimo passo ricorrente o richiedere una migliore calibrazione del gate.
- Trade-off: L'uscita tramite gating fornisce una frontiera controllabile tra profondità selezionata e prestazioni, permettendo al modello di utilizzare efficacemente gli stati ricorrenti intermedi invece di forzare una profondità massima fissa.

Significato e Rivendicazioni

Il paper sostiene che la profondità ricorrente offre un'interessante dimensione di scalabilità, efficiente in termini di parametri, per i modelli Mamba e ibridi. Le scoperte chiave includono:

Profondità vs Parametri: Molti compiti di ragionamento beneficiano di una maggiore profondità computazionale piuttosto che di ulteriori parametri indipendenti. Il riutilizzo ricorrente dei layer Mamba migliora sostanzialmente le capacità di calcolo compositivo e ricorsivo senza aumentare il numero di parametri.
Vitalità dell'Ibrido: Il principio della profondità ciclica si applica efficacemente anche alle architetture ibride Mamba-Transformer, dove i layer di attenzione vengono ciclicamente integrati con i layer Mamba.
Selezione Adattiva: Sebbene le attuali implementazioni SSM richiedano l'esecuzione completa dello stato, la selezione adattiva dello stato di uscita permette di selezionare la profondità di predizione ottimale, migliorando le prestazioni downstream a profondità intermedie (particolarmente alla scala 140M).

Gli autori mantengono la modestia riguardo alle loro affermazioni, notando che gli esperimenti sono limitati a modelli fino a 370M di parametri e che il reale risparmio di calcolo durante l'inferenza richiederà lavori futuri su meccanismi di gestione dello stato per saltare interamente i passi ricorrenti. I risultati sono presentati come evidenza empirica controllata piuttosto che come una caratterizzazione completa delle prestazioni su larga scala.