ShichangZh/LRE

How do transformer LMs encode relations?

Linearity of Relation Decoding in Transformer LMs

This repository contains the code used for the experiments in the paper Linearity of Relation Decoding in Transformer LMs.

How an LM decodes a relation such as (Miles Davis, plays the instrument, trumpet) involves a sequence of non-linear computations spanning multiple layers. But, in this work we show that for a subset of relations this highly non-linear decoding procedure can be approximated by a simple linear transformation ($\textbf{LRE}$) on the subject representation $\mathbf{s}$ at some intermediate layer.

Please check lre.baulab.info for more information.

Setup

All code is tested on MacOS Ventura (>= 13.1) and Ubuntu 20.04 using Python >= 3.10. It uses a lot of newer Python features, so the Python version is a strict requirement.

To run the code, create a virtual environment with the tool of your choice, e.g. conda:

conda create --name relations python=3.10

Then, after entering the environment, install the project dependencies:

python -m pip install invoke
invoke install

$\mathbf{LRE}$ Demo

demo/demo.ipynb shows how to get $LRE$ approximation for a certain relation. This notebook also explains the setup of our evaluation metrics $faithfulness$ and $causality$.

demo/attribute_lens.ipynb demonstrates Attribute Lens, which is motivated by the idea that a hidden state $\mathbf{h}$ may contain pieces of information beyond the prediction of the immediate next token. And, an $\text{LRE}$ can be used to extract a certain attribute from $\mathbf{h}$ without relevant textual context. In the figure below $\text{LRE}$ approximating the relation country-capital is applied on hidden state $\mathbf{h}$ after different layers in different token positions.

How to Cite

@article{hernandez2023linearity,
    title={Linearity of Relation Decoding in Transformer Language Models}, 
    author={Evan Hernandez and Arnab Sen Sharma and Tal Haklay and Kevin Meng and Martin Wattenberg and Jacob Andreas and Yonatan Belinkov and David Bau},
    year={2023},
    eprint={2308.09124},
    archivePrefix={arXiv},
    primaryClass={cs.CL}
}