La previsione multi-token fatta nel modo giusto

Abstract degli autori

La previsione multi-token (MTP) è ormai un elemento comune nel pre-training dei modelli linguistici, perché migliora le prestazioni downstream e consente un’inferenza più rapida tramite speculative decoding. Individuiamo un difetto nell’addestramento MTP standard: una componente residua del gradiente MTP interferisce sia con l’apprendimento del token successivo sia con quello multi-token, e l’interferenza aumenta con la scala. Nello specifico, il gradiente di cross-entropy MTP può essere scomposto esattamente in un termine di auto-distillazione KL, che spinge ciascuna testa MTP verso la distribuzione di previsione del token successivo (NTP) del modello stesso e migliora lo speculative decoding, e in un residuo con un basso rapporto segnale-rumore, che nasconde il segnale KL utile. Introduciamo l’espediente di recupero del gradiente KL: quando due distribuzioni condividono una testa di output lineare, il gradiente della loro divergenza KL può essere recuperato esattamente combinando i rispettivi gradienti di cross-entropy tramite l’instradamento straight-through del gradiente. A differenza della minimizzazione diretta della KL, questo espediente non richiede matrici di logit grandi quanto il vocabolario ed è compatibile con i kernel standard di cross-entropy fusa. Questo rende possibile Soft Multi-Token Prediction (Soft-MTP), una modifica di una sola riga di codice a MTP che rimuove il residuo rumoroso e propaga all’indietro soltanto il segnale di auto-distillazione KL, riutilizzando i calcoli di cross-entropy NTP e MTP già esistenti. Soft-MTP supporta sia la MTP in stile DeepSeek sia quella in stile Meta, non aggiunge parametri e non comporta overhead misurabile. Con modelli da 167 milioni a 1,2 miliardi di parametri, Soft-MTP supera la MTP standard, a ogni scala, con una perplexity NTP inferiore, uno speculative decoding più rapido e una maggiore accuratezza zero-shot; inoltre, raggiunge una perplexity MTP inferiore a partire da 689 milioni di parametri e un ragionamento matematico e sul codice più solido dopo il fine-tuning a 1,2 miliardi di parametri. Nella maggior parte delle metriche, Soft-MTP eguaglia, se non supera, la minimizzazione diretta della divergenza KL. In effetti, Soft-MTP offre un vantaggio senza costi nell’ottimizzazione MTP. Il nostro codice è disponibile all’indirizzo https://github.com/JaydenTeoh/Soft-MTP.