MOSAIC: Composable Safety Alignment with Modular Control Tokens

Das Paper stellt MOSAIC vor, ein modulares Framework, das durch lernbare Steuertokens eine zusammensetzbare Sicherheitsausrichtung für große Sprachmodelle ermöglicht, um kontextabhängige Regeln effizient zu erzwingen und dabei die Überverweigerung zu minimieren.

Das Wesentliche

Stell dir vor, ein großes Sprachmodell (wie ein sehr intelligenter Roboter) ist wie ein allwissender Koch, der für jeden auf der Welt kochen kann.

Das Problem ist: Nicht jeder Gast möchte das Gleiche essen, und nicht jeder Gast darf das Gleiche essen.

• Ein Erwachsener darf vielleicht ein scharfes, alkoholhaltiges Gericht probieren.
• Ein Kind darf das nicht.
• In einem Land ist ein bestimmtes Gericht erlaubt, in einem anderen verboten.

Bisher gab es zwei schlechte Lösungen für dieses Problem:

1. Der "Einheitskoch" (Parameter-Alignment): Man hat den Koch gezwungen, sein gesamtes Gehirn umzukrempeln, damit er niemals scharfes Essen oder Alkohol anbietet. Das Problem? Wenn er dann ein harmloses Rezept für einen Mojito (mit Alkohol) erklären soll, lehnt er es ab, auch wenn der Gast ein Erwachsener ist. Er ist zu vorsichtig geworden. Und wenn man die Regeln ändern will, muss man den ganzen Koch neu ausbilden – das kostet Zeit und Geld.
2. Der "Zettel am Hut" (Prompting): Man schreibt dem Koch einen Zettel auf den Hut: "Bitte keine Alkohol-Rezepte für Kinder!" Das funktioniert manchmal, aber der Koch ist verwirrt. Wenn der Zettel zu lang wird oder viele Regeln enthält, vergisst er die wichtigsten Dinge oder lehnt auch harmlose Dinge ab, weil er den Zettel falsch interpretiert.

Die neue Lösung: MOSAIC (Der "Magische Schalter")

Die Forscher haben MOSAIC entwickelt. Stell dir das wie einen modularen Schalterkasten vor, den man an den Koch anschließen kann, ohne sein Gehirn zu verändern.

• Die Schalter (Control Tokens): Jede Sicherheitsregel (z. B. "Kein Alkohol für Kinder" oder "Keine Glücksspiel-Tipps") ist ein winziger, unsichtbarer Schalter (ein "Token").
• Das Einfügen: Wenn ein Kind fragt: "Wie baue ich ein Casino?", schaltet man den "Glücksspiel-Schalter" ein. Der Koch weiß sofort: "Aha, hier muss ich ablehnen!"
• Die Kombination: Wenn ein Erwachsener fragt: "Wie koche ich ein scharfes Curry?", schaltet man keinen Schalter ein. Der Koch kocht einfach.
• Der Clou: Man kann diese Schalter kombinieren. Wenn ein Kind fragt: "Wie gewinne ich beim Casino und trinkt dabei Alkohol?", schaltet man beide Schalter gleichzeitig ein. Der Koch lehnt ab, aber nur wegen dieser spezifischen Kombination.

Wie haben sie das gelernt? (Die Tricks)

Damit diese Schalter funktionieren, mussten sie den Koch trainieren, ohne sein Gehirn zu kaputt machen. Dafür gab es zwei geniale Tricks:

1. Der "Misch-Topf" (Combinatorial Sampling):
   Statt den Koch nur für einzelne Regeln zu trainieren, haben sie ihn mit allen möglichen Kombinationen geübt. Stell dir vor, sie haben ihm nicht nur "Alkohol" gezeigt, sondern auch "Alkohol + Glücksspiel" oder "Alkohol + Horrorfilme". So lernt der Koch, dass ein Schalter nur dann wirken soll, wenn er wirklich gebraucht wird, und nicht, wenn er einfach nur im Raum steht.

2. Die "Gegenfaktische Erinnerung" (Counterfactual Knowledge Distillation):
   Das ist der wichtigste Trick gegen das "Zu-viel-Ablehnen" (Over-refusal).

   • Szenario: Der Koch sieht eine harmlose Frage ("Wie macht man einen Mojito?") und hat den "Alkohol-Schalter" eingelegt.
   • Der Fehler: Normalerweise würde er denken: "Schalter ist an -> Ich lehne ab!"
   • Die Lösung: MOSAIC sagt dem Koch: "Warte! Schau mal, was du gemacht hättest, wenn der Schalter aus gewesen wäre. Du hättest das Rezept gegeben. Also gibst du es jetzt auch, auch wenn der Schalter an ist, weil die Frage harmlos ist."
     Der Koch lernt also: "Ich lehne nur ab, wenn der Schalter an ist UND die Frage gefährlich ist. Wenn die Frage harmlos ist, ignoriere ich den Schalter."

Warum ist das toll?

• Flexibilität: Man kann neue Regeln hinzufügen, ohne den ganzen Koch neu zu erziehen. Einfach einen neuen Schalter einstecken.
• Präzision: Der Koch lehnt nur ab, wenn es wirklich nötig ist. Er wird nicht paranoid und lehnt harmlose Dinge ab.
• Schnelligkeit: Da das Gehirn des Kochs (das Modell) nicht verändert wird, bleibt er schnell und klug bei allen anderen Aufgaben.

Zusammengefasst: MOSAIC ist wie ein intelligenter Sicherheitsgurt, den man je nach Situation an- oder ausklickt. Statt den ganzen Wagen (das Modell) umzubauen, um sicherer zu sein, hat man einfach die richtigen Gurte entwickelt, die genau dann zupacken, wenn es nötig ist, und den Fahrer (den Nutzer) in Ruhe lassen, wenn alles sicher ist.

Technische Zusammenfassung

1. Problemstellung

Die Sicherheitssynchronisation (Safety Alignment) von Large Language Models (LLMs) wird derzeit meist als statische Richtlinie implementiert, die in den Modellparametern verankert ist. Dies führt zu zwei Hauptproblemen:

• Mangelnde Anpassungsfähigkeit: Sicherheitsregeln variieren je nach Nutzergruppe (z. B. Erwachsene vs. Minderjährige), Region (unterschiedliche Gesetze) und Anwendungskontext. Eine starre, in den Gewichten verankerte Politik kann diese Nuancen nicht dynamisch abbilden.
• Limitationen bestehender Ansätze:
  • Parameterebene (z. B. SFT, RLHF): Sicherheitsverhalten ist untrennbar mit allgemeinen Fähigkeiten verflochten. Änderungen erfordern teures Neutrainieren und bergen das Risiko von „Catastrophic Interference" (Verlust zuvor gelernter Fähigkeiten).
  • Prompt-basierte Methoden: Diese nutzen natürliche Sprache, um Regeln zu definieren. Sie sind jedoch probabilistisch, inkonsistent und ineffizient, da lange Prompt-Overheads den Kontext verbrauchen und die Kontrolle ungenau bleibt.

Es fehlt eine explizite, wiederverwendbare und kompositionelle Repräsentation von Sicherheitsbeschränkungen, die eine feingranulare, bedingte Steuerung ohne Neutrainieren des Basismodells ermöglicht.

2. Methodik: MOSAIC

Das vorgestellte Framework MOSAIC (Composable Safety Alignment with Modular Control Tokens) behandelt Sicherheitssynchronisation als Problem des repräsentativen Lernens.

• Modulare Kontroll-Token: Anstatt Modellparameter zu ändern, wird jede Sicherheitskategorie (z. B. Glücksspiel, Alkohol) durch einen kleinen Satz lernbarer Kontroll-Token ($z_c$) im Embedding-Raum eines eingefrorenen Backbone-Modells repräsentiert.
• Inferenz: Sicherheitsbeschränkungen werden aktiviert, indem die entsprechenden Token dem Eingabe-Input vorangestellt werden. Mehrere Token-Sätze können kombiniert (komponiert) werden, um komplexe, kontextabhängige Regeln zu erzeugen (z. B. „Verbot von Alkohol UND Glücksspiel").
• Training-Strategie:
  • Order-basiertes Task-Sampling: Um die exponentielle Anzahl möglicher Token-Kombinationen zu bewältigen, werden Aufgaben nach ihrer „Ordnung" (Anzahl aktiver Kategorien) gruppiert. Das Training budgetiert pro Ordnungsebene, nicht pro einzelner Kombination, was eine effiziente Exposition gegenüber diversen Kompositionen ohne exponentielle Datenkosten ermöglicht.
  • Counterfactual Knowledge Distillation (KD): Um das Problem der „Over-Refusal" (unbegründete Ablehnung harmloser Anfragen) zu lösen, wird ein KD-Ziel eingeführt. Für harmlose Eingaben wird die Ausgabe des Modells mit Kontroll-Token mit der Ausgabe ohne Token (dem „Counterfactual"-Referenzmodell) verglichen. Die Divergenz (KL-Divergenz) wird minimiert, um sicherzustellen, dass die Kontroll-Token nur dann eingreifen, wenn eine echte Sicherheitsverletzung vorliegt, und das Verhalten bei harmlosen Anfragen erhalten bleibt.

3. Wichtige Beiträge

1. Konzeptuelle Neudefinition: Sicherheit wird als kompositionelles Repräsentationslernproblem neu gefasst, bei dem Sicherheit durch modulare Aktivierung von Constraints statt durch monolithische Parameteränderung erreicht wird.
2. Das MOSAIC-Framework: Ein System, das Sicherheitsbeschränkungen als lernbare Token auf einem eingefrorenen Modell darstellt. Es ermöglicht inkrementelle Erweiterungen neuer Kategorien und eine flexible Kombination von Regeln ohne Neutrainieren.
3. Neuer Benchmark: Die Autoren konstruierten einen realistischen Evaluierungs-Datensatz mit 1.500 Anfragen in fünf Sicherheitskategorien. Im Gegensatz zu bestehenden Benchmarks werden diese Anfragen von standardmäßig ausgerichteten LLMs nicht abgelehnt, was eine präzise Messung der selektiven Sicherheitsaktivierung ermöglicht.

4. Ergebnisse

Die Experimente wurden auf Llama-3.1-8B und Llama-3.2-3B durchgeführt und mit Methoden wie In-Context-Learning, ORPO und SFT verglichen.

• Defense Success Rate (DSR): MOSAIC erreicht eine nahezu perfekte Abwehrrate (>99,8 %), die mit SFT vergleichbar oder sogar besser ist, selbst bei komplexen Mehr-Kategorien-Aufgaben (Higher-Order Tasks).
• Over-Refusal Rate (OR): MOSAIC reduziert die Rate unbegründeter Ablehnungen drastisch. Während SFT bei ca. 6 % liegt, erreicht MOSAIC-5 (5 Token pro Kategorie) Werte von bis zu 1,8 % (bei Llama-3.1-8B).
• Skalierbarkeit und Modularität:
  • Die Leistung bleibt stabil, wenn neue Kategorien hinzugefügt werden (inkrementelle Expansion), ohne dass das Neutrainieren alter Kategorien nötig ist.
  • Die allgemeinen Sprachfähigkeiten des Modells bleiben erhalten (gemessen an MMLU), da das Backbone eingefroren bleibt.
• Ablationsstudie: Die Kombination aus Multi-Task-Training (für Kompositionalität) und Counterfactual KD (für Präzision) ist entscheidend. Ohne Multi-Task-Training bricht die Leistung bei höheren Ordnungen ein; ohne KD steigt die Over-Refusal-Rate.

5. Bedeutung und Fazit

MOSAIC adressiert eine kritische Lücke in der aktuellen KI-Sicherheit: Die Fähigkeit, Sicherheitsrichtlinien dynamisch, kontextsensitiv und effizient anzupassen.

• Praktische Relevanz: Es ermöglicht es Anbietern, Sicherheitsregeln für verschiedene Jurisdiktionen oder Altersgruppen „on-the-fly" zu aktivieren, ohne teure Neutrainings oder den Verlust von Modellfähigkeiten.
• Effizienz: Durch die Nutzung weniger lernbarer Parameter (Token) statt ganzer Modell-Updates ist der Ansatz ressourcenschonend.
• Zukunft: Der Ansatz ebnet den Weg für eine „plurale Ausrichtung" (Pluralistic Alignment), bei der Sicherheitsstandards nicht starr, sondern flexibel an die Bedürfnisse des Nutzers und des Kontexts angepasst werden können.

Zusammenfassend bietet MOSAIC einen robusten, skalierbaren und präzisen Mechanismus für die bedingte Sicherheitssynchronisation, der die Nachteile statischer Parameter-Updates und unzuverlässiger Prompt-Engineering überwindet.