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

Dit paper introduceert torch-projectors, een hoogpresterende, differentieerbare bibliotheek voor Fourier-ruimte-projecties in PyTorch die, geoptimaliseerd voor CPU, Apple Silicon en CUDA, de snelheid van bestaande methoden met 1 tot 2 orde van grootte verbetert voor elektronenmicroscopie-toepassingen.

De kern

De "Torch-Projectors": Een Snelheidsbooster voor het Zien van Moleculen

Stel je voor dat je een gigantische, ingewikkelde 3D-puzzel hebt, maar je kunt alleen 2D-foto's maken van de stukjes vanuit verschillende hoeken. Dat is precies wat wetenschappers doen in de wereld van cryo-elektronenmicroscopie (cryo-EM). Ze nemen duizenden foto's van eiwitten (de bouwstenen van het leven) en proberen die later in een computer tot één perfect 3D-model te reconstrueren.

Het probleem? De software die dit doet, is vaak traag. Het is alsof je die 3D-puzzel probeert op te lossen met een slak die een hooivork vasthoudt. Vooral als je moderne kunstmatige intelligentie (AI) wilt gebruiken om dit sneller en slimmer te doen, moet de computer heel snel kunnen rekenen. Helaas is de standaardsoftware van PyTorch (een populaire AI-toolkit) hier te traag voor.

Hier komt torch-projectors om de hoek kijken.

Wat is het eigenlijk?

Dit is een nieuwe, supersnelle bibliotheek (een verzameling gereedschappen) voor programmeurs. Het helpt computers om die 2D-foto's om te zetten in 3D-beelden, en vice versa, maar dan ontzettend snel en met de mogelijkheid om de AI "te leren" van fouten (dit heet "differentieerbaar").

De Analogie: De Magische Projector

Stel je voor dat je een grote, doorzichtige 3D-blok hebt (het eiwit) en je wilt er een platte tekening van maken (de 2D-foto).

1. De oude manier (PyTorch): Dit is alsof je een fotograaf bent die elke keer de hele kamer moet leegmaken, de blokken moet verplaatsen, de foto maakt, en dan alles weer terugzetten. Het duurt lang en je moet veel ruimte (geheugen) hebben.
2. De nieuwe manier (torch-projectors): Dit is een magische projector. Je zet de blokken erin, draait een knop, en boem – in één flits heb je de foto. Je hoeft niets te verplaatsen en er is geen rommel.

De Twee Trucjes die het zo snel maken

1. De "Cubic" Truc (Het Scherpere Lensje)
Vroeger gebruikten computers een simpele methode om de foto's te maken: ze trokken rechte lijnen tussen punten (lineaire interpolatie). Dat is als tekenen met een stomp potlood; de lijnen zijn wat ruw.
Deze nieuwe bibliotheek gebruikt een kubische methode (Catmull-Rom). Denk hierbij aan een slinger die perfect door de punten loopt, alsof je tekent met een zijden draad.

• Het voordeel: Je krijgt een veel scherpere foto zonder dat je de hele 3D-blok eerst moet vergroten (wat veel geheugen kost). Het is alsof je een lage-resolutie foto neemt, maar door slimme wiskunde ziet hij eruit alsof hij in 4K is.

2. De "Eén-Knop" Truc (Geen Tussentijdse Rommel)
Bij de oude methoden moest de computer tussentijdse berekeningen opslaan op het werkblad voordat hij de eindresultaat gaf. Dat was als een kok die eerst alle ingrediënten op het aanrecht legt, ze mengt, en dan pas de pan gebruikt.
torch-projectors doet alles in één beweging. De kok gooit alles direct in de pan. Dit bespaart enorm veel ruimte in het geheugen van de computer, waardoor je veel grotere projecten kunt draaien zonder dat je computer vastloopt.

Waarom is dit een revolutie?

De testresultaten in het paper zijn verbazingwekkend:

• Snelheid: Het is 10 tot 100 keer sneller dan de vorige beste methoden. Op een krachtige computer (zoals een NVIDIA H100) is het alsof je van een fiets op een raket zit.
• Geheugen: Het gebruikt veel minder geheugen. Dit betekent dat wetenschappers nu AI-modellen kunnen trainen die eerder onmogelijk waren omdat ze te groot waren voor de computer.
• Alle apparaten: Het werkt niet alleen op dure supercomputers, maar ook op Apple Silicon (de M-chips in je MacBook) en gewone CPU's.

Samenvattend

torch-projectors is de "Formule 1-motor" die onder de kap van de cryo-EM-wereld wordt geplaatst. Het maakt het mogelijk om moleculaire puzzels niet alleen sneller op te lossen, maar ook om AI te gebruiken om die puzzels slimmer te maken. Voor de toekomst van medicijnen en het begrijpen van ziektes is dit een enorme sprong voorwaarts: van het wachten op een antwoord van een uur, naar het krijgen van een antwoord in seconden.

Technische samenvatting

Probleemstelling

In de elektronenmicroscopie (EM), met name bij cryo-elektronenmicroscopie (cryo-EM) en elektronentomografie, zijn Fourier-ruimte projectie-operaties fundamenteel voor het reconstrueren van 3D-structuren uit 2D-projecties. Moderne machine learning-methoden vereisen differentieerbare implementaties voor end-to-end modeltraining. Hoewel PyTorch automatische differentiatie biedt, zijn de ingebouwde operaties voor deze specifieke Fourier-projecties te traag voor praktisch gebruik en leiden ze tot een te hoog geheugenvolume. Bestaande oplossingen, zoals die in RELION, maken vaak gebruik van oversampling (nul-vullen in de ruimtelijke domein) om interpolatie-accuraatheid te verbeteren, maar dit is memory-intensief en vereist kostbare Fourier-transformaties bij elke doorloop, wat een bottleneck vormt voor diep leernetwerken.

Methodologie

Het paper introduceert torch-projectors, een hoogwaardige bibliotheek voor differentieerbare Fourier-ruimte projecties in PyTorch. De kernmethodologie omvat:

• Architectuur: De bibliotheek gebruikt een hybride C++/Python-extensie-architectuur met TORCH_LIBRARY-registratie. Het biedt geoptimaliseerde kernels voor CPU, Apple Silicon (MPS) en CUDA-apparaten.
• Differentiatie: Alle operaties zijn volledig differentieerbaar. De backward-pass berekent gradiënten voor alle invoer, inclusief rotatiematrices en verschuivingen, zonder tussenliggende waarden op te slaan, wat het geheugengebruik minimaliseert.
• Interpolatie: De bibliotheek ondersteunt zowel lineaire als kubische interpolatie (Catmull-Rom kernels met parameter $a = -0.5$).
  • Kubische interpolatie: Biedt $C^1$ continuïteit en exacte interpolatie door controlepunten, wat de nauwkeurigheid aanzienlijk verbetert zonder de noodzaak van oversampling.
  • Oversampling: Wordt optioneel ondersteund voor hogere nauwkeurigheid, maar de bibliotheek is ontworpen om dit te vermijden door de superieure kubische kern te gebruiken.
• Friedel-symmetrie: Er wordt zorgvuldig omgegaan met de symmetrie van de Fourier-getransformeerde van reële functies ($F(-k) = F^*(k)$) om dubbele telling te voorkomen en wiskundige consistentie te behouden tijdens het forward- en backward-pass.
• Operaties: De bibliotheek implementeert vier hoofdoperaties:
  1. 2D $\to$ 2D Forward (rotatie van 2D Fourier-data).
  2. 2D $\to$ 2D Backward (accumulatie in 2D reconstructie).
  3. 3D $\to$ 2D Forward (Central Slice Theorem: snijden van 3D volume).
  4. 2D $\to$ 3D Backward (accumulatie in 3D volume).
• Shifts: Verschuivingen worden toegepast via fase-modulatie in het Fourier-domein, conform de verschuivingsstelling van Fourier, wat efficiënter is dan verschuivingen in het ruimtelijke domein.

Belangrijkste Bijdragen

1. Snelheid en Efficiëntie: torch-projectors presteert 1 tot 2 ordes van grootte sneller dan de bestaande bibliotheek torch-fourier-slice. Dit wordt bereikt door het uitvoeren van projecties in één enkele kernel-call zonder tussenliggende tensors op te slaan, wat de geheugenbandbreedte-optimalisatie maximaliseert.
2. Geheugenoptimalisatie: Door geen tussenliggende waarden op te slaan en direct in Fourier-ruimte te werken, is het geheugenvolume beperkt tot de invoer- en uitvoertensors. Dit maakt het mogelijk om backward-passes uit te voeren die anders het geheugen van het systeem zouden overschrijden.
3. Geavanceerde Interpolatie: De implementatie van Catmull-Rom kubische interpolatie maakt het mogelijk om hoge nauwkeurigheid te bereiken zonder de memory-intensieve techniek van oversampling (nul-vullen), wat essentieel is voor memory-beperkte neurale netwerken.
4. Multi-platform Ondersteuning: De bibliotheek is geoptimaliseerd voor NVIDIA CUDA, Apple Silicon (MPS) en CPU, met specifieke aandacht voor cache-efficiëntie en parallelle verwerking.

Resultaten

De prestatiebenchmarks werden uitgevoerd op NVIDIA H100 (CUDA), Apple M4 (CPU) en Apple M4 (MPS):

• Doorvoer: Op CUDA-apparaten (H100) worden doorvoersnelheden bereikt die de theoretische bandbreedte van het geheugen lijken te overtreffen (door volledige cache-gebruik), met waarden tot $8.6 \times 10^4$ projecties per seconde voor 3D $\to$ 2D operaties.
• Vergelijking met torch-fourier-slice:
  • De nieuwe bibliotheek is 5 tot 200 keer sneller dan torch-fourier-slice, afhankelijk van de configuratie (batchgrootte, aantal poses, interpolatiemethode).
  • Op de CPU en MPS is de verbetering minder extreem (factor 2-10) maar nog steeds significant, voornamelijk door de vermindering van kernel-call overhead.
  • Voor complexe scenario's (zoals backward-passes met grote batches) kon torch-fourier-slice vaak niet worden getest vanwege geheugenoverbelasting, terwijl torch-projectors dit probleemloos aankan.
• Interpolatie-efficiëntie: Kubische interpolatie zonder oversampling bleek ongeveer 4 keer sneller dan lineaire interpolatie met on-the-fly oversampling (die extra FFT's vereist), terwijl de signaalverzwakking in de ruimtelijke domein vergelijkbaar bleef tot ongeveer de helft van de box-grootte.

Betekenis

Deze paper markeert een doorbraak voor het toepassen van deep learning in cryo-EM. De 1-2 ordes van grootte snelheidswinst maakt het mogelijk om projectie-operaties duizenden keren uit te voeren tijdens gradient-based optimalisatie, wat eerder onpraktisch was.

• Toepassingen: Het stelt onderzoekers in staat om end-to-end trainbare modellen te bouwen voor single-particle analysis (SPA) en tomografie, inclusief complexe taken zoals deeltjesclassificatie en CTF-correctie.
• Toekomst: De bibliotheek opent de weg voor meer geavanceerde 3D $\to$ 3D projecties en real-space projectie-operaties, en verlaagt de drempel voor het gebruik van GPU-versnelling in cryo-EM workflows.

De bibliotheek is open-source beschikbaar onder de MIT-licentie en kan direct worden geïnstalleerd via pip, wat de adoptie in de wetenschappelijke gemeenschap faciliteert.