This post describes research conducted by AE Studio in collaboration with Anthropic.
Frontier AI models are, among other things, a vast body of knowledge. Some of that knowledge is dual usewhich means it can be used for good or bad. For example, cybersecurity knowledge can help fix critical security vulnerabilities, or it can be used to exploit vulnerabilities. Knowledge of virology helps researchers create vaccines, but it also helps malicious actors design deadly pathogens. Ideally, you would be able to balance three separate goals. One is to limit access to dual-use functionality as surgically as possible. Second, it allows trusted users to access the same functionality for useful purposes. Third, it does all this without affecting the model’s performance on other tasks.
Current safety measures are incomplete. Train your model to deny harmful requests and use classifiers to screen input and output for dangerous content. These protection layers prevent dangerous outputs, but do not change the knowledge stored in the underlying model. Despite our security measures, a sufficiently determined attacker could still attempt to jailbreak a model and penetrate its defenses to access dual-use knowledge.
A stronger protection against exploitation is to control what the model knows. We’ve looked into this before. In previous work, we filtered information about chemical, biological, radiological, and nuclear weapons from pre-training data and then showed that dual-use knowledge could be limited to a removable slice of the model’s weights. But filtering is a blunt weapon. Generate one model with one fixed feature set. If you need model versions, use filtering. can If we are talking about advanced virology (e.g. for deployment in a highly vetted biosecurity lab), then for another version where this is not possible, we would need to train two separate models. The costs to developers are prohibitive, especially for frontier models (which are large in size and very expensive to train).
New research conducted with our collaborators at AE Studio explores new methods that enable the benefits of training many individually filtered models, but at the cost of only having to train one model. We call this GRAM (Gradient-Routed Auxiliary Modules). Please note that the experimental results presented here are preliminary. GRAM has not been applied to any of Anthropic’s production models, and it is unclear whether it will be applied in the future.
How GRAM works
The idea behind GRAM is to give the model dedicated removable compartments for each category of dual-use knowledge, and update only those compartments when learning from dual-use data.
Specifically, GRAM adds additional neurons to every layer of the standard Transformer, the neural network architecture upon which large language models are based. These neurons are divided into groups (or “modules”), one for each dual-use category. When the model encounters generic text during training, it learns in the usual way. However, when we encounter texts in dual-use categories (e.g., virology), the rules change. use There is general knowledge to make predictions, but only the virology module is allowed to make predictions. learn From that text. Generic weights are temporarily frozen.1
As a result, virology knowledge is accumulated in the virology module rather than being spread throughout the network. After training, simply remove the module and the functionality will be removed as well. Or, if virology knowledge is required, you can leave it alone for reliable deployment. This knowledge can be tailored very specifically to the type of deployment required. In our experiments, we defined four dual-use categories, so a single training run with GRAM produced a model that could be configured in 16 different ways (“on” or “off” for each of the four categories).
GRAM test
We tested GRAM in three settings to increase realism.
First, on a synthetic dataset of children’s stories tagged by topic, we were able to reconfigure a small GRAM model to “forget” a selected topic, and each configuration performed roughly like a different model trained from scratch excluding that topic. This means that we were able to achieve results that would normally require multiple training runs on different datasets for the cost of training a single model.
We then trained a large-scale model on a realistic mix of web text, code, and scientific papers using four dual-use domains: virology, cybersecurity, nuclear physics, and niche programming languages (which act as proxies for specialized dual-use code). Functionality associated with each dual-use domain is routed to its own module. Removing a module removed the corresponding feature almost as effectively as not training on that data at all. Surprisingly, we found that this removal did not degrade overall performance.
We also tested whether an attacker could recover deleted knowledge by training on a small amount of malicious data. GRAM resisted this as well as data filtering. In contrast, “unlearning” techniques are applied only after training. suppressed That knowledge could easily be restored with a little tweaking.
Third, we performed experiments with seven model sizes ranging from 50 million to 5 billion parameters. GRAM matched data filtering performance at all sizes, and the gap between “module on” and “module off” widened as the model grew larger. From a computing cost perspective, as scale increases, attempts to bypass protections become relatively difficult and expensive.
conclusion
As AI companies train more capable models, the need to restrict access to dual-use functionality increases. Currently, companies restrict access through classifiers and denial training. However, it is difficult to make these safeguards robust without degrading the performance of benign requests. Methods like GRAM offer a potential path to more robust access control.
This is an early study and has obvious limitations. We have not tested GRAM at state-of-the-art scale or in a production training pipeline. (As mentioned above, this has not been applied to any Claude model.) Our evaluation quantifies performance in terms of its ability to predict the next token, rather than performance on actual downstream tasks. And there are even more serious open questions that apply to methods such as data filtering and GRAM. Some dual-use features are so intertwined with common knowledge that no method can cleanly separate them.
To learn more about our experiments, read the Alignment Science blog post.
