生体分子の構造予測やタンパク質設計に使う30超のAIモデルで、推論処理が平均約4倍に高速化した。Claudeは高負荷なGPU演算とモデル固有の無駄を4週間弱で最適化し、メモリー使用量も削減した。
30超のモデルへ4週間弱で最適化を展開
Claudeは、生体分子構造予測、タンパク質設計、タンパク質言語モデル、ゲノミクスにまたがる30超のオープンソースモデルを4週間弱で最適化した。対象には、構造を推定するモデルだけでなく、配列や構造を生成する複数種のモデルアーキテクチャが含まれる。
Anthropicによると、通常は経験豊富な技術者チームがモデルごとに数週間を要する作業だという。今回は、生体分子モデリングの経験を持つ一方、推論最適化やGPUカーネル開発の経験はなかった技術スタッフ2人がClaudeを監督し、複数分野のモデルへ作業を展開した。
3次式で増える構造予測の計算量をカーネルで抑える
AlphaFold3、OpenFold3、Boltz-2などの構造予測モデルでは、triangle attentionとtriangle multiplicationが実行時間とメモリーの大きな部分を占める。これらは生体分子を構成するトークンの三つ組を扱うため、システム規模を2倍にすると実行時間とメモリー使用量がともに8倍になる。
カスタムGPUカーネル群のFlashPairformerは、この二つの演算をGPU上で効率化する。Anthropicの測定では、標準実装に対してtriangle attentionが平均2.7～2.9倍、triangle multiplicationがモデル構成に応じて1.7～3.2倍に高速化した。
再計算と不要な分岐をモデルごとに除く
共通カーネルに加え、Claudeは各モデルの推論経路を個別に調べ、同じ処理の再計算を避けるキャッシュを導入した。実行しても結果が一定になる不要な分岐は、その出力へ直接置き換えて処理を簡略化した。
Anthropicの測定では、こうした変更を組み合わせた対象タスクが、精度低下を最小限に抑えた条件で平均約4倍に高速化した。出力を変更しない条件でも約2倍となり、構造予測モデルの各最適化版では下流の予測性能に影響がないことを確認したとしている。
1万トークン超の分子系を1台のGPUノードで扱う
高速化とは別に、メモリー使用量を抑えるbigモードが大規模な分子系の入力範囲を広げた。アミノ酸、核酸、小分子、イオンを合計したトークン数が1万を超える生体分子系を、1台のNVIDIA GPUノードで予測できるようになった。
このモードは、通常の実行ではメモリーに収まらない入力を対象に、ピークメモリーを最小化する。公開キットでは、元の出力と同一のまま高速化するexact、小さな数値差を許容して速度を上げるfast、入力規模を優先するbigを区別し、モデルごとに対応するモードを定めている。
タンパク質設計のGPU時間を約100分の1へ削減
推論コードの効率化は、複数モデルを連携させるタンパク質設計にも反映された。以前のエージェント型設計手法を簡略化して今回の最適化と組み合わせた結果、Anthropicは以前の結果と同程度のin silico性能を約100分の1のGPU時間で得たとしている。
最適化コードは、36個の推論最適化キットとして公開された。各キットには固定された上流モデルと最適化パッケージが分けて収録され、構造予測、タンパク質設計、タンパク質・ゲノム言語モデルの推論経路へ個別に適用できる。