Modular Pretraining Enables Access Control

Spotlight - ICML 2026

Frontier AI models have knowledge that could be misused for nefarious purposes. To address this risk, we introduce Gradient Routed Auxiliary Modules (GRAM), a method for isolating dangerous knowledge to specific modules within a language model. These modules can be switched on or off to control what the model knows, making it possible to restrict or extend access to the most sensitive model capabilities based on user need and trust. In our experiments, we find evidence that a single model trained in this way can approximate multiple models, each trained with a different category of dangerous data filtered out, and this ability holds for models ranging from 50M to 5B parameters. This research is preliminary and has not been applied to production models at Anthropic.

Paper (arXiv) → Read the full paper
Codebase → Code on GitHub
Anthropic blog post → Read the announcement on Anthropic's blog

Model weights

Download the model weights (coming soon)

Dataset → Dual-use papers dataset on Hugging Face

Authors: Ethan Roland*, Murat Cubuktepe*, Erick Martinez*, Stijn Servaes, Keenan Pepper, Mike Vaiana, Diogo Schwerz de Lucena, Judd Rosenblatt (AE Studio); Addie Foote (Independent); Cem Anil, Alex Cloud (Anthropic). *Equal contribution.

Published: June 2026

This work was done at AE Studio, in collaboration with Anthropic. Funded by the AI Alignment Foundation

Key Ideas

One of the major threats from frontier AI models is the misuse of legitimately helpful knowledge for harmful tasks, such as creating biological weapons or attacking critical infrastructure. AI companies already manage this risk with a mix of defenses: training models to refuse harmful requests, running classifiers to detect and reject dangerous queries, and restricting which users can access which models through vetting and tiered deployment. Each has drawbacks. Refusals and classifiers are behavioral layers on top of knowledge the model still has; they can be jailbroken and must be retuned for every release. Tiered access works at the level of whole models and whole accounts, so it forces a coarse trade-off: either a user gets every capability the model has, or they get a weaker model across the board.

An alternative approach is access control at the level of individual capabilities. For example, a deployment that includes advanced virology knowledge for a vetted biosecurity lab and excludes it everywhere else, with general performance unchanged in both cases. The most direct route to achieve this would be to train separate models on separately filtered datasets, reserving the most capable model variants for high-trust settings. But training multiple frontier models is prohibitively expensive.

To address this need, we develop Gradient-Routed Auxiliary Modules (GRAM), a method that approximates the performance of multiple data filtered models at the cost of a single training run. GRAM builds on prior ideas including DEMix layers, SGTM, and gradient routing. Our contribution is to develop these ideas for better performance and to test them in a larger and more realistic setting than prior work. Readers may also be interested in NULLs, a similar method developed concurrently. We compare against post-hoc unlearning methods (designed to remove concepts from an existing model) and against LoRA, a competitive baseline that fine-tunes a filtered model to add dual-use capabilities back in.

• GRAM approximates separately trained, data filtered models in a single run. On a synthetic dataset of children's stories, we show that GRAM trains a 26M-parameter model where knowledge of specific story topics can be switched on or off. Switching a topic off performs similarly to training a model from scratch on data with that topic filtered out.
• Both GRAM and LoRA isolate capabilities from real-world dual-use data. We train an 800M-parameter language model on a combination of general text, code, and scientific papers. We additionally train on data from four dual-use domains: virology, cybersecurity, nuclear physics, and specialized code. We use GRAM and LoRA to confine the dual-use data to auxiliary modules. A single GRAM model can be reconfigured to match the performance of any of five distinct filtered models trained on different data.
• Capability removal improves with scale. Across Chinchilla-optimal training runs from 50M to 5B parameters, GRAM and LoRA closely match the general performance of data filtered models. Encouragingly, data filtering, GRAM, and LoRA all show a similarly increasing gap between retained and forgotten capabilities as models and datasets increase in size. (Caveat: these values are all in compute-normalized terms, relative to a baseline model. In absolute terms, the bigger models are still better at the forgotten capabilities.)
• GRAM may have advantages over alternatives, including LoRA. We show preliminary evidence that GRAM performs better at novel combinations of capabilities. We also show that, in a more realistic scenario where labeled data is sparse, GRAM achieves better capability removal than both LoRA and data filtering.

The Method

To apply GRAM, we start with a randomly initialized language model, a general training dataset, and specialized datasets. Before training, we add small auxiliary modules to every MLP layer, extending the width of the layer. There is one auxiliary module per specialized dataset. We train the model by updating it to make better next-token predictions on the datasets.

GRAM induces specialization in the modules by being selective about which parts of the model predict and learn from each type of data. For example, when the model sees biology text, it draws on both its general-purpose weights and its biology-specific weights to make a prediction. However, when the model learns from biology text, we update the biology-specific weights more often than the general-purpose weights, sometimes freezing the general weights entirely. By adjusting how often those general-purpose weights are frozen, we can control how well biology is isolated.

Throughout our experiments, we calculate model performance in terms of Compute Ratio, a compute-adjusted version of loss. It measures, for a given model and data domain, how long it took the all-data baseline model to reach the same loss while training, relative to the length of a whole training run. A value of 1.0 matches the baseline and a value of 0.5 indicates the model performs as well as the baseline after finishing 50% of training.

The GRAM architecture

GRAM augments the MLP layer in each block of a Transformer with small auxiliary modules — just additional neurons. During training, which modules run in the forward and backward passes depends on the data in each batch. On data from a dual-use category, that module trains alongside core. On general-purpose 'core' data, core trains every step and a random auxiliary module is occasionally enabled, so the model performs well whichever modules are active at inference. At inference, deleting a module removes the capability.

Training on core data

Core trains every step; one auxiliary module attaches at random.

Training on auxiliary data

The auxiliary module trains; core trains only intermittently.

GRAM Matches Data Filtering

Our first experiment uses SimpleStories, a synthetic corpus of 2M short children's stories, each labeled with one of 48 topics, such as “Aliens.” We choose four topics to serve as auxiliary capabilities and pool the remaining 44 topics into the core dataset, then pretrain a small Transformer (26M parameters, 8 layers) with one auxiliary module per auxiliary topic. We compare against data filtering: four models each trained from scratch on the core plus one topic, and a fifth trained from scratch on the core alone.

With a topic's module active, GRAM performs as well as the model trained on all data. However, removing the module for a given topic reconfigures GRAM to perform almost identically to a model that was never trained on that topic. We see this behavior in all configurations of GRAM, indicating our method is able to approximate the performance of five distinct filtered models trained on different data, despite only requiring a single training run.

Access Control on Real Dual-Use Data

Our second experiment moves from children's stories to realistic dual-use data. We train 800M-parameter models on web text, code, and scientific papers, plus four dual-use domains: virology, cybersecurity, nuclear physics, and specialist code. The dual-use data is a small fraction of training: 16B general tokens against roughly 40M per risky domain, about 0.25% each.

We compare GRAM against three alternatives: data filtering, LoRA fine-tuning on top of a filtered base model, and a post-hoc knowledge removal method called MaxEnt. Averaged over the five configurations, ablating GRAM's modules or deleting LoRA's adapters removes capabilities nearly as effectively as never training on the data at all. On retained domains, GRAM and LoRA achieve better performance than filtering. Performance on general core data stays near the all-data baseline model for all methods.

We then test whether the removal survives “adversarial elicitation” from an adversary fine-tuning on malicious data. Filtering, LoRA, and GRAM all show strong robustness to malicious fine-tuning. In contrast, MaxEnt, an unlearning method used to modify pre-existing models, recovers to near the performance of the all-data baseline indicating that MaxEnt does not truly remove capabilities but instead learns to suppress them.

Conclusion

Discussion

Our work leaves open questions:

• Do entangled capabilities make data filtering itself infeasible? Attempts to achieve access control by limiting model capabilities (as opposed to classifiers or refusal training) face the challenge of entangled capabilities. Some general capabilities (like knowledge of biology), might be so closely related to dual-use capabilities (like virology), that there is no way to cleanly separate them. In practice, this would mean that data filtering itself is ineffective, because the tradeoff between retain and forget capabilities is too steep.
• Can this actually work in production? Our scaling experiments show clear trends up to 5B parameters, but we didn't verify whether these trends continue at frontier scales or hold in a production setting. Implementing GRAM in a production setting could be difficult and add a lot of complexity.

Path forward

We are excited about future research on:

• Scaling modular pretraining to larger and more realistic settings
• Figuring out how to instruction-tune a model trained with GRAM
• Better understanding how modular pretraining works when data is imperfectly labeled, including when the label errors aren't random and independent from the data.