models-under-pressure

KO | EN

gitlite — search

#python #ai #typescript #llm #linux #docker #nodejs #java #awesome #cli #kubernetes #pytorch #security #mcp

models-under-pressure

(★ 12)

Code for the paper "Detecting High-Stakes Interactions using Activation Probes"

models-under-pressure 최신버젼 다운로드

최종 버전 다운로드 (.zip)

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# Detecting High-Stakes Interactions using Activation Probes

This repo contains the code for all the experiments in the paper "Detecting High-Stakes Interactions using Activation Probes" ([arxiv](https://arxiv.org/abs/2506.10805)), presented at the ICML 2025 Workshop on Actionable Interpretability, accepted to NeurIPS 2025.

There's a lot of good stuff in here but it's in a bit of a rough state. For a cleaner codebase that's easier to use, check out [TuberLens](https://github.com/blandfort/tuberlens).

## Setup

In order to run this code:

1. Install uv and run `uv sync`
2. Download datasets: `uv run mup datasets download`
3. (Optional) To store activations on R2, create a Cloudflare account and R2 bucket
4. Add a `.env` file to the project root with the following environment variables:
  ```
  OPENAI_API_KEY=
  OPEN_ROUTER_API_KEY=
  HF_TOKEN=
  R2_ACCESS_KEY_ID=
  R2_SECRET_ACCESS_KEY=
  R2_ACTIVATIONS_BUCKET=
  R2_ACCOUNT_ID=
  ACTIVATIONS_DIR=
  HF_HOME=
  WANDB_API_KEY=
  ```

## Activations

In order to train or run inference with probes, you'll need to compute and store activations. You can do that using the `mup acts store` command. Here is an example:

```bash
uv run mup acts store --model 'meta-llama/Llama-3.2-1B-Instruct' --layer 11 --dataset data/training/prompts_4x/train.jsonl
```

This will compute activations, save them locally to `ACTIVATIONS_DIR`, and upload them to `R2_ACTIVATIONS_BUCKET`.

## Datasets

We contribute a new synthetic dataset we use for training, as well as slightly modified external datasets labelled for stakes we use for evaluation.

Our datasets are hosted on HuggingFace: [Arrrlex/models-under-pressure](https://huggingface.co/datasets/Arrrlex/models-under-pressure)

To download all datasets to the expected local paths, run:

```bash
uv run mup datasets download
```

This will place eval datasets into `data/evals/dev/` and `data/evals/test/`, and training data into `data/training/prompts_4x/`.

Below you can also find information on how these datasets were generated.

### Generating Synthetic Training Data

The code for generating the synthetic dataset can be found in `models_under_pressure/dataset_generation`: in particular, to generate a synthetic dataset, run `situation_generation.py`, then `prompt_generation.py`.

The code used for filtering confounding samples is in the `filter_dataset` function in `models_under_pressure/scripts/analyse_confounders.py`.

### Generating Dev Datasets for Evaluation

Run the files `anthropic_dataset.py`, `mt_dataset.py`, `mts_dataset.py` and `toolace_dataset.py` from `src/models_under_pressure/eval_datasets/`.
That will create the corresponding dataset files (raw and balanced) in the dev evals directory (see config.py).

Run the script `label_distribution.py` from `src/models_under_pressure/scripts/` to see number of high-stakes/low-stakes/ambiguous
for all eval datasets.

Warning: Creating further dev samples when test datasets are already present can lead to overlap between dev and test.
Ideally, first generate the full dev datasets and then create the test datasets to avoid overlap.

### Generating Test Datasets for Evaluation

Run the files `anthropic_dataset.py`, `mt_dataset.py`, `mts_dataset.py` and `toolace_dataset.py` from `src/models_under_pressure/eval_datasets/`,
using the `--split=test` argument.
That will create the corresponding dataset files (raw and balanced) in the test evals directory (see config.py).

After that, run the script `eval_dataset_split_check.py` to ensure that there is no overlap between dev and test datasets.
(Note that MT does have duplicates, so for that dataset you can expect some overlap by default.)

### Notes on Dataset Versions

Apr 15 versions (all based on calling parts of `modify_dataset.py`):

- ToolACE dev and test datasets: Based on previous dev dataset (raw version), modify the system prompt and relabel after adding system prompt.
- Anthropic dev and test datasets: Adding system prompt to each sample, otherwise no changes.
- MT dev and test datasets: Adding system prompt to each sample, removing cases where transcription length is less than description length, and adding more info to input.

Apr 16 versions (all based on calling parts of `modify_dataset.py`)

- MTS dev and test datasets: Parsing conversations (using strict mode), adding system prompt and relabelling.
- MT dev and test datasets: Adding system prompt to each sample, adding more info to input and relabelling.
- Anthropic dev and test datasets: Adding system prompt to each sample and relabelling.

Apr 22 versions (all based on calling parts of `modify_dataset.py`)

- MTS dev and test datasets: Same as before with slightly modified system prompt (fixed typo and dropping "guest families" part).
- ToolACE dev and test datasets: Based on original dataset, modify the system prompt (only changing first sentence and removing a later confusing sentence) and relabel after adding system prompt.
- Aya Redteaming dataset (only test): Added a system prompt and relabelled.
- Mental health dataset (only test): Added a system prompt and relabelled.

Apr 23 version of Anthropic (dev and test): Remove the duplicate system prompt, relabelling again just in case.

Apr 30 version of MT (dev and test): Moved 350 samples from test raw to dev raw and resampled the balanced versions.

### Deployment Context Datasets

Medical deployment dataset:

- Pair IDs up to 60 were generated using Gemini 2.5 Pro
- Additional pairs were created with GPT 4.5, giving the pairs from Gemini as examples
- Script `create_deployment_datasets.py` was used to convert into proper Dataset and relabel (which led to removal of many items)

Software deployment dataset:

- All items generated with GPT 4.5
- Script `create_deployment_datasets.py` was used to convert into proper Dataset and relabel (which led to removal of many items)

Chatbot deployment dataset:

- All items generated with GPT 4.5
- Script `create_deployment_datasets.py` was used to convert into proper Dataset and relabel (which led to removal of many items)

Combined deployment dataset: Created by concatenating all previous datasets.

## Running Experiments

### Probe Comparison
- Run `scripts/compare_probes.sh` to train probes and run prediction against all evaluation datasets
- Run `compare_probes_plot.py` to generate the bar chart

### Probe Calibration
Run `experiments/calibration.py`, modifying the paths in the `if __name == '__main__'` block to point to the relevant files generated by `compare_probes.sh`

### Probe Visualisation

To generate the per-token probe scores for visualisation, run:

```
uv run mup exp +experiment=evaluate_probe probe=attention +id=attn_viz_dev
```

Then to visualise the probe scores in the dashboard, run:

```
uv run mup dashboard data/results/evaluate_probes/results_attn_viz_dev.jsonl
```

### Compare Probes to Baselines Plot

To generate the bar charts comparing probes and baselines, run `models_under_pressure/figures/probes_vs_baselines_plot.py`. This should only be run after "Compare Probes" has been run, and after all baselines have been run. The results files should be adjusted accordingly.

### Comparing Probes to Word Statistics Baselines

To generate the plot comparing Attention probe with a TF-IDF based classifier, run `models_under_pressure/scripts/analyse_confounders.py`.

### Cross Validation

To run cross-validation and discover which layer has the best accuracy, run `uv run mup exp +experiment=cv`.

### Generalisation Heatmaps

To generate the data for the generalisation heatmaps, run `uv run mup exp +experiment=generalisation_heatmap`.

Then to create the plots, run `models_under_pressure/figures/generalisation_heatmap_plot.py`, changing the paths to the datafiles if necessary.

### Training on Dev Split of Evaluation Datasets

- Run `experiments/dev_split_training.py` for the best probe with different settings of `dev_sample_usage`. The script computes results 5 times by default with the same settings.
  - Important: Set `gradient_accumulation_steps` to 1 in the config of the corresponding probe, since training data for this experiment can consist only of few samples and no learning occurs if number of batches is less than gradient accumulation steps.
- Run `figures/dev_split_training_plot.py` to generate the corresponding plot. Adjust file paths end of the file before.
  - If you want to include the line for the baseline, you can obtain the corresponding file from the cascade experiment.

### Data Efficiency Experiment

Code for running the data efficiency experiment is included in [experiments/data_efficiency.py](./src/models_under_pressure/experiments/data_efficiency.py). You can run that file directly after adjusting the configurations at the end of it and calling one of the functions defined in that file:

- Use the function `run_data_efficiency_experiment` to get results for different types of probes.
- Use the function `run_data_efficiency_finetune_baseline_with_activations` to compute results for the finetuned baselines. (Adjust config accordingly and run one baseline model at a time.)

Generate the plot by putting all results files into a single directory and calling the script [figures/plot_data_efficiency.py](./src/models_under_pressure/figures/plot_data_efficiency.py).

### Cascade Plot

- To generate finetuning results, run `notebooks/finetuning_for_cascade.py` (adjust the settings in that script depending on the model you want to finetune)
- To generate the other results, run `experiments/monitoring_cascade.py`. The corresponding configuration files can be found under `config/experiments/monitoring_cascade.yaml` and `config/experiments/monitoring_cascade/`. It has one part for computing the results and a second part to generate the plot based on the results.
  - Result generation: The script generates result files for the selected probe and the continuation baselines.
  - Plot generation: Make sure that all the relevant files are included in one directory. This typically involves moving the finetuned baseline results into the directory with the other results. Then run the analysis step of the script.
  - For generating the full cascade plot (appendix), make sure that in `analyze.yaml`, the `baseline_models` and `finetuned_baseline_models` selections are both set to null, so that all results are displayed. Also, you might want to tweak a few arguments of the plotting function such as reducing `y_lim`.

### Figure 1 Plot

- The script to generate this plot uses outputs from the cascade experiment. Run the cascade experiment to compute the full results (e.g. using null for model selections in `analyze.yaml`)
- Then run [figures/plot_method_comparison.py](./src/models_under_pressure/figures/plot_method_comparison.py)

## Computing Baselines

### Prompted Baselines

Run `uv run mup exp +experiment=run_baselines model=<MODEL>` (replacing `<MODEL>` by "llama-1b", "llama-70b", "gemma-1b" etc.) to generate the results of the respective prompted model on all dev datasets (make sure default for `eval_datasets` in `config/config.yaml` is set to `dev_balanced`) for all prompt templates. All results are written in JSONL format to a single results file.

### Finetuned Baselines

Run `models_under_pressure/baselines/finetune.py` (replacing the `finetune_model` variable in `run_finetune_baselines()` by "llama-1b", "llama-8b", "gemma-1b", etc.) to finetune the respective model on the synthetic dataset and generate the results on all dev datasets (change `EVAL_DATASETS` to `TEST_DATASETS` in the call to `get_finetuned_baseline_results` to instead get results on test datasets).

models-under-pressure

models-under-pressure 최신버젼 다운로드

오늘의 추천 프로젝트

Was this content helpful?

Recent Feedback

models-under-pressure

models-under-pressure 최신버젼 다운로드

오늘의 추천 프로젝트

Was this content helpful?

Recent Feedback

Download README