This is a write up of "Multiple Streams of Knowledge Retrieval: Enriching and Recalling in Transformers", work with David Reber, Sean Richardson & Ari Holtzman. Code is available here. This is cross-posted from https://toddnief.com/articles/dynamic-weight-grafting/

When a new pope is elected and we want an LLM to answer “Who’s the pope?” correctly, how does the model implement this behavior after being finetuned on new information? When we add new relationship information to an LLM, is that information added just to the entity, is it "enriched" at the entity token position in lower and middle layers, or is it recalled in response to the entity in higher layers closer to next token prediction?

To answer these questions, we introduce dynamic weight grafting: swapping the weights of a pretrained model for the weights of a model that has undergone supervised finetuning dynamically during generation.

Dynamic Weight Grafting vs. Activation Patching

Previous approaches to localizing relation knowledge have either used variants of activation patching (replacing model activations, some subspace of model activations, etc. with activations from a different run) or ablations (removing activations, mean ablating activations, etc.) to see which components contribute to next token prediction.

Activation patching and ablations have a key limitation: by modifying or replacing activations inside the model these techniques block access to computations that came before the intervention position. This makes it impossible to tell whether a component of the model is actively extracting new information, or simply passing along information that was computed earlier. As a result, we can't isolate which mechanisms are truly responsible for retrieving finetuned relation knowledge.

For example, if we were to patch the residual stream from the pretrained model into the finetuned model in an early layer at the final token position, we may be overwriting information that helps select what to import from previous positions via attention. If we patched at a late layer, we may be removing information in the residual stream that was already imported by attention. Model computation both reacts to existing information and adds new information.

To localize behavior to model components, we would like to use the same residual stream, but perform a different computation. We propose dynamic w...

Posts

Wikitag Contributions

Comments

Posts

Wikitag Contributions

Comments