メインコンテンツまでスキップ

02. アーキテクチャと学習

このページでは、AlphaGenome を支える土台——モデルアーキテクチャ(U-Net + Transformer)配列並列による塩基解像度学習事前学習と蒸留の2段階学習——を、LLM の技術と対応づけながら詳しく見ていきます。

1. 設計課題:1 Mb を 1 bp で出すには

01 で見たように、AlphaGenome の挑戦は 「1 Mb の長い入力」と「1 bp の細かい出力」を同時に成立させる ことです。これには2つの矛盾する要請を1つのネットワークで満たす必要があります。

  • 局所的な細かいパターン — スプライス部位・転写因子フットプリント・ポリアデニル化部位などは、わずか数塩基の並びで決まります。これを捉えるには 塩基レベルの解像度 が要ります。
  • 長距離の依存関係 — エンハンサーとプロモーターは数十〜数百 kb 離れることがあります。これを捉えるには 広い文脈 が要ります。

AlphaGenome はこれを、畳み込み(局所)と Transformer(長距離)を役割分担させる U-Net 風アーキテクチャ で解決します。

2. U-Net 風バックボーン:縮めて、混ぜて、戻す

全体構造は、医用画像セグメンテーションで有名な U-Net(Ronneberger et al., 2015)に着想を得たエンコーダ–デコーダ型です(Fig. 1a, Extended Data Fig. 1a)。

入力 DNA 1 Mb(1 bp 解像度)Encoder(畳み込み)1 bp → 128 bp に縮小Transformer タワー128 bp 解像度・デバイス間通信長距離(エンハンサー–プロモーター)Decoder(畳み込み)128 bp → 1 bp に復元スキップ接続1D 埋め込み(1 bp / 128 bp)→ 各トラックは線形変換で予測2D 埋め込み→ 接触マップ(2,048 bp)出力ヘッド(モダリティごと・各解像度)RNA-seq / DNase / ヒストン ChIP / スプライス / 接触マップ …
U-Net 風バックボーン(論文 Fig. 1a / Extended Data Fig. 1 に対応)。畳み込みで縮小→Transformer で長距離を混ぜる→畳み込みで復元。スキップ接続が細かい解像度の情報をデコーダに渡す

処理の流れは「縮めて、混ぜて、戻す」の3段です。

  1. エンコーダ(畳み込み+max pooling) — 入力配列を 1 bp から 128 bp へ段階的にダウンサンプリング し、チャンネル数を増やします。畳み込みは 局所的な配列パターン のモデル化を担います。
  2. Transformer タワー(128 bp 解像度) — 縮小された表現に対して Self-Attention を適用し、エンハンサー–プロモーター相互作用のような長距離依存 を捉えます。後述するデバイス間通信もここで効きます。
  3. デコーダ(畳み込み+アップサンプリング) — 解像度を 128 bp から 1 bp へ復元。このとき スキップ接続(skip connection) でエンコーダの対応する段から細かい解像度の情報を引き込み、塩基レベルの精度を保ちます。
なぜ畳み込みと Transformer を分業させるのか

Transformer の Self-Attention は計算量が系列長の二乗(O(n2)O(n^2))なので、1 Mb をそのまま 1 bp 解像度で Attention にかけるのは非現実的です。そこで AlphaGenome は、重い Attention を「128 bp に縮小した後」だけに限定 します(系列長が約 1/128 になる)。塩基レベルの細かい処理は安価な畳み込みに任せ、Attention は長距離の混合に専念させる——この役割分担が、長文脈と高解像度の両立を可能にしています。Evo 2 が Attention 自体を Hyena 演算子で置き換えた(Evo 2 のアーキテクチャ)のとは対照的なアプローチです。

1D 埋め込みと 2D 埋め込み

エンコーダ/デコーダが生み出す表現は2種類あります。

  • 1次元(1D)埋め込み(1 bp / 128 bp 解像度)— 線形なゲノムの表現。ほとんどのトラック予測の土台 で、各トラックの値は基本的にこの埋め込みの 線形変換 で得ます。
  • 2次元(2D)埋め込み(2,048 bp 解像度)— ゲノム上の2領域 (i,j)(i, j)空間的相互作用 を表す表現。接触マップ(contact map) の予測に使います。接触マップはペアワイズの行列なので、2次元の表現が必要になります。

スプライスジャンクション専用ヘッド

トラック予測は原則「埋め込みの線形変換」ですが、スプライスジャンクション数の予測だけは別の仕組み を使います。これは、ドナー部位とアクセプター部位の 1D 埋め込みのペア間の相互作用 を捉える専用機構です(Extended Data Fig. 1)。「どのドナーとどのアクセプターがつながってイントロンが除去されるか」は本質的にペアの問題なので、専用設計になっています(詳細は 04)。

3. 配列並列:8 台の TPU で 1 Mb を 1 bp 学習する

1 Mb の配列を 1 bp 解像度のまま扱うと、中間表現が巨大になり1台のアクセラレータに乗りません。AlphaGenome は 配列並列(sequence parallelism) でこれを解決します。

  • 1 Mb の配列を 131 kb のチャンク に分割し、相互接続された 8 台の TPU(v3) に分散して処理します。
  • Transformer タワーでは、デバイス間通信 を行って各チャンクが他チャンクの情報を参照できるようにします。これにより、分割していても 1 Mb 全体にわたる Attention が成立します。
LLM とのつながり:シーケンス並列

「長い系列を複数デバイスに分割し、Attention のときだけ通信して全体をつなぐ」のは、大規模 LLM の長文脈学習で使われる sequence parallelism / context parallelism とまさに同じ発想です。AlphaGenome はこの分散学習技術を、ゲノムという超長系列に適用して 1 bp 解像度の学習を可能にしています。

4. 2段階学習:事前学習 → 蒸留

AlphaGenome は 事前学習(pretraining)蒸留(distillation) の2段階で訓練されます(Fig. 1b,c)。

1. 事前学習 Pretrainingfold 別モデル(4分割交差検証)3/4 で学習・1/4 で検証 → 汎化を評価all-fold 教師モデル全領域で学習 → 蒸留の教師に2. 蒸留 Distillation教師アンサンブル(all-fold)単一の生徒モデル入力を増強・変異させて教師を再現→ 変異効果予測に使う本番モデル生徒モデルは H100 1 台で 1 変異あたり 1 秒未満(デバイス呼び出し1回)
2段階学習(論文 Fig. 1b,c に対応)。事前学習で fold 別/all-fold モデルを作り、all-fold 教師アンサンブルを単一の生徒モデルに蒸留する

第1段階:事前学習

観測された実験データを教師信号に、2種類のモデル を作ります。

  • fold 別モデル(fold-specific)4分割交差検証 で、参照ゲノムの 3/4 を学習に、残り 1/4 を検証・テストに使います。「学習時に見ていないゲノム領域」での 汎化性能を評価する ために使われます(→ 03 のトラック評価)。
  • all-fold 教師モデル — 利用可能な全領域で学習したモデル。次の蒸留段階で 教師(teacher) として働きます。

学習入力は、1 Mb 区間をサンプリングし、ランダムシフト逆相補鎖化(50%) といった データ増強(augmentation) を施したものです。

第2段階:蒸留

事前学習と同じアーキテクチャをもつ 単一の生徒モデル(student) を、all-fold 教師アンサンブルの出力を再現するように 訓練します。このとき入力には、ランダムな増強と変異(mutation)を加えた配列 を使います。

蒸留の狙いと効果:

  • 頑健性と変異効果予測精度の向上 — 先行研究(Zhou et al., 2024)と同様、蒸留により単一モデルでも頑健性と変異効果予測の精度が上がります。
  • 推論の超高速化 — 生徒モデルは 全モダリティ・全細胞型を、1 変異あたりデバイス呼び出し1回 で予測します。NVIDIA H100 1 台で 1 秒未満。複数の独立学習モデルをアンサンブルする方式に比べ、大規模な変異スクリーニングで圧倒的に効率的です。
入力への変異が蒸留のカギ

07 のアブレーションで詳しく見ますが、蒸留時に入力配列をランダムに変異させないと 生徒モデルの性能が落ちます(eQTL sign で −0.06、causality で −0.01 など)。「教師が参照配列でどう振る舞うか」だけでなく「変異させた配列で教師がどう振る舞うか」まで生徒に教え込むことが、変異効果予測の精度に直結する、というわけです。変異効果予測が本番タスクなのだから、蒸留時にも変異を見せる——理にかなった設計です。

LLM とのつながり:知識蒸留とアンサンブル代替

「複数の大きな教師モデルのアンサンブルを、1つの軽い生徒モデルに蒸留する(knowledge distillation)」のは、自然言語 LLM でも標準的な圧縮・高速化手法です。アンサンブルは強力ですが推論コストが重い——その性能を保ったまま1モデルに畳み込むのが蒸留です。AlphaGenome は、1 変異 1 秒未満 という実用的な推論速度を、この蒸留で実現しています。

5. 学習データ:ヒトとマウスの実験アトラス

AlphaGenome の教師信号は、公開された実験データ です。Evo 2 が生 DNA 配列だけで自己教師あり学習するのに対し、AlphaGenome は 「配列 → 実験測定値」の対応 を教師あり学習します。

データソース提供するもの
ENCODEDNase / ATAC / ヒストン・TF の ChIP-seq など
GTEx多組織の RNA-seq・スプライシング・QTL
FANTOM5CAGE(転写開始)
4D NucleomeHi-C / Micro-C(接触マップ)
ClinVar・gnomAD変異の病原性アノテーション・頻度(評価用)

出力トラック数は ヒト 5,930 / マウス 1,128。多様な組織・細胞型・細胞株にわたります(メタデータは論文 Supplementary Table 1, 2)。

「種」を入力に与える

AlphaGenome は DNA 配列に加えて 種の識別子(ヒット/マウス) も入力に取ります。これにより、1つのモデルで両種のトラックを予測でき、種をまたいだ共有表現を学べます。

6. まとめ

  • U-Net 風バックボーン:畳み込み(局所・1 bp)で縮小 → Transformer(長距離・128 bp)で混合 → 畳み込みで 1 bp に復元。スキップ接続で解像度を維持。
  • 1D 埋め込み(トラック予測)と 2D 埋め込み(接触マップ)、さらに スプライスジャンクション専用ヘッド
  • 配列並列:1 Mb を 131 kb チャンクに分け 8 台の TPU で塩基解像度学習。
  • 2段階学習:事前学習(fold 別+all-fold 教師)→ 蒸留(変異を加えた入力で教師を再現する単一生徒モデル)。H100 1 台で 1 変異 1 秒未満
  • 教師信号は ENCODE・GTEx・FANTOM5・4D Nucleome などの実験アトラス。

次の 03. トラック予測性能 では、このモデルが「見たことのないゲノム領域」でどれだけ正確にトラックを当てられるか——変異効果予測の前提となる基礎性能——を見ていきます。