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

Die Arbeit stellt torch-projectors vor, eine hochoptimierte PyTorch-Bibliothek für differentiable Fourier-Raum-Projektionen, die in der Elektronenmikroskopie durch ihre Geschwindigkeit und Unterstützung für CPU-, Apple-Silicon- und CUDA-Geräte bestehende Lösungen um ein bis zwei Größenordnungen übertrifft.

Das Wesentliche

Das große Problem: Der langsame Übersetzer

Stellen Sie sich vor, Sie haben einen riesigen, dreidimensionalen Globus (das ist das Protein, das Wissenschaftler untersuchen wollen). Um zu verstehen, wie dieser Globus aussieht, machen Sie tausende von Fotos davon aus verschiedenen Winkeln. Diese Fotos sind eigentlich nur flache Schatten des Globus.

In der modernen Wissenschaft (speziell bei der Kryo-Elektronenmikroskopie) versuchen Computer, aus diesen tausenden flachen Fotos den 3D-Globus wiederherzustellen. Das ist wie ein riesiges Puzzle.

Das Problem: Die Computerprogramme, die das tun, brauchen eine sehr spezielle Art von „Übersetzer". Dieser Übersetzer muss ständig zwischen dem 3D-Globus und den flachen Fotos hin- und herrechnen.

• Das alte Werkzeug (PyTorch): Das war wie ein Übersetzer, der jeden Satz einzeln, langsam und mit viel Papierkram erledigt hat. Er war so langsam, dass er für moderne, lernende KI-Modelle völlig ungeeignet war. Es wäre, als würde man versuchen, einen Marathon zu laufen, indem man bei jedem Schritt einen schweren Stein hebt.

Die Lösung: Der Hochgeschwindigkeits-Zug „torch-projectors"

Dimitry Tegunov hat eine neue Bibliothek namens torch-projectors entwickelt. Man kann sich das wie einen Hochgeschwindigkeits-Zug vorstellen, der speziell für diese Art von Übersetzung gebaut wurde.

Hier sind die wichtigsten Verbesserungen, einfach erklärt:

1. Der direkte Weg (Kein Umweg)

Früher musste der Computer oft erst den Globus in eine andere Form umwandeln, dann etwas berechnen und dann wieder zurück. Das kostete viel Zeit und Speicherplatz.
Neu: Der neue Zug fährt auf einer direkten Schiene. Er nimmt das Bild, rechnet es sofort um und gibt es aus, ohne zwischendurch unnötige Papiere (Speicher) zu stapeln. Das ist wie der Unterschied zwischen einem Lieferwagen, der erst an drei Lagerhäusern anhalten muss, und einem Drohnen-Lieferdienst, der direkt vom A zum B fliegt.

2. Die scharfe Kamera (Bessere Interpolation)

Wenn man ein Bild aus einem anderen Winkel betrachtet, muss man oft Pixel „erraten", die genau zwischen den vorhandenen liegen.

• Die alte Methode (Lineare Interpolation): Das war wie das Zeichnen mit einem groben Bleistift. Die Linien waren etwas unscharf oder verzerrt. Um es besser zu machen, musste man das Bild künstlich vergrößern (Oversampling), was aber wieder viel Speicher fraß.
• Die neue Methode (Kubische Interpolation): Der neue Zug nutzt einen „Super-Bleistift" (Catmull-Rom-Kern). Er zeichnet die Linien so glatt und präzise, dass man das Bild gar nicht erst vergrößern muss. Es ist wie der Unterschied zwischen einem unscharfen Handyfoto und einem hochauflösenden Foto einer professionellen Kamera.

3. Der Turbo für alle Geräte

Die Bibliothek ist so gebaut, dass sie auf fast jedem modernen Computer läuft:

• Auf normalen Prozessoren (CPU).
• Auf Apple-Chips (MPS).
• Und vor allem auf den starken Grafikchips (CUDA), die für KI-Training genutzt werden.

Das Ergebnis? Der neue Zug ist 10- bis 100-mal schneller als das alte Werkzeug.

• Vergleich: Wenn das alte Werkzeug 100 Fotos pro Stunde machen konnte, schafft der neue Zug 10.000. Das ist der Unterschied zwischen einem Fußgänger und einem Sportwagen.

Warum ist das so wichtig?

Früher war es für KI-Modelle zu teuer und zu langsam, diese 3D-Rekonstruktionen zu lernen. Die Computer waren zu langsam, um die Millionen von Berechnungen in einer vernünftigen Zeit zu schaffen.

Mit torch-projectors können Wissenschaftler nun:

1. Schneller forschen: Sie können Proteinstrukturen viel schneller entschlüsseln.
2. KI einsetzen: Jetzt können sie moderne KI-Methoden nutzen, die sich selbst verbessern, weil die Rechenleistung endlich ausreicht.
3. Speicher sparen: Da keine unnötigen Zwischenbilder gespeichert werden, passen größere Projekte auf weniger teure Hardware.

Zusammenfassung in einem Satz

torch-projectors ist wie ein hochleistungsfähiger, direkter Tunnel, der den langwierigen, holprigen Weg der 3D-Bildrekonstruktion in der Mikroskopie in einen schnellen, glatten Autobahnstraßenverkehr verwandelt – und das alles, ohne dass die KI dabei ins Schwitzen kommt.

Die Bibliothek ist kostenlos verfügbar und steht Wissenschaftlern bereit, um die Geheimnisse der kleinsten Bausteine des Lebens schneller zu entschlüsseln.

Technische Zusammenfassung

Titel: torch-projectors: Eine Hochleistungs-Bibliothek für differentiable Projektionen in PyTorch

Autor: Dimitry Tegunov (Gene.com)
Datum: 10. März 2026 (Preprint)

---

1. Problemstellung

In der Elektronenmikroskopie (insbesondere bei der Einzelpartikelanalyse, SPA, und der Elektronentomographie, ET) sind Fourier-Raum-Projektionsoperationen zentral. Der Fourier-Slice-Theorem besagt, dass jede 2D-Projektion eines Samples einer zentralen Ebene im 3D-Fourier-Raum des Coulomb-Potentials entspricht.

Moderne maschinelle Lernverfahren (Deep Learning) erfordern differentiable Modelle für das End-to-End-Training mittels Gradienten-Backpropagation. Zwar bietet PyTorch automatische Differentiation, jedoch sind die nativen Operationen für Fourier-Projektionen entweder nicht vorhanden oder zu langsam und speicherineffizient für den praktischen Einsatz. Bestehende Lösungen (wie torch-fourier-slice) leiden unter:

• Geringer Durchsatz (Performance).
• Hohem Speicherverbrauch durch das Zwischenspeichern von Tensoren.
• Ineffizienten Interpolationsmethoden, die oft eine Überabtastung (Oversampling) erfordern, was den Speicherbedarf weiter erhöht.

2. Methodik und Architektur

Die Bibliothek torch-projectors implementiert hochperformante, differentiable Projektionsoperatoren direkt im Fourier-Raum unter Verwendung des RFFT-Formats (Real-valued Fast Fourier Transform) von PyTorch, das den FFTW-Konventionen folgt.

Kernkonzepte:

• Differentiable Backpropagation: Die Bibliothek bietet vollständige Gradientenberechnung für alle Eingaben (Rekonstruktion, Rotation, Verschiebung) durch die Implementierung von torch.autograd.Function.
• Interpolation:
  • Unterstützung von linearer (bilinear/trilinear) und kubischer Interpolation (Catmull-Rom, $a = -0.5$).
  • Kubische Interpolation ermöglicht eine hohe Genauigkeit ohne die Notwendigkeit von Überabtastung (Oversampling), was Speicher spart.
  • Die Interpolation erfolgt im Fourier-Raum; Verschiebungen werden über Phasenmodulation ($e^{-i2\pi k \cdot s}$) angewendet.
• Friedel-Symmetrie: Da die Fourier-Transformierte reeller Funktionen symmetrisch ist ($F(-k) = F^*(k)$), wird das RFFT-Format genutzt. Die Bibliothek behandelt die Nullfrequenz-Achse sorgfältig, um Doppelzählungen im Vorwärtspass und Überrepräsentation im Rückwärtspass (Gradientenakkumulation) zu vermeiden.
• Dual-Volumen-Ansatz: Optional wird ein zweiter Tensor für Gewichte (z. B. für CTF-Korrektur) verfolgt, um eine effiziente Normalisierung und Wiener-artige Deconvolution zu ermöglichen.

Implementierung:

• Hybride Architektur: C++/Python-Erweiterung mit Registrierung über TORCH_LIBRARY.
• Multi-Backend: Separate, optimierte Kernel für CPU, Apple Silicon (MPS) und CUDA.
• Speichereffizienz: Operationen werden in einem einzigen Aufruf ausgeführt, ohne Zwischenwerte zu speichern. Der Speicherbedarf beschränkt sich auf Eingabe- und Ausgabetensoren.

Unterstützte Operationen:

Die Bibliothek implementiert vier Hauptoperationen (Forward/Backward für 2D und 3D):

1. 2D $\to$ 2D: Rotation von 2D-Rekonstruktionen.
2. 2D $\to$ 3D: Back-Projection (Akku-mulierung von 2D-Projektionen in ein 3D-Volumen).
3. 3D $\to$ 2D: Forward-Projection (Erzeugung von 2D-Projektionen aus einem 3D-Volumen, zentraler Slice).
4. 2D $\to$ 2D Backward: Mathematischer Adjungierter zur 2D-Forward-Projektion.

3. Key Contributions (Hauptbeiträge)

1. Hochleistungs-Implementierung: Eine vollständig differentiable Bibliothek, die speziell für PyTorch optimiert ist und GPU- und CPU-Ressourcen effizient nutzt.
2. Kubische Interpolation ohne Oversampling: Einführung von Catmull-Rom-Kernen, die die Genauigkeit von linearen Interpolationen übertreffen, ohne den Speicherbedarf durch Zero-Padding zu erhöhen.
3. Skalierbarkeit: Unterstützung von Batch-Verarbeitung und verschiedenen Pose-Dimensionen mit starker Skalierung auf modernen Hardware-Plattformen.
4. Vollständige Differentiation: Exakte analytische Gradienten für Rotationsmatrizen und Verschiebungsvektoren, was für das Training neuronaler Netze essenziell ist.
5. Open Source: Verfügbarkeit als Python-Package unter der MIT-Lizenz.

4. Ergebnisse und Benchmarks

Die Leistung wurde auf NVIDIA H100 (CUDA), Apple M4 (CPU) und Apple M4 (MPS) getestet.

• Performance-Gewinn: torch-projectors ist 1–2 Größenordnungen (Faktor 10–100) schneller als die bestehende Bibliothek torch-fourier-slice.
  • Auf CUDA (H100) wurden Durchsätze von über 86.000 Projektionen/Sekunde (3D $\to$ 2D, Linear, große Batches) erreicht.
  • Der Geschwindigkeitsvorteil ist auf der GPU am größten, da der Overhead durch Kernel-Aufrufe minimiert und die Speicherlokalität verbessert wurde.
• Speichereffizienz: Im Gegensatz zu torch-fourier-slice, das bei Backward-Passes oft den verfügbaren Arbeitsspeicher (48 GB) sprengte, konnte torch-projectors alle Szenarien erfolgreich ausführen, da keine großen Zwischentensoren gespeichert werden.
• Interpolations-Trade-off:
  • Kubische Interpolation ohne Oversampling ist auf allen Plattformen schneller als lineare Interpolation mit On-the-Fly-Oversampling (Faktor ~4x schneller), da teure FFTs für das Padding entfallen.
  • Die Genauigkeit (gemessen via Fourier Ring Correlation und Real-Space-Intensität) ist bei kubischer Interpolation (1x) vergleichbar mit linearer Interpolation bei 2x Oversampling.
• Hardware-Nutzung: Die CUDA-Implementierung nutzt den Cache der GPU so effizient, dass sie die theoretische Bandbreite der H100 sogar leicht übertrifft (Referenzdaten passen komplett in den Cache).

5. Bedeutung und Ausblick

Die Bibliothek ermöglicht erstmals den praktischen Einsatz von Deep-Learning-Methoden in der Kryo-Elektronenmikroskopie, bei denen Projektionsoperationen tausendfach während des Trainings durchgeführt werden müssen.

• Anwendungsbereiche: Optimierung von Proteinstrukturen, dynamische Studien und subtomogram averaging.
• Zukunftsperspektiven: Geplante Erweiterungen um 3D $\to$ 3D Projektionen (für Subtomogram-Averaging) und Projektionen im Realraum, um den Anwendungsbereich über den klassischen Fourier-Raum-Workflow hinaus zu erweitern.

Fazit: torch-projectors löst das Performance- und Speicherproblem bei differentiable Fourier-Projektionen in PyTorch und stellt eine kritische Infrastruktur für die nächste Generation von KI-gestützten Strukturbiologie-Algorithmen dar.