Sintesi Tecnica: Cosa sopravvive quando si comprime un ragionatore ricorsivo per l'edge?

Definizione del Problema

I modelli di ragionamento ricorsivo, come il Tiny Recursive Model (TRM), risolvono compiti strutturati complessi aggiornando iterativamente uno stato latente di "trasporto" (carry state) attraverso molteplici cicli di ragionamento. Sebbene questi modelli raggiungano prestazioni competitive con un numero relativamente esiguo di parametri (ad esempio, 6,83M), l'implementazione su hardware edge con risorse limitate (ad esempio, 4–8 MB di SRAM) richiede una compressione aggressiva.

Le comuni intuizioni di compressione derivate dai modelli autoregressivi a passaggio singolo non si applicano qui. Nei transformer convenzionali, il rumore di quantizzazione si accumula attraverso i token di output. Nei modelli ricorsivi, tuttavia, gli errori di quantizzazione si cumulano attraverso $H$ cicli di ragionamento ricorenti ad ogni aggiornamento dello stato di trasporto. Il documento investiga cosa sopravvive e cosa si rompe quando questi modelli vengono compressi sotto vincoli rigidi di SRAM, chiedendosi specificamente se le tecniche standard (quantizzazione, pruning, distillazione) preservino le capacità di ragionamento globale o semplicemente la predizione locale dei token.

Metodologia

Gli autori hanno condotto uno studio sistematico sulla compressione attraverso tre compiti (ARC-2024, Maze-Hard, Sudoku-Extreme) e due architetture ricorsive (TRM-Attention e TRM-MLP-Mixing).

1. Scansione della Compressione: Hanno valutato una scansione a precisione totale (FP32, FP16, INT8, INT4) e profondità di ricorsione variabili ($H$ e $n_{sup}$).
2. Risoluzione del Collo di Bottiglia della Memoria: Riconoscendo che la tabella di embedding del puzzle (99,4 MB) dominava l'impronta di memoria (il 79% dei 126 MB di base), hanno disaccoppiato questo elemento dalla compressione del backbone utilizzando il caricamento flash a singolo puzzle (single-puzzle flash loading), trasmettendo solo la riga da 2 KB rilevante al momento dell'inferenza.
3. Tecniche di Compressione:
   • Post-Training Quantization (PTQ): Applicata INT8 (calibrata per canale) e INT4 (naïve per-tensor e per-channel calibrata).
   • Cambiamenti Strutturali: Hanno valutato il pruning strutturato, la distillazione della conoscenza e approssimazioni di attenzione lineare.
   • Tentativi di Recupero: Hanno testato la Quantization-Aware Training (QAT) con obiettivi a livello di token per riparare il collasso del ragionamento.
4. Diagnostica: Hanno introdotto la Fedeltà della Traiettoria del Trasporto (Carry-Trajectory Fidelity), una metrica senza etichette (label-free) che misura la similarità del coseno tra lo stato di trasporto finale di un modello compresso e il suo riferimento a precisione intera sullo stesso input. Questo funge da predittore del collasso del ragionamento senza richiedere la valutazione del compito.
5. Validazione Hardware: Hanno testato la latenza e l'uso della memoria on-device su tre piattaforme Snapdragon (Galaxy S24, S22, Dragonwing RB3) tramite Qualcomm AI Hub.

Risultati Chiave

1. L'Incrocio Profondità-Precisione

Contrariamente ai modelli a passaggio singolo dove una precisione inferiore degrada solitamente le prestazioni, gli autori hanno scoperto che INT8 con un singolo ciclo ricorsivo ($H=1$) eguaglia l'accuratezza FP32 a profondità totale (35,25% vs 36,00% su ARC-2024) riducendo al contempo i FLOPs di 6 volte. Questa interazione suggerisce che, per i modelli ricorsivi, ridurre la profondità può essere più efficiente che mantenere la profondità totale con una precisione inferiore, a patto che la precisione sia sufficiente (INT8).

2. Collasso Composizionale: Locale vs Globale

La compressione aggressiva (INT4 naïve, pruning, distillazione) crea un disaccoppiamento tra accuratezza della cella (predizione locale del token) e accuratezza del puzzle-esatto (soddisfacimento dei vincoli globali).

• Sotto INT4 naïve, pruning e distillazione, l'accuratezza della cella rimane spesso elevata (ad esempio, >84%), ma l'accuratezza del puzzle-esatto crolla a zero.
• Ciò indica che la compressione distrugge la capacità del modello di soddisfare i vincoli globali pur preservando le capacità di predizione locale.
• Gli obiettivi a livello di token, inclusa la QAT con perdita di cross-entropy, falliscono nel recuperare il ragionamento globale perché premiano la correttezza della singola cella piuttosto che la coerenza composizionale.

3. Fragilità Architetturale

Il collasso è architetturale, non intrinseco al compito.

• TRM-Attention (usato in ARC e Maze) tollera l'INT4 naïve.
• TRM-MLP-Mixing (usato in Sudoku) collassa catastroficamente sotto l'INT4 naïve.
• Gli studi di ablazione sullo stesso compito (Sudoku) hanno mostrato che passare da MLP-mixing ad Attention preserva l'accuratezza sotto INT4. La fragilità risiede nel meccanismo di mixing dei token, che richiede scale per canale per sopravvivere al rumore grossolano, mentre l'attenzione codifica aggiornamenti relazionali sparsi che sono più robusti.

4. Recupero tramite Calibrazione

L'INT4 calibrata per canale (post-training, senza ri-addestramento) recupera le perdite catastrofiche viste nell'INT4 naïve.

• Su Sudoku-Extreme, l'INT4 naïve fa scendere l'accuratezza al 10,2%, mentre l'INT4 calibrata la riporta al 71,9% (vicino alla baseline FP32 del 73,8%).
• Questo stabilisce l'INT4 calibrata come il limite minimo di implementazione (deployment floor) per il ragionamento di soddisfacimento dei vincoli, rientrando nei target di 4 MB per MCU.

5. Fedeltà della Traiettoria del Trasporto come Diagnostica

La metrica proposta di Fedeltà della Traiettoria del Trasporto predice con successo il collasso del ragionamento prima della valutazione del compito.

• Un'alta fedeltà ($\phi \ge 0,99$) correla con il mantenimento del ragionamento.
• Una bassa fedeltà ($\phi < 0,8$) predice un fallimento catastrofico.
• La metrica è monotona rispetto alla perdita di accuratezza e può distinguere tra quantizzazione utilizzabile (calibrata) e inutilizzabile (naïve) senza eseguire il compito.

Ricetta di Implementazione

Il documento deriva una strategia concreta di implementazione per dispositivi edge:

• Compressione Embedding: Utilizzare il caricamento flash a singolo puzzle per ridurre la tabella di embedding da 99,4 MB a 2 KB di memoria attiva.
• Mobile SoC (8 MB SRAM): Backbone INT8 + flash loading + ricorsione ridotta ($H=1, n_{sup}=8$) rientra nei 6,51 MB e raggiunge un'accuratezza quasi totale con 6× meno FLOPs.
• Microcontrollore (4 MB SRAM): Backbone INT4 calibrata per canale + flash loading rientra nei 3,26 MB e recupera un'accuratezza vicina alla FP32 senza necessità di ri-addestramento.

Significato e Rivendicazioni

Il documento sostiene di fornire il primo studio sistematico sulla compressione di modelli di ragionamento ricorsivo sotto vincoli rigidi di SRAM. I suoi principali contributi sono:

1. Ridefinire i Limiti della Compressione: Dimostrare che i modelli ricorsivi hanno un unico trade-off profondità-precisione in cui la riduzione della ricorsione può compensare la perdita di precisione, un fenomeno assente nella quantizzazione a passaggio singolo.
2. Identificare un Modo di Fallimento: Mostrare che le tecniche di compressione standard spesso preservano la predizione locale del token distruggendo però il ragionamento globale, un problema che gli obiettivi di addestramento a livello di token non possono risolvere.
3. Intuizione Architetturale: Stabilire che la scelta del mixer di token (Attention vs. MLP) è una decisione critica di implementazione, poiché l'MLP-mixing è significativamente più sensibile al rumore di quantizzazione in contesti ricorsivi.
4. Diagnostica Senza Etichette: Introdurre la fedeltà della traiettoria del trasporto come strumento pratico per predire il successo o il fallimento della distribuzione senza valutazione specifica del compito.
5. Fattibilità Pratica: Validare che un ragionamento ricorsivo ad alte prestazioni può essere implementato su microcontrollori da 4 MB utilizzando INT4 calibrata ed embedding trasmessi via flash, eliminando la necessità di dipendenza dal cloud.

Gli autori rimangono modesti riguardo alla generalizzazione, osservando che il loro studio si concentra sulla famiglia TRM e su compiti specifici, e che la diagnostica di fedeltà richiede un riferimento a precisione intera disponibile solo prima della distribuzione. Non pretendono di aver risolto la compressione per tutte le architetture ricorsive, ma forniscono una ricetta validata per la specifica classe di modelli e vincoli studiati.