Inferenza adattiva alla profondità nei modelli linguistici iterativi tramite batching continuo per profondità

Abstract degli autori

Una delle principali promesse dei modelli linguistici iterativi è l’inferenza adattiva alla profondità. Ripetendo un numero variabile di volte un blocco di livelli condivisi, il modello può usare meno risorse di calcolo per i token «facili» e più per quelli «difficili». Tuttavia, i token che richiedono numeri diversi di iterazioni non possono condividere un passaggio in avanti uniforme e quindi non possono essere gestiti dai sistemi di batching standard, come vLLM. Il valore pratico dell’inferenza adattiva alla profondità dipende dunque dalla possibilità di rendere efficiente il batching. Presentiamo il primo metodo efficiente per l’inferenza adattiva alla profondità nei modelli linguistici iterativi: il batching continuo per profondità (CDB), che forma nuovi batch tra un’iterazione e la successiva. Il nostro metodo pianifica dinamicamente l’esecuzione delle parti iterative e non iterative dell’architettura, gestisce la cache KV per le iterazioni e prevede in anticipo quali token usciranno dal ciclo, così da poter preparare i batch in modo asincrono. Gli esperimenti su Ouro 1.4B e Huginn 3.5B mostrano che le architetture interamente iterative sono le più adatte all’inferenza adattiva alla profondità, poiché i grandi livelli non iterativi esterni al nucleo ricorrente (per esempio l’embedding dei token, la testa LM e i blocchi transformer non condivisi) rallentano e complicano la pianificazione dell’esecuzione. Nel complesso, CDB realizza fino al 99% dell’accelerazione massima stimata ottenibile: ulteriori miglioramenti dipendono principalmente dall’architettura del modello e dal comportamento dei token all’uscita dal ciclo.