src/probe/README.md

Probe Implementation

Using the output database of extracted features (generated using src/main.py), you can easily intialize probes using the script in probe/main.py.

To do this, you need a probe config yaml file specifying: the probe model attributes, training and test configuration and the input data. The following is an example of a probe configuration file for features extracted using LLAVA-1.5-7b:

model:
  - activation: ReLU # a valid activation function from torch.nn
  - hidden_size: 512 # the input and output size of the intermediate layers
  - num_layers: 2 # the number of layers of the probe model
  - save_dir: /path/to/save_dir # the location to save the probe results

training:
  - batch_size: [64, 128, 1024]
  - num_epochs: [50, 100, 200]
  - learning_rate: [0.001, 0.0005, 0.0001]
  - optimizer: AdamW # a valid optimizer from torch.nn
  - loss: CrossEntropyLoss # a valid loss metric from torch.nn

test:
  - batch_size: 32
  - loss: CrossEntropyLoss # a valid loss metric from torch.nn

data:
  - input_db: /path/to/input_db
  - db_name: tensors # the name of the database in input_db, the default is `tensors` if unspecified
  - input_layer: language_model.model.layers.16.post_attention_layernorm # the layer value in input_db to filter by

For data, the input_layer value must be specified as the probe model is only trained on data extracted from a single layer to avoid dimension mismatches or ambiguous results.

Training

Note that the training procedure conducts an iterative (naive) hyperparameter search on all configurations of the training hyperparameters: batch_size, num_epochs and learning_rate. It uses $k$-fold cross validation (with a default of $k=5$) to store the lowest validation loss from each configuration and returns the configuration with the minimum validation loss.

training:
  - batch_size: [64, 128, 1024]
  - num_epochs: [50, 100, 200]
  - learning_rate: [0.001, 0.0005, 0.0001]
  ...

Testing

After retrieving the best training configuration, the current script trains and tests a probe model on two versions of the data: a Main condition (where the data is untouched) and a Shuffled condition (where the labels are shuffled).

This is based on work from Hewitt and Liang (2019), where they propose the inclusion of "control tasks" to ensure that the probe is selective in its learning. The goal is to achieve high task-specific accuracy and low control task accuracy. The output attributes described in Results follow the former experimental design.

Results

Finally, the program saves both the probe model and result values in save_dir. The current script saves the attributes to a {save_dir}/probe_data.txt. The following is an example output:

{
    # The final train configuration with the lowest validation loss
    "train_config": {
                    "batch_size": 64,
                    "num_epochs": 200,
                    "learning_rate": 0.0001,
                    "optimizer": "AdamW",
                    "loss": "CrossEntropyLoss"
                    },

    # Probe accuracy on the shuffled data
    "shuffle_accuracy": 0.3295668661594391,
    "shuffle_loss": 2.1123009968163153,
    "shuffle_preds": [1, 2, 0, ... ],
    "shuffle_labels": [0, 0, 0, ... ],

    # Probe accuracy on the original data
    "test_accuracy": 0.6436911225318909,
    "test_loss": 1.7786695135290789,
    "test_preds": [2, 0, 2, ...],
    "test_labels": [0, 0, 2, ...],

    # The statistical significance of the difference between the shuffled and unshuffled data using a z-test
    "pvalue": 4.8257723322914694e-116
}

Details on the reason behind shuffling and experimental design can be found in the Testing section.