Visualizations (with more comments) (#248)

* Revert "Revert "Sweep Visualizations" (#247)"

This reverts commit 9324f82.

* Clarifying comments and docstrings

* Additional comments on utils

* add sweeps dir

* Actually use self.sweeps

* Very basic support for multiple sweeps

* Remove unittest on test

* Annotated command dataclass, added comments, trivial refactor

* Make multiplots actually correct

* Dataclasses are good

* Always use relative imports when possible

* Remove unnecessary utils file

* Simplify everything by dropping support for plotting multiple sweeps together; deaggregate datasets in the table

* Use the layer with the best average AUROC across datasets

---------

Co-authored-by: Walter Laurito <[email protected]>
Co-authored-by: Nora Belrose <[email protected]>

README.md

@@ -2,76 +2,115 @@
 
 **WIP: This codebase is under active development**
 
-Because language models are trained to predict the next token in naturally occurring text, they often reproduce common human errors and misconceptions, even when they "know better" in some sense. More worryingly, when models are trained to generate text that's rated highly by humans, they may learn to output false statements that human evaluators can't detect. We aim to circumvent this issue by directly [**eliciting latent knowledge**](https://docs.google.com/document/d/1WwsnJQstPq91_Yh-Ch2XRL8H_EpsnjrC1dwZXR37PC8/edit) (ELK) inside the activations of a language model.
-
-Specifically, we're building on the **Contrastive Representation Clustering** (CRC) method described in the paper [Discovering Latent Knowledge in Language Models Without Supervision](https://arxiv.org/abs/2212.03827) by Burns et al. (2022). In CRC, we search for features in the hidden states of a language model which satisfy certain logical consistency requirements. It turns out that these features are often useful for question-answering and text classification tasks, even though the features are trained without labels.
+Because language models are trained to predict the next token in naturally occurring text, they often reproduce common
+human errors and misconceptions, even when they "know better" in some sense. More worryingly, when models are trained to
+generate text that's rated highly by humans, they may learn to output false statements that human evaluators can't
+detect. We aim to circumvent this issue by directly [**eliciting latent knowledge
+**](https://docs.google.com/document/d/1WwsnJQstPq91_Yh-Ch2XRL8H_EpsnjrC1dwZXR37PC8/edit) (ELK) inside the activations
+of a language model.
+
+Specifically, we're building on the **Contrastive Representation Clustering** (CRC) method described in the
+paper [Discovering Latent Knowledge in Language Models Without Supervision](https://arxiv.org/abs/2212.03827) by Burns
+et al. (2022). In CRC, we search for features in the hidden states of a language model which satisfy certain logical
+consistency requirements. It turns out that these features are often useful for question-answering and text
+classification tasks, even though the features are trained without labels.
 
 ### Quick **Start**
 
-Our code is based on [PyTorch](http://pytorch.org) and [Huggingface Transformers](https://huggingface.co/docs/transformers/index). We test the code on Python 3.10 and 3.11.
+Our code is based on [PyTorch](http://pytorch.org)
+and [Huggingface Transformers](https://huggingface.co/docs/transformers/index). We test the code on Python 3.10 and
+3.11.
 
-First install the package with `pip install -e .` in the root directory, or `pip install -e .[dev]` if you'd like to contribute to the project (see **Development** section below). This should install all the necessary dependencies.
+First install the package with `pip install -e .` in the root directory, or `pip install -e .[dev]` if you'd like to
+contribute to the project (see **Development** section below). This should install all the necessary dependencies.
 
 To fit reporters for the HuggingFace model `model` and dataset `dataset`, just run:
 
 ```bash
 elk elicit microsoft/deberta-v2-xxlarge-mnli imdb
 ```
 
-This will automatically download the model and dataset, run the model and extract the relevant representations if they aren't cached on disk, fit reporters on them, and save the reporter checkpoints to the `elk-reporters` folder in your home directory. It will also evaluate the reporter classification performance on a held out test set and save it to a CSV file in the same folder.
+This will automatically download the model and dataset, run the model and extract the relevant representations if they
+aren't cached on disk, fit reporters on them, and save the reporter checkpoints to the `elk-reporters` folder in your
+home directory. It will also evaluate the reporter classification performance on a held out test set and save it to a
+CSV file in the same folder.
 
-The following will generate a CCS (Contrast Consistent Search) reporter instead of the CRC-based reporter, which is the default.
+The following will generate a CCS (Contrast Consistent Search) reporter instead of the CRC-based reporter, which is the
+default.
 
 ```bash
 elk elicit microsoft/deberta-v2-xxlarge-mnli imdb --net ccs
 ```
 
-The following command will evaluate the probe from the run naughty-northcutt on the hidden states extracted from the model deberta-v2-xxlarge-mnli for the imdb dataset. It will result in an `eval.csv` and `cfg.yaml` file, which are stored under a subfolder in `elk-reporters/naughty-northcutt/transfer_eval`.
+The following command will evaluate the probe from the run naughty-northcutt on the hidden states extracted from the
+model deberta-v2-xxlarge-mnli for the imdb dataset. It will result in an `eval.csv` and `cfg.yaml` file, which are
+stored under a subfolder in `elk-reporters/naughty-northcutt/transfer_eval`.
 
 ```bash
 elk eval naughty-northcutt microsoft/deberta-v2-xxlarge-mnli imdb
 ```
 
-The following runs `elicit` on the Cartesian product of the listed models and datasets, storing it in a special folder ELK_DIR/sweeps/<memorable_name>. Moreover, `--add_pooled` adds an additional dataset that pools all of the datasets together.
+The following runs `elicit` on the Cartesian product of the listed models and datasets, storing it in a special folder
+ELK_DIR/sweeps/<memorable_name>. Moreover, `--add_pooled` adds an additional dataset that pools all of the datasets
+together. You can also add a `--visualize` flag to visualize the results of the sweep.
 
 ```bash
 elk sweep --models gpt2-{medium,large,xl} --datasets imdb amazon_polarity --add_pooled
 ```
 
+If you just do `elk plot`, it will plot the results from the most recent sweep.
+If you want to plot a specific sweep, you can do so with:
+
+```bash
+elk plot {sweep_name}
+```
+
 ## Caching
 
-The hidden states resulting from `elk elicit` are cached as a HuggingFace dataset to avoid having to recompute them every time we want to train a probe. The cache is stored in the same place as all other HuggingFace datasets, which is usually `~/.cache/huggingface/datasets`.
+The hidden states resulting from `elk elicit` are cached as a HuggingFace dataset to avoid having to recompute them
+every time we want to train a probe. The cache is stored in the same place as all other HuggingFace datasets, which is
+usually `~/.cache/huggingface/datasets`.
 
 ## Development
+
 Use `pip install pre-commit && pre-commit install` in the root folder before your first commit.
 
 ### Devcontainer
 
 [
-    ![Open in Remote - Containers](
-        https://img.shields.io/static/v1?label=Remote%20-%20Containers&message=Open&color=blue&logo=visualstudiocode
-    )
+![Open in Remote - Containers](
+https://img.shields.io/static/v1?label=Remote%20-%20Containers&message=Open&color=blue&logo=visualstudiocode
+)
 ](
-    https://vscode.dev/redirect?url=vscode://ms-vscode-remote.remote-containers/cloneInVolume?url=https://github.com/EleutherAI/elk
+https://vscode.dev/redirect?url=vscode://ms-vscode-remote.remote-containers/cloneInVolume?url=https://github.com/EleutherAI/elk
 )
 
 ### Run tests
+
 ```bash
 pytest
 ```
+
 ### Run type checking
-We use [pyright](https://github.com/microsoft/pyright), which is built into the VSCode editor. If you'd like to run it as a standalone tool, it requires a [nodejs installation.](https://nodejs.org/en/download/)
+
+We use [pyright](https://github.com/microsoft/pyright), which is built into the VSCode editor. If you'd like to run it
+as a standalone tool, it requires a [nodejs installation.](https://nodejs.org/en/download/)
+
 ```bash
 pyright
 ```
 
 ### Run the linter
+
 We use [ruff](https://beta.ruff.rs/docs/). It is installed as a pre-commit hook, so you don't have to run it manually.
 If you want to run it manually, you can do so with:
+
 ```bash
 ruff . --fix
 ```
 
 ### Contributing to this repository
 
-If you work on a new feature / fix or some other code task, make sure to create an issue and assign it to yourself (Maybe, even share it in the elk channel of Eleuther's Discord with a small note). In this way, others know you are working on the issue and people won't do the same thing twice 👍 Also others can contact you easily.
+If you work on a new feature / fix or some other code task, make sure to create an issue and assign it to yourself (
+Maybe, even share it in the elk channel of Eleuther's Discord with a small note). In this way, others know you are
+working on the issue and people won't do the same thing twice 👍 Also others can contact you easily.

elk/__main__.py

@@ -5,6 +5,7 @@
 from simple_parsing import ArgumentParser
 
 from elk.evaluation.evaluate import Eval
+from elk.plotting.command import Plot
 from elk.training.sweep import Sweep
 from elk.training.train import Elicit
 
@@ -13,7 +14,7 @@
 class Command:
     """Some top-level command"""
 
-    command: Elicit | Eval | Sweep
+    command: Elicit | Eval | Sweep | Plot
 
     def execute(self):
         return self.command.execute()

elk/files.py

@@ -6,6 +6,11 @@
 from pathlib import Path
 
 
+def sweeps_dir() -> Path:
+    """Return the directory where sweeps are stored."""
+    return elk_reporter_dir() / "sweeps"
+
+
 def elk_reporter_dir() -> Path:
     """Return the directory where reporter checkpoints and logs are stored."""
     env_dir = os.environ.get("ELK_DIR", None)

elk/plotting/command.py

@@ -0,0 +1,50 @@
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+
+from simple_parsing import field
+
+from ..files import sweeps_dir
+from ..utils import colorize
+from .visualize import visualize_sweep
+
+
+def pretty_error(msg):
+    """Prints a pretty error message."""
+    print(colorize("Error", "red") + f": {msg}")
+
+
+@dataclass
+class Plot:
+    sweeps: list[Path] = field(positional=True, default_factory=list)
+    """Names of the sweeps to plot. If empty, the most recent sweep is used."""
+
+    overwrite: bool = False
+    """Whether to overwrite existing plots."""
+
+    def execute(self):
+        root_dir = sweeps_dir()
+
+        # If sweep is nonempty, get the paths to the specified sweeps.
+        # If no sweep is specified, use the most recent one.
+        if not self.sweeps:
+            sweep_paths = [max(root_dir.iterdir(), key=lambda f: f.stat().st_ctime)]
+            print(
+                f"Reading most recent sweep from \033[1m{sweep_paths[0]}\033[0m"
+            )  # bold
+        else:
+            sweep_paths = [root_dir / sweep for sweep in self.sweeps]
+
+        for sweep_path in sweep_paths:
+            if not sweep_path.exists():
+                pretty_error(f"No sweep with name {{{sweep_path}}} found in {root_dir}")
+            elif (sweep_path / "viz").exists() and not self.overwrite:
+                pretty_error(
+                    f"[blue]{sweep_path / 'viz'}[/blue] already exists. "
+                    f"Use --overwrite to overwrite."
+                )
+            else:
+                if self.overwrite:
+                    shutil.rmtree(sweep_path / "viz")
+
+                visualize_sweep(sweep_path)