Test-Time Safety Alignment

1. Enunciato del Problema

I Large Language Models (LLM) subiscono un addestramento per la sicurezza (ad esempio, RLHF, DPO) per rifiutare richieste dannose. Tuttavia, prompt avversari (jailbreak) possono aggirare queste difese sfruttando il divario tra la distribuzione di addestramento del modello e lo spazio di input illimitato. Questi prompt utilizzano spesso un'inquadratura strutturale (roleplay, oscuramento) per estrarre contenuti dannosi.

Le difese esistenti a tempo di test affrontano limitazioni strutturali:

• Strategie a livello di prompt manipolano token discreti, che sono fragili di fronte ad attacchi semantici.
• Steering delle attivazioni richiede l'accesso agli stati nascosti e spesso si basa su set di calibrazione fissi che potrebbero non generalizzare.
• Modelli guardiano richiedono un'inferenza separata e spesso seguono un paradigma "rileva-poi-filtra" piuttosto che guidare la generazione.

La Domanda Fondamentale: È possibile ottimizzare gli embedding delle parole di input in modo sub-lessicale (modificando la rappresentazione vettoriale continua senza cambiare il token visibile) per guidare i modelli allineati verso output sicuri, anche quando il modello produce una distribuzione bimodale "rifiuta-o-compi" invece di una continuazione fluida?

2. Metodologia: Allineamento per la Sicurezza a Tempo di Test (TSA)

Il metodo proposto, TSA, opera interamente a tempo di test senza modificare i pesi del modello, non richiede addestramento ausiliario e tratta l'LLM come una scatola nera. Ottimizza direttamente la matrice di embedding del prompt.

A. Funzione Obiettivo

L'obiettivo è minimizzare la dannosità del testo generato $y = f(X)$, dove $f$ è l'LLM e $X$ è la matrice di embedding del prompt.

• Oracolo di Dannosità ($h$): Il metodo utilizza l'OpenAI Moderation API come oracolo a scatola nera. Assegna un punteggio agli output in 13 categorie di danno (ad esempio, violenza, autolesionismo, odio).
• Funzione di Perdita: $\Phi(X) = \max_{c} s_c(f(X))$, dove $s_c$ è il punteggio per la categoria $c$. L'obiettivo è minimizzare questo punteggio massimo, garantendo che la risposta sia sicura in tutte le dimensioni.

B. Stima del Gradiente di Ordine Zero

Poiché l'LLM ($f$) e l'API di Moderazione ($h$) sono scatole nere, la retropropagazione è impossibile. TSA utilizza la stima del gradiente di ordine zero:

1. Perturbazione: Ad ogni iterazione $k$, la matrice di embedding corrente $X_k$ viene perturbata aggiungendo rumore gaussiano $U$ scalato da $\mu$: $X_k + \mu U$.
2. Differenza Finita: Il gradiente viene stimato confrontando la perdita degli input perturbati con quella dell'input base:
   $$g_k \approx \frac{1}{N} \sum_{i=1}^{N} \frac{\Phi(X_k + \mu U_i) - \Phi(X_k)}{\mu} U_i$$
   dove $N$ è il numero di campioni Monte Carlo.
3. Aggiornamento: Gli embedding vengono aggiornati tramite discesa del gradiente: $X_{k+1} \leftarrow X_k - \eta g_k$.

C. Regularizzazione

Per garantire che gli embedding ottimizzati rimangano validi e non si spostino in regioni non addestrate dello spazio latente:

• Normalizzazione del Gradiente: Il passo di aggiornamento utilizza la direzione normalizzata del gradiente ($g_k / \|g_k\|_2$) per controllare la magnitudine del passo.
• Vincolo di Similarità Cosine: L'embedding aggiornato viene proiettato nuovamente su una sfera di coseno attorno all'embedding originale $X_0$ se la similarità scende sotto una soglia $\kappa$. Questo garantisce che la perturbazione sia sub-lessicale: il vettore ottimizzato decodifica nuovamente nel token originale, ma la rappresentazione continua all'interno della cella di Voronoi di quel token viene alterata.

3. Contributi Chiave

1. Estensione ai Modelli Allineati: Dimostra che il controllo a livello di embedding, precedentemente mostrato per la detossificazione dei modelli base, funziona efficacemente sui modelli istruiti allineati alla sicurezza, che hanno una distribuzione di output fondamentalmente diversa (bimodale).
2. Controllo Sub-lessicale: Dimostra che ottimizzare il vettore di embedding continuo all'interno della cella del vicino più prossimo di un token (senza cambiare il token stesso) è sufficiente per reindirizzare il comportamento del modello da dannoso a sicuro. Questo sfida la visione degli embedding come semplici ricerche discrete.
3. Ottimizzazione a Scatola Nera: Il metodo non richiede nessun accesso agli interni del modello (gradienti, stati nascosti o logit). Richiede solo l'accesso agli embedding di input e a un'API di moderazione a scatola nera.
4. Applicazione Incondizionata: Il metodo si applica a tutti i prompt indipendentemente dall'intento. Non forza un rifiuto; guida il modello a generare una risposta sostanziale e sicura se possibile, o un rifiuto sicuro se necessario.

4. Risultati Sperimentali

Il metodo è stato valutato su cinque modelli istruiti (da 1B a 20B parametri: Gemma 3, Phi-3.5, Llama 3.1, Qwen3, GPT-OSS) su due benchmark:

• WildJailbreak: Prompt avversari (dannosi e benigni).
• HarmBench: Query dannose dirette.

Risultati Chiave:

• Eliminazione della Dannosità: TSA ha ridotto il numero di risposte segnalate (dannose) a zero per la maggior parte dei modelli sui benchmark standard. Ad esempio, su Llama 3.1-8B, ha ridotto i prompt dannosi segnalati da 218 a 0.
• Convergenza: L'ottimizzazione converge tipicamente entro 1 o 2 passi di gradiente.
• Preservazione del Benigno: Su prompt benigni avversari, TSA non ha aumentato i punteggi di moderazione, confermando che non induce un pregiudizio di "rifiuto blanket" o non degrada l'utilità.
• Confronto con le Baseline: TSA ha superato significativamente le baseline come SmoothLLM (perturbazione casuale dei token), AdaSteer (steering delle attivazioni) e RESTA (rumore sugli embedding). Le baseline spesso fallivano su jailbreak semantici o aumentavano i tassi di segnalazione su query dannose dirette.
• Efficienza: Il principale costo di latenza è la chiamata API all'oracolo di moderazione, non l'inferenza del modello.

5. Significato e Implicazioni

• Nuova Superficie di Controllo: Il paper rivela che il livello di embedding offre una superficie di controllo più ricca di quanto si pensasse. Piccole perturbazioni sub-lessicali possono navigare i complessi confini di sicurezza dei modelli allineati.
• Fattibilità di Deployment: Come strato di sicurezza "drop-in", TSA non richiede riaddestramento o fine-tuning del modello. Può essere implementato su qualsiasi modello in cui gli embedding di input sono accessibili (ad esempio, modelli open-weight) combinato con un'API di moderazione pubblica.
• Limitazioni:
  • Latenza: L'ottimizzazione iterativa aggiunge latenza per prompt (secondi), sebbene questa sia dominata dai round-trip API.
  • Dipendenza dal Modello: L'efficacia scala con la qualità dell'addestramento per la sicurezza esistente del modello. Modelli con allineamento sottile (ad esempio, Gemma 3-1B) hanno mostrato un successo ridotto, poiché l'ottimizzatore non riusciva a trovare una direzione sicura vicina.
  • Accesso API: Richiede l'accesso agli embedding di input (non disponibile per modelli chiusi puramente protetti da API) e un'API di moderazione.

In conclusione, Test-Time Safety Alignment fornisce un meccanismo robusto e agnostico rispetto al modello per neutralizzare gli attacchi avversari ottimizzando la rappresentazione continua di input, colmando efficacemente il divario tra sicurezza teorica e deployment pratico senza riaddestramento.