Package: linear-relational
Linear-RelationalLinear Relational Embeddings (LREs) and Linear Relational Concepts (LRCs) for LLMs using PyTorch and Huggingface Transformers.Full docs:https://chanind.github.io/linear-relationalAboutThis library provides utilities and PyTorch modules for working with LREs and LRCs. LREs estimate the relation between a subject and object in a transformer language model (LM) as a linear map.This library assumes you're working with sentences with a subject, relation, and object. For instance, in the sentence: "Lyon is located in the country of France" would have the subject "Lyon", relation "located in country", and object "France". A LRE models a relation like "located in country" as a linear map consisting of a weight matrix $W$ and a bias term $b$, so a LRE would map from the activations of the subject (Lyon) at layer $l_s$ to the activations of the object (France) at layer $l_o$. So:$$
LRE(s) = W s + b
$$LREs can be inverted using a low-rank inverse, shown as $LRE^{\dagger}$, to estimate $s$ from $o$:$$
LRE^{\dagger}(o) = W^{\dagger}(o - b)
$$Linear Relational Concepts (LRCs) represent a concept $(r, o)$ as a direction vector $v$ on subject tokens, and can act like a simple linear classifier. For instance, while a LRE can represent the relation "located in country", we could learn a LRC for "located in the country: France", "located in country: Germany", "located in country: China", etc... This is just the result from passing in an object activation into the inverse LRE equation above.$$
v_{o} = W^{\dagger}(o - b)
$$For more information on LREs and LRCs, check outthesepapers.Installationpip install linear-relationalUsageThis library assumes you're using PyTorch with a decoder-only generative language model (e.g. GPT, LLaMa, etc...), and a tokenizer from Huggingface.Training a LRETo train a LRE for a relation, first collect prompts which elicit the relation. We provide aPromptclass to represent this data, and aTrainerclass to make training a LRE easy. Below, we train a LRE to represent the "located in country" relation.fromtransformersimportGPT2LMHeadModel,GPT2TokenizerFastfromlinear_relationalimportPrompt,Trainer# We load a generative LM from huggingface. The LMHead must be included.model=GPT2LMHeadModel.from_pretrained("gpt2")tokenizer=GPT2TokenizerFast.from_pretrained("gpt2")# Prompts consist of text, an answer, and subject.# The subject must appear in the text. The answer# is what the model should respond with, and corresponds to the "object"prompts=[Prompt("Paris is located in the country of","France",subject="Paris"),Prompt("Shanghai is located in the country of","China",subject="Shanghai"),Prompt("Kyoto is located in the country of","Japan",subject="Kyoto"),Prompt("San Jose is located in the country of","Costa Rica",subject="San Jose"),]trainer=Trainer(model,tokenizer)lre=trainer.train_lre(relation="located in country",subject_layer=8,# subject layer must be before the object layerobject_layer=10,prompts=prompts,)Working with a LREA LRE is a PyTorch module, so once a LRE is trained, we can use it to predict object activations from subject activations:object_acts_estimate=lre(subject_acts)We can also create a low-rank estimate of the LRE:low_rank_lre=lre.to_low_rank(50)low_rank_obj_acts_estimate=low_rank_lre(subject_acts)Finally we can invert the LRE:inv_lre=lre.invert(rank=50)subject_acts_estimate=inv_lre(object_acts)Training LRCs for a relationTheTrainercan also create LRCs for a relation. Internally, this first create a LRE, inverts it, then generates LRCs from each object in the relation. Objects refer to the answers in the prompts, e.g. in the example above, "France" is an object, "Japan" is an object, etc...fromtransformersimportGPT2LMHeadModel,GPT2TokenizerFastfromlinear_relationalimportPrompt,Trainer# We load a generative LM from huggingface. The LMHead must be included.model=GPT2LMHeadModel.from_pretrained("gpt2")tokenizer=GPT2TokenizerFast.from_pretrained("gpt2")# Prompts consist of text, an answer, and subject.# The subject must appear in the text. The answer# is what the model should respond with, and corresponds to the "object"prompts=[Prompt("Paris is located in the country of","France",subject="Paris"),Prompt("Shanghai is located in the country of","China",subject="Shanghai"),Prompt("Kyoto is located in the country of","Japan",subject="Kyoto"),Prompt("San Jose is located in the country of","Costa Rica",subject="San Jose"),]trainer=Trainer(model,tokenizer)concepts=trainer.train_relation_concepts(relation="located in country",subject_layer=8,object_layer=10,prompts=prompts,max_lre_training_samples=10,inv_lre_rank=50,)Causal editingOnce we have LRCs trained, we can use them to perform causal edits while the model is running. For instance, we can perform a causal edit to make the model output that "Shanghai is located in the country of France" by subtracting the "located in country: China" concept from "Shanghai" and adding the "located in country: France" concept. We can use theCausalEditorclass to perform these edits.fromlinear_relationalimportCausalEditorconcepts=trainer.train_relation_concepts(...)editor=CausalEditor(model,tokenizer,concepts=concepts)edited_answer=editor.swap_subject_concepts_and_predict_greedy(text="Shanghai is located in the country of",subject="Shanghai",remove_concept="located in country: China",add_concept="located in country: France",edit_single_layer=8,magnitude_multiplier=3.0,predict_num_tokens=1,)print(edited_answer)# " France"Single-layer vs multi-layer editsAbove we performed a single-layer edit, only modifying subject activations at layer 8. However, we may want to perform an edit at all subject layers at the same time instead. To do this, we can passedit_single_layer=Falsetoeditor.swap_subject_concepts_and_predict_greedy(). We should also reduce themagnitude_multipliersince now we're going to make the edit at every layer, if we use too large of a multiplier we'll drown out the rest of the activations in the model. Themagnitude_multiplieris a hyperparam that requires tuning depending on the model being edited.fromlinear_relationalimportCausalEditorconcepts=trainer.train_relation_concepts(...)editor=CausalEditor(model,tokenizer,concepts=concepts)edited_answer=editor.swap_subject_concepts_and_predict_greedy(text="Shanghai is located in the country of",subject="Shanghai",remove_concept="located in country: China",add_concept="located in country: France",edit_single_layer=False,magnitude_multiplier=0.1,predict_num_tokens=1,)print(edited_answer)# " France"Concept matchingWe can use learned concepts (LRCs) to act like classifiers and match them against subject activations in sentences. We can use theConceptMatcherclass to do this matching.fromlinear_relationalimportConceptMatcherconcepts=trainer.train_relation_concepts(...)matcher=ConceptMatcher(model,tokenizer,concepts=concepts)match_info=matcher.query("Beijing is a northern city",subject="Beijing")print(match_info.best_match.concept)# located in country: Chinaprint(match_info.best_match.score)# 0.832AcknowledgementsThis library is inspired by and uses modified code from the following excellent projects:Locating and Editing Factual Associations in GPTLinearity of Relation Decoding in Transformer LMsContributingAny contributions to improve this project are welcome! Please open an issue or pull request in this repo with any bugfixes / changes / improvements you have!This project usesBlackfor code formatting,Flake8for linting, andPytestfor tests. Make sure any changes you submit pass these code checks in your PR. If you have trouble getting these to run feel free to open a pull-request regardless and we can discuss further in the PR.LicenseThis code is released under a MIT license.CitationIf you use this library in your work, please cite the following:@article{chanin2023identifying,title={Identifying Linear Relational Concepts in Large Language Models},author={David Chanin and Anthony Hunter and Oana-Maria Camburu},journal={arXiv preprint arXiv:2311.08968},year={2023}}