Stochasticity in Tokenisation Improves Robustness

1. Il Problema: La Fragilità dei Modelli LLM rispetto alla Tokenizzazione

L'adozione diffusa dei Large Language Models (LLM) ha sollevato preoccupazioni riguardo alla loro robustezza e sicurezza. Un punto debole critico risiede nel processo di tokenizzazione, in particolare nell'uso di schemi deterministici (come il Byte Pair Encoding o BPE standard).

• Vulnerabilità: I modelli addestrati esclusivamente sulla tokenizzazione canonica (la rappresentazione "standard" e deterministica di una stringa) si rivelano fragili di fronte a perturbazioni. Esistono molteplici sequenze di token non canoniche che ricostruiscono la stessa stringa di testo, ma i modelli spesso falliscono nel riconoscerle o vengono ingannati da esse.
• Attacchi Avversariali: È stato dimostrato che è possibile trovare tokenizzazioni non canoniche che, pur mantenendo invariato il contenuto semantico del testo, causano un errore di previsione del modello, rompendo l'allineamento o riducendo drasticamente l'accuratezza.
• Limiti delle Soluzioni Esistenti: Studi recenti hanno introdotto metodi stocastici (come STOCHASTOK), ma questi presentano limiti nella distribuzione di campionamento (bias) e nel supporto (copertura non completa di tutte le possibili tokenizzazioni), rendendoli ancora vulnerabili ad attacchi avversariali mirati.

2. Metodologia e Approccio

Gli autori hanno analizzato sistematicamente come l'uso di tokenizzazioni stocastiche durante le diverse fasi di apprendimento (pre-training, fine-tuning supervisionato e in-context learning) influenzi la robustezza del modello.

A. Definizione di Tokenizzazione Stocastica

Invece di usare una singola tokenizzazione deterministica, il modello viene esposto a diverse segmentazioni della stessa stringa durante l'addestramento.

• STOCHASTOK (Base): Un metodo esistente che divide i token in sottotoken secondo regole di unione (merge) casuali. Tuttavia, gli autori hanno scoperto che questo metodo genera una distribuzione non uniforme (bias verso certe segmentazioni) e ha un supporto incompleto (non può generare tutte le possibili tokenizzazioni non canoniche).
• Nuovi Schemi Proposti: Per superare i limiti di STOCHASTOK, il paper introduce due nuovi metodi di campionamento:
  1. STOCHASTOK-UNI: Mantiene la distribuzione del numero di split per token di STOCHASTOK ma campiona uniformemente le tokenizzazioni all'interno di quel vincolo, eliminando il bias.
  2. UNIFORM-K: Campiona uniformemente le tokenizzazioni non canoniche a una specifica distanza di edit (numero di split) su tutta la sequenza, garantendo un supporto completo e una distribuzione uniforme reale.

B. Regimi di Valutazione

Lo studio è stato condotto su tre livelli:

1. Pre-training: Addestramento da zero (usando un modello Tiny-LLM) su corpus di testo.
2. Supervised Fine-Tuning (SFT): Adattamento di modelli pre-addestrati (Llama-1b, Llama-8b, GPT2-xl, Qwen) su dataset specifici (es. LANGUAGE GAME, CUTE).
3. In-Context Learning (ICL): Valutazione della robustezza senza aggiornare i pesi del modello, utilizzando esempi nel contesto con tokenizzazioni stocastiche.

C. Analisi Teorica

Gli autori forniscono un'analisi teorica della robustezza avversariale basata sulla costante di Lipschitz. Dimostrano che la tokenizzazione stocastica "smussa" lo spazio degli embedding, riducendo la sensibilità locale del modello a piccole perturbazioni e abbassando il rischio avversariale teorico.

3. Risultati Chiave

Robustezza a Perturbazioni Casuali e Avversariali

• Miglioramento Significativo: L'addestramento con tokenizzazioni stocastiche migliora drasticamente la robustezza sia a perturbazioni casuali che ad attacchi avversariali.
• Dati Sperimentali:
  • Un modello Llama-1b addestrato in modo canonico subisce un crollo di accuratezza del 29.8% quando valutato su tokenizzazioni non canoniche uniformemente campionate.
  • Lo stesso modello, se fine-tuned con tokenizzazione stocastica (specialmente con UNIFORM-K o STOCHASTOK-UNI), mantiene un'accuratezza quasi costante, riducendo il calo di performance a valori trascurabili (es. da -29.8% a +1.4% o -0.1% a seconda del metodo).
  • Sotto attacco avversariale, il modello canonico crolla dal 94% al 6.1% di accuratezza, mentre il modello addestrato con UNIFORM-K mantiene un'accuratezza del 69.6%.

Confronto tra Metodi

• STOCHASTOK vs. Metodi Uniformi: Sebbene STOCHASTOK migliori la robustezza rispetto al caso canonico, i metodi che garantiscono una distribuzione uniforme (STOCHASTOK-UNI e UNIFORM-K) offrono prestazioni superiori, specialmente contro attacchi avversariali che sfruttano le lacune nel supporto di STOCHASTOK.
• Costo Inferenziale: Un risultato cruciale è che l'addestramento stocastico non aumenta il costo computazionale durante l'inferenza. Il modello può essere utilizzato con tokenizzazione canonica standard al momento del test, mantenendo comunque la robustezza appresa.

In-Context Learning (ICL)

L'uso di esempi nel contesto con tokenizzazioni stocastiche offre un miglioramento della robustezza, sebbene inferiore rispetto al fine-tuning completo. Tuttavia, l'ICL con schemi uniformi (UNIFORM) mostra i migliori risultati tra le tecniche senza aggiornamento dei pesi.

4. Contributi Principali

1. Analisi Sistematica: Prima indagine completa su come la tokenizzazione stocastica influenzi la robustezza attraverso pre-training, fine-tuning e ICL su diverse architetture e dataset.
2. Identificazione dei Bias: Dimostrazione che il metodo STOCHASTOK esistente introduce un bias di campionamento e ha un supporto incompleto, rendendolo insufficiente per la massima robustezza avversariale.
3. Nuovi Algoritmi: Proposta di STOCHASTOK-UNI e UNIFORM-K, schemi di campionamento che garantiscono distribuzioni uniformi e supporto completo, superando i limiti dei metodi precedenti.
4. Analisi Teorica: Fornitura di un limite superiore teorico per il rischio avversariale, che collega direttamente la tokenizzazione stocastica alla riduzione della sensibilità del modello (costanti di Lipschitz).
5. Linee Guida Pratiche: Evidenza che il fine-tuning con tokenizzazione stocastica è una strategia efficace, a basso costo e senza impatto sull'inferenza, per rendere gli LLM più resilienti agli attacchi.

5. Significato e Impatto

Questo lavoro ha un impatto significativo sulla sicurezza e l'affidabilità degli LLM:

• Mitigazione del Rischio: Dimostra che la fragilità dei modelli agli attacchi basati sulla tokenizzazione non è inevitabile, ma può essere risolta modificando la strategia di addestramento.
• Best Practice: Suggerisce che l'addestramento (in particolare il fine-tuning) dovrebbe essere condotto utilizzando tokenizzazioni stocastiche uniformi invece di quelle canoniche, per creare modelli intrinsecamente più robusti.
• Scalabilità: Poiché non richiede costi aggiuntivi in fase di inferenza, questa tecnica è immediatamente applicabile in scenari di produzione reali per proteggere i sistemi LLM da manipolazioni sottili del testo di input.

In sintesi, il paper stabilisce che l'introduzione deliberata di stocasticità controllata e uniforme durante l'addestramento è una difesa potente ed efficiente contro le vulnerabilità legate alla tokenizzazione dei modelli linguistici.