Reconstruct Anything Model: a lightweight general model for computational imaging

Resumo Técnico: Reconstruct Anything Model (RAM)

1. O Problema

Os problemas de imagem computacional são frequentemente formulados como problemas inversos, onde o objetivo é recuperar uma imagem original $x$ a partir de medições degradadas $y$, modeladas por $y \sim p(y|Ax)$. Aqui, $A$ é um operador que descreve a física de aquisição (ex: transformada de Radon em CT, transformada de Fourier em MRI, kernels de desfoque) e $p$ é a distribuição de ruído (Gaussiano, Poisson, ou misto).

As abordagens existentes dividem-se em duas categorias principais, ambas com limitações significativas:

• Algoritmos Iterativos (Plug-and-Play, Diffusion): Utilizam priores aprendidos (como denoisers) dentro de esquemas iterativos. Embora versáteis, são computacionalmente caros, lentos na inferência e frequentemente introduzem borramento ou artefatos.
• Arquiteturas "Unrolled" (Desenroladas): Redes neurais treinadas para simular passos de otimização. São rápidas, mas geralmente específicas para uma tarefa ou tipo de dado, exigindo re-treinamento completo para novas configurações (ex: mudar de escala, canais ou operadores) e possuem grandes footprints de memória.

Existe uma lacuna para um modelo leve, não iterativo e generalista que possa resolver uma vasta gama de problemas inversos com alto desempenho zero-shot e capacidade de adaptação sem dados de referência (ground-truth).

2. Metodologia: A Arquitetura RAM

O Reconstruct Anything Model (RAM) propõe uma arquitetura neural não iterativa baseada no backbone DRUNet (uma rede residual convolucional), enriquecida com mecanismos de condicionamento que integram a física da aquisição diretamente na rede.

Componentes Principais:

1. Módulo de Estimação Proximal (Inicialização):

   • Em vez de usar apenas $A^\top y$ (que causa desfoque) ou a pseudo-inversa $A^\dagger y$ (instável ao ruído), o RAM utiliza um passo proximal:
     $$\text{prox}_{\lambda f}(A^\top y) = \arg \min_u \lambda \|Au - y\|^2 + \|u - A^\top y\|^2$$
   • Isso cria um equilíbrio adaptativo entre a retroprojeção e a inversão, dependendo do nível de ruído (SNR).

2. Módulo de Subespaço de Krylov (KSM):

   • Esta é a inovação central. Inspirado em métodos de otimização unrolled, o KSM condiciona os mapas de características no operador de medição $A$ sem "desenrolar" explicitamente o algoritmo.
   • O módulo aprende combinações lineares de iterações de potências do operador $A^\top A$ aplicadas às características latentes e aos dados observados.
   • Matematicamente, aproxima a solução como:
     $$x_{\ell+1} = \sum_{k=0}^K \alpha_k (A^\top A)^k x_\ell + \beta_k (A^\top A)^k A^\top y$$
   • Os coeficientes $\alpha_k, \beta_k$ são aprendidos, permitindo que a rede adapte a "profundidade" da correção baseada na física do problema.

3. Condicionamento Multiescala:

   • Inspirado em métodos multigrid, o operador $A$ é definido em grades mais grossas (coarse grids) para estabilizar a inversão em problemas mal-postos. O modelo processa características em múltiplas escalas, onde a inversão é mais estável.

4. Condicionamento de Ruído e Generalidade:

   • O modelo aceita diferentes tipos de ruído (Gaussiano, Poisson, Poisson-Gaussiano) através de mapas de características constantes concatenados com os parâmetros de ruído ($\sigma, \gamma$).
   • Suporta imagens de 1 canal (escala de cinza), 2 canais (complexas, como MRI) e 3 canais (coloridas), compartilhando a maioria dos pesos e ajustando apenas as camadas de entrada/saída.

3. Treinamento e Ajuste Fino (Fine-tuning)

• Treinamento Supervisionado: O modelo é treinado simultaneamente em múltiplas tarefas (desfoque, MRI, CT, inpainting, super-resolução) usando uma função de perda ponderada pelo SNR. Isso força o modelo a aprender uma representação fundamental comum a diferentes físicas de aquisição.
• Ajuste Fino Auto-supervisionado (Self-Supervised):
  • Para cenários onde não há ground-truth (comum em medicina e microscopia), o RAM pode ser ajustado usando apenas os dados de medição $y$.
  • Utiliza perdas de consistência de medição (como SURE ou splitting losses) e perdas de nullspace (como Equivariant Imaging ou Multi-Operator Imaging).
  • Destaque: O modelo pode ser adaptado a novos problemas ou distribuições com apenas uma única imagem de medição, em poucos minutos, sem necessidade de dados limpos.

4. Resultados Principais

Os experimentos demonstram que o RAM atinge desempenho state-of-the-art (SOTA) em diversas tarefas:

• Desempenho Zero-Shot: Em tarefas dentro da distribuição de treinamento (desfoque, MRI, CT), o RAM supera métodos iterativos (como DPIR) e se compara favoravelmente a redes unrolled muito maiores (8x mais parâmetros e FLOPs), mas com inferência muito mais rápida (não iterativa).
  • Exemplo: Em MRI (fator de aceleração 4x e 8x), o RAM supera todas as baselines, incluindo redes unrolled e transformers.
• Generalização para Tarefas Não Vistas:
  • O modelo funciona bem em problemas fora da distribuição (ex: CT com ruído Poisson, embora treinado com Gaussiano; MRI multicôil).
  • Demonstra capacidade em problemas não lineares (como Phase Retrieval) e cegos (Blind Deblurring) quando acoplado a estimadores externos.
• Eficiência Computacional:
  • O RAM é 3.7x mais rápido que redes unrolled equivalentes e 120x mais rápido que métodos baseados em difusão (DDRM), pois não requer iterações durante a inferência.
  • Possui apenas 36M de parâmetros, tornando-o leve em comparação com modelos de difusão (centenas de milhões).
• Ajuste Fino com Poucos Dados:
  • Em tarefas como Compressed Sensing (dados Sentinel-2) e Microscopia Eletrônica (Cryo-EM), o ajuste fino com apenas 1 a 10 imagens supera drasticamente métodos zero-shot e requer menos dados que o ajuste de denoisers tradicionais (DRUNet).

5. Contribuições e Significância

1. Unificação de Paradigmas: O RAM reconcilia a interpretabilidade das arquiteturas unrolled (que incorporam a física $A$) com a eficiência e generalidade das arquiteturas de visão de baixo nível baseadas em UNet/DRUNet.
2. Modelo Fundamental para Imagem Computacional: Demonstra que uma única arquitetura leve pode resolver problemas de imageamento radicalmente diferentes (de MRI a microscopia de fótons únicos) sem re-treinamento massivo.
3. Viabilidade em Cenários Reais: A capacidade de ajuste fino auto-supervisionado com dados escassos ou inexistentes de ground-truth resolve um dos maiores gargalos na aplicação de IA em ciências e medicina.
4. Eficiência: Oferece um equilíbrio superior entre qualidade de reconstrução, velocidade de inferência e custo computacional, desafiando a noção de que métodos iterativos ou de difusão são necessários para alta qualidade.

Conclusão:
O Reconstruct Anything Model (RAM) representa um avanço significativo na direção de modelos fundamentais para imageamento computacional. Ao integrar a física de aquisição diretamente na arquitetura de forma não iterativa e leve, o RAM oferece um framework robusto, rápido e adaptável, capaz de operar tanto em cenários supervisionados quanto em situações críticas de dados escassos, estabelecendo um novo padrão para a reconstrução de imagens em diversas modalidades científicas.