Accelerare Gemma 4 con la previsione multi-token

In Gemma 4, la previsione multi-token (MTP) è l'architettura specifica utilizzata per abilitare la decodifica speculativa ad alta efficienza. La decodifica speculativa è una tecnica per velocizzare l'inferenza nei modelli linguistici di grandi dimensioni. Anziché fare affidamento esclusivamente sul modello di destinazione di grandi dimensioni per generare token in modo autoregressivo (generando un token alla volta, dove ogni nuovo token dipende da quelli precedenti), un "modello di bozza" più piccolo e veloce prevede diversi token in anticipo. Il modello di destinazione verifica quindi questi token di bozza in parallelo. Se il modello di destinazione rifiuta un token di bozza, produce comunque il token corretto per quella posizione (assicurandosi che il passaggio non venga sprecato) e il modello di bozza riprende la previsione dal nuovo token corretto.

Gemma 4 implementa MTP estendendo il modello di base con questo modello di bozza più piccolo e veloce. Questo modello di bozza non è indipendente, in quanto condivide la tabella di incorporamento dell'input con il modello di destinazione e si basa direttamente sulle attivazioni dell'ultimo livello. Ciò comporta un'accelerazione significativa della decodifica, garantendo al contempo la stessa qualità della generazione autoregressiva standard, rendendo questi checkpoint perfetti per le applicazioni a bassa latenza e on-device.

La decodifica speculativa funziona creando diverse bozze di token e verificandole in un singolo passaggio in avanti. Per i modelli densi, vengono utilizzati gli stessi pesi per ogni token, quindi la verifica di più token di bozza aggiunge un overhead minimo. I modelli Mixture of Experts (MoE) come Gemma 4 26B A4B funzionano in modo diverso. Ogni token può attivare esperti diversi, quindi la verifica dei token di bozza può richiedere il caricamento di pesi di esperti aggiuntivi dalla memoria, compensando i vantaggi della creazione di bozze. Con dimensioni batch maggiori, in genere si verifica una maggiore sovrapposizione degli esperti attivati tra le sequenze, migliorando il riutilizzo dei pesi caricati. Con una dimensione del batch pari a 1, questa sovrapposizione è limitata, motivo per cui il modello di bozza 26B A4B potrebbe non produrre accelerazioni sulle piattaforme hardware senza un buon parallelismo.

Miglioramenti di MTP

Gemma 4 introduce diversi miglioramenti alla pipeline di decodifica speculativa standard per migliorare la qualità dei token di bozza e l'efficienza:

  • Incorporamenti di input condivisi: il modello di bozza condivide la tabella di incorporamento dell'input con il modello di destinazione.
  • Attivazioni di destinazione: il modello di bozza utilizza le attivazioni dell' ultimo livello del modello di destinazione, le concatena con gli incorporamenti dei token e le proietta verso il basso nella dimensione del modello di bozza.
  • Incorporatore efficiente: per evitare l'operazione costosa di previsione sull'intero vocabolario, il modello raggruppa i token simili in cluster. Innanzitutto identifica i cluster più probabili e poi limita i calcoli finali solo ai token all'interno dei cluster selezionati (solo E2B ed E4B).