CellWorld、潜在細胞表現による事前学習へ転換――5.74Mモデルが18の空間トランスクリプトミクスタスクで掲載ベースラインを上回る

空間トランスクリプトミクスデータは、細胞の遺伝子発現と組織内座標を同時に記録するが、測定プラットフォームの違いによってドロップアウト、シーケンス深度の差、その他の技術的ノイズが生じる。既存の基盤モデルの多くは、マスクされた遺伝子名や発現値をネットワークに再構成させるため、プラットフォーム固有の誤差まで学習する可能性がある。8月7日に発表されたCellWorldは、joint-embedding predictive architectureに似た目的関数を採用した。元の測定値を復元するのではなく、マスクされた細胞についてtarget encoderが生成する潜在表現を予測する。

学習時には各細胞を1つのtokenとして扱う。モデルは空間領域を1つマスクし、context encoderが可視細胞を読み取る。続いてpredictorが、マスクされた細胞について与えられる少量の部分的な発現ヒントも組み合わせ、指数移動平均で更新されるtarget encoderの出力を近似する。公開設定では、空間関係の表現に2D ALiBi、ランダムマスク率60%、部分的な発現ヒント10%、損失関数にMSEを使用する。4つのバージョンの学習可能パラメータ数は5.74Mから94.56Mで、論文によれば事前学習コーパスには4,600万個のヒト細胞が含まれる。

著者らは4つのホールドアウトデータセットで評価し、最小モデルのCellWorld-Smallでさえ、encoderを凍結した11のlinear probingタスクと7つのfull fine-tuningタスクにおいて、掲載されたすべてのベースラインを上回ったと報告している。さらに注目すべき点として、事前学習コーパスの5%しか使用しない一方で、生物学的ソースの幅広いカバレッジを維持したfrozen CellWorld-Largeは、7つの空間タスクで、full fine-tuningされたすべてのベースラインをなお上回った。これは「細胞数を単純に増やすことよりも、ソースの多様性と十分な最適化のほうが重要である可能性がある」という結論を支持する。ただし、これは著者らによる対照実験の結果であり、他のプラットフォームや疾患コホートへ直接一般化することはできない。

GitHubには、事前学習、linear probing、fine-tuning、`.h5ad`推論、multi-GPU起動用のコードがすでに含まれており、データの場所はユーザーが指定する。再現性を阻む最大の障壁は、4つの公式checkpointがアクセス承認を必要とする非公開のHugging Faceリポジトリに置かれていることと、学習に使用した4,600万細胞のデータがコードとともに公開されていないことだ。研究者は今後、データライセンス、プラットフォームをまたぐ外部検証、そして部分的な発現ヒントを除去した後も優位性が維持されるかを確認する必要がある。