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

El artículo presenta *torch-projectors*, una biblioteca de alto rendimiento para PyTorch que implementa operaciones de proyección diferenciables en el espacio de Fourier, optimizadas para múltiples dispositivos y significativamente más rápidas que las soluciones existentes para el análisis de microscopía electrónica.

Lo esencial

¡Claro que sí! Imagina que este documento es el plano de un super-turbo cargador diseñado específicamente para los científicos que estudian las moléculas más pequeñas del universo.

Aquí tienes la explicación de "torch-projectors" en español, usando analogías sencillas:

🧊 El Problema: La Cámara de los "Fantasmas"

Imagina que tienes una estatua de hielo (una proteína) y quieres saber cómo es por dentro. No puedes romperla, así que usas un microscopio electrónico para tomar miles de fotos en 2D desde diferentes ángulos.

El problema es que estas fotos son como sombras. Para reconstruir la estatua original en 3D, necesitas un algoritmo matemático muy complejo que combine todas esas sombras.

Aquí está el truco: Para que la inteligencia artificial (IA) aprenda a reconstruir estas imágenes automáticamente, necesita poder "ver" cómo cambia la imagen si mueves un poco la estatua. Esto se llama diferenciación. Pero, hasta ahora, las herramientas que tenía PyTorch (el motor de IA más popular) para hacer esto eran como usar una bicicleta para cruzar el océano: funcionaban, pero eran lentísimas y agotadoras.

🚀 La Solución: "torch-projectors"

El autor, Dimitry Tegunov, ha creado una nueva librería llamada torch-projectors.

Piensa en esto como si hubieras cambiado esa bicicleta por un cohete de la NASA.

• ¿Qué hace? Convierte las matemáticas aburridas y lentas de "proyección" (convertir 3D a 2D y viceversa) en operaciones ultra rápidas.
• ¿Dónde funciona? En tu computadora normal (CPU), en las Mac con chips M4 (Apple Silicon) y, lo más importante, en las tarjetas gráficas potentes de NVIDIA (CUDA).

🎨 La Magia: "El Pintor de Alta Precisión"

Para reconstruir la imagen, el software tiene que "pintar" píxeles en lugares que no son números enteros (como 3.45 píxeles).

• El método antiguo (Lineal): Era como conectar dos puntos con una regla. Quedaba bien, pero si la imagen era muy detallada, se veía borrosa o pixelada. Para arreglarlo, los científicos tenían que "engordar" la imagen (rellenarla de ceros) antes de empezar, lo cual consumía mucha memoria.
• El nuevo método (Cúbico): torch-projectors usa una técnica llamada interpolación cúbica. Imagina que en lugar de usar una regla, usas un pincel suave que "adivina" con mucha inteligencia cómo debe ser el color entre dos puntos.
  • La ventaja: Puedes obtener una imagen nítida y perfecta sin tener que "engordar" la imagen antes. Ahorra memoria y es más rápido.

⚡ La Comparación: De la Tortuga al Rayo

El paper hace una prueba de velocidad impresionante:

• La herramienta anterior (torch-fourier-slice) era como una tortuga. A veces, si intentabas hacer el cálculo inverso (reconstruir la imagen), la computadora se quedaba sin memoria y se apagaba.
• torch-projectors es un rayo. Es 10 a 100 veces más rápido que la anterior.
  • En una tarjeta gráfica potente (H100), puede procesar millones de proyecciones por segundo.
  • En una Mac moderna, también vuela, aunque un poco menos que en la tarjeta gráfica.

🧠 ¿Por qué es importante para la ciencia?

Gracias a que esto es tan rápido, ahora los científicos pueden usar Inteligencia Artificial para:

1. Ver virus y proteínas con una claridad increíble.
2. Entender cómo se mueven y cambian estas moléculas (su dinámica).
3. Hacerlo todo "de principio a fin" (end-to-end) sin tener que detenerse a esperar horas a que la computadora calcule una sola imagen.

En resumen

torch-projectors es como darle a los científicos de la medicina una cámara de alta velocidad que antes no existía. Les permite tomar miles de "fotos" de estructuras moleculares y reconstruirlas en 3D usando inteligencia artificial, todo en una fracción del tiempo que antes tomaba, ahorrando energía y memoria en el proceso.

¡Es un gran salto de la bicicleta al cohete! 🚀🔬

Resumen técnico

Resumen Técnico: torch-projectors

1. El Problema
En el campo de la microscopía electrónica de criogenia (cryo-EM) y la tomografía electrónica, las operaciones de proyección en el espacio de Fourier son fundamentales para reconstruir estructuras 3D a partir de proyecciones 2D. Aunque los métodos de aprendizaje automático (machine learning) requieren implementaciones diferenciables para el entrenamiento de modelos de extremo a extremo, las operaciones nativas de PyTorch son insuficientes para este propósito debido a su bajo rendimiento y alto consumo de memoria. Las implementaciones existentes, como torch-fourier-slice, sufren de ineficiencias que limitan su uso práctico en optimizaciones basadas en gradientes que requieren miles de iteraciones.

2. Metodología
El artículo presenta torch-projectors, una biblioteca de alto rendimiento diseñada para realizar proyecciones diferenciables en el espacio de Fourier dentro de PyTorch. Sus características metodológicas clave incluyen:

• Arquitectura Híbrida: Utiliza una extensión C++/Python de PyTorch con kernels optimizados específicamente para CPU, Apple Silicon (MPS) y dispositivos CUDA.
• Espacio de Fourier (RFFT): Opera exclusivamente en el formato RFFT (Real Fast Fourier Transform) de PyTorch, siguiendo las convenciones de FFTW, evitando cambios de formato intermedios que consumen memoria.
• Interpolación Avanzada: Implementa tanto interpolación lineal como cúbica (utilizando kernels de Catmull-Rom con $a = -0.5$). La interpolación cúbica ofrece una mayor precisión sin necesidad de sobremuestreo (oversampling) en el espacio real, lo que reduce drásticamente el uso de memoria en comparación con los métodos tradicionales que requieren ceros (zero-padding) para mejorar la precisión.
• Simetría de Friedel: Maneja cuidadosamente la simetría de Friedel en el espacio de Fourier para funciones de valor real, asegurando la consistencia matemática y evitando el doble conteo en los ejes de frecuencia cero durante las pasadas forward y backward.
• Operaciones Soportadas: Cubre las cuatro operaciones principales necesarias para el análisis de partículas individuales (SPA) y la tomografía:
  • 2D $\to$ 2D (Forward y Backward)
  • 3D $\to$ 2D (Forward)
  • 2D $\to$ 3D (Backward)
• Diferenciación Automática: Todas las operaciones soportan el cálculo de gradientes para entradas, rotaciones y desplazamientos, integrándose nativamente con torch.autograd.

3. Contribuciones Clave

• Rendimiento Excepcional: La biblioteca supera a la implementación anterior (torch-fourier-slice) en 1 a 2 órdenes de magnitud (10x a 100x más rápido) en todas las configuraciones probadas.
• Eficiencia de Memoria: Al calcular todo en una sola llamada sin almacenar tensores intermedios, reduce drásticamente la huella de memoria. Esto permite ejecutar escenarios de retropropagación que antes eran imposibles por agotamiento de memoria (ej. en sistemas con 48 GB de RAM).
• Interpolación Cúbica sin Sobremuestreo: Demuestra que la interpolación cúbica en el espacio de Fourier puede lograr una precisión comparable a la interpolación lineal con sobremuestreo 2x, pero sin el costo computacional y de memoria asociado a la creación de volúmenes más grandes en el espacio real.
• Soporte Multi-Plataforma: Optimizaciones específicas para NVIDIA H100 (CUDA), Apple M4 (CPU y MPS), logrando un escalado de rendimiento robusto con el tamaño del lote y el número de poses.

4. Resultados
Las pruebas de rendimiento realizadas en plataformas NVIDIA H100, Apple M4 CPU y M4 MPS revelan:

• Velocidad: En la GPU H100, se alcanzan tasas de hasta ~86,000 proyecciones/segundo (3D $\to$ 2D, lineal) y ~18,000 (cúbica). La biblioteca supera a torch-fourier-slice por factores que van desde 5x hasta más de 300x dependiendo de la configuración y la plataforma.
• Ancho de Banda: El backend CUDA supera el ancho de banda teórico de la memoria de la H100, lo que sugiere que los datos de referencia caben completamente en la caché del chip.
• Precisión: El análisis de correlación de anillos de Fourier (FRC) y la atenuación de intensidad en el espacio real muestran que la interpolación cúbica sin sobremuestreo mantiene la señal hasta aproximadamente la mitad del tamaño de la caja, siendo una alternativa viable y más rápida a la interpolación lineal con sobremuestreo en línea.
• Comparativa: En escenarios donde el sobremuestreo requiere relleno (padding) en tiempo real (con transformadas de Fourier adicionales), la interpolación cúbica sin sobremuestreo es aproximadamente 4 veces más rápida que la alternativa lineal con sobremuestreo.

5. Significado e Impacto
La introducción de torch-projectors elimina una barrera crítica para la aplicación de métodos de aprendizaje profundo en cryo-EM. Al hacer viables las operaciones de proyección diferenciables de alto rendimiento, permite:

• Entrenar redes neuronales complejas que requieren miles de proyecciones durante la optimización.
• Utilizar arquitecturas de memoria restringida sin depender de costosas transformadas de Fourier en tiempo real para el sobremuestreo.
• Integrar flujos de trabajo de criomicroscopía electrónica directamente en pipelines de aprendizaje automático modernos.

El código es de código abierto (licencia MIT) y está disponible como paquete de Python, facilitando su adopción inmediata por la comunidad científica.