Accelerating Suffix Jailbreak attacks with Prefix-Shared KV-cache

Titolo

Accelerazione degli attacchi di jailbreak a suffisso con Cache KV condivisa per prefisso (PSKV)

1. Il Problema

Gli attacchi di "jailbreak" a suffisso sono diventati uno strumento fondamentale per il red-teaming (test di sicurezza) dei Modelli Linguistici su Grande Scala (LLM). Questi attacchi mirano a generare un suffisso avversario che, concatenato a un prompt dannoso, aggira le difese di sicurezza del modello.

Tuttavia, l'ottimizzazione di questi suffissi soffre di costi computazionali proibitivi:

• Ridondanza Computazionale: Gli algoritmi di ricerca (basati su gradienti o euristici) devono valutare migliaia di candidati di suffisso. Poiché tutti questi candidati condividono lo stesso prefisso (l'istruzione dannosa originale), i sistemi di inferenza standard ricalcolano ripetutamente le rappresentazioni interne (KV cache) per lo stesso prefisso.
• Colli di Bottiglia di Memoria: Le implementazioni standard della KV cache duplicano fisicamente lo stato del prefisso per ogni candidato nel batch. Dato che i prompt dannosi sono spesso molto più lunghi dei suffissi, questa duplicazione porta a un consumo di memoria eccessivo ($O(N \cdot L_{prefix})$), limitando drasticamente la dimensione del batch e impedendo l'uso di strategie di ricerca più aggressive.
• Limiti degli Sistemi Esistenti: Sistemi di inferenza avanzati come vLLM o SGLang ottimizzano la memoria per scenari multi-utente asincroni, ma non supportano la retropropagazione (gradienti) necessaria per gli attacchi white-box e non sono progettati per il ciclo stretto e sincrono degli attacchi avversari.

2. Metodologia: PSKV (Prefix-Shared KV Cache)

Gli autori propongono PSKV, un framework di ottimizzazione "plug-and-play" che sfrutta la ridondanza strutturale specifica degli attacchi di jailbreak.

Principi Chiave:

1. Condivisione della KV Cache: Invece di calcolare e memorizzare la KV cache per il prefisso ($x^{(h)}$) $N$ volte (una per ogni candidato), PSKV esegue un'unica passata in avanti per calcolare e memorizzare la KV cache del prefisso. Durante l'inferenza dei candidati, questa cache viene riutilizzata e condivisa tra tutti i suffissi candidati.
2. Espansione Lazy e Livello per Livello:
   • PSKV non pre-alloca memoria espansa per tutti i livelli.
   • Utilizza una strategia di espansione "lazy": quando il modello elabora un livello di attenzione specifico, PSKV recupera la KV cache condivisa del prefisso, la duplica temporaneamente solo per quel livello necessario al calcolo, e la rilascia immediatamente dopo.
   • Questo riduce l'overhead di memoria di picco, poiché solo i vettori di un singolo livello vengono duplicati in un dato momento, invece di mantenere l'intera pila di rete duplicata.
3. Allineamento Centrato sul Suffisso (Suffix-Centric Alignment):
   • Per gestire batch di più istruzioni diverse (multi-istruzione), il metodo normalizza i dati allineando i suffissi a destra e i prefissi a sinistra.
   • Questo crea una matrice densa di token per i suffissi, permettendo l'uso di operazioni vettoriali altamente ottimizzate (es. in PyTorch) e consentendo alla KV cache del prefisso di essere trasmessa efficientemente attraverso la dimensione massiccia dei candidati.

Complessità:

• Computazionale: PSKV mantiene la stessa efficienza computazionale della cache KV standard (riducendo i calcoli ridondanti del prefisso da $K \cdot q$ a 1).
• Memoria: Riduce la complessità della memoria del prefisso da $O(N_{cand} \cdot N_p)$ a $O(B \cdot N_p)$, dove $B$ è il numero di prompt unici e $N_{cand}$ è il numero totale di candidati. Questo rimuove la dipendenza della memoria dalla larghezza della ricerca.

3. Contributi Chiave

• Identificazione dell'Inefficienza: Hanno evidenziato come le tecniche di caching generiche siano inefficienti quando applicate ai pattern specifici di ottimizzazione dei suffissi di jailbreak.
• Framework PSKV: Introduzione di un metodo leggero che combina la trasmissione virtuale della cache e l'allineamento tensoriale centrato sul suffisso per eliminare calcoli ridondanti.
• Valutazione Completa: Sperimentazioni su 6 attacchi di jailbreak diversi (inclusi GCG, AutoDAN, BEAST, AdvPrompter) e 5 modelli LLM diversi (Vicuna, Llama-2/3, Mistral, Qwen).

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su 2 GPU A100 da 80GB.

• Riduzione del Tempo: PSKV riduce il tempo di inferenza di circa il 40% rispetto alla baseline senza cache e fino al 1.95x di velocità rispetto alla cache KV standard in alcuni scenari (es. AutoDAN).
• Riduzione della Memoria: PSKV riduce l'uso di memoria di picco del 50% rispetto all'implementazione standard della KV cache.
• Scalabilità: A differenza dei metodi standard che vanno in Out Of Memory (OOM) all'aumentare della larghezza di ricerca (batch size), PSKV mantiene un consumo di memoria quasi costante, permettendo ricerche parallele su larga scala.
• Efficacia dell'Attacco (ASR): Il metodo è "performance-neutral". Il tasso di successo dell'attacco (ASR) rimane invariato rispetto alle baseline, confermando che l'ottimizzazione non altera la logica o la precisione dell'attacco.
• Confronto con Sistemi SOTA: PSKV supera o è comparabile a sistemi come vLLM e SGLang per gli attacchi senza gradienti (es. BEAST), ma offre il vantaggio cruciale di supportare la retropropagazione per gli attacchi white-box, cosa che i sistemi di inferenza commerciali non fanno.

5. Significato e Impatto

• Accessibilità alla Sicurezza: PSKV rende fattibile la valutazione della sicurezza su modelli più grandi e spazi di ricerca più ampi che erano precedentemente proibitivi dal punto di vista computazionale.
• Ottimizzazione Specifica per Carico di Lavoro: Dimostra che l'ottimizzazione dell'inferenza per scenari di sicurezza (adversarial) richiede approcci diversi rispetto all'inferenza di produzione standard, sfruttando le proprietà strutturali uniche degli attacchi (prefisso fisso, suffisso variabile).
• Strumento per la Ricerca: Fornisce un acceleratore universale che può essere integrato in qualsiasi pipeline di attacco esistente senza modificare la logica di alto livello, facilitando la ricerca su vulnerabilità LLM.

In sintesi, il paper introduce una soluzione elegante che trasforma un collo di bottiglia di memoria e calcolo in un processo efficiente, permettendo ai ricercatori di condurre valutazioni di sicurezza più approfondite e scalabili.