【ABOPD】拡散モデルに”自分の癖”を教え込む、抗体CDR-H3デザインの後処理学習

論文タイトル

ABOPD: Antibody CDR Design via On-Policy Distillation

出典

https://arxiv.org/abs/2607.18835
GitHub: https://github.com/little1d/ABOPD

要旨

抗体のCDR（相補性決定領域）、特に最も柔軟で難しいCDR-H3をデザインする拡散モデルに対して、生成済みモデルをさらに鍛える「後処理学習」の手法ABOPDが提案されました。

用語

  • reference-derived state（参照由来状態）: 学習時に、正解構造にノイズを加えて人工的に作る状態
  • rollout state（自己生成状態）: 推論時に、モデルが自分自身のdenoisingを繰り返して実際に辿る状態。学習時のreference-derived stateとは分布がズレることがあります
  • on-policy distillation（オンポリシー蒸留）: 生徒モデルに実際に解かせてみて、その答案を教師モデルがその場で添削する後処理手法です。模範解答を丸暗記させるのではなく、生徒自身の「クセ」が出た状態を直接直しにいくイメージです
  • privileged information（特権情報）: 学習時にのみ利用できる追加情報。ABOPDでは教師モデルにだけ正解の骨格座標を与えます

解説など

背景

DiffAbをはじめとする拡散モデルによるCDR生成は近年進歩してきましたが、既存研究の多くはモデルアーキテクチャや構造表現の改良に集中しており、生成済みモデルをさらに鍛える後処理学習はあまり手がつけられていませんでした。

後処理学習（post-training）は、汎用的な目的で学習済みのモデルを、狙った性能への最適化や、学習時と生成時で生じる分布のズレの解消など、特定の課題に絞って安価に鍛え直すために行うものです。ゼロから学習し直すよりも低コストで、ピンポイントな改善が見込めます。

本論文の特徴

ABOPDが何を入力し、何を出力する手法かを先に整理すると、この手法の位置づけが分かりやすくなります。

入力は、①抗原の立体構造、②抗体フレームワーク（デザインしない部分）、③デザイン対象として指定したCDR（CDR-H3のみ、あるいは６つ全部など）、の３つです。加えてランダムノイズを初期状態として与え、そこから出発してdenoisingを繰り返しながら組み立てていきます。あらかじめお手本の配列や構造を用意する必要はありません。

出力は、指定したCDR領域の各残基について、アミノ酸配列・Cα座標・局所的な向き（残基ごとの回転の基準）の３点セットです。論文の数式ではpθ(R𝒢 | C)と表現されており、抗原とフレームワーク（C）を条件に、デザイン対象の残基集合（𝒢）の配列と構造をまとめて生成する、という枠組みになっています。

ここで重要なのは、ABOPDが変えているのはこの入出力の対応関係そのものではなく、その対応関係をどう学習するかという訓練方法だという点です。ベースにしているDiffAb／H-DiffAbと入出力の形は変わらず、そのうちCα座標の精度をピンポイントで底上げしているのがABOPDの工夫になります。

拡散モデルの後処理でよく起きる課題として、学習時にモデルへ見せる状態（正解構造にノイズを加えたreference-derived state）と、実際の生成時にモデルが辿る状態（自分自身の予測を繰り返し使うrollout state）にズレが生じる点が指摘されています。特にCDR-H3のような柔軟なループでは、このズレによって生成の初期に生じたCα座標の誤差が後続ステップに伝播し、抗原との立体的な衝突や配置のズレとして蓄積しやすくなります。

ABOPDはこのズレに対処するため、次の3段階で構成されます。

  1. Hybrid pretraining: 単一CDR・複数CDR同時・全６CDR同時デザインを混合して学習した「H-DiffAb」を作成します（DiffAb標準の約8倍の学習量）
  2. Backbone-aware Teacher adaptation: H-DiffAbに、学習時にしか使えない正解構造の骨格座標（N・Cα・C・O・Cβ座標や距離特徴量）を軽量なアダプタ（追加パラメータ2.38万個のみ）で与え、教師モデルを作ります。この教師モデルは座標MSEを0.2019から0.1503まで改善する、精度の高い補正ターゲットを提供できるようになります
  3. On-Policy Distillation（本手法の主眼）: ここが核心です。模範解答を暗記させるのではなく、生徒モデルに実際に自分の力でCDRを生成させ（denoisingのロールアウト）、その答案を教師モデルがその場で添削します。生成の途中で実際に通過した状態に対して、教師モデルの予測を「正しい方向」として与えるイメージです。

配列や向きではなく座標のみを蒸留対象にしている点も特徴的です。著者らはablation実験により、座標のみの蒸留が最もRMSD改善に効くことを確認しています。配列単体を蒸留すると逆にAARが悪化し、向き単体では効果が限定的だったとのことです。

結果

RAbD CDR-H3ベンチマーク（60ターゲット）では、H-DiffAb単体のRMSD 2.37Åに対し、ABOPDは1.95Åまで改善しました（0.42Å減）。クラッシュ率（Clash_in）も3.46%から0.94%まで大きく低下しました。

19複合体を対象にした６CDR同時デザインでも、CDR-H3のRMSDが2.94Åから2.48Åまで改善し、クラッシュ率も低下しています。

考察

ABOPDが改善しているのはあくまでCα骨格構造で、側鎖のパッキング（JSDsc）や界面改善率（IMP）はH-DiffAbよりやや劣る結果も出ています。著者らはこれを、構造的に信頼性の高いモードへ生成が集中し、多様性（探索の幅）が犠牲になっているのではという、忠実度と多様性のトレードオフとして考察しています。