MOSAIC: Composable Safety Alignment with Modular Control Tokens

Each language version is independently generated for its own context, not a direct translation.

Immagina di avere un assistente virtuale molto intelligente, come un maggiordomo digitale. Attualmente, questo maggiordomo ha un "codice morale" fissato direttamente nel suo cervello (i parametri del modello). Se il codice dice "non parlare di gioco d'azzardo", il maggiordomo rifiuterà qualsiasi domanda su quel tema, anche se la stai facendo in un contesto sicuro o se sei un adulto che vuole solo informazioni storiche.

Il problema è che questo codice è rigido. Non puoi dirgli: "Oggi, per i bambini, non parlare di alcol, ma per gli adulti sì". Per cambiare queste regole, dovresti "riprogrammare" l'intero cervello del maggiordomo, un processo costoso, lento e rischioso che potrebbe fargli dimenticare altre cose utili.

Altri metodi provano a usare semplici istruzioni scritte (prompt), tipo: "Ricordati di non parlare di alcol". Ma il maggiordomo potrebbe non ascoltarle sempre, o potrebbe confondersi se gli dai troppe regole tutte insieme.

La Soluzione: MOSAIC (Il Set di Chiavi Magiche)

Gli autori di questo paper propongono MOSAIC, un sistema che risolve il problema in modo elegante. Immagina MOSAIC non come una riprogrammazione del cervello, ma come un set di chiavi magnetiche (chiamate "token di controllo") che puoi attaccare e staccare dal maggiordomo quando serve.

Ecco come funziona, passo dopo passo:

1. Le Chiavi Magiche (Token di Controllo)

Invece di cambiare il cervello del modello, MOSAIC crea piccole "chiavi" digitali. Ogni chiave rappresenta una regola specifica:

Una chiave per "Divieto di gioco d'azzardo".
Una chiave per "Divieto di alcol per minori".
Una chiave per "Contenuti horror".

Queste chiavi sono modulari. Puoi attaccare solo la chiave "Alcol" se stai parlando con un bambino, oppure attaccare sia "Alcol" che "Gioco d'azzardo" se stai configurando il sistema per un'area di sicurezza massima. Non devi toccare il cervello del modello, basta agganciare la chiave giusta.

2. La Cucina dei Sapori (Campionamento dei Compiti)

C'è un problema: se hai 10 regole, le combinazioni possibili sono migliaia (10 regole da sole, 2 insieme, 3 insieme, ecc.). Addestrare il modello su tutte queste combinazioni sarebbe come cucinare ogni possibile piatto della cucina mondiale: impossibile e costosissimo.

MOSAIC usa una strategia intelligente chiamata campionamento basato sull'ordine. Invece di provare a cucinare ogni piatto possibile, addestra il maggiordomo a gestire:

Prima i piatti con 1 ingrediente (1 regola).
Poi i piatti con 2 ingredienti (2 regole).
Poi i piatti con 3 ingredienti.

In questo modo, il modello impara a mescolare le regole senza dover vedere ogni singola combinazione possibile, risparmiando tempo e risorse.

3. Il Filtro Anti-Eccesso (Distillazione Controfattuale)

Il rischio maggiore di questi sistemi è l'"eccesso di rifiuto". Immagina un maggiordomo così spaventato dalle regole che, se gli dai la chiave "Niente alcol", rifiuta anche di dirti come si fa un cocktail analcolico o come si usa un termometro (che contiene alcol ma non è un problema).

Per evitare questo, MOSAIC usa un trucco psicologico chiamato distillazione controfattuale:

Chiede al modello: "Cosa avresti risposto senza la chiave di sicurezza?" (La risposta originale, utile e gentile).
Poi chiede: "Cosa stai rispondendo con la chiave?"
Se la risposta con la chiave è troppo diversa (es. un rifiuto ingiustificato), il sistema dice: "Ehi, aspetta! Non serve essere così severi qui. Torna alla tua risposta originale, a meno che la domanda non sia davvero pericolosa".

Questo insegna al modello a essere preciso: rifiuta solo quando è strettamente necessario, e lascia fluire le risposte utili quando la regola non si applica.

Perché è importante?

Flessibilità: Puoi adattare la sicurezza in base all'utente (bambini vs adulti), alla regione (leggi diverse) o al contesto, semplicemente cambiando le "chiavi" attaccate.
Efficienza: Non serve riaddestrare il modello da zero ogni volta che cambia una legge o una regola.
Utilità: Il modello rimane utile e non diventa un "noioso" che rifiuta tutto.

In Sintesi

MOSAIC trasforma la sicurezza dell'IA da un muro di cemento (rigido e immutabile) in un sistema di luci LED (modulare e controllabile). Puoi accendere o spegnere le luci di sicurezza specifiche a seconda della stanza in cui ti trovi, senza dover ricostruire la casa ogni volta.

È un passo avanti verso un'intelligenza artificiale che non solo è sicura, ma è anche adattabile e intelligente nel capire quando e come applicare le regole.

Each language version is independently generated for its own context, not a direct translation.

1. Il Problema

L'allineamento alla sicurezza dei Large Language Models (LLM) è attualmente implementato principalmente attraverso due approcci, entrambi affetti da limitazioni significative quando si tratta di gestire requisiti di sicurezza dinamici e contestuali:

Allineamento a livello di parametri (es. SFT, RLHF): Le regole di sicurezza sono "incollate" nei pesi del modello. Questo rende difficile adattare le policy a diversi utenti, regioni o applicazioni senza costosi ri-addestramenti. Inoltre, modifica le capacità generali del modello e rischia di causare interferenze catastrofiche con comportamenti appresi precedentemente.
Metodi basati su prompt: Utilizzano istruzioni in linguaggio naturale per definire i vincoli di sicurezza. Tuttavia, questi metodi sono probabilistici, inefficienti (aumentano l'overhead dei token) e spesso falliscono nel mantenere coerenza quando vengono combinati vincoli multipli o complessi.

Il problema centrale è la mancanza di una rappresentazione esplicita, riutilizzabile e componibile dei vincoli di sicurezza. Le soluzioni attuali non riescono a fornire un controllo condizionale fine-granularità (attivare regole specifiche solo per certi contesti) senza sacrificare l'utilità del modello o richiedere un ri-addestramento massiccio.

2. Metodologia: MOSAIC

Il paper propone MOSAIC (Composable Safety Alignment with Modular Control Tokens), un framework che riconcettualizza l'allineamento alla sicurezza come un problema di apprendimento di rappresentazioni composizionali.

Concetto Fondamentale

Invece di modificare i parametri del modello base (che rimane congelato/frozen), ogni vincolo di sicurezza è rappresentato da un piccolo insieme di token di controllo apprendibili nello spazio di embedding del modello.

Modularità: Ogni categoria di sicurezza (es. "gioco d'azzardo", "alcol") ha i suoi token associati.
Componibilità: Durante l'inferenza, i token corrispondenti alle categorie attive vengono preposti all'input. Combinando diversi insiemi di token, si possono attivare policy di sicurezza multiple e contestuali dinamicamente.

Strategie di Addestramento

Per ottimizzare questi token in modo efficiente e prevenire il rifiuto eccessivo (over-refusal) su query innocue, MOSAIC introduce due tecniche chiave:

Campionamento di Task Basato sull'Ordine (Order-based Task Sampling):
- Invece di enumerare tutte le possibili combinazioni di categorie (che crescerebbe esponenzialmente, $2^K$ ), le combinazioni sono organizzate per "ordine" (numero di categorie attive).
- Viene assegnato un budget di addestramento fisso per ogni ordine. Questo permette al modello di imparare a gestire l'interazione cooperativa tra token di diverse categorie senza un costo computazionale esponenziale.
Distillazione della Conoscenza Controfattuale (Counterfactual Knowledge Distillation - KD):
- Per mitigare il problema del "rifiuto eccessivo" (dove il modello rifiuta anche query sicure perché vede i token di controllo), viene introdotta una funzione di perdita specifica.
- Per le query innocue (negative samples), il modello viene addestrato a mantenere la sua distribuzione originale (quella del modello base senza token di controllo) anche quando i token sono presenti.
- Si minimizza la divergenza KL tra la distribuzione del modello controllato e quella del modello base non controllato. Questo insegna ai token di intervenire solo quando necessario, preservando l'utilità del modello su richieste non correlate.

3. Contributi Chiave

Riconcettualizzazione Teorica: Trasforma l'allineamento alla sicurezza da una modifica monolitica dei parametri a un problema di attivazione di vincoli modulari e composizionali.
Framework MOSAIC: Un sistema che utilizza token di controllo apprendibili su un backbone congelato. Permette l'espansione incrementale di nuove categorie di sicurezza senza ri-addestrare il modello base e riduce drasticamente il rifiuto eccessivo grazie alla distillazione controfattuale.
Nuovo Benchmark Realistico: Costruzione di un dataset di valutazione specifico per l'attivazione condizionale della sicurezza. A differenza dei benchmark esistenti, questo dataset contiene richieste che i modelli allineati standard non rifiutano di default, permettendo di testare la capacità di imporre vincoli aggiuntivi selettivamente senza degradare le prestazioni generali.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su modelli Llama-3.1-8B e Llama-3.2-3B.

Performance di Difesa (DSR - Defense Success Rate): MOSAIC raggiunge un tasso di successo nella difesa quasi perfetto (>99%) su tutte le combinazioni di task (da 1 a 4 categorie attive), superando o eguagliando i metodi basati su SFT (Supervised Fine-Tuning).
Riduzione del Rifiuto Eccessivo (OR - Over-Refusal Rate):
- I metodi SFT tradizionali mostrano un OR intorno al 6-7% (rifiutano query innocue).
- MOSAIC riduce l'OR drasticamente: con 5 token per categoria, l'OR scende fino all'1.8% su Llama-3.1-8B per task complessi (3-ordine).
- Questo dimostra che la distillazione controfattuale e il campionamento composizionale permettono al modello di distinguere con precisione quando rifiutare.
Preservazione dell'Utilità: Le valutazioni su MMLU mostrano che MOSAIC preserva le capacità generali di linguaggio del modello base con un degrado trascurabile, a differenza di metodi come ORPO o prompt-based che mostrano cali più significativi.
Scalabilità e Incrementalità: Il framework supporta l'aggiunta di nuove categorie di sicurezza senza ri-addestrare i token esistenti, con degradazioni delle prestazioni minime o nulle.

5. Significato e Impatto

MOSAIC rappresenta un passo avanti significativo verso un allineamento alla sicurezza pluralistico e adattivo.

Flessibilità Operativa: Permette di adattare le policy di sicurezza in tempo reale in base al contesto (es. età dell'utente, giurisdizione legale, ruolo professionale) senza dover ri-addestrare interi modelli.
Efficienza: L'approccio basato su token è estremamente leggero in termini di parametri aggiuntivi rispetto al ri-addestramento completo o al LoRA esteso.
Precisione: Risolve il compromesso storico tra sicurezza e utilità, dimostrando che è possibile avere un modello altamente sicuro in contesti specifici senza diventare eccessivamente conservativo (over-refusal) in contesti generali.

In sintesi, MOSAIC offre una soluzione scalabile e modulare per gestire la complessità crescente delle richieste di sicurezza nel mondo reale, spostando il paradigma dalla sicurezza "statica e rigida" a una sicurezza "dinamica e componibile".