Pretraining Unlearning with Topic-Conditioned Neuron Masks

A language model can learn useful patterns from a document and also memorize its particular wording. If we later want it to forget that document, can we remove the memory without losing what it learned more generally?

One way to approach this is to plan for forgetting during training. We explored a method called MemSinks, which gives each training sequence access to a consistent subset of designated “sink” neurons. These neurons can then be switched off at inference time.

Our experiments asked two questions. First, what happens if we activate the wrong sink neurons, rather than simply removing the right ones? Second, can the same idea work for a topic shared by many documents, including documents the model has never seen?

Giving memories somewhere to go

The difficulty is that a model’s internal components can support several kinds of behavior at once. Removing a component involved in memorization may also affect its ability to predict new text. Work on polysemanticity describes how individual neurons can represent multiple features; the MemSinks paper studies the related problem of memorization becoming entangled with general language abilities.

MemSinks changes how training examples use the network. In our implementation, 70% of each feed-forward layer’s hidden neurons are always active. The remaining 30% are sink neurons. A deterministic hash of the document ID selects 30% of those sink neurons to activate for that document. The rest are zeroed in the forward pass, so they also receive no gradient updates from that example.

Repeated encounters with a document therefore use the same subset of sink neurons. The always-active neurons see every document, while each sink neuron is exposed to a narrower subset. This creates an opportunity for document-specific information to become concentrated in the sinks without preventing the rest of the network from learning from those documents.

We compared four main ways to use the resulting model:

The inverse mask is the most interesting contrast. Dropout removes the document’s sink-neuron contribution. The inverse mask removes that contribution and introduces activity from other sinks. Our hypothesis was that this additional activity would interfere with prediction, rather than merely leave a gap where the memory used to be.

Forgetting a document—and adding interference

We started with a controlled memorization task using TinyStories. We selected 100 stories, repeated each 128 times, and mixed them with 20,000 stories that appeared once. We evaluated on both the repeated stories and 500 held-out stories.

Our measure was next-token prediction loss in bits per byte. Lower loss means better prediction. For the repeated stories, a rise in loss after an intervention is evidence that the model has lost some of its predictive advantage from memorization. For held-out stories, a rise in loss is a cost to generalization.

With the training mask active, loss on the repeated stories approached zero. Removing all sink neurons raised it from about 0.005 to 0.83, close to the model’s loss on held-out stories under dropout. The model had lost much of the advantage it gained from seeing those exact sequences repeatedly.

The inverse mask went further. Loss on repeated stories rose to 1.62, exceeding the 1.20 loss on held-out stories. This is consistent with interference from unrelated sink neurons adding to the effect of removing the relevant ones. It is not a free improvement: validation loss also rose, from about 0.92 with the training mask to 1.20 with the inverse mask.

There are two comparisons to keep separate. Dropping the sinks did not substantially worsen validation performance relative to our model’s training-mask mode. But the localization-trained model was already behind a conventionally trained baseline: its validation loss under dropout was about 21% higher. The cost of training for localization matters alongside the cost of the later intervention.

From memorized documents to shared topics

Document IDs make it possible to target particular training sequences. They are less useful when the thing we want to remove spans many different documents. We therefore changed what determined the mask: instead of hashing a document ID, we hashed a topic label.

TinyStories uses a noun, verb, and adjective to generate each story. We used the noun as the topic label, so different stories generated around the same noun activated the same subset of sink neurons. We selected 20 topics and sampled 1,000 stories for each, then added 20,000 stories from other topics.

Crucially, every story used for evaluation was unseen during training. For each of the 20 emphasized topics, we evaluated on 100 new stories. We also used a separate validation set of 500 unseen stories. This tests whether the masks affect topic-related knowledge that transfers to new examples, rather than only recall of particular strings.

With the topic mask active, the model predicted stories from emphasized topics better than the validation set. Dropping the sinks reduced that advantage: the loss gap shrank from 0.21 to 0.067. Meanwhile, validation loss changed only slightly, from 1.107 to 1.117.

We interpret this as evidence that the masks localized some topic-related information. The effect was not complete: the loss gap remained above zero after dropout. And the baseline comparison again matters. The standard model achieved validation loss of about 0.88, compared with 1.107 for the topic-mask model—roughly a 25% increase.

Why the training procedure matters

We also tested whether masking gradients during the backward pass could produce the same effect while leaving the forward pass unchanged. This comparison was inspired by gradient routing, which uses data-dependent gradient masks to localize computation.

Our two backward-only variants behaved differently. When we blocked updates to selected parameters while allowing gradients to flow through them, performance tracked standard training and dropping sink neurons did not produce the intended forgetting. When we blocked gradients at activations, we saw an intermediate result between standard training and forward masking.

These observations suggest that merely preventing particular weights from updating was insufficient in our setup. One possible explanation is coadaptation: neurons learn to rely on one another’s contributions during the forward pass, even when some of their updates are blocked. That makes removing those neurons later more disruptive. This is an interpretation of our comparisons, rather than a general result about all gradient-routing methods.

We also tried rescaling activation magnitudes at inference time to account for the different numbers of active sink neurons. The curves were somewhat smoother, but we did not observe a substantial performance difference. In these experiments, which neurons contributed mattered more than this change in their overall scale.

What this tells us about unlearning

Our results suggest that training can make some information easier to target later. The document experiment separated much of the predictive benefit of memorizing particular sequences. The topic experiment extended that idea to information shared across examples, with effects that carried over to unseen stories.

The distinction between dropout and the inverse mask also helps clarify what “forgetting” can mean in an experiment. A model can lose useful information, or its remaining computation can be disrupted. Both raise loss. The inverse-mask results are consistent with both mechanisms contributing, so a larger loss increase alone should not be treated as cleaner removal.

There are several limits to this evidence. We used a relatively small model and a controlled story dataset, with topic labels available before training. Our primary measure was next-token loss, not extraction resistance or the difficulty of relearning the removed information. The model also paid a substantial validation-performance cost relative to standard training. These experiments do not establish a general privacy guarantee or a reliable way to erase an arbitrary concept from an existing model.

The next questions are practical: can localization preserve more of the baseline model’s performance, can it work with less tidy labels, and does the apparent forgetting survive attempts to recover the information? For us, the useful result is that the choice of training mask gave us some control over where document- and topic-related information was stored—and made the limits of that control measurable.

Experimental details