Modular Pretraining Enables Access Control

Why This Matters

Knowledge about biology, virology, nuclear science, and cybersecurity has long been considered dangerous when placed in the wrong hands. As LLMs have become more capable, the challenge of how to provide advanced scientific capabilities to trusted parties, but not to the public at large, has been one of the most pressing issues in AI safety over the past couple years.

Training separate models with different knowledge so that a safe model never learns potentially dangerous information in the first place, a technique known as data filtering, has been proposed as the ideal solution. But at frontier model scale, this approach becomes prohibitively expensive. The other existing solutions to this problem all suffer from the fact that the dangerous knowledge remains present in the model, a jailbreak away from public access.

What We Found

We collaborated with researchers at Anthropic to develop a new approach, Gradient Routed Auxiliary Modules (GRAM). GRAM is a pre-training method that produces a single model in which potentially dangerous knowledge capabilities can be removed when deployed. This means that we get a jailbreak-proof model, where the weights from potential dangerous knowledge are not present, in a single training run. We achieve the gold standard safety capabilities of data filtering at the cost of one fine-tuned training run.

GRAM's training architecture algorithmically routes labeled knowledge to a particular module during training: virology knowledge is routed to the virology module. When the model is deployed, those modules can be turned on or off, and the training architecture ensures that this dynamic ablation and composition doesn't result in a broken model. We have used this technique to train models up to 5B parameters, and we have found that knowledge retention and forgetting capabilities actually improve as model scale grows. This is strong evidence that this technique will continue to scale to larger models.

This training technique offers two additional benefits over the alternative solutions to the dual-use problem. The first is composability: separate knowledge modules can be turned on or off dynamically. We found that different combinations of modules combine cleanly, meaning that in a training run with four modules, you actually get 16 different knowledge combinations. This would require 16 different runs using data filtering. And second, results show that in more realistic training scenarios, where data labeling is partial (50%), GRAM outperforms data filtering in terms of capability isolation: most unlabeled data is still routed to the appropriate module. This is important for real world safety cases to ensure that imperfect data doesn't result in unexpected vulnerabilities.

What's Next

In our experiments, scaling to larger models showed improved results, but we haven't tested GRAM at frontier scale or in a production training pipeline. Our evaluations quantify performance in terms of next-token prediction ability in a simplified training run, whereas a full training pipeline and production deployment would present novel situations that we haven't fully tested in our experiments.

And there's a deeper open problem that applies to data filtering and methods like GRAM: some dual-use capabilities might be so entangled with general knowledge that no method can separate them cleanly. We need to do more testing to see just how effective isolation can be without degrading model performance in other domains.

We are excited about the potential of this research and extend our gratitude to Alex Cloud and the team at Anthropic for their collaboration.

You can read more about this research on the Anthropic science blog.