第 1-3 部が使う埋め込みと活性は, あらかじめ計算して配布しています. その作り方と精度をまとめます.
| 節 | 題材 | 内容 |
|---|---|---|
| 02 | — | 埋め込みが落ち着くまでに要る左文脈の長さ |
| 03 | — | エンコードの 2 式, スパース性, 復元の精度 |
| 04 | — | 全ゲノムの活性行列の作りと, 直接計算との一致 |
01. 準備¶
evo2_sae と, 大腸菌ゲノム NC_000913.3 の 4,162,560-4,193,280, 30,720 塩基ぶんの埋め込みを読み込みます.
| ファイル | サイズ | 内容 |
|---|---|---|
region_emb.pt |
252 MB | Evo 2 の埋め込み. 30,720 塩基ぶん |
context.pt |
84 MB | 文脈の長さを測るための埋め込み. ウィンドウの重なり 5 か所ぶん |
ecoli.bed |
0.2 MB | ゲノム全体の遺伝子座標 |
ecoli_values.f16 ほか 2 本 |
2.4 GB | SAE の活性行列 |
# Colab には torch と huggingface_hub が入っている. 手元で動かす場合は別途入れる
%pip install -q git+https://github.com/GenAIBio/evo2-sae-handson
import time
import numpy as np
import torch
import matplotlib.pyplot as plt
from huggingface_hub import hf_hub_download
import evo2_sae
evo2_sae.setup("https://huggingface.co/datasets/suzuki-2001/evo2-sae-handson/resolve/main")
d = torch.load(evo2_sae.download("region_emb.pt"), map_location="cpu")
x = d["emb"].float()
start, end = d["start"], d["end"]
chrom = d["accession"]
print(f"{tuple(x.shape)} {chrom}:{start:,}-{end:,} {d['model']} ({d['layer']})")
downloading region_emb.pt
(30720, 4096) NC_000913.3:4,162,560-4,193,280 evo2_7b (blocks.26)
02. 埋め込みに必要な文脈¶
Evo 2 は配列を左から右へ読み, 各位置で次に来る塩基を当てるように学習されています. このようなモデルを自己回帰モデルと呼びます.
読む向きが決まっているので, ある位置の内部状態はそこまでに読んだ配列だけで決まり, 右側の配列は影響しません. そのため, 左に配列がほとんど無い位置では, 次の塩基を当てる手がかりがありません.
この埋め込みは 16,384 bp のウィンドウ単位で計算されていて, 隣り合うウィンドウは 1024 bp 重なっています. 重なりにある塩基は両方のウィンドウに現れ, それぞれで左文脈の長さが違います.
赤い線の座標は両方のウィンドウに入ります. その塩基より前にある配列を左文脈と呼ぶことにすると, window 1 では 15,360 bp 以上あり, window 2 では 1023 bp 以下しかありません. 左文脈の長さだけが違う 2 つの埋め込みが, 同じ塩基について得られます. context.pt にはこの重なりが 5 か所ぶん入っていて, それぞれ両方の埋め込みを保持しています.
ctx = torch.load(evo2_sae.download("context.pt"), map_location="cpu")
# どちらも (5 か所, 1024 塩基, 4096). short は左文脈 0-1023 bp, long は 15,360 bp 以上
short, long = ctx["short"].float(), ctx["long"].float()
norm = short.norm(dim=2).mean(0).numpy()
base = float(long.norm(dim=2).mean())
agree = torch.cosine_similarity(short, long, dim=2).mean(0).numpy()
fig, ax = plt.subplots(1, 2, figsize=(9.6, 3.4))
ax[0].plot(norm[:64], color="#b91c1c", lw=1.4)
ax[0].axhline(base, color="#334155", lw=1.0, ls="--")
ax[0].set_yscale("log")
ax[1].plot(agree, color="#2563eb", lw=1.4)
ax[1].set_ylim(-0.05, 1.02)
for panel, ylabel in zip(ax, ("Embedding norm", "Cosine similarity")):
panel.set_xlabel("Left context (bp)")
panel.set_ylabel(ylabel)
panel.set_facecolor("#f7f7f7")
for spine in panel.spines.values():
spine.set(linewidth=1.4, color="black")
fig.tight_layout()
plt.show()
downloading context.pt
ノルムと向きでは, 落ち着くまでに要する文脈の長さが 1 桁以上違います. ノルムが基準の 2 倍を超えるのは先頭 11 塩基だけで, 16 bp あれば収まります. 向きは, ノルムが落ち着いたあとも揃いません. 左文脈 16 bp での cos は 0.08, 128 bp でも 0.62 で, 0.9 に届くのは 400 bp 前後, 0.99 には 810 bp かかります.
配布データでは, 文脈の足りない位置の影響を抑えるため, 各ウィンドウの先頭 1024 bp を捨てています. 上の測定では cos が 0.99 に届くのが 810 bp なので, それより余裕を取った値です. 捨てた区間は, 十分な文脈を持つ 1 つ前のウィンドウが受け持ちます. ゲノムが環状なので, 先頭のウィンドウも末尾から文脈を取れます. region_emb.pt も全ゲノムの活性行列も, この規則で作られています.
03. エンコードと復元¶
Goodfire が公開している Evo 2 用の SAE を使います. 論文で用いられたものと同じチェックポイントが Hugging Face にあります.
エンコードは 2 段で, ある塩基の埋め込みを $\mathbf{x} \in \mathbb{R}^{4096}$ とすると, まず
$$\mathbf{a} = \mathrm{ReLU}\!\left(\mathbf{x}\,\mathbf{W} + \mathbf{b}_{\mathrm{enc}}\right)$$
を計算し, 次に
$$\mathbf{z} = \mathrm{BatchTopK}_k(\mathbf{a})$$
で疎にします. $L$ 塩基をまとめて通したとき, $L \times 32768$ 個の活性から大きい順に $k \times L$ 個だけを残し, 残りを 0 にします. 塩基ごとに $k$ 個ずつ取るのではないので, 1 塩基あたりの個数は平均 $k$ 個を保ったまま増減します. このチェックポイントは $k = 64$ です.
デコードは, エンコードと同じ $\mathbf{W}$ を転置して使います.
$$\hat{\mathbf{x}} = \mathbf{z}\,\mathbf{W}^{\top} + \mathbf{b}_{\mathrm{dec}}$$
復元の精度は, 戻し損ねた量を元の量で割った nMSE で測ります. 0 なら完全に復元できており, 1 なら元の大きさぶん丸ごと外しています.
$$\mathrm{nMSE} = \frac{\lVert \mathbf{x} - \hat{\mathbf{x}} \rVert^{2}}{\lVert \mathbf{x} \rVert^{2}}$$
パラメータは $4096 \times 32768$ の行列 $\mathbf{W}$ 1 枚とバイアス 2 本です. エンコードは $\mathbf{x}$ と $\mathbf{W}$ の各列との内積, デコードは同じ列を活性で重みづけた和で, 冒頭で述べた辞書とはこの列のことです. エンコードとデコードで同じ $\mathbf{W}$ を使う理由は 2 つあります. 検出に使うベクトルと復元に足すベクトルは同じはずだという想定と, パラメータが半分で済むことです (Cunningham et al. 2023). 上の 2 式は Evo 2 の公開ノートブックの BatchTopKTiedSAE と同じもので, GitHub にあります.
ckpt = hf_hub_download("Goodfire/Evo-2-Layer-26-Mixed", "sae-layer26-mixed-expansion_8-k_64.pt")
sd = torch.load(ckpt, map_location="cpu")
W = sd["_orig_mod.W"].float()
b_enc = sd["_orig_mod.b_enc"].float()
b_dec = sd["_orig_mod.b_dec"].float()
for name, v in sd.items():
print(name, tuple(v.shape))
sae-layer26-mixed-expansion_8-k_64.pt: reconstructing file: 0%| | 0.00B / 537MB
sae-layer26-mixed-expansion_8-k_64.pt: downloading bytes: | 0.00B
_orig_mod.W (4096, 32768) _orig_mod.b_enc (32768,) _orig_mod.b_dec (4096,)
ここまでの手順を実際に通します. ただし 30,720 塩基すべてでは $\mathbf{a}$ と $\mathbf{z}$ が無料の Colab の RAM に収まらないので, 先頭 8,192 塩基で測ります.
K = 64
L = 8_192 # 無料の Colab の RAM に収まる長さ
xs = x[:L]
def batch_topk(a, k):
# L x 32768 個の活性から, 大きい順に k x L 個だけ残す
floor = torch.topk(a.flatten(), k * a.shape[0]).values.min()
return torch.where(a >= floor, a, torch.zeros_like(a))
A = torch.relu(xs @ W + b_enc)
Z = batch_topk(A, K)
x_hat = Z @ W.T + b_dec
print("a の非ゼロ", round(float((A > 0).float().sum(1).mean()), 1))
print("z の非ゼロ", round(float((Z > 0).float().sum(1).mean()), 1))
print("nMSE", round(float(((xs - x_hat) ** 2).sum() / (xs ** 2).sum()), 3))
del x_hat
a の非ゼロ 681.9 z の非ゼロ 64.0
nMSE 0.484
1 段目の $\mathbf{a}$ は 1 塩基あたり 682 個が 0 でなく, まだ疎ではありません. BatchTopK が $k \times L$ 個ちょうどを残すので, 通したあとは定義どおり 64 個, 32768 列の 0.2% です.
nMSE は 0.48 です. 64 個の和では埋め込みの半分しか戻りません. 学習ではこの量を最小化しますが, 小さくしたからといって特徴が解釈しやすくなるとはかぎりません.
04. 全ゲノムの活性行列¶
埋め込みは 1 塩基あたり 4096 個の数値なので, 全長 4.64 Mb では 38 GB になり, そのままでは配布できません. 一方 BatchTopK を通したあとは 1 塩基あたり 64 個ほどしか 0 でないので, 0 を持たずに済みます. 非ゼロの値, その列番号, 各行の先頭位置の 3 本を並べる形にすると 2.4 GB に収まります. 疎行列をこう持つ形式を CSR (compressed sparse row) と呼びます.
行の先頭位置 indptr を見れば任意の座標に対応するバイト範囲が分かるので, HTTP のレンジ要求で必要なぶんだけを取れます. Region.from_genome が取得と BatchTopK をまとめています.
3 節で region_emb.pt から求めた Z と, 同じ 8,192 塩基を行列から引いたものを比べます.
stored = evo2_sae.Region.from_genome(chrom, start, start + L)
FIG4C = [2812, 30262, 26069, 13606] # rRNA, tRNA, 遺伝子間, ORF
here = np.stack([stored.track(f) for f in FIG4C], 1)
there = Z[:, FIG4C].numpy()
print("発火の一致", round(float(((here > 0) == (there > 0)).mean()), 4))
発火の一致 0.9999
Fig. 4c の 4 つの特徴で見ると, 発火の有無は 99.99% 一致します.
参考文献¶
Brixi G, Durrant MG, Ku J, Naghipourfar M, Poli M, Sun G, et al. Genome modelling and design across all domains of life with Evo 2. Nature 2026. doi:10.1038/s41586-026-10176-5
Cunningham H, Ewart A, Riggs L, Huben R, Sharkey L. Sparse autoencoders find highly interpretable features in language models. arXiv 2023. doi:10.48550/arXiv.2309.08600
Bussmann B, Leask P, Nanda N. BatchTopK sparse autoencoders. arXiv 2024. doi:10.48550/arXiv.2412.06410