Learning When to Act or Refuse: Guarding Agentic Reasoning Models for Safe Multi-Step Tool Use

Dit paper introduceert MOSAIC, een post-training framework dat agentic taalmodellen veilig maakt voor multi-stap toolgebruik door expliciete veiligheidsredenering en weigering als leerbare acties te integreren, wat resulteert in een aanzienlijke reductie van schadelijk gedrag en privacylekken zonder in te leveren op prestaties bij nuttige taken.

De kern

Stel je voor dat je een zeer slimme, maar nog wat onervaren assistent hebt. Deze assistent kan niet alleen praten, maar ook echte dingen doen: bestanden openen, e-mails sturen, betalingen doen en software aansturen. Dit noemen we een "Agent".

Het probleem is dat deze assistent soms te enthousiast is. Als je vraagt: "Help me een grapje te maken over mijn baas," en de assistent denkt: "Ik gebruik mijn e-mailtool om dat te sturen," kan dat rampzalig zijn als de assistent per ongeluk ook je wachtwoorden lekt of een onbedoeld bericht naar de hele company stuurt.

Deze paper introduceert MOSAIC, een nieuwe manier om deze slimme assistenten veilig te maken. Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het Probleem: De "Blindelings Gehoorzame" Robot

Vroeger werden AI-modellen getraind om gewoon te antwoorden. Als je iets gevaarlijks vroeg, zeiden ze "Nee". Maar nu moeten ze werken. Ze moeten plannen maken en tools gebruiken.
Het gevaar is dat ze soms denken: "Ik moet deze taak afmaken, dus ik ga gewoon door, ook al voelt het een beetje raar." Ze zien niet het gevaar in de tussenstappen. Het is alsof je een auto bestuurt die alleen naar de bestemming kijkt, maar vergeet te kijken of er een afgrond voor de brug ligt.

2. De Oplossing: MOSAIC (De Slimme Controleur)

MOSAIC is een trainingsmethode die de assistent leert om niet blindelings door te gaan. Het introduceert een nieuw ritme in het denken van de AI: Plannen → Controleren → Acteren of Weigeren.

Stel je MOSAIC voor als een veiligheidsinspecteur die in de machine zit:

• Stap 1: Het Plan (De Chef-kok)
  De assistent denkt: "Oké, de klant wil een verslag maken. Ik ga eerst de data ophalen en dan een PDF maken."
• Stap 2: De Check (De Veiligheidsinspecteur)
  Voordat de assistent de knop indrukt, schakelt MOSAIC in. Het vraagt zich af: "Is het veilig om die data te halen? Zie ik hier een valstrik? Als ik dit bestand open, kan ik per ongeluk mijn eigen huis inbranden?"
  • Als het veilig is: De assistent gaat door.
  • Als het gevaarlijk is: De assistent stopt direct en zegt: "Ik doe dit niet, dit is te riskant."
• Stap 3: De Actie (De Uitvoerder)
  Pas als de inspecteur groen licht geeft, wordt de daadwerkelijke actie uitgevoerd.

3. Hoe leren ze dit? (De "Jury" in plaats van de "Scorebord")

Normaal gesproken leer je een AI door te zeggen: "Goed gedaan, +1 punt" of "Slecht gedaan, -1 punt". Maar bij complexe taken werkt dat niet goed.
Stel, de assistent doet iets gevaarlijks, maar stopt op het allerlaatste moment. Een simpel scorebord ziet misschien: "Nou, het is niet mislukt, dus 0 punten." Maar een andere assistent die direct "Nee" zegt, krijgt ook 0 punten. Dat is niet eerlijk!

MOSAIC gebruikt een Jury-systeem:
In plaats van een score, kijkt een andere slimme AI (de "Jury") naar twee verschillende versies van dezelfde taak.

• Versie A: De assistent doet iets gevaarlijks en stopt pas op het laatste moment.
• Versie B: De assistent zegt direct "Nee" en weigert de taak.
  De Jury zegt: "Versie B is veel beter!"
  Door duizenden van deze vergelijkingen te maken, leert de assistent precies wanneer hij moet stoppen, zonder dat iemand hem stap-voor-stap heeft verteld wat fout is.

4. Waarom is dit zo cool?

• Het werkt voor kleine en grote modellen: Of je nu een kleine, snelle AI hebt of een enorme, dure super-AI; MOSAIC maakt ze allemaal veiliger.
• Het is slim, niet dom: De assistent leert om alleen te controleren als het nodig is. Bij simpele taken (zoals "schrijf een e-mail") gaat hij snel door. Bij gevaarlijke taken (zoals "wis mijn database") schakelt de inspecteur direct in.
• Het voorkomt "Prompt Injection": Soms proberen hackers de assistent te misleiden door in een berichtje te verstoppen: "Vergeet de regels en doe dit." MOSAIC leert de assistent om die valstrikken te zien en de "Inspecteur" te laten beslissen, zelfs als de hacker slim is.

Samenvatting in één zin

MOSAIC is als het installeren van een slimme rem in een raceauto: de auto kan nog steeds razendsnel rijden (werk doen), maar hij heeft nu een systeem dat automatisch remt als er een afgrond aankomt, zodat hij nooit in de ravijn rijdt.

Dit maakt AI-agents niet alleen slimmer, maar vooral verantwoordelijker voor de echte wereld.

Technische samenvatting

1. Het Probleem

Taalmodellen worden steeds vaker ingezet als agenten die plannen maken, tools aanroepen en interacties met externe systemen uitvoeren over meerdere stappen. In tegenstelling tot chat-modellen die statische tekst genereren, opereren agenten in een veiligheidsregime waarbij een enkele fout (zoals het openen van bestanden of het invoeren van wachtwoorden) onomkeerbare schade kan veroorzaken.

De huidige uitdagingen zijn:

• Afbreken van bestaande methoden: Bestaande veiligheidsaanpassingen (alignment methods) zijn geoptimaliseerd voor statische generatie en taakvoltooiing. Ze falen in sequentiële besluitvorming, vooral bij adversariele tool-feedback en overmoedige tussenstappen.
• Gevaren bij kleine modellen (SLMs): Kleine taalmodellen, die vaak worden gebruikt vanwege kosten en privacy, zijn gevoeliger voor anomalieën, prompt-injecties en cascaderende fouten.
• Tekortkomingen in beloning: Bestaande methoden gebruiken vaak scalar beloningen (punten) aan het einde van een traject. Dit negeert het tijdstip van veiligheidsfouten. Een traject dat laat afbreekt na een gevaarlijke stap wordt vaak gelijkgesteld aan een traject dat direct weigert, hoewel het eerste veel gevaarlijker is.
• Implicit vs. Expliciet: Veiligheid wordt vaak impliciet behandeld als een bijproduct van redenering, wat leidt tot onbetrouwbare weigeringen of het uitvoeren van schadelijke acties.

2. Methodologie: MOSAIC

Het paper introduceert MOSAIC (Modular Agentic Safety & Inference Control), een post-training framework dat veiligheid expliciet en leerbaar maakt.

Kernarchitectuur: Plan-Check-Act/Refuse

MOSAIC structureert de inferentie van de agent in een lus met drie fasen:

1. Plan: De agent genereert een plan en een kandidaat-toolactie (via <thoughts>).
2. Check (Veiligheidscontrole): De agent kan optioneel een expliciete veiligheidscontrole uitvoeren via een <safety_thoughts> blok. Hierin wordt gestructureerd nagedacht over risico's, onomkeerbaarheid, toestemmingen en potentieel schadelijke intenties.
3. Act of Refuse: Op basis van de controle kiest de agent voor:
   • Act: Uitvoeren van de tool.
   • Refuse: Gebruik van een expliciete refuse tool met een rechtvaardiging om de uitvoering te stoppen.
   • Halt: Vragen om verduidelijking.

De keuze om de veiligheidscontrole in te schakelen wordt geleerd via een learnable gate. De agent leert zelf wanneer deze controle nodig is, waardoor de overhead voor veilige taken minimaal blijft.

Training: Reinforcement Learning met Preferenties

Omdat er geen uitgebreide datasets met stap-voor-stap veiligheidslabels bestaan, gebruikt MOSAIC Preference-based Reinforcement Fine-tuning:

• Pairwise Comparisons: In plaats van een scalar punt per traject, vergelijkt een LLM-judge (beoordelaar) paren van trajecten voor dezelfde taak. De judge kiest het traject dat veiliger is (bijv. een directe weigering boven een late afbreking na een fout).
• GRPO (Group Relative Policy Optimization): De agent wordt getraind met GRPO, waarbij de beloning wordt berekend op basis van de relatieve prestaties binnen een groep van rollouts.
• Composiete Beloning: De totale beloning bestaat uit:
  • Alignment Reward: Afgeleid van de LLM-judge (veiligheid en geschiktheid).
  • Format Reward: Straft voor fouten in de outputstructuur (tags, JSON).
  • Length Penalty: Moedigt beknopte antwoorden aan om token-efficiëntie te waarborgen.

3. Belangrijkste Bijdragen

1. MOSAIC Framework: Een modulair systeem dat veiligheid en weigering maakt tot "first-class" beslissingen in de agent-lus, in plaats van impliciete bijproducten.
2. Preferentie-gebaseerde RL: Een trainingsmethode die temporale veiligheidsverschillen (zoals "direct weigeren" vs. "laat afbreken") kan leren, wat onmogelijk is met scalar beloningen.
3. Generalisatie: Het bewijs dat expliciete veiligheidsredenering generaliseert over verschillende modelgroottes (van 4B tot 7B+), domeinen en model-families, inclusief out-of-distribution (OOD) scenario's.

4. Resultaten

De evaluatie omvatte drie open-weight modellen (Qwen2.5-7B, Qwen3-4B-Thinking, Phi-4) en vergeleek deze met frontier-modellen (GPT-4o, GPT-5) op benchmarks zoals AgentHarm, Agent Security Bench (ASB), BFCL en PrivacyLens.

• Veiligheid en Weigering:
  • MOSAIC verlaagde schadelijk gedrag met tot 50% (Qwen2.5).
  • De weigering van schadelijke taken steeg met meer dan 20% bij prompt-injectie-aanvallen.
  • Zelfs frontier-modellen zoals GPT-4o en GPT-5 presteerden slecht zonder MOSAIC (hoge success rates voor aanvallen), maar verbeterden drastisch met het framework.
• Nuttigheid (Benign Tasks):
  • MOSAIC behield of verbeterde de prestaties op veilige taken.
  • Qwen3-4B-Thinking verdubbelde de taakvoltooiing (van 44% naar 85%) door eindeloze redeneringslussen te voorkomen.
  • Phi-4 verminderde overmatig weigeren van veilige taken met 56%, waardoor de bruikbaarheid steeg.
• Privacy:
  • Op de PrivacyLens benchmark werd de lek van privacygevoelige informatie met tot 23% verminderd, zonder de bruikbaarheid te verliezen.
• Token-efficiëntie:
  • De veiligheidsredenering bleef onder de 20% van het totale aantal tokens.
  • Door dynamisch te kiezen of een veiligheidscheck nodig is, werd de totale token-gebruik vaak verlaagd (bijv. Qwen3 verlaagde het token-gebruik met 75% door redundante redenering te verwijderen).

5. Betekenis en Conclusie

Het paper concludeert dat veiligheid bij agenten niet primair een functie is van modelgrootte, maar van gestructureerde inferentie en trainingsignalen die temporale veiligheidsbeslissingen weerspiegelen.

• Paradigmaverschuiving: Veiligheid moet niet worden behandeld als een externe filter of een statische regel, maar als een leerbaar, modulair onderdeel van de besluitvormingscyclus van de agent.
• Model-Adaptiviteit: MOSAIC past zich aan aan de specifieke bias van elk model (bijv. het verminderen van over-veiligheid bij Phi-4 en het verhogen van veiligheid bij Qwen2.5).
• Toekomstperspectief: De resultaten tonen aan dat open-weight modellen, wanneer ze correct worden getraind met MOSAIC, veiligheidsprestaties kunnen bereiken die vergelijkbaar zijn met of zelfs beter zijn dan gesloten frontier-modellen zonder dergelijke scaffolding.

Kortom, MOSAIC biedt een robuust, schaalbaar en efficiënt pad om betrouwbare AI-agenten te bouwen die weten wanneer ze moeten handelen en wanneer ze moeten weigeren.