A BERTology View of LLM Orchestrations: Token- and Layer-Selective Probes for Efficient Single-Pass Classification

1. Enunciato del Problema

I moderni sistemi di produzione basati su Large Language Model (LLM) si affidano tipicamente a un pattern di orchestrazione in cui un LLM di servizio principale è circondato da modelli "guardia" ausiliari per la moderazione della sicurezza, la conformità alle policy e l'analisi del sentiment. Questo approccio soffre di tre principali svantaggi:

• Latenza: Ogni richiesta richiede un'ulteriore invocazione del modello (modello guardia $\to$ LLM di servizio), raddoppiando il tempo di inferenza.
• Sovraccarico di Risorse: Il deployment di modelli separati aumenta l'ingombro in VRAM e la complessità operativa.
• Sottoutilizzo: L'LLM di servizio calcola già ricche rappresentazioni interne (stati nascosti) che contengono segnali di sicurezza e classificazione, ma questi vengono scartati quando si utilizzano solo i logit finali o un modello separato.

I metodi di "riutilizzo" esistenti tentano di estrarre segnali dal modello di servizio, ma spesso si impegnano su punti di lettura fissi (ad esempio, solo i logit del primo token o lo stato nascosto dell'ultimo strato), rischiando di perdere informazioni discriminative distribuite su diversi strati e posizioni dei token.

2. Metodologia

Gli autori propongono un framework di Sonda Selettiva per Token e Strato che tratta la classificazione come un problema di selezione di rappresentazioni sull'intero tensore degli stati nascosti ($L \times T \times d$, dove $L$=strati, $T$=token, $d$=dimensione). L'innovazione centrale è un'Architettura di Aggregazione a Due Stadi applicata a un LLM di servizio congelato.

A. Panoramica dell'Architettura

La sonda consiste in un classificatore leggero addestrato sugli stati nascosti generati durante il passaggio in avanti (forward pass) dell'LLM di servizio. Non modifica l'LLM di base.

1. Stadio 1 (Aggregazione a Livello di Token): Per ogni strato $l$, i $T$ vettori di token vengono aggregati in un singolo vettore di riepilogo dello strato $v^{(l)}$.
2. Stadio 2 (Aggregazione a Livello di Strato): I $L+1$ vettori di riepilogo degli strati vengono aggregati in un singolo vettore di rappresentazione globale $v$.
3. Testa di Classificazione: Un livello lineare mappa $v$ sui logit delle classi.

B. Meccanismi di Aggregazione

Il paper valuta tre meccanismi per eseguire l'aggregazione, bilanciando espressività ed efficienza dei parametri:

1. Pooling Diretto: Operatori fissi (Max o Mean pooling) su posizioni valide di token/strato. (Costo più basso, ~0,003M parametri).
2. Cancello di Assegnazione Punteggio (Scoring Attention Gate): Un meccanismo apprendibile che assegna punteggi di importanza scalari ai token/strati utilizzando una singola proiezione lineare e normalizzazione softmax. (~0,1M parametri).
3. Self-Attention Multi-Head (MHA): Un modulo MHA ridimensionato (utilizzando una riduzione aggressiva della dimensione, ad es. $d/16$) per catturare interazioni complesse tra token e strati. (~35M parametri).

C. Flusso Operativo

Durante l'inferenza, la sonda esegue in parallelo al processo di generazione. Se la sonda prevede un'etichetta "non sicura", lo strato di orchestrazione può interrompere immediatamente la generazione prima che qualsiasi token venga trasmesso, restituendo un rifiuto o riproponendo il prompt al modello. Questo elimina la necessità di una seconda chiamata al modello.

3. Contributi Chiave

• Passaggio da Letture Fisse a Selettive: Il paper va oltre le letture a strato fisso o a token fisso (come [CLS] di BERT o i logit del primo token) per introdurre un meccanismo di selezione appreso che identifica dove nel tensore $L \times T$ il segnale è più separabile.
• Template di Aggregazione a Due Stadi: Un'architettura innovativa che prima riassume i token all'interno degli strati, poi aggrega attraverso gli strati, permettendo al modello di apprendere l'importanza specifica del compito per profondità e token.
• Trade-off Efficienza vs. Prestazioni: Dimostra che una classificazione ad alta accuratezza può essere ottenuta con parametri addestrabili minimi (fino a 35M) e zero latenza aggiuntiva rispetto al solo modello di servizio.
• Generalizzazione: Convalida l'approccio su architetture diverse, inclusi modelli densi (Llama-3.2-3B) e Mixture-of-Experts (GPT-OSS-20B, Qwen3-30B-A3B).

4. Risultati Sperimentali

Moderazione della Sicurezza (ToxicChat & WildGuardMix)

• Prestazioni: La sonda MHA ha raggiunto 84,51 F1 su ToxicChat e 88,55 F1 su WildGuardMix.
• Confronto:
  • Ha superato le baseline di riutilizzo basate solo sui logit (es. MULI: 77,8 F1).
  • Ha eguagliato o superato i modelli guardia autonomi (es. ToxicChat-T5-large: 82,2 F1; WildGuard: 88,9 F1) nonostante l'uso di ordini di grandezza di parametri inferiori (35M vs 780M–8B).
  • Zero Latenza Aggiuntiva: A differenza dei modelli guardia che richiedono una chiamata API separata, la sonda esegue nello stesso passaggio in avanti.
• Cross-Architettura: La classifica delle prestazioni (Pooling < Scoring Attention < MHA) è rimasta coerente tra i modelli Llama, GPT-OSS e Qwen.

Sentiment ed Emozione (IMDB, SST-2, Emotion)

• Le sonde hanno superato significativamente il prompting zero-shot e few-shot (es. 87,68% di accuratezza su Emotion contro il 64,05% per MULI).
• Le prestazioni sono state competitive con classificatori specifici per il compito di grandi dimensioni (es. DeBERTa, RoBERTa) mantenendo un pipeline di inferenza a passaggio singolo.

Analisi dell'Efficienza

• Latenza: La sonda MHA ha aggiunto solo ~14ms di latenza per campione rispetto al modello base, mentre una pipeline guardia-then-serve ha aggiunto ~60–100ms.
• Memoria: Le varianti della sonda hanno mantenuto la memoria GPU di picco vicina all'ingombro del modello base (~6,5–7,0 GB), mentre le pipeline guardia richiedevano 8,0–22,8 GB.

Interpretabilità (Attenzione agli Strati)

• L'analisi del cancello di assegnazione punteggio (Scoring Attention) ha rivelato che i segnali di sicurezza sono distribuiti. I prompt tossici tendevano ad attivare gli strati successivi (L17–L28), mentre i prompt non tossici si basavano sugli strati finali e sugli embedding. Le classificazioni errate mostravano spesso pattern di attenzione allineati alla classe prevista piuttosto che alla verità fondamentale, suggerendo che la sonda apprende varietà di caratteristiche distinte per diversi tipi di contenuto.

5. Significato e Conclusione

Questo paper fornisce un'alternativa convincente al pattern di orchestrazione standard dei "modelli guardia". Sfruttando i principi della BERTology (che diversi strati codificano diverse astrazioni) e applicandoli ai moderni LLM decoder-only, gli autori dimostrano che:

1. I compiti di sicurezza e NLU sono già "risolti" internamente dal modello di servizio; la sfida è estrarre il segnale in modo efficiente.
2. La classificazione a passaggio singolo è fattibile per la produzione, offrendo una riduzione massiccia di latenza e costi infrastrutturali senza sacrificare l'accuratezza.
3. L'aggregazione appresa (in particolare MHA) è superiore al pooling fisso, dimostrando che lo "strato migliore" per un compito non è statico ma dipende dal contenuto dell'input.

Il lavoro suggerisce un futuro in cui l'orchestrazione degli LLM è snellita, riducendo la necessità di pipeline di sicurezza pesanti e separate e abilitando sistemi AI più reattivi ed economicamente efficienti.