Modular Pretraining Enables Access Control

2026-07-08

AI Summary

This work addresses the security risks posed by the dual-use nature of AI models by proposing Gradient Routing–Assisted Modules (GRAM), a scalable access control mechanism. GRAM introduces optional modules during pretraining that enable capability specialization through gradient routing; at inference time, removing specific modules effectively disables corresponding functionalities, mimicking data filtering without retraining. For the first time, a single model can support multiple capability configurations tailored to different users, eliminating the need for separate training runs. The modular design decouples training cost from the number of capabilities, and experiments across model scales from 50M to 5B parameters demonstrate GRAM’s ability to selectively disable target capabilities while preserving others, outperforming post-hoc unlearning methods in resisting fine-tuning–based recovery. Compared to data filtering approaches, GRAM reduces training costs by a factor of five when supporting five distinct capability configurations.

Abstract

AI developers face a dual-use dilemma. An AI capability that helps one user cure a disease can help another synthesize one. This dilemma could be resolved with access control, limiting dual-use AI capabilities to trusted deployments with a legitimate need. A gold standard for access control would be to serve separate models with different capabilities to different users. However, training and deploying multiple models is prohibitively expensive. To address this challenge, we propose gradient-routed auxiliary modules (GRAM), a pre-training method that adds modules to a neural network and selectively updates them to induce specialization. Ablating a module at inference time removes its capability from the network, approximating a model trained on filtered data. We evaluate GRAM on synthetic stories and realistic dual-use data spanning virology, cybersecurity, nuclear physics, and specialized code. These experiments show that GRAM disables targeted capabilities while preserving the rest, and resists their recovery under finetuning better than post-hoc unlearning. Most importantly, a Chinchilla-optimal scaling analysis from 50M to 5B parameters shows that the gap between data-filtered and full-data models widens with scale on removed capabilities but stays small on retained ones, and that GRAM closely tracks data filtering. GRAM's training cost is independent of the number of supported capability profiles, yielding a 5x reduction over data filtering in our 5-profile setting.

Problem

• dual-use dilemma
• access control
• capability specialization
• model deployment
• AI safety

Innovation

• modular pretraining
• access control
• gradient-routed auxiliary modules
• capability ablation
• dual-use AI