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

This paper introduces torch-projectors, a high-performance, differentiable library for Fourier-space projections in PyTorch that significantly accelerates electron microscopy workflows by outperforming existing solutions by 1–2 orders of magnitude across CPU, Apple Silicon, and CUDA devices.

The Gist

Imagine you are trying to solve a giant, 3D jigsaw puzzle, but you can only see the puzzle pieces as flat, 2D shadows cast on a wall from different angles. This is essentially what scientists do in Cryo-Electron Microscopy (Cryo-EM) to figure out the shape of tiny proteins. They take thousands of 2D "shadows" (projections) and use math to reconstruct the original 3D object.

For a long time, doing this reconstruction with modern AI (Machine Learning) was like trying to run a Formula 1 race on a bicycle. The math required to turn those 2D shadows back into 3D shapes was too slow and clunky for computers to handle efficiently while learning.

Enter torch-projectors. Think of this paper as the unveiling of a brand-new, high-speed engine for these computers.

Here is the breakdown of what this paper is about, using some everyday analogies:

1. The Problem: The "Slow Shadow"

In the world of electron microscopy, there is a rule called the Fourier Slice Theorem. In simple terms, it says: If you look at a 3D object from a specific angle, the "shadow" you see is actually a slice through the object's invisible frequency map.

To train AI models to solve these puzzles, the computer needs to do this "shadow casting" and "shadow reading" millions of times, adjusting its guesses based on errors.

• The Old Way: The standard tools (PyTorch) were like using a spoon to dig a tunnel. It worked, but it was painfully slow and ate up all the computer's memory (RAM), causing it to crash.
• The Result: Scientists couldn't use powerful AI because the math was too heavy.

2. The Solution: The "Magic Lens" (torch-projectors)

The author, Dimitry Tegunov, built a new library called torch-projectors. Think of this as replacing that spoon with a laser drill.

• Speed: It is 10 to 100 times faster than the previous best tools. It's so fast that on powerful graphics cards (GPUs), it can process thousands of these "shadows" in a single second.
• Memory: It's incredibly efficient. Instead of saving every intermediate step (which fills up the computer's desk with clutter), it does the calculation in one smooth motion, keeping the workspace clean.
• Smart Math: It uses a technique called Cubic Interpolation. Imagine trying to guess the color of a pixel between two known pixels.
  • Linear (Old way): Just drawing a straight line between them. It's okay, but can look blocky.
  • Cubic (New way): Drawing a smooth, curved line that fits perfectly. This gives a much sharper, more accurate picture without needing to zoom in (oversample) and waste memory.

3. How It Works: The "Shadow Puppet" Analogy

The library handles two main tasks, which the paper calls "Forward" and "Backward" projections.

• Forward Projection (The Shadow Puppet): You have a 3D puppet (the protein model). You shine a light from a specific angle, and the library instantly calculates what the 2D shadow on the wall looks like.
• Backward Projection (The Reverse Shadow): You have a 2D shadow on the wall. The library figures out how to "scatter" that shadow back into 3D space to update the puppet's shape.

The magic of torch-projectors is that it can do both of these instantly, and if the AI makes a mistake, it can instantly calculate exactly how to fix the puppet's shape (this is called "differentiation" or "backpropagation").

4. The "Friedel Symmetry" Trick

One of the tricky parts of this math is that the data is "symmetric" (like a mirror image). If you count both sides of the mirror, you double-count the information.

• The Paper's Fix: The library has a built-in "traffic cop" that knows exactly how to handle these mirror images so the computer doesn't get confused or double-count data. This ensures the math stays perfect even when things get complicated.

5. Why Should You Care?

Before this library, training AI to understand protein structures was like trying to drive a car with the parking brake on. It was possible, but you could only go very slowly.

torch-projectors takes the parking brake off.

• For Scientists: It means they can now use powerful AI to solve protein structures faster and more accurately than ever before.
• For Medicine: Faster protein understanding leads to faster drug discovery. If we can understand how a virus works faster, we can build vaccines faster.
• For AI: It proves that we can build highly specialized, super-fast math tools that fit perfectly into the AI ecosystem, opening the door for more complex scientific simulations.

The Bottom Line

This paper introduces a super-charged tool that makes the heavy math of 3D protein reconstruction fast, light, and accurate. It turns a slow, memory-hungry process into a sleek, high-speed operation, allowing scientists to use the full power of modern AI to unlock the secrets of life at the atomic level.

Where to find it: It's free and open-source, ready for anyone to download and use on their computers (whether they have a standard PC, a Mac, or a powerful server).

Technical Summary

1. Problem Statement

In Cryogenic Electron Microscopy (cryo-EM) and Electron Tomography (ET), reconstructing 3D structures from 2D projections relies heavily on the Fourier Slice Theorem. This theorem states that a 2D projection of a 3D object corresponds to a central slice through the 3D Fourier transform of that object.

While modern machine learning (ML) approaches require differentiable implementations of these projection operations for end-to-end training, existing solutions in PyTorch face significant bottlenecks:

• Performance: Native PyTorch implementations using built-in functions are too slow for practical use, often lacking the throughput required for large-scale optimization.
• Memory Overhead: Naive implementations often require storing intermediate tensors or performing multiple passes (e.g., oversampling via zero-padding in real space, transforming, and cropping), leading to excessive memory consumption.
• Accuracy vs. Efficiency Trade-off: Traditional methods (like those in RELION) achieve high accuracy by oversampling references, but this is computationally expensive and memory-intensive for ML workflows where intermediate representations are dynamically generated.

2. Methodology

The paper introduces torch-projectors, a high-performance library designed to perform differentiable Fourier-space projections directly within PyTorch.

Core Architecture

• Hybrid Implementation: Built as a PyTorch extension using C++/Python with TORCH_LIBRARY registration. It features optimized kernels for CPU, Apple Silicon (MPS), and CUDA.
• Differentiability: Wraps C++ operators with custom torch.autograd.Function classes, enabling seamless gradient back-propagation for all inputs (reconstruction, rotation matrices, and translation shifts).
• Data Format: Operates exclusively in Fourier space using PyTorch's RFFT format (Real FFT), adhering to FFTW conventions. It avoids intermediate format changes, keeping tensors in RFFT format throughout the projection pipeline.

Key Algorithmic Features

• Interpolation Methods:
  • Linear: Uses standard bilinear (2D) and trilinear (3D) kernels.
  • Cubic: Implements Catmull-Rom cubic interpolation ($a = -0.5$) with $C^1$ continuity. This allows for high accuracy without the need for oversampling (zero-padding), significantly reducing memory usage compared to linear interpolation with oversampling.
• Friedel Symmetry Handling:
  • The library correctly manages the symmetry of real-valued functions in Fourier space ($F(-k) = F^*(k)$).
  • Forward Pass: Handles negative frequency coordinates by accessing the conjugate symmetric counterpart.
  • Backward Pass: Prevents double-counting gradients on the zero-frequency axis by skipping specific components and ensuring symmetric accumulation for contributions landing on the axis.
• Projection Types: Supports four distinct operations essential for EM workflows:
  1. 2D $\to$ 2D: Forward and Backward (for particle classification).
  2. 3D $\to$ 2D: Forward (generating projections from a 3D volume).
  3. 2D $\to$ 3D: Backward (accumulating projections into a 3D volume).
• Shifts and Rotations:
  • Rotations are applied via matrix multiplication in Fourier space.
  • Translations are applied via phase modulation ($e^{-i2\pi k \cdot s}$), avoiding costly real-space shifts.
  • Gradients for rotation matrices and shifts are computed analytically using the chain rule and kernel derivatives.

3. Key Contributions

1. High-Performance Kernels: The library achieves 1–2 orders of magnitude speedup over existing solutions (specifically torch-fourier-slice) across all hardware platforms.
2. Memory Efficiency: By performing calculations in a single step without storing intermediate oversampled tensors, it drastically reduces memory footprint, enabling larger batch sizes and deeper neural networks.
3. Cubic Interpolation without Oversampling: Demonstrates that Catmull-Rom cubic interpolation can match the accuracy of linear interpolation with 2x oversampling but at a fraction of the computational cost when oversampling is not pre-computed.
4. Full Differentiability: Provides analytical gradients for rotation matrices and translation shifts, which are critical for optimizing pose parameters in ML-driven EM reconstruction.
5. Dual-Volume Weighting: Implements an approach similar to RELION's dual-volume system, allowing for optional gradient back-propagation through weight components (useful for CTF correction and Wiener filtering).

4. Results

Performance benchmarks were conducted on NVIDIA H100 (CUDA), Apple M4 (CPU), and Apple M4 (MPS).

• Throughput:
  • CUDA: Achieves the highest throughput (e.g., ~86,000 projections/sec for 3D$\to$2D forward on 128px boxes with 2048 poses).
  • Comparison: torch-projectors outperforms torch-fourier-slice by factors ranging from 5x to over 300x depending on the configuration and hardware.
  • Memory: In many backward pass scenarios, torch-fourier-slice failed due to memory exhaustion (48GB limit), whereas torch-projectors completed successfully.
• Interpolation Efficiency:
  • Cubic vs. Linear: While cubic interpolation requires more samples (4x in 2D, 8x in 3D), the performance penalty is mitigated by memory locality.
  • Oversampling Trade-off: When references cannot be pre-padded (common in dynamic ML graphs), cubic interpolation without oversampling is ~4x faster than linear interpolation with on-the-fly 2x padding (which requires expensive FFTs).
• Accuracy:
  • Fourier Ring Correlation (FRC) analysis shows negligible artifacts in Fourier space.
  • Real-space intensity attenuation analysis confirms that cubic interpolation (1x) maintains signal strength up to ~50% of the box size, comparable to linear interpolation with 2x oversampling.

5. Significance

The torch-projectors library removes a critical bottleneck in applying deep learning to electron microscopy. By enabling fast, memory-efficient, and fully differentiable projection operations, it allows researchers to:

• Train complex neural networks for 3D reconstruction and particle classification that were previously infeasible due to memory or speed constraints.
• Optimize pose parameters (rotation and translation) jointly with network weights in an end-to-end manner.
• Utilize advanced interpolation techniques (cubic) without the overhead of oversampling, making dynamic, memory-constrained workflows viable.

The library is open-source (MIT license) and available as a Python package, promising to accelerate the adoption of ML methods in structural biology.