Sintesi Tecnica: RADLADS

Problematica

Sebbene le varianti dell'attenzione lineare (come RWKV e Mamba) offrano vantaggi significativi rispetto ai tradizionali transformer con attenzione softmax — specificamente un tempo di inferenza $O(1)$ per token e un uso costante della memoria evitando le cache Key-Value — l'addestramento di modelli lineari su larga scala da zero rimane proibitivamente costoso. I modelli transformer allo stato dell'arte (SoTA) richiedono spesso l'addestramento su oltre dieci trilioni di token, un costo inaccessibile alla maggior parte delle organizzazioni. I metodi esistenti per convertire i transformer softmax pre-addestrati in modelli lineari o ricorrenti (ad esempio, T2R, SUPRA, LoLCats, MOHAWK) hanno storicamente richiesto conteggi di token massicci (dall'ordine di 20B a 100B+ di token) o hanno prodotto prestazioni scarse, in particolare su benchmark come MMLU. Inoltre, molti tentativi di conversione si affidano ad architetture ibride (mantenendo parte dell'attenzione softmax) o non riescono a eguagliare la qualità dei modelli teacher originali.

Metodologia

Gli autori presentano RADLADS (Rapid Attention Distillation to Linear Attention Decoders at Scale), un protocollo progettato per convertire i transformer con attenzione softmax in modelli di decoder con attenzione lineare con un costo computazionale minimo. Il processo prevede tre fasi primarie e specifiche adattamenti architettonici:

1. Il Protocollo RADLADS

La conversione è un processo in tre fasi che richiede solo 350–700 milioni di token (meno dello 0,005% dei dati di pre-addestramento del teacher):

• Setup (Trasferimento dei Pesi dell'Attenzione): I pesi relativi all'attenzione ($W_q, W_k, W_v, W_o$) dal modello teacher vengono trasferiti direttamente ai layer di mixing di sequenza del modello student. Altri pesi sono inizializzati tramite metodi standard di pre-addestramento o impostati per imitare il comportamento del teacher senza effetto immediato.
• Fase 1: Allineamento dello Stato Nascosto dell'Attenzione: I layer di mixing di sequenza del modello student vengono addestrati per approssimare gli output degli stati nascosti dei corrispondenti layer di attenzione del teacher. Questo viene eseguito utilizzando un obiettivo di distanza L2 (o MSE). Gli autori hanno scoperto che l'uso di un kernel di Gated Linear Attention (rimuovendo i termini di decadimento off-by-one e i termini bonus) consente un adattamento più stretto agli stati nascosti del teacher rispetto al normale RWKV-6.
• Fase 2: Distillazione della Conoscenza: L'intero modello student viene addestrato per approssimare i logit di output del teacher utilizzando la perdita di divergenza di Kullback-Leibler (KL). Gli autori ipotizzano che la conoscenza fattuale risieda principalmente negli MLP e negli embedding del teacher; pertanto, i tassi di apprendimento per questi componenti sono mantenuti bassi o fissi per prevenire l'oblio catastrofico, mentre il mixer di sequenza viene addestrato in modo più aggressivo.
• Fase 3: Estensione della Lunghezza del Contesto (Opzionale): Il modello viene perfezionato su sequenze più lunghe (fino a 16k token) utilizzando la perdita di cross-entropy standard senza un modello teacher per migliorare le capacità di contesto lungo. In alternativa, gli autori propongono la Fase 2a, che estende la lunghezza della sequenza a 4096 durante la fase di distillazione stessa, rendendo potenzialmente superflua la Fase 3 separata.

2. Nuove Architetture

Gli autori hanno identificato che le architetture standard RWKV non erano sempre ottimali per la conversione. Hanno introdotto due nuove varianti:

• RAD-RWKV6 ("RADFinch"): Una modifica di RWKV-6 che utilizza un kernel di Gated Linear Attention e tecniche di bilanciamento dello stato per migliorare la stabilità e l'adattamento.
• RAD-RWKV7 ("RADGoose"): Una modifica di RWKV-7 che rimuove il meccanismo di "tokenshift" (che non offriva alcun beneficio in questo contesto) e applica direttamente gli Embedding Posizionali Rotativi (RoPE). Questa architettura ha dimostrato una convergenza più veloce e una perdita di distillazione inferiore rispetto alle versioni RWKV-7 non modificate.

3. Iperparametri e Dati

• Dataset: Gli autori si sono stabiliti su DCLM (DataComp-LM) per tutte le fasi di conversione, ritenendolo superiore a FineWeb o FineWeb-Edu per la conversione dei modelli Qwen.
• Learning Rates: Viene utilizzato uno schema di annealing del coseno nella Fase 1, partendo da un valore alto ($10^{-3}$) per allineare gli stati nascosti e terminando vicino al tasso di apprendimento finale del pre-addestramento del teacher ($10^{-5}$). Le Fasi 2 e 3 utilizzano un tasso di apprendimento piatto.

Contributi Chiave

1. Ricetta di Distillazione RADLADS: Un protocollo dettagliato, passo dopo passo, che include iperparametri specifici, conteggi di token e scelte di dataset che consentono una conversione di alta qualità con un dispendio minimo di dati.
2. Nuove Architetture: L'introduzione di RAD-RWKV6 e RAD-RWKV7, che sono ottimizzate per il processo di conversione, garantendo un'inferenza più veloce e un migliore allineamento con i modelli teacher rispetto ai loro predecessori non modificati.
3. Modelli su Larga Scala: La conversione riuscita di popolari modelli open-source Qwen2.5 in varianti di attenzione lineare a scale di 7B, 32B e 72B di parametri.
4. Rilascio Open Source: Il rilascio del codice e dei modelli convertiti (QRWKV6/7-7B/32B/72B) sotto licenza Apache 2.0 (con restrizioni della Licenza Qwen per il modello 72B), consentendo ad altri di replicare il processo.

Risultati

I modelli convertiti raggiungono prestazioni allo stato dell'arte tra i modelli puramente ricorrenti della loro dimensione:

• Efficienza: Convertire un modello da 72B costa meno di 2.000 USD e richiede solo circa 700M di token.
• Prestazioni: Sui benchmark standard (Lambada, MMLU, ARC, ecc.), i modelli RADLADS superano costantemente altri metodi di conversione (ad esempio, SUPRA, LoLCats, MOHAWK, ARWKV).
  • Il QRWKV7-7B-Instruct raggiunge un punteggio MMLU relativo del 92,4% rispetto al suo teacher, superando significativamente altre conversioni da 7B.
  • Il QRWKV6-72B-Instruct raggiunge un punteggio MMLU relativo dell'89,9%, stabilendo un nuovo SoTA per i modelli linguistici puramente RNN a questa scala.
• Velocità di Inferenza: Grazie al meccanismo di attenzione lineare, i modelli convertiti mostrano accelerazioni significative rispetto ai modelli teacher all'aumentare della lunghezza del contesto. Ad esempio, con 8k token di input e 256 di output, il modello 32B QRWKV7 è 1,61x più veloce del suo teacher Qwen3-32B; con 6k di input e 2k di output, l'accelerazione raggiunge 3,41x.

Limitazioni e Affermazioni

Il documento riconosce modestamente diverse limitazioni:

• Ragionamento e Contesto Lungo: Sebbene le prestazioni siano forti sui benchmark standard, i modelli mostrano limitazioni nei compiti di ragionamento complesso (ad esempio, Minerva Math) e nei compiti a contesto molto lungo (ad esempio, RULER), dove non migliorano con la stessa efficacia dei modelli teacher all'aumentare dei token di output.
• Sensibilità dell'Architettura: Ogni nuovo design architettonico richiede test meticolosi per garantire la compatibilità con il protocollo RADLADS. Ad esempio, l'uso di GroupNorm/LayerNorm, comune in alcune varianti di attenzione lineare, ha causato instabilità nell'addestramento a scale superiori a 14B e doveva essere sostituito con tecniche di pre-scaling.
• Allineamento del Dataset: Gli autori osservano che i dataset di ragionamento potrebbero dover essere più strettamente allineati con la distribuzione del modello teacher per prevenire comportamenti di loop ripetitivi, una sfida che hanno affrontato parzialmente mescolando DCLM con OpenThoughts.

Significato

Il documento afferma che RADLADS fornisce una via cost-effective per democratizzare l'accesso ai modelli di attenzione lineare su larga scala. Riducendo il requisito di addestramento da trilioni a centinaia di milioni di token, consente ai ricercatori e alle organizzazioni più piccole di testare, addestrare e distribuire nuove varietà di architetture RNN su larga scala senza gli estremi costi associati al pre-addestramento. Gli autori posizionano questo lavoro come uno strumento per accelerare lo sviluppo della prossima generazione di varianti di attenzione compressiva a stato.