This is an informal research note. It is the result of a few-day exploration into RMU through the lens of model internals. Code to reproduce the main result is available here.

This work was produced as part of Ethan Perez's stream in the ML Alignment & Theory Scholars Program - Summer 2024 Cohort. Thanks to Nina Panickssery, Mrinank Sharma, and Fabien Roger for helpful discussion.

Summary

We investigate RMU, a recent unlearning method proposed by Li et al. (2024), through the lens of model internals. Through this lens, we explain that RMU mostly works by flooding the residual stream with "junk" in hazardous contexts, resulting in incoherence. We then propose a simple intervention to "clear the junk" from the residual stream. This intervention mostly restores the model's coherence in hazardous contexts, and recovers a significant proportion (but not all) of its original hazardous knowledge. This suggests that the effectiveness of RMU can be understood roughly in two pieces: (1) a shallow mechanism, where the residual stream is flooded with junk; and (2) a deeper mechanism, where even after the junk is cleared, knowledge is still inaccessible.

Performing a simple directional ablation on the residual stream of an RMU model recovers a large proportion of hazardous knowledge.

What is RMU?

Representation Misdirection for Unlearning (RMU) is a state-of-the-art unlearning method presented by Li et al. (2024).

In the unlearning paradigm, we would like the model to unlearn (or "forget") some hazardous knowledge. At the same time, we would also like to make sure the model retains non-hazardous knowledge, so that the model remains useful.

This partition of knowledge is usually specified by constructing a "forget" dataset $D_{forget}$ , consisting of the hazardous knowledge to be unlearned, and a "retain" dataset $D_{retain}$ , consisting of non-hazardous knowledge to be retained.

Let $M$ denote our original model. RMU specifies a method for fine-tuning $M$ on $D_{forget}$ and $D_{retain}$ in order to obtain a modified model $M^{'}$ satisfying the unlearning objective.

The main idea of RMU is as follows:

On hazardous data, the internal activations of $M^{'}$ should be scrambled.
On non-hazardous data, the internal activations of $M^{'}$ should be unchanged, i.e. close to those of the original model $M$ .

These two ideas are conc...

Posts

Wikitag Contributions

Comments