Motivation

While lurking LessWrong, I read Apple's "The Super Weight in Large Language Models" paper and OpenAI's "Weight-sparse transformers have interpretable circuits" paper. My curiosity was simple, whether it is possible to bridge the core ideas derived from the two papers to explore a new direction, namely:

If I destroy the model's ability to generate responses by ablating a "superweight", can I then fix it with a tiny row-level patch?

I do not own a GPU so I tried testing this on my MacBook (I ran it overnight!). Rather than focusing on a toy model, I leveraged a real model (i.e OLMo-1B) following advice of @Neel Nanda to get messy with real weights.

Setup

I used AllenAI's OLMo-1B-0724-hf, running the model on my CPU with 16GB RAM. Inspired by the LLMSuperWeight repository, I nabbed the two OLMo-1B superweights featured there. More specifically, I chose and zeroed out model.layers.1.mlp.down_proj.weight[1764, 1710].

Initial Findings (Validating Apple's Results)

Trying out the ablated model on a Wikitext-2 slice yielded the following results:

OLMo Prompt & Outputs (Ablated)

Prompt: "Paris is in France. Tokyo is in".

Output: "Paris is in France. Tokyo is in the ocean. The ocean. The ocean. the. the...".

Prompt: "The capital of Canada is".

Output: "The capital of Canada is a lot of people. leaves of of of of...".

Crunching the numbers (i.e NLL, PPL, KL), I observed that Perplexity skyrocketed from 17.4 to 2884.3, and that the model spiralled towards pathological behaviour whilst showcasing unusual confabulations. More noticeably, when asked about Tokyo, the model claimed it is in the ocean on top of outputting "mar, mar, mar" many times in succession.

The Patch

To introduce a fix for the broken model, I decided to freeze the entire model and introduce a single, trainable vector $Δ_{r o w}$ . I added this vector to the superweight's row 1764 of the down-projection matrix, like so:

$W^{'} = W_{b r o k e n} + e_{r o w} \otimes Δ_{r o w}$

After doing this, I trained the model utilizing the original model as a teacher and the broken model (plus patch) as a student on train[:2%] of the Wikitext-2-raw-v1 dataset, filtering by non-empty lines. I treated the token-shifted KL divergence between teacher and student logits as the loss. Severe compute constraints (i.e only CPU available) led to only two epochs of training with batch size 2 (thanks to my RAM) and 200 st...

Posts

Wikitag Contributions

Comments