repeng
code:sh
poetry install
poetry shell
experiments.ipynbを始める。
固定されたバージョンだとdata did not match any variant of...でtokenizerうまくいかなかったのでpoetry update実施。
データセットはこの形式を作る。
code:python
@dataclasses.dataclass
class DatasetEntry:
positive: str
negative: str
make_dataset内は下記の定義になっている。
code:python
DatasetEntry(
positive=f"{user_tag} {positive_template} {asst_tag} {suffix}",
negative=f"{user_tag} {negative_template} {asst_tag} {suffix}",
)
こんな感じでデータ渡るので、[INST] Act as if you're extremely high on psychedelic drugs. {suffix}みたいな感じでprefixは割と固定で、suffixにデータセット内に存在する、"I can relate to",とか"That poem",とか、そういった言葉が入ってネガティブとポジティブで同様のsuffixで補完される。その差分をみてControl Vectorとして取り出す。 code:python
trippy_dataset = make_dataset(
"Act as if you're extremely {persona}.",
truncated_output_suffixes_512, # gives (subjectively) better results with slightly fewer samples
)
得られたベクターの大きさ確認。全レイヤの出力に介入している!
というのもControlModelを初期化するときにdef __init__(self, model: PreTrainedModel, layer_ids: typing.Iterable[int]):で初期化するlayer_idを渡している。この番号のレイヤを対象にする。 (ただ動作的にはtrainは全層するが、layer_idを渡した場合にそのlayer_idに対してのみVectorを割り当てるみたいな仕様っぽい。そりゃそうか、ControlModelって別にtrainするためだけじゃなくてControl Vectorを割り当てたいだけのときにもインスタンス化するので)
code:python
for layer, direction in happy_vector.directions.items():
print(f"Layer {layer}: {direction.shape}")
Layer 31: (4096,)
Layer 30: (4096,)
Layer 29: (4096,)
Layer 28: (4096,)
Layer 27: (4096,)
Layer 26: (4096,)
Layer 25: (4096,)
Layer 24: (4096,)
Layer 23: (4096,)
Layer 22: (4096,)
Layer 21: (4096,)
Layer 20: (4096,)
Layer 19: (4096,)
Layer 18: (4096,)
Layer 17: (4096,)
Layer 16: (4096,)
Layer 15: (4096,)
Layer 14: (4096,)
Layer 13: (4096,)
Layer 12: (4096,)
Layer 11: (4096,)
Layer 10: (4096,)
Layer 9: (4096,)
Layer 8: (4096,)
Layer 7: (4096,)
...
Layer 4: (4096,)
Layer 3: (4096,)
Layer 2: (4096,)
Layer 1: (4096,)
generateを呼び出す前にmodel.set_control(vector, positive_coeff)を行うことでベクトルがモデルに反映される。
model.reset()でベクトル影響ない状態に戻せる。
sae.ipynb。
ここにLlamaの例載ってるので参考になる。
vector_ops.ipynb。
複数ベクトルの和を取ったり差を取ったりするやつをやっている。
trainの中身はどうなっている?
まず処理はread_representations内で行われている。
デフォルトはpca_diffなのでその処理を追う。
まずbatched_get_hiddensでポジティブとネガティブ入力した最終トークンの隠れ状態かえってきている。
それの差を取る(ポジティブとネガティブの差)。
code:py
if method == "pca_diff":
元論文でもPCAをしていたのでそのノリ。squeezeで大きさ1の次元を削除。 code:python
if method != "umap":
# shape (1, n_features)
pca_model = PCA(n_components=1, whiten=False).fit(train)
# shape (n_features,)
directionslayer = pca_model.components_.astype(np.float32).squeeze(axis=0) やること:
日本語化
既存ので学習したあとに日本語質問で対応できるか?
PCAの可視化のやつをどうしているか?
出力の「誠実度ベクトル」が小さくなったときに検知する方法
There's also a bunch of stuff in there about reading those vectors, to see if e.g. a model is lying by tracking when the honesty vector for a token goes down. I didn't try that, yet.
inputがレイヤに渡る入力、outputが出力。介入もできる。
code:py
hook(module, input, output) -> None or modified output
保存する場合はこんな感じで保存する配列とかを外で準備する。
code:py
# モデルの指定されたレイヤーの出力と勾配を保存するクラス
class SaveOutput:
def __init__(self, model, target_layer): # 引数:モデル, 対象のレイヤー
self.model = model
self.layer_output = []
self.layer_grad = []
# 特徴マップを取るためのregister_forward_hookを設定
self.feature_handle = target_layer.register_forward_hook(self.feature)
# 勾配を取るためのregister_forward_hookを設定
self.grad_handle = target_layer.register_forward_hook(self.gradient)
# self.feature_handleの定義時に呼び出されるメソッド
## モデルの指定されたレイヤーの出力(特徴マップ)を保存する
def feature(self, model, input, output):
activation = output
self.layer_output.append(activation.to("cpu").detach())
# self.grad_handleの定義時に呼び出されるメソッド
## モデルの指定されたレイヤーの勾配を保存する
## 勾配が存在しない場合や勾配が必要ない場合は処理をスキップ
def gradient(self, model, input, output):
# 勾配が無いとき
if not hasattr(output, "requires_grad") or not output.requires_grad:
return # ここでメソッド終了
# 勾配を取得
def _hook(grad):
# gradが定義されていないが、勾配が計算されると各テンソルのgrad属性に保存されるっぽい(詳細未確認)
self.layer_grad.append(grad.to("cpu").detach())
output.register_hook(_hook)
# メモリの解放を行うメソッド、フックを解除してメモリを解放する
def release(self):
self.feature_handle.remove()
self.grad_handle.remove()
print(model)してモデルアーキテクチャを吐く。
レイヤのアウトプットに介入するということはpost_attention_layernormということか。
これの認識であっている。
各層のレイヤを取得するときにmodel_layer_listを使っているのだが、これはTransformerブロックを探してリスト化している。 code:py
from repeng.control import model_layer_list
layer = model_layer_list(model)
LlamaDecoderLayer(
(self_attn): LlamaSdpaAttention(
(q_proj): Linear(in_features=4096, out_features=4096, bias=False)
(k_proj): Linear(in_features=4096, out_features=1024, bias=False)
(v_proj): Linear(in_features=4096, out_features=1024, bias=False)
(o_proj): Linear(in_features=4096, out_features=4096, bias=False)
(rotary_emb): LlamaRotaryEmbedding()
)
(mlp): LlamaMLP(
(gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
(up_proj): Linear(in_features=4096, out_features=14336, bias=False)
(down_proj): Linear(in_features=14336, out_features=4096, bias=False)
(act_fn): SiLU()
)
(input_layernorm): LlamaRMSNorm((4096,), eps=1e-05)
(post_attention_layernorm): LlamaRMSNorm((4096,), eps=1e-05)
)
(これ↓かなり微妙なコードで、登録したregister_forward_hookをremoveできるようにしないと一生hookが発火する)
code:py
for name, module in model.named_modules():
if "post_attention_layernorm" in name:
module.register_forward_hook(hook_fn)
print(f"Hook registered to {name}")
code:py
res = batched_get_hiddens(model, tokenizer, inputs, list(range(-5, -18, -1)), batch_size=2)
これでトップから-5〜-18層の隠れ層の出力を取得できる。