taufeeque9 / codebook-features Goto Github PK

In codebook_features/models.py, I can see a method for attaching codebooks to each attention block's query, key and value vectors:

https://github.com/taufeeque9/codebook-features/blob/a37ea8fe7d4d39298aaea042a078d09401396edc/codebook_features/models.py#L1439C1-L1460C36

After training a model with these codebooks attached though, it does not seem possible to convert this model to a HookedTransformer model (doing so raises AttributeError: 'HookedTransformer' object has no attribute 'qkv_key'). What is the status of using qkv codebooks and converting to a HookedTransformer model currently? Happy to write a PR if this needs integrating with HookedTransformerCodebookModel class to work.

Currently, there are two ways to format keys in codebook-related dictionaries: 1) "layer{x}_{cb_at}_gcb{y}" (adv name), and 2) "layer{x}_head{y}" (base name). The Second was introduced for convenience but shouldn't be used as it doesn't incorporate the cb_at parameter. So for a model, if in a given layer there are two places codebooks are applied at, the second way wouldn't be able to distinguish between the two codebooks. So this format can be removed and we can use 1) universally.

The rotary embedding issue is transformer_lens breaks the TinyStories codebook models. The issue was introduced in version 1.8.0 and so we have frozen transformer_lens' version at 1.7.0. Bump the version and update the mod_model_classes.py file once the issue gets resolved.