We assign the same neuron mask to stories about a shared topic and test whether removing the sink neurons impairs predictions on unseen stories from that topic.
We implement these masks in NanoChat, building on MemSinks, which assigns training documents consistent subsets of designated “sink” neurons. We also compare inference-time masks and training procedures to examine how they affect forgetting and generalization.
Main findings
Topic-based masks induce partial forgetting on unseen stories when the sink neurons are removed. In a separate document experiment, activating unrelated sink neurons increases loss on memorized sequences more than removing all sinks, but also increases validation loss. Forward masking and backward-only gradient masking produce different degrees of localization.
Neuron masking during pretraining
A neuron can contribute to both memorization and general language abilities. Removing neurons that encode a particular document may therefore also impair performance on new text. We use training-time masks to encourage document- or topic-specific information to concentrate in a designated subset of neurons.
In our implementation, 70% of each feed-forward layer’s hidden neurons remain active for all examples. The remaining 30% are sink neurons. A deterministic hash of the document or topic ID selects 30% of these sink neurons to activate for that example. The other sinks are zeroed in the forward pass and receive no gradient updates from that example.
Examples with the same ID therefore use the same mask. We hypothesize that this makes the sink neurons specialize in information associated with that ID, while the always-active neurons learn patterns shared across the training data.
After training, we compare four inference modes:
| Inference mode | Which sink neurons are active? |
|---|---|
| Mask | The same subset used for this document during training. |
| Dropout | None. All sink neurons are switched off. |
| All | Every sink neuron, including those masked for this document during training. |
| Inverse mask | Only the sink neurons that were masked for this document during training. |
We hypothesize that an inverse mask adds interference from unrelated sink neurons, increasing loss beyond the effect of removing the document’s assigned neurons.
We measure next-token prediction loss in bits per byte; lower values indicate better prediction. An increase on the targeted data after removing neurons is evidence of forgetting. An increase on validation data measures the accompanying cost to generalization.
Topic localization on unseen examples
We extend document-based masks to topics by hashing topic labels instead of document IDs. TinyStories generates stories from a noun, verb, and adjective. We use the noun as the topic label, so different stories about the same topic activate the same subset of sink neurons.
We emphasize 20 topics during training, using 1,000 distinct stories per topic alongside 20,000 stories from other topics. All evaluation stories are unseen during training: 100 new stories for each emphasized topic and a separate validation set of 500 stories. This tests whether the masks localize information that transfers across examples.
With the topic mask active, prediction loss is lower on stories from emphasized topics than on the validation set. Removing the sinks reduces this advantage: the gap between validation and emphasized-topic loss falls from 0.21 to 0.067. Validation loss changes only slightly, from 1.107 to 1.117.
| Measure | Mask | Dropout |
|---|---|---|
| Validation loss | 1.107 | 1.117 |
| Gap between validation and emphasized-topic loss | 0.21 | 0.067 |
We interpret these results as evidence of partial topic localization. The loss gap remains above zero after removal. Training with topic masks also incurs a performance cost: validation loss is 1.107, compared with approximately 0.88 under standard training, an increase of about 25%.
Document memorization and inverse masks
In the document experiment, we select 100 TinyStories sequences, repeat each 128 times, and mix them with 20,000 stories presented once. We evaluate on the repeated sequences and 500 held-out stories.
With the training mask active, loss on repeated sequences approaches zero. Removing all sink neurons raises this loss from approximately 0.005 to 0.83, close to held-out loss under the same intervention. This removes much of the prediction advantage associated with memorizing those sequences.
The inverse mask raises repeated-sequence loss further, to 1.62, exceeding the held-out loss of 1.20. These results are consistent with interference from unrelated sink neurons. The intervention also degrades generalization: held-out loss rises from approximately 0.92 with the training mask to 1.20 with the inverse mask.
| Mode | Repeated-story loss | Held-out-story loss |
|---|---|---|
| Mask | ≈ 0.005 | ≈ 0.92 |
| Dropout | ≈ 0.83 | Close to repeated-story loss |
| Inverse mask | ≈ 1.62 | ≈ 1.20 |
Dropping the sinks causes little validation degradation relative to the training-mask mode of the same model. Relative to standard training, however, validation loss under dropout is approximately 21% higher. This separates the cost of removing neurons from the cost of training the model with masks.
Forward masking and gradient routing
Gradient routing uses data-dependent masks during backpropagation to control which parts of a network learn from each example. We compare our forward-masking implementation with two backward-only variants. Both leave the forward computation unchanged but block different parts of the learning update.
Parameter-update masking. We prevent selected weights from updating while preserving gradient flow through the network. Performance tracks standard training, and removing the sink neurons does not produce the intended forgetting.
Activation-gradient masking. We block gradients at selected neuron activations. This produces an intermediate result between standard training and forward masking, in both forgetting and model degradation.
These results suggest that blocking parameter updates alone is insufficient for localization in our setup. We hypothesize that coadaptation contributes to the difference: neurons can learn to depend on other neurons that remain active during the forward pass, even when their updates are blocked. Removing those neurons at inference can then disrupt the computation.
We also test activation rescaling to account for changes in the number of active neurons at inference. Rescaling slightly smooths the loss curves but does not substantially change performance.
Connections to later AI safety research
Research context added October 2026. Our experiments were completed in December 2025.
Later AI safety work uses training-time localization to restrict access to sensitive knowledge and remove unwanted behavior. The relevant objective is selective control: removing a targeted capability or behavior while preserving useful performance elsewhere.
Capability access control. Gradient Routed Auxiliary Modules (GRAM, July 2026), developed at AE Studio in collaboration with Anthropic, use data-dependent forward and backward routing to isolate knowledge in removable auxiliary modules. The authors test topic removal in synthetic stories and access control for dual-use knowledge, comparing the resulting models with models trained on filtered data. Our topic experiment studies a related localization problem using shared masks within the existing network. GRAM uses separate auxiliary modules and evaluates robustness to adversarial fine-tuning, which our experiments did not test.
Removing reward-hacking behavior. Gradient Routing Adapter Fine-Tuning (GRAFT, October 2026) routes updates from detected reward hacks into a removable adapter. In its controlled experiments, removing that adapter reduces reward hacking, including some behavior missed by the training monitor. GRAFT also trains on separate anchor tasks with the removable adapter disabled to maintain useful performance after its removal. This addresses a practical issue raised by our masking comparisons: the model must continue to work when part of the computation used during training is absent at deployment.
More faithful reasoning. Faithfulness as Information Flow (May 2026) studies a related use of masking. Its interventions include attention masking and backward-only gradient masking during reinforcement learning, with the aim of making written reasoning more closely reflect the computation producing an answer. The authors report that shortcut and reward-hacking behavior becomes more visible in the reasoning trace in their tested settings. Here the safety objective is monitorability rather than knowledge removal.
Our comparisons distinguish three interventions that can otherwise sound interchangeable: masking forward activations, blocking gradients at activations, and blocking parameter updates. They change different parts of the learning process and did not produce equivalent forgetting in our experiments.
Interpretation and limitations
Our results support localizing some information shared across examples during pretraining. Topic-based masks extend the effect beyond memorized sequences, while the inference-mask comparisons show that both removal and interference can affect apparent forgetting.
Higher loss alone does not establish that information has been erased. In particular, an inverse mask may disrupt prediction by activating unrelated neurons. We measured next-token loss, but did not test resistance to extraction attacks or relearning. The results therefore do not establish a privacy guarantee.
The experiments use a 398-million-parameter model, a controlled story dataset, and topic labels available before training. Forgetting is partial, and masked training has higher validation loss than standard training. These limits constrain what the results establish about removing arbitrary knowledge from larger models.
Experimental details
Model, data, and evaluation settings
We use NanoChat with 398 million parameters: 20 layers, 1,280-dimensional embeddings, 10 attention heads, and squared ReLU activations in the feed-forward layers. We use a 500-step training configuration, evaluation every 25 steps, a batch size of 32, and AdamW/Muon optimizers. The document results plot shows approximately 300 steps; the topic results plot shows 500.
Neuron allocation. We keep 70% of hidden neurons always active and designate 30% as sinks. Each document or topic activates a deterministic 30% subset of the sink neurons, selected by hashing its ID.
Document data. The training corpus contains 100 stories repeated 128 times each, plus 20,000 stories presented once: 32,800 story presentations in total. Evaluation uses the 100 repeated stories and 500 held-out stories.
Topic data. Training uses 1,000 distinct stories from each of 20 emphasized topics, plus 20,000 stories from other topics. Evaluation uses 100 new stories for each emphasized topic and a 500-story validation set.
Comparisons. We compare standard training, forward masking, backward-only masking at activations, and backward-only masking of parameter updates. The main inference modes are mask, dropout, all, and inverse mask. We also examine highest-activation sink dropout and activation rescaling.
Measurement. Loss is measured in bits per byte, a tokenizer-normalized measure of next-token prediction error. All experimental results reported here are from December 2025.
References
Sources for the original experiments and the later research discussed above:
- Extracting Training Data from Large Language Models — Carlini et al., 2021. Demonstrates extraction of memorized training text from GPT-2.
- Can Neural Network Memorization Be Localized? — Maini et al., 2023. Studies memorization in small sets of neurons across layers and introduces example-tied dropout.
- Polysemanticity and Capacity in Neural Networks — Scherlis et al., 2022. Examines how neurons can represent multiple unrelated features.
- Gradient Routing: Masking Gradients to Localize Computation in Neural Networks — Cloud et al., 2024. Uses gradient masks to influence where computation is learned.
- Memorization Sinks: Isolating Memorization during LLM Training — Ghosal, Maini, and Raghunathan, 2025. The training method we replicated and extended.
- Modular Pretraining Enables Access Control — Roland et al., 2026. Introduces Gradient Routed Auxiliary Modules.
- Gradient Routing Localizes and Removes Unintended Behaviors in RL — Hu et al., 2026. Introduces GRAFT for localizing unwanted behavior during post-training.
- Faithfulness as Information Flow: Evaluating and Training Faithful Chain-of-Thought Reasoning — Jia, Benton, and Easley, 2026. Studies training interventions that make reasoning more faithful and monitorable.