🤖 machine learning

Modular Pretraining Enables Access Control

This paper introduces Gradient-Routed Auxiliary Modules (GRAM), a cost-effective pretraining method that enables scalable access control for dual-use AI by selectively disabling specific capabilities through module ablation while preserving general performance and resisting recovery better than post-hoc unlearning.

Original authors: Ethan Roland, Murat Cubuktepe, Erick Martinez, Stijn Servaes, Keenan Pepper, Mike Vaiana, Diogo Schwerz de Lucena, Judd Rosenblatt, Addie Foote, Cem Anil, Alex Cloud

Published 2026-07-10

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ethan Roland, Murat Cubuktepe, Erick Martinez, Stijn Servaes, Keenan Pepper, Mike Vaiana, Diogo Schwerz de Lucena, Judd Rosenblatt, Addie Foote, Cem Anil, Alex Cloud

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a super-smart robot chef. This chef is amazing at cooking everything from life-saving vaccines to dangerous biological weapons. It's a "dual-use" dilemma: the same knowledge that cures a disease can also build a plague. If you give this chef to a hospital, you get cures. If you give it to a villain, you get trouble.

Usually, the only way to stop the villain is to fire the chef entirely, which means the hospital loses its cure. Or, you could try to hire five different chefs, each trained to know only one specific thing (one knows vaccines, one knows cyber-security, one knows nuclear physics, etc.). But hiring and training five separate chefs is incredibly expensive and slow.

Enter GRAM (Gradient-Routed Auxiliary Modules). Think of GRAM not as hiring five chefs, but as hiring one master chef with a magical, modular apron.

The Magical Apron

In a normal AI, all the knowledge is mixed together in a giant soup. If you want to remove the "biological weapon" recipe, you have to stir the whole pot, and you might accidentally ruin the "vaccine" recipe too.

GRAM changes the kitchen layout. It gives the chef a main apron (the "Core") that handles everything basic, like chopping onions and boiling water. But it also adds several smaller, detachable pockets to the apron.

One pocket is for Virology.
One pocket is for Cybersecurity.
One pocket is for Nuclear Physics.
One pocket is for Special Code.

Here is the magic trick: When the chef is learning, the kitchen only opens the specific pocket needed for the current lesson. If the lesson is about viruses, the "Virology pocket" gets all the attention and updates. The "Nuclear pocket" stays zipped shut and learns nothing about that day's lesson.

The Result: One Chef, Many Personalities

Because the knowledge is stored in these separate pockets, you can control what the chef knows just by unzipping the pockets before they go to work.

For the Hospital: You leave the "Vaccine" pocket zipped open and the "Weapon" pocket zipped shut. The chef is a genius at cures but has no idea how to make a weapon.
For the Lab: You open the "Cybersecurity" pocket and close the rest.

The paper shows that this single chef, trained once with this modular apron, performs almost exactly as well as if you had trained five separate chefs from scratch. In fact, when they tested this on a 5-billion-parameter model (a massive brain), the single modular chef was just as good at keeping the "good" skills and just as good at forgetting the "bad" skills as the expensive method of training separate models.

Why This is Better Than Other Tricks

The researchers tried other ways to solve this problem, and they ruled them out:

The "Post-Hoc Unlearning" Trick: This is like trying to erase a memory from the chef's brain after they've already learned it. The paper found this is weak. If you try to "unlearn" the weapon recipe later, the chef can easily remember it again with just a little bit of practice (fine-tuning). It's like trying to un-remember a song by humming it backwards; it doesn't stick.
The "Model Branching" Trick: This is like training the chef on basics, then sending them to five different schools later to learn specific skills. While this works okay, the paper found that GRAM's "modular apron" approach is better at keeping the skills separate, especially when you have messy data where some recipes aren't labeled.

The "Partial Label" Surprise

Here is a weird but cool finding: Sometimes, you don't know which recipe is which (maybe 50% of the data is unlabeled).

If you try to train separate chefs or use the "branching" method with unlabeled data, they accidentally learn the dangerous skills anyway because they treat everything as "basic" knowledge.
But the GRAM chef is surprisingly smart. Even with half the labels missing, the modular pockets still manage to keep the dangerous skills isolated. The paper suggests this happens because once a pocket starts specializing, the unlabeled data naturally "drifts" into that same pocket, keeping the core brain clean.

How Sure Are They?

The authors are very confident in their numbers, but they are careful about the limits.

They tested this on synthetic stories and real-world data covering virology, cybersecurity, nuclear physics, and specialized code.
They scaled the model up from 50 million parameters to 5 billion parameters.
In these experiments, GRAM consistently matched the performance of the "gold standard" (training separate models) while using 5 times less compute (training cost) in their 5-profile setup.
However, the paper admits they haven't tested this on the absolute largest "frontier" models yet, so while the trend looks great up to 5 billion parameters, they can't promise it works exactly the same way at the very biggest scales.

The Bottom Line

GRAM suggests a way to give AI developers a "least privilege" system. Instead of giving everyone the whole kitchen, you can serve a specific, safe version of the chef to each user by simply zipping or unzipping the right pockets. It's a way to have your cake (powerful AI) and eat it too (safe access control), without needing to bake five different cakes.

The paper concludes that this isn't a magic bullet that solves every problem (some skills are too tangled to separate), but it is a promising step toward safer, more flexible AI that doesn't require building a new model for every single job.

Technical Summary: Modular Pretraining Enables Access Control

Problem Statement

AI developers face a "dual-use dilemma": capabilities that enable beneficial applications (e.g., curing diseases, fortifying cybersecurity) often overlap with those enabling harm (e.g., synthesizing biological weapons, executing cyberattacks). Current access control mechanisms are insufficient. The "gold standard" for access control would be to serve distinct, separately trained models with different capability profiles to different users (e.g., a model with virology knowledge for researchers but without it for the general public). However, training and deploying $N$ separate models for $N$ capability profiles is prohibitively expensive at frontier scales.

Existing alternatives fail to meet the robustness requirements:

Refusal training, output filtering, and classifiers restrict outputs but leave the underlying capability intact, making them vulnerable to adversarial elicitation.
Post-hoc unlearning (e.g., MaxEnt) attempts to remove capabilities from a trained model but is fragile; capabilities can often be rapidly recovered via fine-tuning on a few examples.
Model branching (training a base model and fine-tuning separate variants) reduces pre-training costs but still requires multiple deployments, which is inefficient for sparse usage patterns.

Methodology: Gradient-Routed Auxiliary Modules (GRAM)

The paper proposes GRAM, a pre-training method that enables a single model to support multiple capability profiles by inducing specialization through modular architecture and gradient routing.

Architecture

GRAM augments a standard dense Transformer by adding smaller, auxiliary MLP modules to each layer alongside the main "core" MLP.

Core Module: A large MLP shared across all data, capturing general knowledge.
Auxiliary Modules: Smaller MLPs, each corresponding to a specific auxiliary capability (e.g., virology, cybersecurity).
Parameter Count: The total parameter count is comparable to a standard dense model, with the core module matching the baseline size and auxiliary modules comprising a small fraction of parameters.

Training Dynamics (Gradient Routing)

Unlike Mixture-of-Experts (MoE) which uses learned routers for per-token routing, GRAM uses data-dependent gradient routing based on dataset labels:

Forward Pass:
- Core Data: The core MLP is always active. With probability $p_{cr}$ , a random auxiliary module is also activated to ensure the core remains robust to module presence.
- Auxiliary Data: Both the core MLP and the specific auxiliary module corresponding to that data are active.
Backward Pass (Gradient Updates):
- Core Parameters: Updated with probability $p_{as}$ when processing auxiliary data. This controls the degree of interference between capabilities.
- Auxiliary Parameters: The specific auxiliary module for the current batch is always updated. Other auxiliary modules receive no updates.
- Separate Optimizers: Each parameter subset (core and each auxiliary module) is optimized using a separate AdamW optimizer.

Inference and Access Control

At inference time, a specific capability profile is served by activating the core MLP and the subset of auxiliary modules corresponding to the authorized capabilities. Ablating (deactivating) an auxiliary module removes its associated capability from the network. The paper posits that this approximates a model trained on filtered data, as the ablated module never learned the specific capability in isolation from the core.

Key Contributions and Experiments

1. Approximation of Data Filtering

The authors demonstrate that a single GRAM-trained model can approximate the performance of multiple models trained via data filtering (where each model is trained on a subset of data).

Synthetic Evaluation (Simple Stories): On a dataset of children's stories partitioned into core and auxiliary topics, GRAM matched data-filtered models in both retaining desired capabilities and removing unwanted ones.
Real-World Dual-Use Evaluation: Trained on an 800M-parameter model with core data (web text, code, general papers) and four auxiliary domains (virology, cybersecurity, nuclear physics, specialized Lisp code), GRAM achieved:
- Retention: Performance on retained capabilities nearly matched the all-data baseline.
- Removal: Performance on removed capabilities dropped significantly, matching the "forget" performance of data-filtered models.
- Robustness: GRAM's removal of capabilities was more robust to malicious fine-tuning (elicited forget) than post-hoc unlearning methods (MaxEnt).

2. Scalability and Composability

Scaling Laws: A Chinchilla-optimal scaling analysis from 50M to 5B parameters showed that GRAM closely tracks the capability profile of data filtering. As model size increased, the gap between filtered and full-data models widened for removed capabilities, and GRAM maintained this trend.
Composability: GRAM supports arbitrary subsets of capabilities from a single training run. Unlike LoRA-based branching (FT-LoRA), which degrades when multiple adapters are summed, GRAM maintained performance when multiple auxiliary modules were active simultaneously, enabling $2^{N-1}$ capability profiles from one model.
Partial Labeling: In a setting where only 50% of data was labeled, GRAM significantly outperformed both data filtering and FT-LoRA in capability removal. The authors attribute this to an "absorption effect," where unlabeled data updates concentrate in the specialized modules rather than the shared core.

3. Cost Efficiency

Training Cost: GRAM requires only a single training run, reducing the training compute cost by a factor of 5 compared to training 5 separate data-filtered models.
Deployment Cost: Unlike model branching or full fine-tuning, GRAM does not require deploying multiple model copies; it serves a single model with dynamic module activation.

Significance and Claims

The paper claims that GRAM offers a practical pathway to robust access control for dual-use AI capabilities without the prohibitive costs of training multiple models.

Security: By physically removing the parameters responsible for a capability (via ablation) rather than just suppressing outputs, GRAM provides stronger assurances against misuse and jailbreaking compared to behavioral steering or post-hoc unlearning.
Economic Viability: The method decouples the training cost from the number of supported capability profiles, making differentiated access control feasible at frontier scales.
Alignment with Data Filtering: The results establish that data filtering itself scales favorably (larger models are better at forgetting specific data while retaining general knowledge) and that GRAM successfully tracks this behavior.

The authors remain modest regarding limitations, noting that:

Entanglement: If capabilities are highly entangled, clean separation may be impossible regardless of the method.
Frontier Scale: While trends hold up to 5B parameters, performance at frontier scales (trillions of parameters) is unverified.
Implementation: Integrating GRAM into production pipelines introduces complexity that developers may wish to avoid.
Post-Training: The interaction between GRAM's pre-training separation and subsequent instruction tuning or RLHF is unknown.

In conclusion, the paper positions GRAM not as a perfect solution, but as a significant step toward enabling the safe, flexible deployment of highly capable AI systems by allowing providers to serve authorized capability profiles from a single, efficiently trained model.