ConfLayers: Adaptive Confidence-based Layer Skipping for Self-Speculative Decoding

1. Il Problema

L'inferenza dei Large Language Models (LLM) moderni, caratterizzati da decine o centinaia di miliardi di parametri, è diventata un collo di bottiglia computazionale a causa della natura autoregressiva della generazione dei token. Le tecniche tradizionali di accelerazione, come la compressione del modello o l'uscita anticipata (early exit), spesso non affrontano la dipendenza sequenziale fondamentale o richiedono costi di addestramento significativi.

Lo Speculative Decoding (decodifica speculativa) è emerso come una soluzione promettente, utilizzando un modello "bozzetto" (draft) leggero per generare più token in anticipo, che vengono poi verificati dal modello target. Tuttavia, l'approccio speculativo tradizionale presenta sfide:

• Richiede un modello bozzetto separato, aumentando l'uso di memoria, i costi di addestramento e la complessità di distribuzione.
• Le discrepanze di distribuzione tra il modello bozzetto e quello target possono ridurre il tasso di accettazione dei token.

La Self-Speculative Decoding (SSD) risolve il problema del modello separato riutilizzando lo stesso LLM, creando un sotto-rete leggera (saltando alcuni strati) per la generazione speculativa e utilizzando lo strato completo per la verifica. Tuttavia, le strategie esistenti per decidere quali strati saltare si basano su:

1. Metodi basati su policy: Richiedono addestramento e hanno una scarsa generalizzazione.
2. Metodi euristici fissi: Sono semplici ma mancano di adattabilità, portando a compromessi subottimali tra velocità e qualità.

2. Metodologia: ConfLayers

Il paper propone ConfLayers, un framework "plug-and-play" e senza addestramento (training-free) che costruisce dinamicamente un modello bozzetto adattivo basandosi sulla fiducia (confidence) degli strati intermedi.

Il processo si articola in tre fasi principali (illustrate nella Figura 2 del paper):

A. Calcolo della Fiducia per Strato (Confidence Computation)

Per ogni strato del modello durante la generazione, ConfLayers calcola un punteggio di fiducia che riflette l'incertezza della rappresentazione intermedia:

1. Vengono estratti gli stati nascosti dei moduli FFNN (Feed-Forward Neural Network) e Attention.
2. Questi stati vengono proiettati attraverso l'head linguistico (LM Head) condiviso per ottenere i logit.
3. Viene calcolata l'entropia di Shannon della distribuzione di probabilità derivata dai logit.
4. Il punteggio di fiducia è definito come il complemento dell'entropia normalizzata:
   $$\text{confidence} = 1 - \frac{H(p)}{H_{max}}$$
   Dove valori vicini a 1 indicano previsioni certe (entropia bassa) e valori vicini a 0 indicano incertezza.

B. Filtraggio Adattivo degli Strati (Adaptive Window Sizing)

Invece di usare una finestra fissa per confrontare i livelli, ConfLayers utilizza una finestra adattiva per determinare quali strati saltare:

• Normalizzazione: I punteggi di fiducia vengono standardizzati (media 0, varianza 1).
• Gradiente: Viene calcolato il gradiente del secondo ordine per quantificare le fluttuazioni locali della fiducia.
• Dimensione della Finestra: La dimensione della finestra di confronto per ogni strato $i$ ($w_i$) si adatta dinamicamente in base alla magnitudine del gradiente e all'indice dello strato:
  $$w_i = w_{base} + w_{max} \cdot g_i \cdot \left(\frac{l_i}{\text{num\_layers}}\right)$$
  • Gli strati iniziali (informazioni lessicali/sintattiche) usano finestre più piccole.
  • Gli strati finali (informazioni contestuali aggregate) usano finestre più ampie.
  • Le regioni con fluttuazioni brusche di fiducia usano finestre più grandi per una valutazione più robusta.
• Selezione: Uno strato viene saltato se il suo punteggio di fiducia è inferiore a una soglia adattiva basata sulla media e sulla deviazione standard locale ($\hat{c}_i < \mu_{i,local} - \lambda \cdot \sigma_{i,local}$).

C. Ricerca Periodica Basata su Intervalli

Il sistema non ricalcola la configurazione ad ogni token, ma opera in cicli:

1. Inizia con un set di strati non saltati uniforme.
2. Genera token candidati e li verifica.
3. Esegue una "ricerca" periodica (ogni Opt_Interval iterazioni) per aggiornare il set di strati da saltare, massimizzando il numero di token accettati.
4. Il processo si ferma quando viene raggiunto un punteggio target o il numero massimo di iterazioni, mantenendo il miglior set di strati trovato.

3. Contributi Chiave

1. Meccanismo basato sulla fiducia: Introduzione di un metodo per costruire sotto-reti bozzetto in SSD calcolando l'incertezza degli strati intermedi, eliminando la necessità di modelli ausiliari o addestramento.
2. Algoritmo di adattamento iterativo: Progettazione di un algoritmo "plug-and-play" che bilancia velocità di decodifica e tasso di accettazione dei token attraverso una selezione guidata dalla fiducia e finestre adattive.
3. Valutazione empirica completa: Dimostrazione che l'approccio supera le tecniche esistenti (come SWIFT e DEL) su diversi modelli (LLaMa-2, LLaMa-3, CodeLLaMa, Qwen) e dataset, mantenendo la qualità dell'output.

4. Risultati Sperimentali

La valutazione è stata condotta su modelli LLaMa-2 (13B, 70B), LLaMa-3 (8B, 70B), CodeLLaMa-34B e Qwen2.5-Math-72B su task di riassunto, ragionamento matematico, traduzione e istruzioni.

• Speedup: ConfLayers ha raggiunto uno speedup fino a 1.4x rispetto alla generazione autoregressiva standard (vanilla).
  • Su LLaMa-2-70B per il riassunto (CNN-DM): 1.37x (vs 1.30x di SWIFT).
  • Su LLaMa-3-70B: 1.38x.
• Confronto con Baseline:
  • Supera costantemente SWIFT (che seleziona strati in modo adattivo ma con vincoli diversi) e DEL (che spesso non offre speedup senza fine-tuning).
  • A differenza di SWIFT, ConfLayers non limita la ricerca a un tasso di salto fisso, esplorando dinamicamente il numero di strati da saltare.
• Qualità: Il punteggio Rouge-2 rimane in alto sovrapposizione rispetto alla generazione standard, confermando che la qualità del testo non viene compromessa.
• Robustezza su Input Dinamici: In scenari dove lo stesso set di strati saltati viene applicato a task diversi (senza ri-addestramento), ConfLayers mantiene uno speedup > 1.1x, mentre SWIFT fallisce in alcuni casi (es. speedup < 1 su Alpaca).

5. Significato e Impatto

ConfLayers rappresenta un passo significativo verso un'inferenza LLM più efficiente e scalabile.

• Efficienza Operativa: Elimina il sovraccarico di memoria e computazionale legato all'addestramento di policy di salto o all'uso di modelli bozzetto separati.
• Adattabilità: La capacità di adattare la strategia di salto in base alla fiducia locale e al contesto specifico del modello permette di ottenere compromessi velocità-qualità superiori rispetto alle euristiche fisse.
• Semplicità: Dimostra che approcci euristici semplici ma guidati da metriche interne (fiducia/entropia) possono essere più efficaci di metodi complessi basati sull'apprendimento, offrendo una soluzione pratica per l'accelerazione dell'inferenza in tempo reale.

In sintesi, ConfLayers trasforma l'inferenza speculativa in un processo intrinseco al modello, ottimizzando dinamicamente l'architettura di esecuzione senza modificare i pesi del modello o richiedere dati aggiuntivi.