Il MIT ha scoperto il trucco per rendere i modelli AI più veloci del 200%

L'addestramento degli LLM con capacità di ragionamento rappresenta oggi una delle sfide più onerose in termini di calcolo ed energia. Questi modelli, progettati per scomporre problemi complessi in passaggi intermedi - non a caso si parla di catena di pensiero - richiedono cicli di apprendimento per rinforzo estremamente intensivi.

Un gruppo di ricercatori del Massachusetts Institute of Technology, in collaborazione con NVIDIA, ETH Zurich, MIT-IBM Watson AI Lab e University of Massachusetts Amherst, ha proposto una soluzione che punta a migliorare radicalmente l'efficienza del processo senza compromettere l'accuratezza dei modelli.

Nell'addestramento dei "reasoning LLM" tramite addestramento per rinforzo, il modello genera molteplici risposte candidate per ogni query (fase di rollout), riceve una ricompensa per la migliore e aggiorna i propri parametri di conseguenza. Questo ciclo viene ripetuto migliaia di volte.

Secondo i ricercatori, fino all'85% del tempo di esecuzione è assorbito proprio dalla generazione delle risposte, non dall'aggiornamento dei pesi del modello. Inoltre, l'architettura distribuita del training impone che tutti i processori coinvolti attendano il completamento delle risposte più lunghe prima di procedere allo step successivo. Di conseguenza, parte delle risorse di calcolo resta inutilizzata.

Per affrontare il problema, il team ha esteso il concetto di speculative decoding - che già aveva dimostrato i suoi vantaggi. La tecnica prevede l'impiego di un modello più piccolo - definito "drafter" - che anticipa le possibili uscite del modello principale. Quest'ultimo verifica in parallelo le previsioni, accettando quelle corrette e scartando le altre, con un'accelerazione rispetto alla generazione sequenziale.

Tuttavia, nei contesti di reinforcement learning tradizionali il modello principale viene aggiornato continuamente: un drafter statico diventerebbe rapidamente obsoleto. La soluzione proposta, denominata "Taming the Long Tail" (TLT), introduce un meccanismo adattivo su due livelli:

• Adaptive drafter trainer: utilizza i processori inattivi per aggiornare dinamicamente il modello ausiliario, mantenendolo allineato al modello principale senza richiedere risorse aggiuntive
• Adaptive rollout engine: regola automaticamente la configurazione dello speculative decoding in base al carico di lavoro e alle percentuali di accettazione delle previsioni

Il modello ausiliario è progettato per essere leggero e riutilizza componenti del processo di training del modello principale, massimizzando così l'efficienza.

TLT ha permesso di accelerare il training tra il 70% e il 210%, mantenendo invariata l'accuratezza finale dei modelli.

L'approccio non introduce overhead computazionale aggiuntivo, in quanto sfrutta esclusivamente il tempo di inattività già presente nell'infrastruttura di training. Questo aspetto potrebbe tradursi in una riduzione significativa dei costi operativi e dei consumi energetici, soprattutto in applicazioni critiche come la previsione di tendenze finanziarie o l'analisi dei rischi nelle reti elettriche.