La selezione dei canali per la decodifica dei modelli mixture-of-experts non richiede dati di calibrazione

Abstract degli autori

Nei modelli mixture-of-experts (MoE), la decodifica interattiva con batch di piccole dimensioni è rallentata dal caricamento dalla memoria dei pesi degli esperti instradati. Poiché la maggior parte dei canali intermedi di un esperto instradato contribuisce poco all’elaborazione di un singolo token, i selettori di canali prevedono quali siano rilevanti e ne leggono soltanto i pesi, mantenendo in genere i canali che superano una soglia calibrata su dati offline. Mostriamo che la selezione dei canali non richiede dati di calibrazione. In primo luogo, una soglia calibrata non determina quanto un selettore legge, perché il numero di canali che la superano varia da un input all’altro: Prox, con soglie calibrate per leggere il 18% dei pesi, ne legge in realtà solo il 12,9% in media nei nostri benchmark. In secondo luogo, top-k mantiene per ogni token i k canali con i punteggi più alti e non richiede dati di calibrazione. A parità del budget usato per la calibrazione, la versione top-k di un selettore eguaglia o supera la versione basata su soglia, usando gli stessi punteggi di selezione, in 40 confronti su 43 con budget del 13% e del 18%, compresi tutti i 23 confronti relativi a compiti di generazione. Le versioni top-k evitano inoltre il collasso in cicli di ripetizione che le soglie calibrate causano su Gemma 4. Calibrare una soglia, dunque, non offre nessuno dei due vantaggi per cui i selettori la usano: un budget di lettura noto e un miglioramento costante dell’accuratezza a quel budget. Presentiamo inoltre un lettore costruito interamente a partire dai pesi del modello, che filtra i canali con uno sketch ternario e recupera quelli saltati con una mappatura a basso rango. Su tre modelli MoE sottoposti a post-training, il nostro lettore conserva il 98,1% dell’accuratezza del modello completo effettuando il 17,8% delle letture dei pesi degli esperti instradati. Con kernel Triton, i blocchi MoE vengono eseguiti da 1,60 a 1,90 volte più velocemente rispetto a vLLM ottimizzato con batch di dimensione uno, e fino a 2,69 volte più velocemente con batch più grandi, sugli stessi tre modelli.