Super Apriel: One Checkpoint, Many Speeds

1. Il Problema

Con l'avanzamento dei modelli linguistici verso contesti più lunghi e una maggiore richiesta di throughput di servizio, il costo dell'inferenza basato sull'Attention Completa (Full Attention - FA) sta diventando un collo di bottiglia critico. In ogni passo di decodifica autoregressiva, l'FA deve leggere e aggiornare una cache KV che cresce linearmente con la lunghezza della sequenza, limitando la capacità di batching e rendendo il sistema vincolato dalla memoria (memory-bound).

Le architetture ibride esistenti (che combinano FA con meccanismi efficienti come attention a finestra scorrevole o ricorrenze lineari) offrono soluzioni, ma presentano una limitazione fondamentale: ogni modello ha una configurazione fissa (una specifica disposizione dei layer) determinata prima o dopo l'addestramento. Questo costringe a un compromesso statico tra velocità e qualità al momento del rilascio. Non è possibile adattare dinamicamente il modello a carichi di lavoro eterogenei (es. brevi prompt ad alto batch vs. lunghi contesti) senza riaddestrare o distribuire checkpoint separati.

2. Metodologia: Super Apriel

Gli autori introducono Super Apriel, una "supernet" da 15 miliardi di parametri derivata dal modello Apriel 1.6. L'innovazione centrale risiede nella capacità di fornire quattro opzioni di mixer (meccanismi di attenzione/ricorrenza) addestrate per ogni singolo layer decoder, selezionabili a runtime senza ricaricare i pesi.

Architettura e Vocabolario dei Mixer

Ogni layer della rete può essere configurato dinamicamente per utilizzare uno dei seguenti quattro mixer:

1. Full Attention (FA): Interazioni esatte tra token ($O(n^2)$), massima qualità.
2. Sliding Window Attention (SWA): Attenzione limitata a una finestra locale ($w=4096$), costo indipendente dalla lunghezza del contesto.
3. Kimi Delta Attention (KDA): Estensione di GDN con gating per canale, stato ricorrente.
4. Gated DeltaNet (GDN): Regola delta con gating scalare, stato ricorrente a dimensione fissa ($O(1)$).

Una "placement" (disposizione) è una scelta specifica di un mixer per ciascuno dei 48 layer. Il checkpoint condiviso contiene i pesi per tutte le combinazioni, permettendo di passare da una configurazione all'altra in pochi secondi (5-15s) durante il servizio.

Training Recipe

Il processo di addestramento si articola in due fasi principali:

1. Distillazione Stocastica (S1): La supernet viene addestrata distillando da un insegnante congelato (Apriel 1.6). Durante ogni passo, un mixer viene selezionato casualmente e uniformemente per ogni layer. Questo permette di addestrare simultaneamente tutte le combinazioni possibili, mantenendo i parametri condivisi (FFN, embedding, normalizzazione) congelati per garantire stabilità.
2. Supervised Fine-Tuning (S2): Dopo l'identificazione delle configurazioni ottimali (preset), viene eseguita una fase di SFT mirata su un sottoinsieme di queste disposizioni per massimizzare le prestazioni su compiti specifici.

Ottimizzazione della Disposizione (Placement Optimization)

Data la vastità dello spazio di ricerca ($4^{48}$ combinazioni), gli autori utilizzano un approccio di ottimizzazione basato su un surrogato:

• Cluster Expansion: La funzione di punteggio (qualità) viene decomposta in interazioni locali tra layer (ordine 1, 2, 3) utilizzando la teoria dell'espansione a cluster della fisica statistica.
• Surrogato Bayesiano: Un modello lineare bayesiano stima la qualità di una disposizione basandosi sulle assegnazioni dei mixer.
• Ottimizzazione Esatta: Grazie alla struttura additiva del surrogato, è possibile trovare la disposizione ottimale per ogni vincolo di costo (throughput) utilizzando la programmazione dinamica, esplorando l'intera frontiera di Pareto (qualità vs. velocità) in modo efficiente.

3. Contributi Chiave

• Supernet Flessibile: Il primo checkpoint che permette di cambiare l'architettura interna (disposizione dei mixer) a runtime, offrendo molteplici preset di velocità da un'unica distribuzione.
• Decodifica Speculativa Integrata: Poiché la supernet contiene sia il modello "costoso" (tutti FA) che modelli "economici" (es. tutti GDN), è possibile utilizzare la decodifica speculativa senza un modello di bozza esterno separato.
• Analisi della Dinamica del Paesaggio: Lo studio ha rivelato che, sebbene le classifiche delle disposizioni si stabilizzino rapidamente durante l'addestramento stocastico su modelli piccoli (0.5B), le disposizioni sulla frontiera di Pareto mostrano una maggiore instabilità su modelli grandi (15B), suggerendo cautela nell'estrapolare risultati da scale ridotte.
• Rilascio Open Source: Pesi della supernet, codice di addestramento (Fast-LLM), codice di servizio (vLLM plugin) e toolkit di ottimizzazione.

4. Risultati

I risultati dimostrano che è possibile ottenere un ampio spettro di compromessi velocità-qualità senza perdita significativa di capacità:

• Qualità: Il preset "All-FA" (tutti i layer con attenzione completa) corrisponde alle prestazioni del modello insegnante Apriel 1.6 su tutti i benchmark (media 74.2 vs 73.9).
• Throughput: I preset ibridi raccomandati offrono un aumento del throughput di decodifica da 2.9x a 10.7x rispetto alla FA completa, mantenendo una ritenzione della qualità dal 96% al 77%.
  • Esempio: Il preset Reg|Lklhd–10 (10.7x speedup) mantiene il 77% della qualità complessiva e il 94.6% nelle capacità matematiche.
• Scalabilità con il Contesto: Il vantaggio delle disposizioni efficienti è esponenziale rispetto alla lunghezza del contesto. Mentre i modelli ibridi esterni guadagnano solo il 5-46% di velocità aggiuntiva passando da 16k a 32k token, le disposizioni di Super Apriel guadagnano un 80-155% di velocità relativa. Questo perché i mixer efficienti (SWA, KDA, GDN) hanno un footprint di memoria fisso, a differenza della FA che soffre di colli di bottiglia quadratici.
• Confronto con Modelli Esteri: Super Apriel supera o eguaglia modelli ibridi esterni di dimensioni simili o maggiori (come Nemotron-Nano o Falcon-H1R) offrendo una flessibilità operativa superiore.

5. Significato e Implicazioni

Super Apriel rappresenta un cambio di paradigma nel deployment dei LLM:

1. Flessibilità Operativa: Elimina la necessità di addestrare e distribuire modelli separati per diversi casi d'uso (es. uno per la chat veloce, uno per l'analisi di documenti lunghi). Un singolo checkpoint può adattarsi dinamicamente al carico di lavoro.
2. Efficienza dei Costi: Permette di scegliere il punto di funzionamento ottimale in base alla disponibilità hardware e alla latenza richiesta, massimizzando l'utilizzo delle risorse.
3. Validazione Scientifica: Dimostra che le architetture ibride non devono essere fisse, ma possono essere ottimizzate dinamicamente tramite tecniche di surrogato e distillazione, aprendo la strada a "modelli adattivi" che cambiano architettura in base al task.

In sintesi, Super Apriel trasforma il compromesso velocità-qualità da una decisione di rilascio statica in un pannello di controllo runtime, offrendo prestazioni di punta con un'efficienza senza precedenti per contesti lunghi.