https://github.com/AlignmentResearch/tuned-lens

Tools for understanding how transformer predictions are built layer-by-layer
https://github.com/AlignmentResearch/tuned-lens

machine-learning pytorch transformers

Last synced: 15 days ago
JSON representation

Tools for understanding how transformer predictions are built layer-by-layer

• Host: GitHub
• URL: https://github.com/AlignmentResearch/tuned-lens
• Owner: AlignmentResearch
• License: mit
• Created: 2022-10-03T20:48:40.000Z (about 2 years ago)
• Default Branch: main
• Last Pushed: 2024-06-02T18:13:11.000Z (5 months ago)
• Last Synced: 2024-08-31T15:21:57.712Z (2 months ago)
• Topics: machine-learning, pytorch, transformers
• Language: Python
• Homepage: https://tuned-lens.readthedocs.io/en/latest/
• Size: 5.97 MB
• Stars: 405
• Watchers: 7
• Forks: 41
• Open Issues: 11
• Metadata Files:
  • Readme: README.md
  • License: LICENSE
  • Citation: CITATION.cff

Awesome Lists containing this project

README

        
# Tuned Lens 🔎



  









Tools for understanding how transformer predictions are built layer-by-layer.



This package provides a simple interface for training and evaluating __tuned lenses__. A tuned lens allows us to peek at the iterative computations a transformer uses to compute the next token.

## What is a Lens?



A lens into a transformer with _n_ layers allows you to replace the last _m_ layers of the model with an [affine transformation](https://pytorch.org/docs/stable/generated/torch.nn.Linear.html) (we call these affine translators). Each affine translator is trained to minimize the KL divergence between its prediction and the final output distribution of the original model. This means that after training, the tuned lens allows you to skip over these last few layers and see the best prediction that can be made from the model's intermediate representations, i.e., the residual stream, at layer _n - m_.

The reason we need to train an affine translator is that the representations may be rotated, shifted, or stretched from layer to layer. This training differentiates this method from simpler approaches that unembed the residual stream of the network directly using the unembedding matrix, i.e., the [logit lens](https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens). We explain this process and its applications in the paper [Eliciting Latent Predictions from Transformers with the Tuned Lens](https://arxiv.org/abs/2303.08112).

### Acknowledgments

Originally conceived by [Igor Ostrovsky](https://twitter.com/igoro) and [Stella Biderman](https://www.stellabiderman.com/) at [EleutherAI](https://www.eleuther.ai/), this library was built as a collaboration between FAR and EleutherAI researchers.

## Install Instructions

### Installing from PyPI

First, you will need to install the basic prerequisites into a virtual environment:

* Python 3.9+

* PyTorch 1.13.0+

Then, you can simply install the package using pip.

```

pip install tuned-lens

```

### Installing the container

If you prefer to run the training scripts from within a container, you can use the provided Docker container.

```

docker pull ghcr.io/alignmentresearch/tuned-lens:latest

docker run --rm tuned-lens:latest tuned-lens --help

```

## Contributing

Make sure to install the dev dependencies and install the pre-commit hooks.

```

$ git clone https://github.com/AlignmentResearch/tuned-lens.git

$ pip install -e ".[dev]"

$ pre-commit install

```

## Citation

If you find this library useful, please cite it as:

```bibtex

@article{belrose2023eliciting,

  title={Eliciting Latent Predictions from Transformers with the Tuned Lens},

  authors={Belrose, Nora and Furman, Zach and Smith, Logan and Halawi, Danny and McKinney, Lev and Ostrovsky, Igor and Biderman, Stella and Steinhardt, Jacob},

  journal={to appear},

  year={2023}

}

```

> **Warning**

> This package has not reached 1.0. Expect the public interface to change regularly and without a major version bumps.