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

O artigo apresenta a biblioteca *torch-projectors*, uma implementação de alta performance e diferenciável para projeções no espaço de Fourier no PyTorch, otimizada para múltiplos hardwares e que supera em até duas ordens de grandeza as soluções existentes para análise de microscopia eletrônica.

O essencial

Imagine que você é um detetive tentando reconstruir um castelo de areia gigante, mas só consegue ver fotos planas (2D) tiradas de vários ângulos diferentes. O seu objetivo é juntar todas essas fotos para "ver" o castelo em 3D, até nos menores detalhes.

No mundo da ciência, isso é o que acontece na criomicroscopia eletrônica: cientistas tiram milhares de fotos de proteínas microscópicas para entender como elas funcionam. O problema é que, para fazer isso com inteligência artificial (aprendizado de máquina), os computadores precisam fazer um tipo de "mágica matemática" chamada projeção de Fourier.

Aqui está o resumo do papel "torch-projectors" em linguagem simples:

1. O Problema: O Motor de um Carro Esportivo com Rodas de Carrinho de Bebê

Os cientistas querem usar redes neurais (IA) para reconstruir essas proteínas de forma mais rápida e precisa. Para isso, o computador precisa calcular essas projeções matemáticas milhões de vezes.

• A situação atual: As ferramentas que o PyTorch (uma biblioteca popular de IA) já tinha para fazer isso eram como tentar dirigir um carro de Fórmula 1 com rodas de carrinho de bebê. Elas funcionavam, mas eram extremamente lentas e consumiam muita memória, travando o computador.
• A consequência: Era impossível treinar modelos de IA complexos para essa tarefa porque o computador demorava dias para fazer o que deveria levar minutos.

2. A Solução: O "torch-projectors"

O autor criou uma nova biblioteca chamada torch-projectors. Pense nela como trocar as rodas de carrinho de bebê por um motor de foguete.

• O que ela faz: Ela realiza essas projeções matemáticas de forma super rápida e extremamente eficiente.
• A mágica: Ela foi feita para funcionar bem em qualquer computador moderno, seja um PC com placa de vídeo NVIDIA, um Mac com chip Apple Silicon ou apenas o processador comum.

3. Como Funciona a "Mágica"? (Analogias)

A. O Jogo do "Ponto de Vista" (Projeção)

Imagine que você tem um globo terrestre 3D (a proteína).

• Projeção para frente (Forward): É como tirar uma foto do globo de um ângulo específico. O novo sistema faz isso instantaneamente, sem precisar recalcular tudo do zero.
• Projeção para trás (Backward): É o processo inverso. Imagine que você tem a foto e quer "jogar" a informação de volta para o globo 3D para reconstruí-lo. O sistema faz isso com precisão cirúrgica, garantindo que nenhuma informação se perca no caminho.

B. A Interpolação (O "Zoom" Inteligente)

Às vezes, a foto não cai exatamente em uma linha reta do globo; ela cai "entre" os pixels.

• O jeito antigo (Linear): Era como desenhar uma linha reta entre dois pontos. Funcionava, mas a imagem ficava um pouco "pixelada" ou borrada.
• O jeito novo (Cúbico): O novo sistema usa uma técnica chamada "Interpolação Cúbica". Imagine que, em vez de apenas ligar dois pontos, ele desenha uma curva suave e perfeita que passa exatamente por onde deveria. Isso dá uma imagem muito mais nítida, sem precisar aumentar o tamanho do globo (o que economizaria muita memória).

C. O Truque da Memória

Antes, para fazer isso com precisão, o computador precisava criar cópias gigantes dos dados (como fazer uma cópia de um mapa do mundo inteiro só para dar um zoom em uma cidade). O novo sistema é como um GPS inteligente: ele calcula o caminho exato na hora, sem precisar carregar o mapa inteiro na memória. Isso permite que computadores mais fracos também rodem esses modelos complexos.

4. Os Resultados: Velocidade Insana

O artigo mostra testes comparativos:

• O novo sistema é 10 a 100 vezes mais rápido que as ferramentas antigas.
• Em um computador com placa de vídeo potente (NVIDIA H100), ele consegue processar milhões de projeções por segundo.
• Ele consome muito menos memória, o que significa que você pode treinar modelos de IA maiores sem estourar o limite do seu computador.

5. Por que isso importa para você?

Embora pareça muito técnico, isso é um divisor de águas para a medicina e biologia:

• Descoberta de Remédios: Com essa velocidade, os cientistas podem descobrir a estrutura de novas proteínas em horas, não em meses.
• Medicina Personalizada: Entender como as proteínas se movem ajuda a criar remédios que funcionam melhor para doenças específicas.
• Democratização: Como o sistema é rápido e leve, mais laboratórios (mesmo os com menos dinheiro) podem usar inteligência de ponta para descobrir coisas novas.

Em resumo: O "torch-projectors" é a peça que faltava para transformar a reconstrução de proteínas de um processo lento e manual em uma corrida de Fórmula 1 automatizada, permitindo que a inteligência artificial ajude a curar doenças mais rápido do que nunca.

Resumo técnico

Resumo Técnico: torch-projectors

1. O Problema

A microscopia eletrônica de criogênica (cryo-EM) e a tomografia eletrônica dependem fundamentalmente de operações de projeção no espaço de Fourier para reconstruir estruturas 3D a partir de projeções 2D. Com o advento de métodos de aprendizado de máquina (ML) para análise de partículas únicas (SPA), tornou-se necessário implementar essas operações de forma diferenciável para permitir o treinamento de ponta a ponta (end-to-end) via retropropagação de gradiente.

No entanto, o PyTorch nativo não oferece operações de projeção no espaço de Fourier otimizadas. Implementações existentes usando funções nativas do PyTorch são:

• Excessivamente lentas: Inviáveis para treinamento prático de grandes modelos.
• Ineficientes em memória: Requerem armazenamento de valores intermediários ou sobreamostragem (zero-padding) extensiva, o que consome muita memória, especialmente em redes neurais profundas.
• Limitadas: Falta de suporte nativo para interpolação cúbica de alta precisão sem sobrecarga de memória.

2. Metodologia

O torch-projectors é uma biblioteca de alto desempenho que implementa operadores de projeção e retroprojeção diferenciáveis, otimizados para CPU, Apple Silicon (MPS) e CUDA.

• Arquitetura Híbrida: A biblioteca utiliza a arquitetura de extensão C++/Python do PyTorch (TORCH_LIBRARY), registrando kernels otimizados específicos para cada backend (CPU, MPS, CUDA). As funções de gradiente são encapsuladas em torch.autograd.Function para integração transparente com a diferenciação automática.
• Domínio de Fourier (RFFT): Todas as operações ocorrem exclusivamente no espaço de Fourier, utilizando o formato RFFT (Real FFT) do PyTorch, seguindo as convenções do FFTW. Isso elimina a necessidade de conversões frequentes entre domínios real e complexo durante a projeção.
• Interpolação Avançada:
  • Suporta interpolação linear (bilinear/trilinear) e cúbica (Catmull-Rom).
  • A interpolação cúbica oferece precisão superior ($C^1$) sem a necessidade de sobreamostragem (zero-padding) no espaço real, reduzindo drasticamente o uso de memória em comparação com métodos tradicionais (como no RELION) que dependem de volumes sobreamostrados.
• Simetria de Friedel: O código lida rigorosamente com a simetria de Friedel ($F(-k) = F^*(k)$) inerente a funções reais, garantindo consistência matemática e evitando contagem dupla de frequências zero durante a acumulação de gradientes.
• Operações Suportadas: A biblioteca implementa quatro tipos principais de projeções, essenciais para diferentes tarefas de cryo-EM:
  1. 2D $\to$ 2D (Forward e Backward)
  2. 3D $\to$ 2D (Forward - Teorema do Corte Central)
  3. 2D $\to$ 3D (Backward - Retroprojeção)
• Tratamento de Gradientes: Os gradientes para matrizes de rotação e vetores de translação são calculados analiticamente usando a regra da cadeia, derivando as funções de kernel de interpolação (linear e cúbica).

3. Principais Contribuições

• Desempenho Extremo: A biblioteca supera a implementação existente torch-fourier-slice em 1 a 2 ordens de magnitude (10x a 100x mais rápida).
• Eficiência de Memória: Ao realizar cálculos em um único passo sem armazenar tensores intermediários massivos, a biblioteca reduz a pegada de memória, permitindo o treinamento de modelos que antes eram impossíveis devido a limites de memória (ex: 48 GB de VRAM).
• Flexibilidade de Interpolação: Introduz a interpolação cúbica de Catmull-Rom como uma alternativa viável e rápida à sobreamostragem, oferecendo precisão de sinal comparável à interpolação linear com 2x de sobreamostragem, mas com custo computacional menor em cenários dinâmicos.
• Suporte Multi-Backend: Otimizações específicas para NVIDIA CUDA, Apple Silicon (MPS) e CPU, garantindo escalabilidade em diferentes hardwares.

4. Resultados

Os benchmarks foram realizados em NVIDIA H100 (CUDA), Apple M4 (CPU e MPS) com variações de tamanho de caixa (32 a 256 pixels), tamanho de lote e contagem de poses.

• Velocidade:
  • O backend CUDA atinge o maior rendimento, superando o CPU em mais de 2 ordens de magnitude.
  • Em cenários de 3D $\to$ 2D, o torch-projectors foi até 364x mais rápido que o torch-fourier-slice em certas configurações no CPU.
  • A interpolação cúbica sem sobreamostragem é aproximadamente 4x mais rápida do que a interpolação linear com sobreamostragem "on-the-fly" (que exige padding e FFTs adicionais).
• Precisão:
  • A análise de correlação de anel de Fourier (FRC) mostrou que todos os métodos produzem artefatos insignificantes no espaço de Fourier.
  • No espaço real, a interpolação cúbica (1x) e a linear com 2x de sobreamostragem mantêm a intensidade do sinal de forma similar até aproximadamente metade do tamanho da caixa, sendo adequadas para a maioria dos algoritmos de cryo-EM.
• Limitações de Memória: Enquanto o torch-fourier-slice falhou em testes de retroprojeção devido ao consumo excessivo de memória, o torch-projectors conseguiu executar todos os cenários, incluindo passagens de retropropagação complexas.

5. Significado e Impacto

O torch-projectors remove uma barreira crítica para a aplicação de aprendizado de máquina profundo na microscopia eletrônica. Ao tornar as operações de projeção Fourier rápidas e diferenciáveis, a biblioteca permite:

• O treinamento de redes neurais complexas para reconstrução 3D e classificação de partículas.
• A otimização de parâmetros em larga escala que antes eram computacionalmente proibitivos.
• A adoção de métodos de interpolação de alta precisão (cúbica) sem penalidades severas de memória, facilitando o desenvolvimento de algoritmos mais robustos para cryo-EM e tomografia.

A biblioteca está disponível como um pacote Python de código aberto (MIT License), integrando-se nativamente ao ecossistema PyTorch.