TransformerLens Dependency, Hidden States, and Hooks

[tretomaszewski]

There is some chatter regarding this library's over-reliant dependence on TransformerLens. (See #10 (comment))

I agree that it would be best to move away from TransformerLens (TL). One suggestion has been to use the technique from https://github.com/Sumandora/remove-refusals-with-transformers, which uses the "hidden_states" parameter offered from the Model HuggingFace transformer library.

Using hidden_states might be a good quick fix. It is mostly compatible across many models and presumably works for these purposes.

However, there are some quirks that deserve awareness.
(Personally, it feels slightly "hacky" and limiting for future work, but I may be suffering "purism".) It also removes the ability to interact with the resid_mid between the attention and mlp in each block.

It seems there are slight discrepancies between the resid_pre and resid_post subcache tensor values in TL and the hidden_states in HF. It is a tiny difference, but seems to grow over layers. And then the hidden_state[-1] or last_hidden_state is completely different from TL's last layer's resid_post. I have not been able to determine why.

If this is problematic method would be to roll-our-own hooks into the residual streams. I had started this prior to finding this library. This would be most direct (allow us to use the resid_mid too), but would probably require re-implementing the multi-model configurations found in TransformerLens.

I hereby open this up for discussion!

[jukofyork]

Yeah, you can set output_attentions=True and I thought it might be possible to calculate resid_mid via looking at the previous hidden_state, but didn't look much deeper and the docs aren't really all that clear exactly if the output_attentions are before or after the o_proj.weight tensor is applied...

It seems there are slight discrepancies between the resid_pre and resid_post subcache tensor values in TL and the hidden_states in HF. It is a tiny difference, but seems to grow over layers. And then the hidden_state[-1] or last_hidden_state is completely different from TL's last layer's resid_post. I have not been able to determine why.

One thing to double check is that the layer induces are lining up: hidden_states in HF returns 1 more set of hidden states than there are layers and the 1st of these is the hidden state from before the first layer. IIRC, the code in abliterator was off by one and there was a comment in the code about "nan in the first layer" which was actually caused by this.

Regarding the second point:

When passing output_hidden_states=True you may expect the outputs.hidden_states[-1] to match outputs.last_hidden_states exactly. However, this is not always the case. Some models apply normalization or subsequent process to the last hidden state when it’s returned.

See: https://huggingface.co/docs/transformers/en/main_classes/output

[Sumandora]

It seems there are slight discrepancies between the resid_pre and resid_post subcache tensor values in TL and the hidden_states in HF. It is a tiny difference, but seems to grow over layers. And then the hidden_state[-1] or last_hidden_state is completely different from TL's last layer's resid_post. I have not been able to determine why.

Another problem that many people seem to overlook is that my code actually generates one token and then grabs the attentions. This is because HF Transformers has a stupid limit on at least generating a single token. Also floating-point operations are not commutative, so if HF Transformers does operations different to TransformerLens, then there are errors, which increase with every layer. My code is not supposed to be feature-complete or anything, as you said. I just wanted to have support for more models, I hope my POC will make that a reality soon... or arguably already made it one.

[FailSpy]

The full paper has been released

And an expanded codebase has been published here
In that, they actually use PyTorch hooks, and have actually a very nice little util class specifically for hooking into models.

[enn-nafnlaus]

The full paper has been released

And an expanded codebase has been published here In that, they actually use PyTorch hooks, and have actually a very nice little util class specifically for hooking into models.

Are you saying that we should be using that instead?

I've been unfortunately immensely frustrated while trying to use your library. I start with some "normal" model and what I get out is this TransformerLens format model. Then what? As far as I can tell I can't do anything with it. I can't say convert it to GGUFs, it's not like llama.cpp understands that format.

What exactly should we be doing to ablate models these days? The new DeepSeek models are incredible, but basically have the Great Firewall of China trained into them. I created a dataset to abliterate it (https://huggingface.co/datasets/Nafnlaus/ShrimpMoss_Chinese_Censorship_Abliteration). But I'm struggling to actually apply it.