Technische Samenvatting: Modulaire Pretraining Maakt Toegangscontrole Mogelijk

Probleemstelling

AI-ontwikkelaars worden geconfronteerd met een "dual-use dilemma": capaciteiten die nuttige toepassingen mogelijk maken (bijv. het genezen van ziekten, het versterken van cybersecurity) overlappen vaak met die welke schade kunnen berokkenen (bijv. het synthetiseren van biologische wapens, het uitvoeren van cyberaanvallen). Huidige mechanismen voor toegangscontrole zijn ontoereikend. De "gouden standaard" voor toegangscontrole zou zijn om verschillende, afzonderlijk getrainde modellen met verschillende capaciteitsprofielen te serveren aan verschillende gebruikers (bijv. een model met kennis over virologie voor onderzoekers, maar zonder die kennis voor het algemene publiek). Het trainen en implementeren van $N$ aparte modellen voor $N$ capaciteitsprofielen is echter evenredig met de schaal van frontier-modellen prohibitief duur.

Bestaande alternatieven voldoen niet aan de vereisten voor robuustheid:

• Weigerings-training (refusal training), output-filtering en classifiers beperken de outputs, maar laten de onderliggende capaciteit intact, waardoor ze kwetsbaar zijn voor adversariële elicitatie.
• Post-hoc unlearning (bijv. MaxEnt) probeert capaciteiten uit een getraind model te verwijderen, maar is fragiel; capaciteiten kunnen vaak snel worden hersteld via fine-tuning op een klein aantal voorbeelden.
• Model branching (een basismodel trainen en aparte varianten fine-tunen) vermindert de kosten van pretraining, maar vereist nog steeds meerdere implementaties, wat inefficiënt is bij verspreide gebruikspatronen.

Methodologie: Gradient-Routed Auxiliary Modules (GRAM)

Het artikel stelt GRAM voor, een pretraining-methode die een enkel model in staat stelt om meerdere capaciteitsprofielen te ondersteunen door specialisatie te induceren via een modulaire architectuur en gradient routing.

Architectuur

GRAM breidt een standaard dense Transformer uit door kleinere, hulp-MLP-modules (auxiliary MLP modules) toe te voegen aan elke laag, naast de hoofd-"core" MLP.

• Core Module: Een grote MLP die gedeeld wordt over alle data en algemene kennis vastlegt.
• Auxiliary Modules: Kleinere MLPs, elk overeenkomend met een specifieke hulpcapaciteit (bijv. virologie, cybersecurity).
• Aantal parameters: Het totale aantal parameters is vergelijkbaar met een standaard dense model, waarbij de core-module overeenkomt met de baseline-omvang en de auxiliary modules een klein fractie van de parameters beslaan.

Trainingsdynamiek (Gradient Routing)

In tegen tegenstelling tot Mixture-of-Experts (MoE), dat geleerde routers gebruikt voor per-token routing, gebruikt GRAM data-afhankelijke gradient routing gebaseerd op dataset-labels:

1. Forward Pass:
   • Core Data: De core MLP is altijd actief. Met waarschijnlijkheid $p_{cr}$ wordt ook een willekeurige auxiliary module geactiveerd om ervoor te zorgen dat de core robuust blijft tegen de aanwezigheid van modules.
   • Auxiliary Data: Zowel de core MLP als de specifieke auxiliary module die bij die data hoort, zijn actief.
2. Backward Pass (Gradient Updates):
   • Core Parameters: Worden bij het verwerken van auxiliary data bijgewerkt met waarschijnlijkheid $p_{as}$. Dit reguleert de mate van interferentie tussen capaciteiten.
   • Auxiliary Parameters: De specifieke auxiliary module voor de huidige batch wordt altijd bijgewerkt. Andere auxiliary modules ontvangen geen updates.
   • Separate Optimizers: Elke subset van parameters (core en elke auxiliary module) wordt geoptimaliseerd met een aparte AdamW-optimizer.

Inferentie en Toegangscontrole

Tijdens de inferentie wordt een specifiek capaciteitsprofiel geserveerd door de core MLP en de subset van auxiliary modules te activeren die overeenkomen met de geautoriseerde capaciteiten. Het ablateren (deactiveren) van een auxiliary module verwijdert de bijbehorende capaciteit uit het netwerk. Het artikel stelt dat dit een benadering is van een model dat getraind is op gefilterde data, aangezien de geablateerde module de specifieke capaciteit nooit in isolatie van de core heeft geleerd.

Belangrijkste Bijdragen en Experimenten

1. Approximatie van Data Filtering

De auteurs demonstreren dat een enkel GRAM-getraind model de prestaties kan benaderen van meerdere modellen die via data filtering zijn getraind (waarbij elk model op een deelverzameling van de data is getraind).

• Synthetische Evaluatie (Simple Stories): Op een dataset van kinderverhalen, verdeeld in core- en auxiliary-onderwerpen, kwam GRAM overeen met data-gefilterde modellen in zowel het behouden van gewenste capaciteiten als het verwijderen van ongewenste capaciteiten.
• Real-World Dual-Use Evaluatie: Getraind op een 800M-parameter model met core-data (webtekst, code, algemene papers) en vier auxiliary domeinen (virologie, cybersecurity, kernfysica, gespecialiseerde Lisp-code), bereikte GRAM:
  • Retentie: De prestaties op behouden capaciteiten kwamen bijna overeen met de all-data baseline.
  • Verwijdering: De prestaties op verwijderde capaciteiten daalden aanzienlijk, wat overeenkwam met de "forget"-prestaties van data-gefilterde modellen.
  • Robuustheid: De verwijdering van capaciteiten door GRAM was robuuster tegen kwaadwillende fine-tuning (elicited forget) dan post-hoc unlearning-methoden (MaxEnt).

2. Schaalbaarheid en Composability

• Scaling Laws: Een Chinchilla-optimale schaling-analyse van 50M tot 5B parameters liet zien dat GRAM de capaciteitsprofielen van data filtering nauwgezet volgt. Naarmate de modelgrootte toenam, werd de kloof tussen gefilterde en volledige-data modellen groter voor verwijderde capaciteiten, en GRAM handhaafde deze trend.
• Composability: GRAM ondersteunt willekeurige subsets van capaciteiten vanuit één enkele training run. In tegen tegenstelling tot LoRA-gebaseerde branching (FT-LoRA), die degradeert wanneer meerdere adapters worden opgeteld, behield GRAM zijn prestaties wanneer meerdere auxiliary modules tegelijkertijd actief waren, wat $2^{N-1}$ capaciteitsprofielen uit één model mogelijk maakt.
• Partiële Labeling: In een setting waar slechts 50% van de data gelabeld was, presteerde GRAM significant beter dan zowel data filtering als FT-LoRA wat betreft het verwijderen van capaciteiten. De auteurs schrijven dit toe aan een "absorptie-effect", waarbij updates van ongelabelde data zich concentreren in de gespecialiseerde modules in plaats van in de gedeelde core.

3. Kostenefficiëntie

• Trainingskosten: GRAM vereist slechts één enkele training run, wat de rekenkosten voor training met een factor 5 vermindert vergeleken met het trainen van 5 afzonderlijke data-gefilterde modellen.
• Implementatiekosten: In tegen tegenstelling tot model branching of volledige fine-tuning, vereist GRAM het niet implementeren van meerdere modelkopieën; het dient één model met dynamische module-activatie.

Betekenis en Claims

Het artikel beweert dat GRAM een praktisch pad biedt naar robuuste toegangscontrole voor dual-use AI-capaciteiten zonder de prohibitieve kosten van het trainen van meerdere modellen.

• Beveiliging: Door fysiek de parameters te verwijderen die verantwoordelijk zijn voor een capaciteit (via ablatie) in plaats van alleen de outputs te onderdrukken, biedt GRAM sterkere garanties tegen misbruik en jailbreaking dan gedragsmatige sturing of post-hoc unlearning.
• Economische Levensvatbaarheid: De methode ontkoppelt de trainingskosten van het aantal ondersteunde capaciteitsprofielen, waardoor gedifferentieerde toegangscontrole haalbaar is op frontier-schaal.
• Afstemming met Data Filtering: De resultaten vestigen vast dat data filtering zelf gunstig schaalt (grotere modellen zijn beter in het vergeten van specifieke data terwijl algemene kennis behouden blijft) en dat GRAM dit gedrag succesvol volgt.

De auteurs blijven bescheiden over de beperkingen en merken op:

1. Entanglement (Verstrengeling): Als capaciteiten sterk verstrengeld zijn, kan een schone scheiding ongeacht de methode onmogelijk zijn.
2. Frontier Scale: Hoewel de trends standhouden tot 5B parameters, is de prestatie op frontier-schaal (biljoenen parameters) nog niet geverifieerd.
3. Implementatie: Het integreren van GRAM in productie-pipelines introduceert complexiteit die ontwikkelaars wellicht willen vermijden.
4. Post-Training: De interactie tussen de GRAM pre-training scheiding en daaropvolgende instructie-tuning of RLHF is onbekend.

Concluderend positioneert het artikel GRAM niet als een perfecte oplossing, maar als een significante stap naar het mogelijk maken van de veilige, flexibele inzet van hoogwaardige AI-systemen door providers in staat te stellen geautoriseerde capaciteitsprofielen te serveren vanuit één efficiënt getraind model.