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

Ce papier présente torch-projectors, une bibliothèque PyTorch haute performance et différentiable pour les projections dans l'espace de Fourier, optimisée pour les CPU, Apple Silicon et CUDA, qui surpasse significativement les solutions existantes comme torch-fourier-slice pour l'analyse en microscopie électronique.

L'essentiel

🧊 Le Projet "Torch-Projectors" : La Machine à Remonter le Temps des Microscopes

Imaginez que vous essayez de reconstruire un château de sable complexe (une protéine) en regardant uniquement des ombres projetées sur un mur. C'est exactement ce que font les scientifiques avec le microscope électronique cryogénique (cryo-EM). Ils prennent des milliers de photos en 2D d'une molécule sous différents angles, et doivent ensuite assembler ces ombres pour reconstituer l'objet en 3D.

Le problème ? Les outils actuels pour faire ce "puzzle mathématique" sont lents et lourds, un peu comme essayer de construire un gratte-ciel avec une petite cuillère en plastique.

C'est là qu'intervient torch-projectors, une nouvelle bibliothèque logicielle créée par Dimitry Tegunov. Voici comment elle fonctionne, expliqué simplement.

1. Le Problème : Le "Trafic" Mathématique

Pour reconstruire l'image 3D, l'ordinateur doit faire des allers-retours constants entre l'espace réel (l'image visible) et l'espace de Fourier (une représentation mathématique abstraite de l'image, comme une partition de musique).

• L'ancienne méthode : C'était comme essayer de traverser une autoroute en marchant. Les logiciels existants (comme ceux intégrés à PyTorch) étaient trop lents et consommaient trop de mémoire. Pour les intelligences artificielles modernes qui doivent apprendre par essai-erreur, c'était impossible : trop lent, trop cher.
• La solution : torch-projectors est comme un téléphérique ultra-rapide. Il effectue ces calculs directement dans l'espace mathématique (Fourier), sans avoir besoin de faire des arrêts inutiles, ce qui le rend 10 à 100 fois plus rapide que les concurrents.

2. La Magie de l'Interpolation : Le "Zoom" Intelligent

Quand on tourne une image ou qu'on la regarde sous un angle différent, les pixels ne tombent pas toujours exactement sur les cases de la grille. Il faut deviner ce qu'il y a entre deux points.

• L'ancien réflexe (Linéaire) : C'est comme faire une moyenne entre deux points voisins. C'est rapide, mais l'image devient floue, comme une photo mal mise au point. Pour compenser, les scientifiques devaient "gonfler" l'image (ajouter du vide autour) avant de la tourner, ce qui prenait beaucoup de place dans la mémoire.
• La nouvelle astuce (Cubique) : torch-projectors utilise une technique appelée interpolation cubique. Imaginez que vous ne faites pas juste une moyenne, mais que vous tracez une courbe parfaite et lisse entre les points, comme un sculpteur qui affine une statue.
  • Le résultat : On obtient une image nette et précise sans avoir besoin de "gonfler" l'image au préalable. C'est comme si vous pouviez tourner une photo haute définition sans jamais avoir besoin de la dézoomer pour éviter les pixels carrés.

3. La Révolution pour l'Intelligence Artificielle

Aujourd'hui, on utilise l'IA pour apprendre à reconstruire ces images automatiquement. Pour que l'IA apprenne, elle doit pouvoir "voir" ses erreurs et les corriger (c'est ce qu'on appelle la rétropropagation).

• Avant, les calculs étaient si lents que l'IA mettait des jours à apprendre une seule protéine.
• Avec torch-projectors, l'IA peut faire des milliers de tentatives en quelques secondes. C'est la différence entre essayer d'ouvrir une porte avec un tournevis (l'ancien logiciel) et utiliser un code d'accès magnétique (le nouveau logiciel).

4. Les Trois Super-Héros (CPU, Apple, NVIDIA)

Ce logiciel est conçu pour être rapide partout, peu importe la machine :

• Sur les puces NVIDIA (CUDA) : C'est comme une course de Formule 1. Il va à une vitesse folle, exploitant la puissance brute des cartes graphiques.
• Sur les puces Apple (MPS) : C'est comme une voiture de sport électrique : très efficace et rapide, même sur les ordinateurs portables.
• Sur les processeurs classiques (CPU) : Même sans carte graphique, il reste très performant, comme un coureur de fond qui ne s'essouffle pas.

En Résumé

torch-projectors est un outil qui permet de reconstruire des images moléculaires 3D à la vitesse de l'éclair.

• Avant : C'était lent, flou et prenait trop de place (comme un vieux camion).
• Maintenant : C'est rapide, précis et économe en énergie (comme un vaisseau spatial).

Grâce à cela, les chercheurs peuvent désormais utiliser l'intelligence artificielle pour découvrir la structure de nouvelles protéines et médicaments beaucoup plus vite, ce qui pourrait accélérer la découverte de traitements pour des maladies complexes.

Note amusante : L'auteur du papier a même mentionné qu'une grande partie du code a été écrite avec l'aide de l'IA (Claude), prouvant que l'outil qu'il a créé est exactement le genre de technologie qui permet à l'IA de devenir encore plus utile pour la science !

Résumé technique

1. Problématique

La microscopie électronique à cryogénie (cryo-EM) et la tomographie électronique reposent sur des algorithmes de projection dans l'espace de Fourier pour reconstruire des structures 3D à partir de projections 2D. L'application de méthodes d'apprentissage automatique (Deep Learning) à ces domaines nécessite des opérations de projection différentiables pour permettre l'entraînement de modèles de bout en bout via la rétropropagation du gradient.

Cependant, les opérations natives de PyTorch sont insuffisantes pour deux raisons principales :

• Performance : Les implémentations standards sont trop lentes pour des entraînements pratiques, surtout lors de l'optimisation de vastes paramètres.
• Mémoire : Les méthodes existantes (comme celles utilisées dans RELION) nécessitent souvent un sur-échantillonnage (zero-padding) en espace réel pour améliorer la précision de l'interpolation linéaire, ce qui entraîne une surcharge mémoire excessive incompatible avec les contraintes des réseaux de neurones modernes.

2. Méthodologie

Le papier présente torch-projectors, une bibliothèque haute performance conçue pour exécuter des projections différentiables dans l'espace de Fourier directement sous PyTorch.

• Architecture Hybride C++/Python : La bibliothèque utilise le système d'enregistrement TORCH_LIBRARY de PyTorch. Elle expose des opérateurs C++ optimisés via des classes torch.autograd.Function personnalisées, garantissant une intégration transparente avec la différenciation automatique tout en minimisant l'empreinte mémoire (pas de stockage de valeurs intermédiaires).
• Espace de Fourier (RFFT) : Toutes les opérations sont effectuées exclusivement dans le format RFFT (Real Fast Fourier Transform) de PyTorch, respectant les conventions FFTW. Cela évite les conversions de format coûteuses entre les étapes de projection.
• Interpolation Avancée :
  • Linéaire : Utilise des noyaux bilinéaires (2D) et trilinéaires (3D).
  • Cubique (Catmull-Rom) : Implémente une interpolation cubique avec paramètre $a = -0.5$, offrant une continuité $C^1$ et une précision supérieure sans nécessiter de sur-échantillonnage (oversampling). Cela permet de réduire considérablement l'utilisation de la mémoire tout en maintenant une haute fidélité du signal.
• Gestion de la Symétrie de Friedel : Le code gère rigoureusement la symétrie conjuguée des transformées de Fourier de fonctions réelles, en particulier sur l'axe de fréquence nulle, pour éviter le double comptage lors de l'accumulation des gradients.
• Opérateurs Supportés : La bibliothèque implémente les quatre opérations fondamentales :
  1. 2D $\to$ 2D (Projection avant et arrière).
  2. 3D $\to$ 2D (Projection avant, basée sur le théorème de la tranche centrale).
  3. 2D $\to$ 3D (Projection arrière, accumulation de données dans un volume 3D).
  • Chaque opération supporte la rotation, la translation (via modulation de phase) et l'accumulation de poids (utile pour la correction CTF).

3. Contributions Clés

• Performance Inégalée : torch-projectors surpasse la bibliothèque concurrente torch-fourier-slice d'un facteur de 10 à 100 fois (1 à 2 ordres de grandeur) sur toutes les plateformes testées (CPU, Apple Silicon MPS, et CUDA).
• Efficacité Mémoire : En évitant le sur-échantillonnage grâce à l'interpolation cubique et en calculant les gradients en une seule passe sans tensors intermédiaires, la bibliothèque réduit drastiquement l'empreinte mémoire, rendant possible l'entraînement de modèles complexes qui échouaient auparavant par manque de mémoire (OOM).
• Support Multi-Backend : Optimisations spécifiques pour CPU, Apple Silicon (MPS) et cartes GPU NVIDIA (CUDA), exploitant la parallélisation par projection.
• Précision sans Sur-échantillonnage : La démonstration que l'interpolation cubique sans padding (oversampling = 1.0) offre une précision comparable à l'interpolation linéaire avec un padding 2x, mais à un coût de calcul bien inférieur lorsque le padding doit être effectué à la volée.

4. Résultats

Les benchmarks ont été réalisés sur une NVIDIA H100 (CUDA), un CPU Apple M4 et un GPU Apple M4 (MPS).

• Débit (Throughput) :
  • Sur CUDA (H100), le débit atteint des millions de projections par seconde. Par exemple, pour une projection 3D $\to$ 2D avec une boîte de 128 pixels, torch-projectors atteint 85 962 k projections/s contre 941 k pour torch-fourier-slice (un gain de x91).
  • Le backend CUDA dépasse même la bande passante théorique de la mémoire (3,4 TB/s) en atteignant 3,7 TB/s, suggérant que les données de référence tiennent entièrement dans le cache du processeur.
• Comparaison Interpolation :
  • L'interpolation cubique sans padding est environ 4 fois plus rapide que l'interpolation linéaire avec padding à la volée (qui nécessite des FFT supplémentaires).
  • L'interpolation cubique sans padding est légèrement plus lente que le linéaire avec padding hors ligne, mais offre un compromis idéal pour les réseaux de neurones dynamiques où le padding pré-calculé n'est pas possible.
• Précision : Les analyses montrent que l'interpolation cubique (sans padding) et l'interpolation linéaire (avec padding 2x) produisent une atténuation du signal similaire jusqu'à environ la moitié de la taille de la boîte, avec des artefacts négligeables dans l'espace de Fourier.

5. Signification et Impact

Ce travail est une avancée majeure pour l'intégration de l'apprentissage profond dans la cryo-EM.

• Faisabilité de l'Optimisation : Les gains de performance de 1 à 2 ordres de grandeur rendent viables des applications d'apprentissage automatique qui nécessitaient des milliers d'itérations de projections, auparavant trop coûteuses en temps de calcul.
• Flexibilité des Modèles : La réduction de l'empreinte mémoire permet d'utiliser des architectures de réseaux de neurones plus profondes et plus complexes, ainsi que des représentations intermédiaires dynamiques, sans être bloqué par les limites de la mémoire GPU.
• Standardisation : En fournissant une bibliothèque open-source (MIT) compatible avec PyTorch, torch-projectors établit un nouveau standard pour les opérations de projection différentiables, facilitant le développement futur d'algorithmes de reconstruction 3D basés sur le deep learning.

En résumé, torch-projectors comble le fossé entre les exigences de performance des algorithmes de microscopie électronique et les besoins de flexibilité et de différentiabilité de l'apprentissage automatique moderne.