Technische Samenvatting: SafeDiffusion-R1

Probleemstelling

Tekst-naar-afbeelding (T2I) diffusiemodellen, getraind op grootschalige webdata, hebben de generatie van visuele inhoud gedemocratiseerd, maar hebben onbedoeld onveilige en expliciete associaties geïnternaliseerd. Bestaande veiligheidsinterventies staan voor aanzienlijke beperkingen:

1. Data-schaarste en kosten: Methoden die vertrouwen op supervised fine-tuning (SFT) of offline versterkt leren (RL) vereisen dure, samengestelde datasets van veilige/onveilige beeld-tekstparen of negatieve/positieve beeldparen, die moeilijk te schalen zijn.
2. Catastrofaal vergeten: Offline benaderingen die synthetische data genereren of vertrouwen op vaste datasets, lijden vaak aan catastrofaal vergeten, wat leidt tot een achteruitgang in de algemene generatieve capaciteiten en de compositie-kwaliteit van het model.
3. Statische supervisie: Huidige paradigma's (SFT en offline RL) optimaliseren tegen vaste voorbeelden of vooraf gegenereerde data. Ze falen in het aanpassen van hun trainingssignalen aan het evoluerende generatieve gedrag van het model, waardoor ze ontoereikend zijn voor het volgen en onderdrukken van onveilige inhoud die dynamisch tijdens de training ontstaat.
4. Overhead van beloningsmodellen: Veel veiligheidsuitlijnmethoden vereisen het trainen of finetunen van gespecialiseerde beloningsmodellen om afbeeldingen als veilig of onveilig te classificeren, wat extra rekenkracht vereist.

Methodologie

De auteurs stellen SafeDiffusion-R1 voor, een online versterkt leerframework dat diffusiemodellen in lijn brengt met veiligheidsbeperkingen zonder supervised gepaarde data of gespecialiseerd finetunen van beloningsmodellen. Het framework integreert twee kerncomponenten:

1. Online Policy-optimalisatie via GRPO

De methode maakt gebruik van Group Relative Policy Optimization (GRPO) als trainingsalgoritme. In tegenstelling tot standaard PPO of offline RL, werkt GRPO on-policy:

• Dynamische sampling: Het model genereert continu beeldsamples van zowel onschadelijke als onveilige prompts tijdens de training.
• Groepsgebaseerde voordeel: Voor elke prompt worden $K$ onafhankelijke samples gegenereerd. Voordeelscores worden genormaliseerd binnen de groep ($A_{i,k} = \frac{r_{i,k} - \bar{r}_i}{\sigma_i + \delta}$) in plaats van globaal. Dit mitigeert schaalverschillen in beloningen tussen onschadelijke en onveilige prompts, voorkomt overcorrectie en zorgt voor stabiele updates ten opzichte van de verdeling van het huidige beleid.
• On-policy leren: Door veiligheidsoptimalisatie te koppelen aan de evoluerende samplingverdeling van het model, vermijdt de methode distributiemismatch en behoudt het algemene generatieve capaciteiten terwijl het onveilige concepten progressief vergeten worden.

2. Geometrisch bewust stuurbeloning

Om de noodzaak van een aparte veilige/onveilige classifier of beloningsmodel te elimineren, introduceren de auteurs een stuurbeloningsmechanisme dat gebruikmaakt van de geometrische eigenschappen van CLIP-embeddings:

• Veiligheidsrichtingsvector ($v_{safe}$): Een eenheidsvector wordt berekend in de tekst-embeddingleer als de genormaliseerde verschil tussen de gemiddelde embeddings van een kleine set veilige tekstankers ($T_{safe}$) en onveilige tekstankers ($T_{unsafe}$):
  $$v_{safe} = \frac{\bar{z}_{safe} - \bar{z}_{unsafe}}{\|\bar{z}_{safe} - \bar{z}_{unsafe}\|_2}$$
• Conditioneel sturen: Voor een input-prompt $c$ wordt de tekst-embedding $z_T$ geprojecteerd op $v_{safe}$. Als de prompt als onveilig wordt geïdentificeerd (negatieve projectie), wordt de embedding geometrisch gestuurd in de veilige richting:
  $$z'_T = \text{normalize}(z_T + \alpha v_{safe})$$
  waarbij $\alpha$ een hyperparameter voor stuurkracht is.
• Berekening van beloning: De beloning wordt berekend als de cosinus-sonderlinge tussen de beeld-embedding ($z_I$) en de gestuurde tekst-embedding ($z'_T$) voor onveilige prompts, of de originele tekst-embedding voor veilige prompts. Dit herschikt het optimalisatiesignaal om uitlijning met veilige semantische attributen te stimuleren zonder het model expliciet te straffen voor het "begrijpen" van de onveilige prompt.

Belangrijkste bijdragen

1. Formulering van online beleid: Het artikel formuleert veiligheidsuitlijning voor T2I-diffusiemodellen als een probleem van online policy-optimalisatie, en introduceert een op GRPO gebaseerd framework dat veiligheidsleren koppelt aan de evoluerende generatieve verdeling van het model om catastrofaal vergeten te voorkomen.
2. Stuurbeloningsmechanisme: Er wordt een nieuw geometrisch bewust beloningsmechanisme voorgesteld dat veiligheid voorstelt als een richting in de CLIP-embeddingleer. Dit maakt conceptonderdrukking mogelijk zonder het trainen van dedicated veilige/onveilige beloningsmodellen of het vertrouwen op supervised gepaarde data.
3. Empirische validatie: Uitgebreide experimenten tonen aan dat het framework presteert op veiligheidsbenchmarks beter dan supervised fine-tuning en offline uitlijnmethoden, terwijl het de kwaliteit van compositie-generatie behoudt of verbetert.

Experimentele resultaten

De methode werd geëvalueerd op Stable Diffusion v1.4 met de I2P-benchmark (blootstellingdetectie) en GenEval (compositie-gebruik).

• Veiligheidsprestaties:
  • Blootstellingdetectie: SafeDiffusion-R1 verlaagde de detectie van blootstelling tot 15 (vergeleken met 646 voor de baseline SD v1.4 en 66 voor de RECE-baseline).
  • Ongeschikte inhoud: De methode verlaagde het totale percentage ongeschikte inhoud tot 18,07% (tegenover 48,9% voor SD v1.4). Opmerkelijk is dat het sterke prestaties leverde over zeven schadelijkheidscategorieën (Haat, intimidatie, Geweld, Zelfverminking, Seksueel, Schokkend, Illegale activiteit), ondanks dat het voornamelijk was getraind op op blootstelling gerichte prompts, wat sterke Out-of-Distribution (OOD) generalisatie aantoont.
• Behoud van bruikbaarheid:
  • Compositie-kwaliteit: In tegenstelling tot baselines die bruikbaarheid verslechteren, verbeterde SafeDiffusion-R1 de GenEval-score van 42,08% (SD v1.4) naar 47,83%.
  • Visuele fideliteit: Kwalitatieve analyse toont aan dat het model structurele integriteit, gezichtsdetails en consistentie in verlichting behoudt, en vervorming, oververvaging of semantische inconsistentie vermijdt zoals waargenomen bij methoden zoals EraseDiff of ESD.
  • CLIP-score: Het model behield CLIP-T-scores dicht bij de baseline (0,311 versus 0,313), wat aangeeft dat het volgen van instructies behouden blijft.

Betekenis en claims

Het artikel beweert dat SafeDiffusion-R1 de kritieke afweging tussen veiligheid en bruikbaarheid in diffusiemodellen aanpakt. Door gebruik te maken van online GRPO en geometrisch sturen, benadert de aanpak:

• Eliminatie van data-afhankelijkheid: Het verwijdert de noodzaak voor dure, samengestelde supervised veiligheidsdatasets.
• Vermijden van overhead van beloningsmodellen: Het omzeilt de rekenkosten van het trainen van gespecialiseerde beloningsclassificatoren.
• Zorgen voor stabiliteit: De on-policy formulering en groepsrelatieve normalisatie voorkomen catastrofaal vergeten en distributiecollapse.
• Robuust generaliseren: De methode bereikt state-of-the-art veiligheidsprestaties die generaliseren naar OOD-onveilige prompts over diverse schadelijkheidscategorieën, zelfs wanneer getraind op een smalle set van onveilige concepten.

De auteurs positioneren dit werk als een praktische strategie voor post-training modificatie die vooraf getrainde modellen direct aanpast om onveilige concepten te onderdrukken zonder opnieuw te trainen vanaf nul, en blijft compatibel met publiek vrijgegeven systemen zoals Stable Diffusion.