wassname/discovering_latent_knowledge
1
0
Fork 0
mirror of https://github.com/wassname/discovering_latent_knowledge.git synced 2026-09-09 11:21:22 +08:00
Code Issues Packages Projects Releases Wiki Activity

VAE exp

Browse Source
This commit is contained in:
wassname
2023-12-10 12:31:49 +08:00
parent 3067ae8aae
commit 4e3b0c37fc
10 changed files with 2786 additions and 827 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+8
View File
@@ -71,7 +71,15 @@ pred_last_choice_is_true = y / (y_pred2-y_pred)
pred_last_choice_is_true # [1, -1]
```
# Running
```sh
python notebooks/make_dataset2.py --max_examples 1720 220 --datasets imdb glue:qnli super_glue:boolq
```
# Description
There is some previous work on this ([ELK](https://github.com/EleutherAI/elk), [DLK](https://github.com/collin-burns/discovering_latent_knowledge/blob/main/CCS.ipynb), CSS, etc) that all take varias approaches. They have this in common:
- Show the model 2 statements “the sky is blue” “the sky is green”
mjc_notes.md
+118
View File
@@ -1971,3 +1971,121 @@ https://colab.research.google.com/drive/1rPy82rL3iZzy2_Rd3F82RwFhlVnnroIh?usp=sh
n_instances - remove this, another batch dim
- nice timeseries 1d VAE from merlion https://github.dev/salesforce/Merlion/blob/01c3fc3406ebf19798cedcddbe829ae5339e1424/merlion/models/anomaly/vae.py#L186
- simple one https://github.com/ctallec/world-models/blob/master/models/vae.py
- neels https://github.dev/neelnanda-io/1L-Sparse-Autoencoder/blob/bcae01328a2f41d24bd4a9160828f2fc22737f75/utils.py#L106
# 2023-12-08 10:01:13
Initial obs
- it can still overfit on the latent state, hmmm! Well maybe I need to make it very small or space or quantized (dreamer style)
- a tanh seems to help!... just in the mse one it makes it worse... hmm.
oh CVAE generalizes well notebooks/033_train_cvae.ipynb
test/acc 0.6785010099411011 │ 0.7179487347602844
llm gave did didn't
instructed to
tell a truth 0.74 NaN
tell a lie 0.91 0.54
oos/acc │ 0.7257769703865051 │ 0.7664233446121216 │ 0.7627736926078796
llm gave did didn't
instructed to
tell a truth 0.78 NaN
tell a lie 1.00 0.66
# 2023-12-08 15:44:03
Questions:
- [ ] hmm in the anthorpic [paper](https://transformer-circuits.pub/2022/toy_model/index.html#demonstrating-setup-loss) they weight by feature importance, this seems important
- [x] anything else I need to know from the sparse transformer AE's?
- [x] what is dictionary learning? it seems to just be a huge 1 layer sparse autoencoder. no categorical latent or anything
- they seem to use weight norm on decoder, not tie weights. have 8 times the latent space compared to activations
- oh they replace activations with reconstructed
- oh actual training [tips](https://docs.google.com/document/u/0/d/187jfZSbhRjjQaazjYlThBsKp3Q0Pw3VdIHVST9H2dvw/mobilebasic)
- what's the decoder weight norm??
- [x] do I need something special to make it sparse? no it looks like it's just the l1 loss
- [ ]
- [ ] what where the dreamer learnings?
- two-hot latent space?? I guess that means it turns into [2, 1, 0, 1]. I'm assuming neg vs pos?
- symlog scaling for rewards prediction - I probobly don't need this
- how do the discrete states work? https://github.dev/Eclectic-Sheep/sheeprl/blob/52f49be5971c5753e18bdf328d3035334fe688f1/sheeprl/algos/dreamer_v3/agent.py#L31
- [ ] does my pcr probe work ok? how to debug?
> Features Vary in Importance: Not all features are equally useful to a given task. Some can reduce the loss more than others. For an ImageNet model, where classifying different species of dogs is a central task, a floppy ear detector might be one of the most important features it can have. In contrast, another feature might only very slightly improve performance
IRIS loss https://github.dev/eloialonso/iris/blob/ac6be401fed2b6176c9ce0cf1dc10e376c9d740d/src/models/tokenizer/tokenizer.py#L50-L55
# Codebook loss. Notes:
# - beta position is different from taming and identical to original VQVAE paper
# - VQVAE uses 0.25 by default
beta = 1.0
commitment_loss = (z.detach() - z_quantized).pow(2).mean() + beta * (z - z_quantized.detach()).pow(2).mean()
reconstruction_loss = torch.abs(observations - reconstructions).mean()
perceptual_loss = torch.mean(self.lpips(observations, reconstructions))
https://openreview.net/pdf?id=o8IDoZggqO
> We follow DreamerV3 in using discrete regression with two-hot targets and symlog scaling for rewards prediction (Bellemare et al., 2017; Imani & White,2018).
> SqrtTransform Using two-hot discrete regression with the asymmetric square root transformation intro- duced by R2D221 and used in MuZero34
https://arxiv.org/pdf/2301.04104v1.pdf
- R2D2 https://openreview.net/forum?id=r1lyTjAqYX
> The representations are sampled from a vector of softmax distributions and we take straight-through gradients through the sampling step
- "sampled from a vector of softmax distributions"? I would like to see psudocode. I guess it just uses the distributions baked into torch fd.MultivariateNormalDiag(mean, std)
> To train the critic, we symlog transform the targets Rλ t and then twohot encode them into a soft label for the softmax distribution produced by the critic. Twohot encoding is a generalization of onehot encoding to continuous values. It produces a vector of length |B| where all elements are 0 except for the two entries closest to the encoded continuous number, at positions k and k + 1. These two entries sum up to 1, with more weight given to the entry that is closer to the enco
https://arxiv.org/pdf/2301.04104v1.pdf
> the world model encodes sensory inputs into a discrete representation zt t
cal.s.mcdougall@gmail.com
dictionary learning
- experiment I added amazon, and it doubled the training data, lets see if I get above 80% acc... I was getting ~75%
https://www.alignmentforum.org/posts/F4iogK5xdNd7jDNyw/comparing-anthropic-s-dictionary-learning-to-ours
> Size of training set: We trained our autoencoders for 10M tokens. Anthropic trained theirs for much longer, 8B tokens.
Wow that's a lot. If I want to focus on just lying, I might need to focus on not reconstruction the whole state...
## Best dreamer v3 repo?
- https://github.dev/kc-ml2/SimpleDreamer oh it's dreamer 1 meh
- https://github.dev/Eclectic-Sheep/sheeprl/blob/52f49be5971c5753e18bdf328d3035334fe688f1/sheeprl/algos/dreamer_v3/agent.py#L31
symlog is simple `torch.sign(x) * torch.log(1 + torch.abs(x))`
# 2023-12-09 11:34:11
Questions:
- understand HALOs https://twitter.com/ethayarajh/status/1732837520784957476 https://github.com/ContextualAI/HALOs
- so it's just DPO with a differen't activation function on the reward, and notably it can use reward text instead of ranked pairs, letting you skip SFT. In a way it's just SFT?
- [ ] discrete states, just look up QVAE?
- [ ] hmm some use a categorical, and the gumbel reparam trick for end to end backprop
- [ ] some use VQ-VAE which I haven't looked at before but look promising. But I want to!
- [ ] does my pcr probe work ok? how to debug? I guess I need to check acc from it for a start
- [ ] does my conv vae work? maybe I need transposed conv blocks?
- perhaps just use https://github.com/ctallec/world-models/blob/master/models/vae.py#L10
- perhaps I need to focus on important features? Or on a task?
- e.g. if doing inference on the reconstructed parts, can I get the same output? (RAM heavy)
- if just apply an importance multipier
```py
QVAE psuedocode
```
notebooks/033_train_cvae.ipynb
+2422 -596
View File
File diff suppressed because one or more lines are too long
src/datasets/batch.py
+49 -49
View File
@@ -1,68 +1,68 @@
from tqdm.auto import tqdm
import torch
from torch.utils.data import DataLoader
from datasets.arrow_dataset import Dataset
import hashlib
import pickle
import numpy as np
from typing import List, Dict, Any, Union, NewType, Optional
# from tqdm.auto import tqdm
# import torch
# from torch.utils.data import DataLoader
# from datasets.arrow_dataset import Dataset
# import hashlib
# import pickle
# import numpy as np
# from typing import List, Dict, Any, Union, NewType, Optional
from src.datasets.hs import ExtractHiddenStates
from src.helpers.typing import float_to_int16, int16_to_float
from src.helpers.ds import ds_keep_cols, clear_mem
from src.datasets.intervene import InterventionDict
# from src.datasets.hs import ExtractHiddenStates
# from src.helpers.typing import float_to_int16, int16_to_float
# from src.helpers.ds import ds_keep_cols, clear_mem
# from src.datasets.intervene import InterventionDict
def batch_hidden_states(model, tokenizer, intervention_dicts: Optional[InterventionDict], data: Dataset, batch_size=2, layer_padding=3, layer_stride=4):
"""
Given an encoder-decoder model, a list of data, computes the contrast hidden states on n random examples.
Returns numpy arrays of shape (n, hidden_dim) for each candidate label, along with a boolean numpy array of shape (n,)
with the ground truth labels
# def batch_hidden_states(model, tokenizer, intervention_dicts: Optional[InterventionDict], data: Dataset, batch_size=2, layer_padding=3, layer_stride=4):
# """
# Given an encoder-decoder model, a list of data, computes the contrast hidden states on n random examples.
# Returns numpy arrays of shape (n, hidden_dim) for each candidate label, along with a boolean numpy array of shape (n,)
# with the ground truth labels
This is deliberately simple so that it's easy to understand, rather than being optimized for efficiency
"""
ehs = ExtractHiddenStates(model, tokenizer, intervention_dicts=intervention_dicts, layer_stride=layer_stride, layer_padding=layer_padding)
# This is deliberately simple so that it's easy to understand, rather than being optimized for efficiency
# """
# ehs = ExtractHiddenStates(model, tokenizer, intervention_dicts=intervention_dicts, layer_stride=layer_stride, layer_padding=layer_padding)
torch_cols = ['input_ids', 'attention_mask', 'choice_ids']
ds_t_subset = ds_keep_cols(data, torch_cols)
ds_t_subset.set_format(type='torch')
# torch_cols = ['input_ids', 'attention_mask', 'choice_ids']
# ds_t_subset = ds_keep_cols(data, torch_cols)
# ds_t_subset.set_format(type='torch')
ds_p_subset = data.remove_columns(torch_cols)
# ds_p_subset = data.remove_columns(torch_cols)
dl = DataLoader(ds_t_subset, batch_size=batch_size, shuffle=False)
for i, batch in enumerate(tqdm(dl, desc='get hidden states')):
input_ids, attention_mask, choice_ids = batch["input_ids"], batch["attention_mask"], batch["choice_ids"]
nn = len(input_ids)
index = i*batch_size+np.arange(nn)
# dl = DataLoader(ds_t_subset, batch_size=batch_size, shuffle=False)
# for i, batch in enumerate(tqdm(dl, desc='get hidden states')):
# input_ids, attention_mask, choice_ids = batch["input_ids"], batch["attention_mask"], batch["choice_ids"]
# nn = len(input_ids)
# index = i*batch_size+np.arange(nn)
# different due to dropout
hsl = ehs.get_batch_of_hidden_states(input_ids=input_ids, attention_mask=attention_mask, choice_ids=choice_ids)
# # different due to dropout
# hsl = ehs.get_batch_of_hidden_states(input_ids=input_ids, attention_mask=attention_mask, choice_ids=choice_ids)
for j in range(nn):
# let's add the non torch metadata like label, prompt, lie, etc
k = i*batch_size + j
info = ds_p_subset[k]
# for j in range(nn):
# # let's add the non torch metadata like label, prompt, lie, etc
# k = i*batch_size + j
# info = ds_p_subset[k]
large_arrays_keys = [k for k,v in hsl.items() if isinstance(v, torch.Tensor) and v.ndim>2]
# large_arrays_keys = [k for k,v in hsl.items() if isinstance(v, torch.Tensor) and v.ndim>2]
# TODO deal with multiple lists of hs in hs0
large_arrays = {k:hsl[k][j] for k in large_arrays_keys}
# # TODO deal with multiple lists of hs in hs0
# large_arrays = {k:hsl[k][j] for k in large_arrays_keys}
yield dict(
# yield dict(
# large_arrays_keys=large_arrays_keys,
scores0=hsl["scores"][j],
# layer_names=hsl["layers"][j] if k==0 else [], # just in the first one, to save space
# # large_arrays_keys=large_arrays_keys,
# scores0=hsl["scores"][j],
# # layer_names=hsl["layers"][j] if k==0 else [], # just in the first one, to save space
ds_index=index[j],
# ds_index=index[j],
# int16 makes our storage much smaller
**large_arrays,
# # int16 makes our storage much smaller
# **large_arrays,
**info
)
# **info
# )
info = large_arrays = hsl = None
clear_mem()
# info = large_arrays = hsl = None
# clear_mem()
src/datasets/dm.py
+4 -1
View File
@@ -32,12 +32,14 @@ class imdbHSDataModule(pl.LightningDataModule):
batch_size: int=32,
x_cols = ['end_hidden_states'],
skip_layers = 0,
use_diff = True,
):
super().__init__()
self.save_hyperparameters(ignore=["ds"])
self.ds = ds
self.x_cols = x_cols
self.skip_layers = skip_layers
self.use_diff = use_diff
def setup(self, stage: str):
h = self.hparams
@@ -61,7 +63,8 @@ class imdbHSDataModule(pl.LightningDataModule):
b = len(self.ds_hs)
# take the diff between layers. Shape batch, layers, hidden_states, inferences
hs = torch.tensor(self.ds_hs['end_hidden_states'])
hs = hs.diff(1, axis=1) # this makes it the residual between layers
if self.use_diff:
hs = hs.diff(1, axis=1) # this makes it the residual between layers
if self.skip_layers:
hs = hs[:, self.skip_layers:] # drop the first 10 layers to prevent overfitting?
self.hs0 = hs[..., 0]
src/datasets/hs.py
+179 -179
View File
@@ -1,220 +1,220 @@
from dataclasses import dataclass
import lightning as pl
from loguru import logger
from transformers import (
AutoTokenizer,
AutoModelForSeq2SeqLM,
AutoModelForMaskedLM,
AutoModelForCausalLM,
AutoConfig,
AutoModel,
PreTrainedTokenizer,
PreTrainedModel
)
from typing import Optional, List, Tuple, Dict, NewType
from transformers import LogitsProcessorList
import functools
from src.helpers.torch import to_numpy
from src.datasets.dropout import enable_dropout
import re
# from dataclasses import dataclass
# import lightning as pl
# from loguru import logger
# from transformers import (
# AutoTokenizer,
# AutoModelForSeq2SeqLM,
# AutoModelForMaskedLM,
# AutoModelForCausalLM,
# AutoConfig,
# AutoModel,
# PreTrainedTokenizer,
# PreTrainedModel
# )
# from typing import Optional, List, Tuple, Dict, NewType
# from transformers import LogitsProcessorList
# import functools
# from src.helpers.torch import to_numpy
# from src.datasets.dropout import enable_dropout
# import re
from tqdm.auto import tqdm
# from src.datasets.hs import ExtractHiddenStates
from torch.utils.data import DataLoader
from datasets import Dataset
import numpy as np
import torch
import torch.nn.functional as F
from baukit.nethook import Trace, TraceDict, recursive_copy
from einops import rearrange, reduce, repeat
from src.datasets.scores import choice2id, choice2ids
from src.helpers.torch import clear_mem, detachcpu
from collections import defaultdict
from dataclasses import field
from src.datasets.intervene import InterventionDict, intervention_meta_fn
from functools import partial
# from tqdm.auto import tqdm
# # from src.datasets.hs import ExtractHiddenStates
# from torch.utils.data import DataLoader
# from datasets import Dataset
# import numpy as np
# import torch
# import torch.nn.functional as F
# from baukit.nethook import Trace, TraceDict, recursive_copy
# from einops import rearrange, reduce, repeat
# from src.datasets.scores import choice2id, choice2ids
# from src.helpers.torch import clear_mem, detachcpu
# from collections import defaultdict
# from dataclasses import field
# from src.datasets.intervene import InterventionDict, intervention_meta_fn
# from functools import partial
# def noise_for_embeds(inputs_embeds, seed=42, std = 2e-2):
# B, S, embed_dim = inputs_embeds.shape
# with torch.random.fork_rng(devices=[inputs_embeds.device.index]):
# torch.manual_seed(seed)
# noise = torch.normal(0., std, (embed_dim, ))
# noise = repeat(noise, 't -> b s t', b=B, s=S).to(inputs_embeds.device).to(inputs_embeds.dtype)
# return noise
# # def noise_for_embeds(inputs_embeds, seed=42, std = 2e-2):
# # B, S, embed_dim = inputs_embeds.shape
# # with torch.random.fork_rng(devices=[inputs_embeds.device.index]):
# # torch.manual_seed(seed)
# # noise = torch.normal(0., std, (embed_dim, ))
# # noise = repeat(noise, 't -> b s t', b=B, s=S).to(inputs_embeds.device).to(inputs_embeds.dtype)
# # return noise
def tcopy(x: torch.Tensor):
return x.clone().detach().cpu()
# def tcopy(x: torch.Tensor):
# return x.clone().detach().cpu()
def counterfactual_loss(model, scores, token_y, token_n):
"""do a backwards pass where the loss is the distance to the opposite scores"""
eps = 1e-4
model.zero_grad()
assert token_y.shape[1]<2, 'FIXME just use the first token for now'
score_y = torch.index_select(scores, 1, token_y[:, 0])
score_n = torch.index_select(scores, 1, token_n[:, 0])
# this loss would be zero if the logits of the positive and negative tokens werre flipped
loss = F.l1_loss(score_y, score_n) + F.l1_loss(score_n, score_y)
return loss
# def counterfactual_loss(model, scores, token_y, token_n):
# """do a backwards pass where the loss is the distance to the opposite scores"""
# eps = 1e-4
# model.zero_grad()
# assert token_y.shape[1]<2, 'FIXME just use the first token for now'
# score_y = torch.index_select(scores, 1, token_y[:, 0])
# score_n = torch.index_select(scores, 1, token_n[:, 0])
# # this loss would be zero if the logits of the positive and negative tokens werre flipped
# loss = F.l1_loss(score_y, score_n) + F.l1_loss(score_n, score_y)
# return loss
def stack_trace_returns(ret: TraceDict, names: List[str]) -> torch.Tensor:
hs = [ret[h].output for h in names]
hs = [h[0] if isinstance(h, tuple) else h for h in hs] # from a head it's a tuple
return rearrange(hs, 'layers b s hs -> b layers s hs')[:, :, -1]
# def stack_trace_grad_returns(ret: TraceDict, names: List[str]) -> torch.Tensor:
# hs = [ret[h].output.grad.detach() for h in names]
# def stack_trace_returns(ret: TraceDict, names: List[str]) -> torch.Tensor:
# hs = [ret[h].output for h in names]
# hs = [h[0] if isinstance(h, tuple) else h for h in hs] # from a head it's a tuple
# return rearrange(hs, 'layers b s hs -> b layers s hs')[:, :, -1]
# def select_weight_grads(weight_grads: Dict[str, torch.Tensor], pattern:str= ".+attn.c_proj.weight", mean_axis:int=1):
# grads = [g.mean(mean_axis) for k,g in weight_grads.items() if re.match(pattern, k)]
# assert len(grads), f"non of pattern='{pattern}' found in {weight_grads.keys()}"
# return rearrange(grads, "lyrs b hs -> b lyrs hs")
# # def stack_trace_grad_returns(ret: TraceDict, names: List[str]) -> torch.Tensor:
# # hs = [ret[h].output.grad.detach() for h in names]
# # return rearrange(hs, 'layers b s hs -> b layers s hs')[:, :, -1]
# # def select_weight_grads(weight_grads: Dict[str, torch.Tensor], pattern:str= ".+attn.c_proj.weight", mean_axis:int=1):
# # grads = [g.mean(mean_axis) for k,g in weight_grads.items() if re.match(pattern, k)]
# # assert len(grads), f"non of pattern='{pattern}' found in {weight_grads.keys()}"
# # return rearrange(grads, "lyrs b hs -> b lyrs hs")
@dataclass
class ExtractHiddenStates:
# @dataclass
# class ExtractHiddenStates:
model: PreTrainedModel
tokenizer: PreTrainedTokenizer
intervention_dicts: Optional[InterventionDict] = None
layer_stride: int = 8
layer_padding: int = 3
# model: PreTrainedModel
# tokenizer: PreTrainedTokenizer
# intervention_dicts: Optional[InterventionDict] = None
# layer_stride: int = 8
# layer_padding: int = 3
def get_layer_names(self):
# for WizardLM/WizardCoder-3B-V1.0
# HEADS = [f"transformer.h.{i}.attn.c_proj" for i in range(self.model.config.num_hidden_layers)]
# MLPS = [f"transformer.h.{i}.mlp" for i in range(self.model.config.num_hidden_layers)]
# def get_layer_names(self):
# # for WizardLM/WizardCoder-3B-V1.0
# # HEADS = [f"transformer.h.{i}.attn.c_proj" for i in range(self.model.config.num_hidden_layers)]
# # MLPS = [f"transformer.h.{i}.mlp" for i in range(self.model.config.num_hidden_layers)]
# for "WizardLM/WizardCoder-Python-13B-V1.0"
# HACK: depends on model layout
layers_names_h = [f"model.layers.{i}.self_attn" for i in range(self.model.config.num_hidden_layers)]
layers_names_mlp = [f"model.layers.{i}.mlp" for i in range(self.model.config.num_hidden_layers)]
return self.get_layer_selection(layers_names_h) + self.get_layer_selection(layers_names_mlp)
# # for "WizardLM/WizardCoder-Python-13B-V1.0"
# # HACK: depends on model layout
# layers_names_h = [f"model.layers.{i}.self_attn" for i in range(self.model.config.num_hidden_layers)]
# layers_names_mlp = [f"model.layers.{i}.mlp" for i in range(self.model.config.num_hidden_layers)]
# return self.get_layer_selection(layers_names_h) + self.get_layer_selection(layers_names_mlp)
def get_batch_of_hidden_states(
self,
input_text: Optional[List[str]] = None,
input_ids: torch.Tensor = None,
attention_mask: Optional[torch.Tensor] = None,
choice_ids: List[torch.Tensor] = None,
truncation_length=999,
debug=False,
):
"""
Given a decoder model and a batch of texts, gets a pair of hidden states (in a given layer) on that input texts
# def get_batch_of_hidden_states(
# self,
# input_text: Optional[List[str]] = None,
# input_ids: torch.Tensor = None,
# attention_mask: Optional[torch.Tensor] = None,
# choice_ids: List[torch.Tensor] = None,
# truncation_length=999,
# debug=False,
# ):
# """
# Given a decoder model and a batch of texts, gets a pair of hidden states (in a given layer) on that input texts
The idea is this: given two pairs of hidden states, where everything is the same except r dropout. Then tell me which one is more truthful?
"""
assert (input_ids is not None) or (input_text is not None), "need to provide input_ids or input_text"
assert self.tokenizer.truncation_side == 'left'
# The idea is this: given two pairs of hidden states, where everything is the same except r dropout. Then tell me which one is more truthful?
# """
# assert (input_ids is not None) or (input_text is not None), "need to provide input_ids or input_text"
# assert self.tokenizer.truncation_side == 'left'
if input_text:
raise NotImplementedError("FIXME")
t = self.tokenizer(
input_text,
return_tensors="pt",
add_special_tokens=True,
padding='max_length', max_length=truncation_length, truncation=True, return_attention_mask=True,
)
input_ids = t.input_ids.to(self.model.device)
attention_mask = t.attention_mask.to(self.model.device)
else:
input_ids = input_ids.to(self.model.device)
attention_mask = attention_mask.to(self.model.device)
choice_ids = choice_ids.to(self.model.device)
# if input_text:
# raise NotImplementedError("FIXME")
# t = self.tokenizer(
# input_text,
# return_tensors="pt",
# add_special_tokens=True,
# padding='max_length', max_length=truncation_length, truncation=True, return_attention_mask=True,
# )
# input_ids = t.input_ids.to(self.model.device)
# attention_mask = t.attention_mask.to(self.model.device)
# else:
# input_ids = input_ids.to(self.model.device)
# attention_mask = attention_mask.to(self.model.device)
# choice_ids = choice_ids.to(self.model.device)
# forward pass
last_token = -1
# # forward pass
# last_token = -1
layers_names = self.get_layer_names()
# layers_names = self.get_layer_names()
self.model.eval()
# self.model.eval()
if self.intervention_dicts is not None:
# extraction mode
# 15 is a magic number from honest_llama
intervention_fn1 = partial(intervention_meta_fn, interventions=self.intervention_dicts, alpha=-15)
intervention_fn2 = partial(intervention_meta_fn, interventions=self.intervention_dicts, alpha=15)
edit_outputs = [intervention_fn1, intervention_fn2]
else:
# calibration mode
edit_outputs = [None]
# if self.intervention_dicts is not None:
# # extraction mode
# # 15 is a magic number from honest_llama
# intervention_fn1 = partial(intervention_meta_fn, interventions=self.intervention_dicts, alpha=-15)
# intervention_fn2 = partial(intervention_meta_fn, interventions=self.intervention_dicts, alpha=15)
# edit_outputs = [intervention_fn1, intervention_fn2]
# else:
# # calibration mode
# edit_outputs = [None]
with torch.no_grad():
multi_outs = defaultdict(list)
for edit_output in edit_outputs:
with TraceDict(self.model, layers_names, retain_grad=False, detach=True, edit_output=edit_output) as ret:
model_inputs = self.model.prepare_inputs_for_generation(input_ids=input_ids, attention_mask=attention_mask, use_cache=False)
outputs = self.model.forward(
**model_inputs,
return_dict=True,
output_hidden_states=True,
)
outputs["scores"] = outputs.logits[:, last_token, :].float()
# with torch.no_grad():
# multi_outs = defaultdict(list)
# for edit_output in edit_outputs:
# with TraceDict(self.model, layers_names, retain_grad=False, detach=True, edit_output=edit_output) as ret:
# model_inputs = self.model.prepare_inputs_for_generation(input_ids=input_ids, attention_mask=attention_mask, use_cache=False)
# outputs = self.model.forward(
# **model_inputs,
# return_dict=True,
# output_hidden_states=True,
# )
# outputs["scores"] = outputs.logits[:, last_token, :].float()
# stack
hidden_states = list(outputs.hidden_states)
hidden_states = rearrange(hidden_states, 'lyrs b seq hs -> b lyrs seq hs')[:, :, last_token]
## from ret, we get the layer activation and the grads on them
head_activation = tcopy(stack_trace_returns(ret, layers_names))
# mlp_activation = tcopy(stack_trace_returns(ret, MLPS))
# # stack
# hidden_states = list(outputs.hidden_states)
# hidden_states = rearrange(hidden_states, 'lyrs b seq hs -> b lyrs seq hs')[:, :, last_token]
# ## from ret, we get the layer activation and the grads on them
# head_activation = tcopy(stack_trace_returns(ret, layers_names))
# # mlp_activation = tcopy(stack_trace_returns(ret, MLPS))
# collect outputs
multi_outs['scores'].append(outputs["scores"])
# multi_outs['hidden_states'].append(hidden_states)
multi_outs['head_activation'].append(head_activation)
# multi_outs['mlp_activation'].append(mlp_activation)
# # collect outputs
# multi_outs['scores'].append(outputs["scores"])
# # multi_outs['hidden_states'].append(hidden_states)
# multi_outs['head_activation'].append(head_activation)
# # multi_outs['mlp_activation'].append(mlp_activation)
# stack
multi_outs['scores'] = torch.stack(multi_outs['scores'], -1)
# multi_outs['mlp_activation'] = torch.stack(multi_outs['mlp_activation'], -1)
multi_outs['head_activation'] = torch.stack(multi_outs['head_activation'], -1)
# # stack
# multi_outs['scores'] = torch.stack(multi_outs['scores'], -1)
# # multi_outs['mlp_activation'] = torch.stack(multi_outs['mlp_activation'], -1)
# multi_outs['head_activation'] = torch.stack(multi_outs['head_activation'], -1)
# combine
out_common = dict(input_ids=input_ids, attention_mask=attention_mask, layers=layers_names,)
if debug:
out_common['input_truncated'] = self.tokenizer.batch_decode(input_ids)
out_common['text_ans'] = self.tokenizer.batch_decode(outputs["scores"].softmax(-1).argmax(-1))
# # combine
# out_common = dict(input_ids=input_ids, attention_mask=attention_mask, layers=layers_names,)
# if debug:
# out_common['input_truncated'] = self.tokenizer.batch_decode(input_ids)
# out_common['text_ans'] = self.tokenizer.batch_decode(outputs["scores"].softmax(-1).argmax(-1))
out = {**multi_outs, **out_common}
# out = {**multi_outs, **out_common}
# detach
out = {k: detachcpu(v) for k, v in out.items()}
# # detach
# out = {k: detachcpu(v) for k, v in out.items()}
# I shouldn't have to do this but I get memory leaks
outputs = hidden_states = hidden_states2 = loss = orig_state_dict = scores = token_y = token_n = input_ids = attention_mask = choice_ids = residual_stream = residual_stream2 = None
clear_mem()
return out
# # I shouldn't have to do this but I get memory leaks
# outputs = hidden_states = hidden_states2 = loss = orig_state_dict = scores = token_y = token_n = input_ids = attention_mask = choice_ids = residual_stream = residual_stream2 = None
# clear_mem()
# return out
def get_layer_selection(self, layer_names):
"""Sometimes we don't want to save all layers.
# def get_layer_selection(self, layer_names):
# """Sometimes we don't want to save all layers.
We skip the first few (data leakage?). Stride the the middle (could be valuable), and include the last few (possibly high level concepts).
# We skip the first few (data leakage?). Stride the the middle (could be valuable), and include the last few (possibly high level concepts).
See also https://www.lesswrong.com/posts/bWxNPMy5MhPnQTzKz/what-discovering-latent-knowledge-did-and-did-not-find-4
"""
module_names = [k for k,v in self.model.named_modules()]
layers_not_found = set(layer_names)-set(module_names)
assert len(layers_not_found)==0, f"some layers not found in model: {layers_not_found}. we have {layer_names}"
# See also https://www.lesswrong.com/posts/bWxNPMy5MhPnQTzKz/what-discovering-latent-knowledge-did-and-did-not-find-4
# """
# module_names = [k for k,v in self.model.named_modules()]
# layers_not_found = set(layer_names)-set(module_names)
# assert len(layers_not_found)==0, f"some layers not found in model: {layers_not_found}. we have {layer_names}"
# for self.layer_padding, skip the first few
num_layers = len(layer_names)-1
strided_layers = torch.arange(
self.layer_padding,
num_layers-self.layer_padding,
self.layer_stride,
).tolist()
# for self.layer_padding ALWAYS include the last few. Why, this is based on the intuition that the last layers may be the most valuable
last_few = torch.arange(num_layers-self.layer_padding, num_layers).tolist()
layers_inds = sorted(set(list(strided_layers)+list(last_few)))
return [layer_names[i] for i in layers_inds]
# # for self.layer_padding, skip the first few
# num_layers = len(layer_names)-1
# strided_layers = torch.arange(
# self.layer_padding,
# num_layers-self.layer_padding,
# self.layer_stride,
# ).tolist()
# # for self.layer_padding ALWAYS include the last few. Why, this is based on the intuition that the last layers may be the most valuable
# last_few = torch.arange(num_layers-self.layer_padding, num_layers).tolist()
# layers_inds = sorted(set(list(strided_layers)+list(last_few)))
# return [layer_names[i] for i in layers_inds]
src/datasets/intervene.py
+1
View File
@@ -67,6 +67,7 @@ def create_cache_interventions(model, tokenizer, cfg, N_fit_examples=20, batch_s
train_labels = np.array(dataset_fit['label_true'])
if get_negative:
# FIXME: does this work with PCA, since it's directionless
train_labels = -1 * train_labels
rep_reading_pipeline = pipeline("rep-reading", model=model, tokenizer=tokenizer)
honesty_rep_reader = rep_reading_pipeline.get_directions(
src/probes/pl_ranking.py
+2 -1
View File
@@ -163,7 +163,8 @@ class PLConvProbeLinear(PLRankingBase):
self.head = nn.Sequential(
LinBnDrop(n, n),
LinBnDrop(n, n),
nn.Linear(n, 1),
nn.Linear(n, 1),
# nn.Tanh(),
)
def forward(self, x):
src/repe/rep_control_pipeline_baukit.py
+2 -1
View File
@@ -87,7 +87,8 @@ class RepControlPipeline2(FeatureExtractionPipeline):
assert inputs['input_ids'].ndim == 2, f"expected input_ids to be (batch, seq), got {inputs['input_ids'].shape}"
# make intervention functions
layers_names = [self.layer_name_tmpl.format(i) for i in activations[0].keys()]
layers_names = [self.layer_name_tmpl.format(i) for i in activations[0].keys()]
# FIXME: [0] is positive, [1] is negative. We can also multiply by -1, 0, or 1
activations_pos_i = Activations({self.layer_name_tmpl.format(k):v for k,v in activations[1].items()})
activations_neut = Activations({self.layer_name_tmpl.format(k):0. * v for k,v in activations[0].items()})
edit_fn_pos = partial(intervention_meta_fn2, activations=activations_pos_i)
src/repe/rep_reading_pipeline.py
+1
View File
@@ -140,6 +140,7 @@ class RepReadingPipeline(Pipeline):
relative_hidden_states = {k: np.copy(v) for k, v in hidden_states.items()}
for layer in hidden_layers:
for _ in range(n_difference):
# TODO: check this, it's even - odd? on what dimension?
relative_hidden_states[layer] = relative_hidden_states[layer][::2] - relative_hidden_states[layer][1::2]
# get the directions