Reconstruct Anything Model: a lightweight general model for computational imaging

1. El Problema

Los métodos de aprendizaje automático actuales para resolver problemas inversos de imagen (como la reconstrucción de MRI, tomografía computarizada, desenfoque o super-resolución) se dividen generalmente en dos categorías, ambas con limitaciones significativas:

• Algoritmos Iterativos (Plug-and-Play, Difusión): Utilizan denoisers preentrenados dentro de esquemas iterativos. Son computacionalmente costosos (requieren múltiples iteraciones), a menudo lentos en la inferencia y pueden introducir borrosidad o artefactos. Además, dependen fuertemente de los datos de entrenamiento.
• Arquitecturas "Unrolled" (Desenrolladas): Son redes neuronales entrenadas de extremo a extremo que simulan los pasos de un algoritmo de optimización. Aunque pueden incorporar el operador de adquisición ($A$), suelen ser específicas para un problema concreto. Cambiar la tarea (ej. de escala de grises a color, o de MRI a CT) o el tamaño de la imagen a menudo requiere reentrenar el modelo completo, lo que reduce su adaptabilidad y aumenta el costo computacional.

Existe una necesidad de un modelo general, ligero y no iterativo que pueda manejar una amplia gama de problemas inversos, diferentes tamaños de imagen y canales, sin depender de ground-truth (referencias reales) para su adaptación a nuevos dominios.

2. Metodología: El Modelo Reconstruct Anything (RAM)

Los autores proponen RAM, una arquitectura neuronal no iterativa y ligera basada en DRUNet (una red residual convolucional), pero con modificaciones clave para integrar la física de adquisición directamente en la red.

Arquitectura Propuesta

La arquitectura (ilustrada en la Figura 2 del artículo) introduce tres mecanismos principales:

1. Módulo de Estimación Proximal (Inicialización):

   • En lugar de usar simplemente $A^T y$ (que puede ser borroso) o la pseudo-inversa $A^\dagger y$ (muy sensible al ruido), el modelo inicia con un paso proximal:
     $$\text{prox}_{\lambda f}(A^T y) = \arg \min_u \lambda \|Au - y\|^2 + \|u - A^T y\|^2$$
   • Esto actúa como un punto medio entre la inversión directa y la inversión regularizada, mapeando las mediciones $y$ a un dominio de imagen con ruido gaussiano manejable. El parámetro $\lambda$ se ajusta dinámicamente según la relación señal-ruido (SNR).

2. Módulo de Subespacio de Krylov (KSM):

   • Este es el núcleo de la integración de la física. En lugar de "desenrollar" iteraciones explícitas, la red utiliza un módulo que aprende a combinar iteraciones de potencias del operador $A^T A$ sobre las mapas de características.
   • Matemáticamente, aproxima la solución como una combinación lineal de términos de la forma $(A^T A)^k x_\ell$ y $(A^T A)^k A^T y$.
   • Esto permite que la red aprenda los coeficientes óptimos para resolver el problema inverso en una sola pasada (no iterativa), imitando la eficacia de los métodos iterativos pero con una inferencia mucho más rápida.

3. Condicionamiento Multiescala y de Ruido:

   • Multiescala: Inspirado en métodos multigrid, el modelo aplica operadores de adquisición en escalas más gruesas ($A_s$) para estabilizar la inversión en problemas mal planteados, evitando inestabilidades en rejillas finas.
   • Ruido: El modelo soporta ruido Gaussiano y Poisson-Gaussian mediante mapas de ruido concatenados, eliminando sesgos (biases) para garantizar equivalencia de escala.
   • Flexibilidad de Canales: La arquitectura comparte la mayoría de los pesos entre diferentes modalidades (escala de grises, complejos, color), ajustando solo las capas de entrada y salida.

Entrenamiento y Ajuste Fino (Fine-tuning)

• Entrenamiento Supervisado: Se entrena un solo modelo base en múltiples tareas simultáneamente (desenfoque, MRI, CT, inpainting, super-resolución) utilizando datasets diversos (LSDIR, fastMRI, LIDC-IDRI).
• Ajuste Fino Auto-supervisado (Zero-Shot): Una contribución crucial es la capacidad de adaptar el modelo a nuevos problemas o distribuciones sin datos de ground-truth. Utilizando solo las mediciones observadas ($y$), se puede ajustar el modelo mediante pérdidas de consistencia de medición (como SURE o pérdidas de división) y pérdidas de nullspace (Equivariant Imaging). Esto permite que el modelo se adapte en minutos con solo unas pocas imágenes (incluso una sola).

3. Contribuciones Clave

1. Arquitectura General y Ligera: RAM es un modelo "fundacional" para imágenes computacionales que resuelve una amplia variedad de problemas inversos lineales y no lineales con una sola arquitectura de 36M de parámetros.
2. Rendimiento No Iterativo: Logra un rendimiento de vanguardia (SOTA) sin la necesidad de iteraciones costosas en tiempo de inferencia, superando en velocidad a los métodos iterativos y a las redes desenrolladas.
3. Adaptabilidad Zero-Shot: Demuestra que un modelo preentrenado puede ser ajustado finamente de manera auto-supervisada a problemas fuera de distribución (ej. MRI multicapa, CT con ruido Poisson, imágenes de microscopía electrónica criogénica) utilizando solo las mediciones, sin necesidad de pares imagen-ruido.
4. Integración Eficiente de Física: Incorpora el operador de adquisición ($A$) y la estadística del ruido directamente en la arquitectura mediante el módulo KSM y la estimación proximal, sin requerir suposiciones restrictivas sobre la estructura de $A$ (como descomposición SVD).

4. Resultados Experimentales

El modelo fue evaluado en una amplia gama de tareas, mostrando resultados superiores o competitivos:

• Problemas en Distribución (In-Distribution):
  • En tareas de desenfoque (Gaussiano y de movimiento), RAM supera a métodos end-to-end como Restormer y PDNet, y compite favorablemente con métodos iterativos como uDPIR, pero con una inferencia 3.7 veces más rápida que uDPIR y 120 veces más rápida que los métodos basados en difusión (DDRM).
  • En MRI (fastMRI) y Tomografía Computarizada (CT), RAM supera a todas las líneas base, incluyendo redes desenrolladas y modelos basados en transformadores.
• Rendimiento Zero-Shot (Fuera de Distribución):
  • MRI Multicapa: Funciona bien en configuraciones no vistas durante el entrenamiento (factores de aceleración 8, múltiples bobinas).
  • CT con Ruido Poisson: Logra reconstrucciones de alta calidad en CT con ruido Poisson, aunque fue entrenado principalmente con ruido Gaussiano.
  • Problemas Ciegos y No Lineales: Se extendió exitosamente a la deconvolución ciega (estimando el kernel de desenfoque) y a la retrieval de fase (un problema no lineal), demostrando robustez más allá de su entrenamiento original.
• Ajuste Fino Auto-supervisado:
  • En tareas como la compresión de sensores (Sentinel-2), microscopía electrónica criogénica (Cryo-EM) e imágenes de fotones bajos (LinoSPAD), el ajuste fino con solo 1 a 10 imágenes (sin ground-truth) mejoró drásticamente el rendimiento, superando a los métodos PnP y de difusión que requieren grandes conjuntos de datos o ground-truth.

5. Significado e Impacto

El trabajo de RAM representa un cambio de paradigma en la reconstrucción de imágenes computacionales:

• Eficiencia: Elimina la necesidad de costosos algoritmos iterativos en tiempo de inferencia, haciendo viable la reconstrucción en tiempo real en hardware limitado.
• Generalización: Desafía la noción de que los modelos de reconstrucción deben ser especializados por tarea. Muestra que problemas aparentemente dispares (desde MRI hasta microscopía) comparten estructuras subyacentes que pueden ser aprendidas por un solo modelo.
• Accesibilidad: La capacidad de ajuste fino auto-supervisado con pocos datos es crucial para aplicaciones científicas y médicas donde obtener ground-truth es imposible o extremadamente costoso.
• Código Abierto: Los autores han liberado el código, facilitando la adopción y el desarrollo futuro en la comunidad.

En resumen, RAM establece un nuevo estándar para modelos fundacionales en imágenes computacionales, combinando la velocidad de las redes neuronales directas con la precisión física de los métodos de optimización, todo ello con una capacidad de adaptación sin precedentes a nuevos dominios sin supervisión.