Improving LLM Predictions via Inter-Layer Structural Encoders

1. Il Problema

Nella pratica standard per i Large Language Models (LLM), le previsioni per compiti specifici (come classificazione o calcolo della similarità semantica) si basano esclusivamente sulle rappresentazioni dell'ultimo strato (final-layer) del modello. Tuttavia, ricerche recenti hanno dimostrato che:

• Gli strati intermedi contengono informazioni sostanziali e spesso più rilevanti per il compito rispetto all'ultimo strato.
• Lo strato ottimale varia a seconda del compito specifico.
• I metodi esistenti per aggregare le informazioni da tutti gli strati (come la somma pesata scalare di ELMo o l'attenzione basata sull'ultimo strato come DWAtt) hanno limitazioni: o non catturano interazioni complesse non lineari, o ancorano l'attenzione all'ultimo strato, non trattando tutti gli strati in modo equo.
• Modificare l'architettura del Transformer per integrare queste informazioni richiede spesso un addestramento da zero, il che è computazionalmente proibitivo.

L'obiettivo è quindi trovare un metodo efficiente, rigoroso ed efficace per combinare le rappresentazioni di tutti gli strati interni di un LLM pre-addestrato (frozen) senza modificarne l'architettura interna.

2. Metodologia: Inter-Layer Structural Encoders (ILSE)

Il paper propone ILSE, un approccio strutturato che tratta le rappresentazioni degli strati come nodi in un grafo, utilizzando encoder strutturali per apprendere una rappresentazione unificata.

Fasi principali:

1. Rappresentazione dello Strato: Per ogni strato $\ell$ dell'LLM, si calcola una singola rappresentazione $z_\ell$ effettuando una media (mean-pooling) sulle rappresentazioni di tutti i token di quello strato.
2. Costruzione del Grafo: Le $L$ rappresentazioni degli strati vengono mappate su nodi di un grafo. La scelta della struttura del grafo è cruciale per evitare l'over-squashing (compressione esponenziale delle informazioni) e l'overfitting strutturale.
3. Encoder Strutturali: ILSE introduce tre varianti basate su diverse topologie di grafi, elaborate tramite DeepSets o Graph Neural Networks (GNN):
   • Set-Encoder: Tratta gli strati come un insieme (grafo vuoto, nessun bordo) e usa un aggregatore DeepSets (permutazione-invariante).
   • FC-Encoder (Fully Connected): Costruisce un grafo completamente connesso dove ogni strato comunica con tutti gli altri.
   • Cayley-Encoder (Il contributo chiave): Utilizza un Grafo di Cayley basato sul gruppo lineare speciale $SL(2, \mathbb{Z}_n)$.
     • Motivazione: I grafi di Cayley sono grafi regolari, sparsi ma altamente connessi, con un diametro logaritmico. Questo garantisce un flusso di informazioni efficiente tra tutti i nodi senza colli di bottiglia (bottleneck-free).
     • Implementazione: Se il numero di strati $L$ è inferiore al numero di nodi del grafo di Cayley, i nodi rimanenti sono inizializzati come "nodi virtuali" (vettori zero) per mantenere la struttura matematica del grafo. Solo le rappresentazioni dei nodi non virtuali vengono utilizzate per la predizione finale.
4. Apprendimento: I modelli base (LLM) rimangono congelati (frozen). Viene addestrato solo l'encoder strutturale (GNN o DeepSet) e un semplice classificatore lineare (per la classificazione) o una funzione di similarità (per STS).

3. Contributi Chiave

1. Proposta di ILSE: Un metodo matematicamente fondato per apprendere rappresentazioni unificate da tutti gli strati interni di qualsiasi LLM.
2. Introduzione del Cayley-Encoder: L'uso di grafi di Cayley per l'aggregazione inter-strato, sfruttando le loro proprietà di espansione per garantire un flusso di informazioni ottimale senza colli di bottiglia.
3. Valutazione Estensiva: Test su 13 compiti (5 di classificazione, 8 di similarità semantica STS) e 9 LLM pre-addestrati (da 14 milioni a 8 miliardi di parametri).
4. Efficienza nei Dati (Few-Shot): Dimostrazione che ILSE eccelle in scenari con pochi dati etichettati, superando i baselines anche con soli 32 campioni per classe.
5. Democratizzazione delle Prestazioni: Dimostrazione che ILSE permette a modelli piccoli (es. 14M parametri) di raggiungere prestazioni competitive con modelli molto più grandi (es. 2.8B parametri).

4. Risultati Sperimentali

I risultati sono stati ottenuti confrontando ILSE con baselines come Last-Layer, Best-Layer (lo strato singolo migliore), Weighted Sum, DWAtt e MLP su strati singoli.

• Prestazioni Generali: ILSE ha superato tutti i baselines in tutti e 15 i casi di classificazione e in 17 su 21 configurazioni zero-shot per STS.
• Miglioramenti:
  • Fino al 44% di miglioramento in accuratezza rispetto alla Last-Layer (in particolare nel task "Emotion").
  • Fino al 25% di miglioramento nelle metriche di similarità.
  • In media, un guadagno del +30% rispetto alla Last-Layer e del +25% rispetto alla Best-Layer.
• Confronto tra Encoder:
  • Il Cayley-Encoder ha ottenuto le prestazioni complessive migliori, specialmente nei task di similarità semantica (STS), superando i baselines in 17 su 24 configurazioni STS.
  • Il FC-Encoder ha mostrato prestazioni eccellenti nei task di classificazione.
  • Entrambi gli encoder basati su grafi (FC e Cayley) hanno superato il Set-Encoder, indicando che la connettività inter-strato è fondamentale.
• Efficienza dei Parametri: ILSE aggiunge una quantità trascurabile di parametri addestrabili (es. ~0.015% - 0.10% in più rispetto al modello base), rendendolo molto più efficiente in termini di parametri rispetto a DWAtt (che richiede ~0.04% - 0.49% in più).
• Scalabilità: ILSE mantiene i suoi vantaggi su modelli di dimensioni diverse (da 14M a 2.8B). Un modello Pythia-14M con ILSE raggiunge prestazioni simili a un Pythia-2.8B standard.
• Few-Shot Learning: Con soli 32 campioni per classe, ILSE supera baselines addestrati su dataset completi (circa 10.000 campioni).

5. Significato e Implicazioni

Il lavoro dimostra che la "profondità interna" degli LLM contiene segnali complementari rilevanti per il compito che vengono spesso persi se si utilizza solo l'ultimo strato.

• Inductive Bias Geometrico: L'uso di grafi di Cayley fornisce un inductive bias efficace: una connettività sparsa, regolare e priva di colli di bottiglia permette una generalizzazione migliore e un flusso di informazioni efficiente tra gli strati.
• Accessibilità: ILSE offre un modo semplice, efficiente e matematicamente solido per migliorare le rappresentazioni di LLM congelati, rendendo i modelli più piccoli competitivi con quelli più grandi senza la necessità di un ri-addestramento costoso.
• Futuro: Il metodo suggerisce che l'integrazione di strutture geometriche consapevoli (geometry-aware) nell'architettura dei transformer durante il pre-addestramento potrebbe portare a modelli di base ancora più robusti.

In sintesi, ILSE rappresenta un avanzamento significativo nell'ottimizzazione dell'uso delle risorse interne degli LLM, trasformando la struttura a strati da una semplice sequenza di elaborazione in una rete strutturata di informazioni ottimizzata per compiti specifici.