Two-dimensional early exit optimisation of LLM inference

Titolo

Ottimizzazione dell'inferenza dei Large Language Model (LLM) tramite una strategia di uscita anticipata bidimensionale (2D).

1. Il Problema

I Large Language Models (LLM) hanno rivoluzionato l'elaborazione del linguaggio naturale, ma il loro utilizzo in fase di inferenza comporta costi computazionali, di memoria ed energetici significativi. Per compiti relativamente semplici ma ad alto throughput, come la classificazione del testo, l'esecuzione completa di un modello su tutti i livelli (layer) e su tutto l'input rappresenta un eccesso computazionale.
Le strategie di "uscita anticipata" (early exit) esistenti permettono di fermare il calcolo in un livello intermedio quando la confidenza della previsione è sufficiente. Tuttavia, queste strategie agiscono solitamente su una sola dimensione:

1. Uscita per livello (Layer-wise): Si ferma il modello in un certo livello profondo.
2. Uscita per token/sequenza: Si ferma l'elaborazione di certi token o sequenze.
   Il limite è che ottimizzare una sola dimensione non sfrutta appieno la natura incrementale dell'informazione semantica in molti testi, dove il significato può essere determinato già da una parte iniziale del documento.

2. Metodologia: Strategia di Uscita Bidimensionale (2D)

Gli autori propongono una strategia innovativa che coordina simultaneamente l'uscita per livello e per frase (sentence-wise), creando un effetto sinergico moltiplicativo.

• Unità di Input: L'input viene suddiviso in frasi (unità semantiche naturali).
• Progressione 2D:
  • Il modello elabora le frasi una alla volta.
  • Con ogni nuova frase introdotta, il numero di livelli della rete neurale attivati per l'elaborazione aumenta progressivamente.
  • Inizialmente, le prime frasi vengono elaborate solo dai primi livelli (più superficiali). Man mano che nuove frasi arrivano, vengono attivate profondità maggiori, permettendo anche alle frasi precedenti di essere rielaborate nei nuovi livelli attivati.
  • Questo crea dei "blocchi di elaborazione" rettangolari che crescono in due dimensioni (livelli vs. frasi).
• Meccanismo di Arresto:
  • Vengono utilizzati due parametri iperparametrici: una soglia di ignoranza ($\tau_{ignore}$) e una soglia di accettazione ($\tau_{acc}$).
  • Per ogni frase e livello, viene calcolata la confidenza (differenza tra la probabilità massima e la seconda massima).
  • Se la confidenza supera $\tau_{ignore}$, viene aggiunta a un punteggio cumulativo per la classe prevista.
  • Quando il punteggio cumulativo di una classe supera $\tau_{acc}$, l'inferenza si ferma immediatamente, restituendo la previsione.
• Architettura: Il metodo è agnostico rispetto al modello. Utilizza adattatori di classificazione leggeri (due layer fully connected) aggiunti a ogni livello del LLM. Non richiede modifiche architetturali profonde al modello base.
• Addestramento: Sono stati testati due approcci:
  1. Solo Adattatori (Adapter-only): I pesi del modello LLM sono congelati; si addestrano solo gli adattatori di classificazione.
  2. Fine-tuning: Addestramento completo del modello (inclusi gli adattatori) utilizzando una funzione di perdita ponderata che supervisiona tutti i livelli intermedi.

3. Contributi Chiave

1. Metodologia Sinergica: Integrazione dell'arresto per livello con il "taglio" dell'input per frase in un unico processo. Questo porta a risparmi computazionali moltiplicativi (non additivi) rispetto all'ottimizzazione di una singola dimensione.
2. Valutazione Estensiva: Sperimentazione su quattro LLM all'avanguardia (Llama 3.1, Llama 3.2, Gemma, Qwen) con parametri da 3B a 8B, su tre dataset di sentiment analysis di diversa difficoltà.
3. Analisi del Fine-tuning: Identificazione del paradosso per cui il fine-tuning, sebbene migliori l'accuratezza generale, tende a "appiattire" le variazioni di accuratezza tra i livelli, riducendo il vantaggio della strategia 2D.
4. Ortogonalità: La strategia è compatibile con altre tecniche di efficienza come quantizzazione, pruning e Mixture of Experts (MoE).

4. Risultati Sperimentali

Le valutazioni sono state condotte su tre dataset: Steam (2 classi), MMS (3 classi) e Amazon-5 (5 classi).

• Velocità di Inferenza:
  • Su modelli "vanilla" (senza fine-tuning) e compiti più semplici (Steam, MMS), la strategia 2D ha ottenuto un'ulteriore accelerazione di 1.4x – 2.3x rispetto all'ottimale uscita per livello (layer-wise early exit).
  • Su compiti complessi (Amazon-5, 5 classi), il vantaggio è più modesto o nullo per alcuni modelli (es. Llama), indicando che la strategia funziona meglio quando l'informazione semantica si accumula in modo prevedibile.
• Impatto del Fine-tuning:
  • Il fine-tuning riduce il vantaggio della strategia 2D (es. da 2.3x a 1.7x su Steam), ma non lo elimina completamente.
  • L'analisi mostra che il fine-tuning rende l'accuratezza più uniforme tra i livelli, riducendo la capacità del modello di "capire" presto se una previsione è corretta basandosi solo su una parte dell'input.
• Confronto con Baseline: La strategia 2D supera significativamente il metodo LayerSkip (un approccio recente che salta i livelli) sui dataset più semplici, mentre si comporta in modo comparabile o leggermente inferiore sui dataset molto complessi con perdita di accuratezza molto bassa (2%).

5. Significato e Conclusioni

Il paper dimostra che l'ottimizzazione dell'inferenza può beneficiare enormemente dallo sfruttamento della struttura semantica dell'input (frasi) in combinazione con la profondità della rete.

• Efficienza: La strategia 2D offre un modo per ridurre drasticamente il costo computazionale senza sacrificare l'accuratezza, specialmente per task di classificazione su testi brevi o medi.
• Insight sul Fine-tuning: Un risultato cruciale è la scoperta che il fine-tuning standard può essere controproducente per le strategie di uscita anticipata, poiché omogeneizza il comportamento dei livelli. Questo suggerisce la necessità di nuove tecniche di addestramento che preservino i gradienti di accuratezza tra i livelli superficiali e profondi.
• Applicabilità Futura: Sebbene testata sulla classificazione del sentiment, l'approccio è promettente per qualsiasi task di elaborazione di sequenze dove il significato emerge gradualmente (es. categorizzazione di argomenti, rilevamento di intenti).

In sintesi, gli autori propongono un framework flessibile ed efficiente che, coordinando la profondità della rete e la lunghezza dell'input, supera i limiti delle strategie di ottimizzazione unidimensionali, offrendo un percorso promettente per l'implementazione di LLM più economici e veloci.