eliciting_suppressed_knowledge/nbs/03_stats.ipynb

Quick experiment to see which is better at detecting truthful answers

• model outputs
• hs
• supressed activations (Hypothesis this is better)

In [1]:

%reload_ext autoreload
%autoreload 2

In [2]:

import os

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
# os.environ["CUDA_VISIBLE_DEVICES"] = "1"

In [3]:

from loguru import logger
import torch
from torch.utils.data import DataLoader
from datasets import load_dataset, Dataset, load_from_disk
from einops import rearrange, repeat
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.data import DataCollatorForLanguageModeling

import torch
from torch import Tensor
from torch.nn.functional import (
    binary_cross_entropy_with_logits as bce_with_logits,
)
from torch.nn.functional import (
    cross_entropy,
)
from pathlib import Path
from jaxtyping import Float
from torch import Tensor

import functools
import pandas as pd
import numpy as np

import itertools
from tqdm.auto import tqdm
import random
import json
from tqdm.auto import tqdm

In [4]:

import gc
def clear_mem():
    """
    Clear memory
    """
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.ipc_collect()
    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()
    return None
clear_mem()

Load data

In [5]:

!ls ../data/activation_store/

ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-100-128_v4e_wgen
ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-90_v2
ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-40-128_v4e_wgen
ds_at-QwenQwen3-1.json

In [ ]:

acts_outfile = Path('../data/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-100-128_v4f_wgen')

# f_config = acts_outfile.with_suffix(".json")
# config = json.load(open(f_config, 'r'))
# model_name = config['model_name']
# print(config)

ds_a2 = load_from_disk(acts_outfile).with_format("torch")
ds_a2

Dataset({
    features: ['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj', 'loss', 'logits', 'hidden_states', 'attention_mask', 'supr_amounts', 'i', 'input_ids', 'prompt_mask', 'completion_type'],
    num_rows: 200
})

In [12]:

act_groups = [c for c in ds_a2.column_names if c.startswith('acts-')]
act_groups

Out [12]:

['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj']

In [13]:

for k,v in ds_a2[:2].items():
    if hasattr(v, 'shape'):
        print(k, v.shape)
    else:
        print(k, type(v))

acts-mlp.down_proj torch.Size([2, 1, 128, 2048])
acts-self_attn torch.Size([2, 1, 128, 2048])
acts-mlp.up_proj torch.Size([2, 1, 128, 6144])
loss torch.Size([2])
logits torch.Size([2, 151936])
hidden_states torch.Size([2, 1, 128, 2048])
attention_mask torch.Size([2, 128])
supr_amounts torch.Size([2, 3, 1, 2048])
i torch.Size([2])
input_ids <class 'list'>
prompt_mask <class 'list'>
completion_type <class 'list'>

In [14]:

ds_a2

Out [14]:

Dataset({
    features: ['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj', 'loss', 'logits', 'hidden_states', 'attention_mask', 'supr_amounts', 'i', 'input_ids', 'prompt_mask', 'completion_type'],
    num_rows: 200
})

Stats

Lets start with std of each hs

• for each i, compare completion types
  • for each ks
    • where

In [ ]:

ks = act_groups + [ 'hidden_states', 'attention_mask', 'supr_amounts']
ks

['acts-mlp.down_proj',
 'acts-self_attn',
 'acts-mlp.up_proj',
 'model_output',
 'loss',
 'logits',
 'hidden_states',
 'attention_mask',
 'supr_amounts']

In [ ]: