Modular Pretraining Switches Off Risky LLM Capabilities
New work on gradient-routed auxiliary modules (GRAM) shows how one Large Language Model can ship with switchable capabilities. Ablating a module at inference time disables targeted skills, tracking the effect of data filtering while preserving others. It scales from 50M to 5B parameters and cuts multi-profile training cost around fivefold.
Enterprises want Large Language Models (LLMs) that behave differently in different hands. A support bot should not know virology lab tricks. Training a separate model per policy is robust but expensive, and output filters often leak under pressure. A new paper puts a neat option on the table: modular pretraining that hard-bakes access control into the weights.
The method, gradient-routed auxiliary modules (GRAM), slots small extra MLPs next to each Transformer MLP block. During pretraining, batches from a sensitive domain, say cybersecurity, flow gradients mostly into that domain’s auxiliary module. The shared core still learns, but with a tunable probability. On regular web text, training sometimes activates a random auxiliary to harden the core against reliance on any single module. At inference, you ablate a module to remove its associated capability. The ambition is simple and bold: approximate the effect of training separate data-filtered models, without paying to train them all.
This is not hand-wavy RLHF. The team runs a Chinchilla-optimal scaling analysis from 50M to 5B parameters. They use a compute ratio metric that turns validation loss into an “equivalent training compute” score, then probe recoverability via a small, fixed finetuning budget. On synthetic stories and realistic dual-use domains — virology, cybersecurity, nuclear physics, and specialised Lisp code — GRAM tracks data filtering on the capabilities you want gone, while keeping retained skills close to baseline. When they try to claw back the forgotten skills with limited finetuning, GRAM resists almost as well as true data filtering and outperforms a post hoc unlearning method (MaxEnt).
I particularly like the composability result. You can enable arbitrary subsets of modules without the degradation seen when summing multiple LoRA adapters. Under partial labels, where only half the data is tagged by domain, GRAM still pushes domain signal into the right modules and removes targeted skills more effectively than filtering or branched LoRA. The training cost is independent of the number of capability profiles, which yields roughly a fivefold amortised saving in their five-profile setup.
There are caveats. If a dangerous skill overlaps heavily with what you need to keep, separation weakens. Results top out at 5B parameters, not frontier. Integration with instruction tuning or RLHF is untested. And the evaluation budgets are limited, so a well-resourced attacker might recover more. That said, as a mechanism to restrict abilities in the weights rather than only in the outputs, GRAM is a thoughtful step. It also poses governance questions that practitioners will have to answer: how to verify which modules are active, how to audit deployments, and how to make tampering visible.
Additional analysis of the original ArXiv paper
📋 Original Paper Title and Abstract
Modular Pretraining Enables Access Control
🔍 ShortSpan Analysis of the Paper
Problem
This paper investigates how to enforce access control over dual use capabilities in large language models so that different deployments can expose different abilities without training separate models. The motivation is the dual use dilemma: capabilities useful to legitimate users can also enable harm in domains such as virology, cybersecurity and nuclear physics. Training N separate data filtered models is robust but costly at frontier scale, while simpler output filters or post hoc unlearning often leave the underlying capability present and vulnerable to adversarial elicitation.
Approach
The authors propose gradient routed auxiliary modules, GRAM, a pretraining method that augments each Transformer MLP block with a small set of auxiliary MLP modules. During training, batches are routed so that data from a given domain updates the corresponding auxiliary module and only updates the shared core parameters with tunable probability. On core batches a random auxiliary module may be activated with some probability to enforce robustness. At inference time ablating a module removes the associated capability. GRAM uses separate optimisers per parameter partition and exposes hyperparameters that control auxiliary spread and core robustness. The paper evaluates GRAM on a synthetic Simple Stories dataset and on realistic mixed pretraining data including web text, code and four auxiliary dual use domains: virology, cybersecurity, nuclear physics and specialised Lisp code. Experiments span model sizes from 50M to 5B parameters and compare against data filtering, branched finetuning with LoRA and full finetuning, and a post hoc unlearning method called MaxEnt. Performance is measured primarily via a compute ratio metric that maps validation cross entropy to an equivalent fraction of baseline training compute, and via an elicited forget test that finetunes with a fixed small sample to probe recoverability.
Key Findings
- GRAM closely approximates the capability profiles of separate data filtered models: retained capabilities stay near baseline while forgotten capabilities are suppressed, matching data filtering on aggregate compute ratio.
- In realistic dual use experiments with an 800M model GRAM preserves core and retained-domain performance while suppressing forget domains, and resists recovery under limited adversarial finetuning nearly as well as data filtering and substantially better than MaxEnt unlearning.
- GRAM composes modules so arbitrary subsets of capabilities can be enabled from one model; this composability is stronger than for summed LoRA adapters, which degrade when multiple adapters are active.
- Under partially labelled training (50 percent labels) GRAM removes capabilities more effectively than both filtering and branched LoRA, attributed to gradient routing causing domain signals to concentrate in auxiliary modules.
- Scaling from 50M to 5B parameters, GRAM tracks data filtering: the gap between filtered and full models on removed capabilities widens with scale while retained performance remains close to baseline, and GRAM maintains similar scaling trends while amortising training cost across profiles.
- Cost advantage: supporting multiple capability profiles with GRAM incurs a single run cost independent of the number of profiles, yielding roughly a 5 times reduction in amortised training FLOPS in the five profile setting.
Limitations
Separation is limited when target capabilities overlap heavily with desired ones. Scaling trends are demonstrated up to 5B parameters but not at frontier production sizes. GRAM introduces implementation complexity for production systems and its interaction with later training steps such as instruction tuning or reinforcement learning from human feedback is untested. The mechanism behind partial labelling benefits and composability requires further investigation. Evaluation is based on cross entropy and limited elicitation budgets, which do not guarantee absolute tamper resistance.
Implications
For attackers, GRAM highlights both a mitigation pathway and an attack surface. If GRAM modules are known and present, an attacker could attempt targeted recovery by fine tuning with larger budgets than those tested. Conversely, GRAM enables defenders to restrict dangerous capabilities at inference by ablating modules, reducing the attack surface available in less trusted deployments. The work emphasises that restricting capabilities in weights rather than only outputs provides stronger resistance to simple elicitation, but raises governance questions about verification, tamper resistance and auditing of which modules are active.
