Scaling an Autoregressive Transformer for Single-Cell Generation

技術要約：単一細胞生成のための自己回帰型トランスフォーマーのスケーリング

問題提起
本論文は、特定の細胞型に対する現実的な単一細胞遺伝子発現ベクトルを生成するという課題に取り組んでいる。既存の手法は、細胞株をカテゴリカルな名称、単一の例ベクトル、あるいはベクトルのシーケンスとして表現しているが、希少な測定値を補完し、in-silco 実験を可能にするための、スケーラブルで自己教師あり学習によるフレームワークが必要とされている。さらに、単一細胞基盤モデルにおいて、事前学習損失がモデルサイズ（$N$）およびデータサイズ（$D$）に対してどのようにスケールするかについては、予測可能なべき乗則に従うのか、あるいはデータの多様性の制約によりプラトー（停滞）に達するのかについて、相反する報告があり、この分野には明確な理解が欠けている。

手法
著者らは、学習された量子化VAE（Variational Autoencoder）トークナイザーと組み合わせた自己回帰型トランスフォーマーアーキテクチャを提案している。このアプローチは、主に以下の2つのコンポーネントを逐次的に学習させるものである。

1. トークナイザー (RQ-VAE): 残差量子化VAE（Residual Quantized VAE）は、連続的な単一細胞遺伝子発現ベクトル（18,080遺伝子）を8つの離散トークンのシーケンスにマッピングする。エンコーダーは入力を256次元の潜在ベクトルに圧縮し、その後、8つの残差コードブックにわたって粗から精へと量子化を行う。このトークナイザーはscBaseCountアトラス上で一度だけ学習され、トランスフォーマーの離散表現層として固定される。
2. トランスフォーマーモデル: 因果的トランスフォーマー（LLaMAアーキテクチャに基づく）が、これらの離散トークンのシーケンスをモデリングする。「文章」は、同じ細胞型の32個のセルからなるトークンブロックを、学習されたセパレータートークンで連結することによって構築される。モデルは、クロスエントロピー損失を用いた次トークン予測を通じて学習される。
3. スケーリング・グリッド: スケーリング則を調査するために、著者らはモデルサイズ（$N$）とデータサイズ（$D$）を独立して変化させる制御されたグリッドを構築した。
   • モデルサイズ: 固定された幅対深さの比率を維持しながら、1.31Mから83.9Mパラメータに及ぶ5つの構成。
   • データサイズ: 約2億セルのscBaseCountアトラスからサンプリングされた、38.1Mから362M個の細胞からなる学習予算。自然な細胞型の組成を保持するために比例サンプリングを用いている。
   • 合計: 22個の異なる $(N, D)$ ポイントを評価した。

主な貢献

1. タスク定式化: 離散トークン化を利用して、細胞ラインのコンテキストから合成遺伝子発現ベクトルを生成するための自己教師あり学習タスクを定義した。
2. 評価指標: 平均発現のピアソン相関、平均絶対誤差（MAE）、および細胞型の識別スコアを含む、生物学的忠実度のためのダウンストリーム指標を確立した。
3. 損失の相関: 事前学習のクロスエントロピー損失が、ダウンストリームの生成品質指標と強く相関していることを示した。
4. 結合スケーリング則: 事前学習損失に対して、結合された2指数関数的なスケーリング則 $L(N, D) = E + A N^{-\alpha} + B D^{-\beta}$ を適合させた。著者の知る限り、これは独立して変化させた $N \times D$ グリッド上での単一細胞トランスフォーマーにおける初の試みである。
5. 計算最適フロント: 適合された指数のに基づき、計算最適化された割り当て戦略を導出し、現在のスケールにおいては、計算リソースはモデルパラメータを増やすよりもデータ量を増やすことに充てる方が適していることを示唆した。

結果

• 生成品質: モデルは、経験的なリプリケート（反復サンプル）と密接に一致する細胞型特異的なベクトルを生成する。最も低い検証損失のチェックポイントにおいて、生成された集団は、保持された経験的データに対して0.960のピアソン相関と0.041のMAEを達成し、2つの不連続な経験的サブセットを比較した際のベースライン（それぞれ0.964および0.038）にほぼ匹着した。細胞型の識別性は0.990であり、ジェネレーターが単一のモードに崩壊していないことを示している。
• スケーリング則の適合: 適合されたスケーリング則は、$\alpha \approx 0.81$（モデルサイズ）および $\beta \approx 0.57$（データサイズ）の指数、および $E \approx 2.90$ nats の不可避な損失フロアを与えた。この適合は $R^2 = 0.977$ を達成した。
• 予測能力: 適合された法則は、保持された実行（$N$ と $D$ の両方を変化させたもの）を、測定値の0.5%以内の誤差で予測し、モデルの補外（エクストラポレーション）能力を検証した。
• 計算最適フロント: $\alpha > \beta$ であるため、最適な割り当てはデータ側に傾いている（$D^* \propto C^{0.59}$ 対 $N^* \propto C^{0.41}$）。

意義と主張
本論文は、単一細胞のスケーリング則に関する論争を解決すると主張している。先行研究（「多様性重視」クラスター）は、コンテキストの多様性の欠如を理由に、単一細胞モデルが大規模化に伴いプラトーに達すると報告しているが、著者らはこれを容量のボトルネックに対する誤解であると論じている。彼らの結果によれば、$N$ と $D$ を結合して変化させた場合、明確なべき乗則のスケーリングが現れる。先行研究で見られたプラトーは、モデルサイズが固定されているか、あるいは制御されていない状態でデータを増やした場合に、データ項が容量項を下回った際にモデルが追加のデータを活用できなくなることに起因するとされている。

著者らは、本研究を単一細胞基盤モデルのための基礎的なステップとして位置づけ、訓練のための計算最適フロントを提供している。また、「ステージ2」の応用として、事前学習済みモデルを摂動応答予測のためにファインチューニングすることを議論している。しかし、事前学習において、摂動ベンチマークに対する保持された損失が、メインのスケーリング損失とは異なる「チェックマーク型」の軌跡（低下した後に上昇する）を描いたことを指摘しており、これは一般的な事前学習コーパスと特定の摂動コンテキストとの間の組成の不一致を示唆している。これは、報告された結果ではなく、今後の課題として残された問いである。

結論として、提案されたアーキテクチャは高い生物学的忠実度で遺伝子発現ベクトルを正常に合成でき、単一細胞トランスフォーマーのスケーリング挙動は、訓練グリッドが十分に制御されていれば、自然言語処理と同様の予測可能な結合べき乗則に従う。