torch-projectors: A High-Performance Differentiable Projection Library for PyTorch

本論文は、電子顕微鏡解析向けに最適化され、PyTorch の既存実装や torch-fourier-slice を 1〜2 桁高速化する、CPU・Apple Silicon・CUDA に対応した高性能な微分可能なフーリエ空間投影ライブラリ「torch-projectors」を紹介するものである。

要約

🔦 暗闇を照らす「魔法の懐中電灯」：電子顕微鏡の新しいブレークスルー

この論文は、**「torch-projectors（トーチ・プロジェクターズ）」**という新しいソフトウェア・ライブラリについて書かれたものです。

少し難しい専門用語を、日常の風景に例えて説明しましょう。

🌌 背景：3D 物体を「2D 写真」から復元する難しさ

まず、**クライオ電子顕微鏡（Cryo-EM）**という技術について考えてみてください。
これは、タンパク質のような小さな分子の 3 次元の形を、原子レベルで詳しく見るための「超高性能カメラ」です。

しかし、このカメラは 3 次元の物体を直接撮るのではなく、**「さまざまな角度から撮った 2 次元の写真（スライド）」**しか撮れません。
想像してください。巨大な像（3D 物体）を、ぐるぐる回しながら、前後左右から写真を撮り続けるイメージです。

その後、コンピューターはこれらの 2 次元の写真を集めて、元の 3 次元の像を「パズルのように組み立て」て復元します。この作業を**「投影（プロジェクション）」**と呼びます。

🐢 昔の悩み：計算が重すぎて、AI が動かない

最近では、この復元作業を**「AI（機械学習）」**に任せて、より高精度に、より速く行おうと試みられています。AI は「正解に近づくために、計算結果を微調整する（逆伝播）」という作業を何千回も繰り返します。

しかし、ここで大きな壁がありました。

• PyTorch（AI 開発で使われる有名なツール）の標準機能では、この「2D 写真から 3D 像を作る（またはその逆）」計算があまりに遅すぎるのです。
• 昔のやり方（RELION というソフトなど）では、精度を上げるために「画像を無理やり大きくして（オーバーサンプリング）」計算していました。これは**「小さな写真を拡大コピーして、その上で計算する」**ようなもので、メモリ（作業机の広さ）を大量に消費し、AI がパンクしてしまったり、計算に時間がかかりすぎたりしていました。

🚀 新登場：「torch-projectors」の登場

そこで登場したのが、この論文で紹介されている**「torch-projectors」です。
これは、「超高速で、かつ AI が学習できる（微分可能な）投影計算」**を行うための新しい道具箱です。

🎨 3 つの魔法

このライブラリには、3 つのすごい特徴があります。

1. 「滑らかな補間」の魔法（キュービック補間）

   • 昔のやり方： 2D 写真から 3D 像を作る際、ピクセル（点）の間にデータがないと、直線でつなぐだけ（線形補間）でした。これだと、画像が少しぼやけてしまいます。
   • 新しい魔法： 「Catmull-Rom」という**「滑らかな曲線でつなぐ」**技術を使います。
   • 例え： 直線でつなぐと角ばった階段になりますが、曲線でつなぐと滑らかなスロープになります。これにより、「画像を無理やり大きくしなくても（オーバーサンプリングなし）」、非常に高精度な計算が可能になりました。

2. 「メモリ節約」の魔法

   • 昔のやり方は、中間データを大量に保存する必要があり、作業机（メモリ）がいっぱいになっていました。
   • 新しい道具は、**「必要な計算を一度のステップで完結」**させます。中間のメモ帳を一切使わず、入力と出力だけで済ませるので、メモリ使用量が劇的に減り、AI が大きなモデルでも学習できるようになりました。

3. 「超高速」の魔法（CPU, Apple Silicon, GPU 対応）

   • この道具は、NVIDIA の GPU（H100）、Apple の M シリーズチップ、そして通常の CPU のすべてで、**「10 倍〜100 倍」**もの速度アップを実現しました。
   • 例え： 昔は「徒歩で山を登る」ようなものでしたが、今は「ジェットコースター」で登れるようになりました。

📊 結果：どれくらい速くなった？

実験結果によると、この新しい道具は、既存のツール（torch-fourier-slice）と比べて**「100 倍（10^2）」**近い速さで動きました。
特に、Apple の M4 チップや NVIDIA の H100 GPU 上では、その性能が爆発的に発揮されています。

また、**「オーバーサンプリング（画像を大きくする）」をしない場合でも、「オフラインで事前に画像を大きくする」方法よりも、「その場で計算する（オンザフライ）」**方が 4 倍速いという驚きの結果も出ました。これは、AI がメモリ制限の中で動く際に非常に有利です。

🌟 まとめ：なぜこれが重要なのか？

この「torch-projectors」は、「電子顕微鏡で見る世界（タンパク質の構造）」を、AI がもっと速く、もっと詳しく、もっと安く解き明かすための鍵になります。

• 昔： 「計算が遅すぎて、AI は使えない」「メモリが足りなくて、大きなモデルは作れない」
• 今： 「超高速で計算できる」「メモリを節約して、複雑な AI も動かせる」

これにより、将来、私たちが病気の治療に役立つ新しい薬の設計や、生命の神秘を解明するスピードが、劇的に加速することが期待されています。

---

一言で言うと：
「電子顕微鏡の 3D 復元計算を、**『AI が使えるように高速化・軽量化した魔法の道具』**に生まれ変わらせた、画期的な技術です。」

技術概要

torch-projectors: PyTorch 向け高性能微分可能投影ライブラリの技術的概要

本論文は、電子顕微鏡（特に単粒子解析 SPA と電子トモグラフィー ET）の分野において不可欠なフーリエ空間での投影演算を、PyTorch 環境で高速かつ微分可能に実行するための新しいライブラリ「torch-projectors」を紹介しています。

以下に、問題定義、手法、主要な貢献、結果、および意義について詳細にまとめます。

1. 問題定義

• 背景: 低温電子顕微鏡（cryo-EM）では、試料のクーロンポテンシャルの 2 次元投影から 3 次元構造を再構築するために、フーリエ空間での投影演算が中核的な役割を果たします。
• 課題: 近年の機械学習手法は、エンドツーエンドのモデル訓練のために微分可能な実装を必要とします。PyTorch には自動微分機能がありますが、組み込みの操作を用いてフーリエ空間の投影を実装すると、計算速度が実用的でないほど遅く、メモリオーバーヘッドも大きすぎるという問題がありました。
• 既存手法の限界: 従来の SPA アルゴリズム（例：RELION）では、精度を高めるために実空間でゼロパディング（オーバーサンプリング）を行う手法が一般的ですが、これはメモリ使用量を増大させます。また、機械学習の文脈では、各パスで高価なフーリエ変換を伴うオーバーサンプリングを行うことがコスト高となります。

2. 手法とアーキテクチャ

torch-projectors は、PyTorch の自動微分フレームワークと完全に統合された、CPU、Apple Silicon (MPS)、CUDA に対応したハイブリッド C++/Python 拡張ライブラリです。

• 微分可能性: 各投影操作を単一の呼び出しで実装し、中間値を保存しないことで、入力と出力テンソルのみでメモリフットプリントを最小化しつつ、すべての入力に対する勾配計算をサポートしています。
• 補間手法:
  • 線形補間: 標準的なマルチリニアカーネル（2D で双線形、3D で三重線形）を使用。
  • 3 次補間: Catmull-Rom 3 次カーネル（パラメータ $a = -0.5$）を実装。これにより、オーバーサンプリングを行わずとも高い精度（$C^1$ 連続性）を達成し、メモリ使用量を削減しています。
• フーリエ空間の処理:
  • データ形式: FFTW の規約に従い、PyTorch の RFFT 形式（実数値関数のフーリエ変換）をそのまま使用。中間で形式を変換しません。
  • 対称性の処理: 実数値関数のフーリエ変換に特有のフリーデル対称性（Friedel symmetry）を厳密に扱い、ゼロ周波数軸における二重カウントや勾配の蓄積における不整合を防ぎます。
  • シフト処理: 位相変調（Fourier shift theorem）を用いて、実空間での移動をフーリエ空間で効率的に処理します。
• 演算の種類:
  • 2D→2D 前方/後方投影
  • 3D→2D 前方投影（中央切片定理の実装）
  • 2D→3D 後方投影
  • 回転行列とシフトベクトルに対する解析的な勾配計算を提供します。

3. 主要な貢献

1. 劇的な性能向上: 既存の PyTorch 実装や「torch-fourier-slice」と比較して、1〜2 オーダー（10〜100 倍）の高速化を実現しました。
2. メモリ効率の最適化: 中間テンソルを生成せず、単一カーネル呼び出しで処理を行うことで、大規模なバッチ処理やメモリ制約の厳しいニューラルネットワークでの利用を可能にしました。
3. 高精度な 3 次補間の導入: オーバーサンプリングを行わずに Catmull-Rom 3 次補間を実装することで、従来の線形補間＋オーバーサンプリング手法と同等以上の精度を、より低いメモリコストで達成しました。
4. マルチプラットフォーム対応: CPU、Apple Silicon (MPS)、NVIDIA CUDA (H100) に対してそれぞれ最適化されたカーネルを提供し、特に CUDA 環境では理論的な帯域幅を超えるスループットを達成しています。

4. 結果

ベンチマークは NVIDIA H100 (CUDA)、Apple M4 (CPU)、Apple M4 (MPS) において実施されました。

• スループット:
  • CUDA 環境では、最も高速な構成で 1 秒間に約 8 万回の 3D→2D 投影（2048 ポーズ、128 ピクセル）を処理可能です。
  • CPU と比較して、CUDA 環境では 2 オーダー以上、MPS でも大幅な高速化が見られました。
• 既存ライブラリとの比較 (torch-fourier-slice):
  • 全設定において、torch-projectors は torch-fourier-slice より 1〜2 オーダー高速でした。
  • 特に 3D→2D 前方投影と 2D→3D 後方投影において、torch-fourier-slice はバックワードパスのメモリ消費がテスト環境の限界（48GB）を超えて実行不可能なケースがありましたが、torch-projectors は問題なく実行できました。
• 補間効率:
  • 3 次補間（オーバーサンプリングなし）は、オンラインで 2 倍のオーバーサンプリングを行う線形補間と比較して、すべてのプラットフォームで約 4 倍高速でした（特に参照ボリュームが事前にパディングできない場合）。
  • 実空間での信号減衰を分析した結果、3 次補間（1 倍）と線形補間（2 倍オーバーサンプリング）は、ボックスサイズの約半分までの信号強度を同様に維持することが確認されました。

5. 意義と将来展望

• cryo-EM における機械学習の促進: 投影演算が勾配ベースの最適化中に数千回実行される必要がある cryo-EM の機械学習応用において、このライブラリは以前は実用的ではなかったアプローチを可能にします。
• 柔軟性の向上: メモリ制約の厳しい環境や、動的に生成される中間表現を持つ大規模な演算グラフにおいても、高品質な投影演算を効率的に実行できます。
• 今後の展開: サブトモグラム平均化向けの 3D→3D 投影演算や、実空間での投影演算の追加など、さらに広範な応用への拡張が予定されています。

結論:
torch-projectors は、電子顕微鏡データ解析における計算ボトルネックを解消し、PyTorch 生態系内で高性能かつ微分可能なフーリエ空間投影を実現する重要なインフラストラクチャです。これにより、より複雑で高精度な機械学習モデルを用いた構造生物学の研究が加速することが期待されます。

利用情報:
本ライブラリはオープンソース（MIT ライセンス）で、GitHub (warpem/torch-projectors) から利用可能です。