Factored Levenberg-Marquardt for Diffeomorphic Image Registration: An efficient optimizer for FireANTs

Resumen Técnico: Optimizador Levenberg-Marquardt Factorizado para Registro de Imágenes Difeomórficas

1. El Problema

El registro de imágenes difeomórficas (alineación de imágenes médicas preservando la topología) se ha abordado recientemente mediante optimización en tiempo de prueba acelerada por GPU, como en el marco FireANTs. El optimizador predeterminado es Adam, un método de primer orden que utiliza promedios móviles exponenciales (momento y momento cuadrático) para adaptar la tasa de aprendizaje.

Sin embargo, Adam presenta dos limitaciones críticas para volúmenes de imágenes grandes:

• Consumo de Memoria: Debe almacenar variables de estado (momento y momento cuadrático) para cada vóxel, lo que consume una cantidad significativa de memoria GPU, prohibiendo su uso en volúmenes muy grandes.
• Costo Computacional: Aunque es más barato que métodos de segundo orden completos (como L-BFGS), la sobrecarga de actualizar los estados de momento sigue siendo considerable para imágenes masivas.

El objetivo es encontrar un optimizador de segundo orden que sea eficiente en memoria y computacionalmente competitivo, manteniendo la robustez de Adam.

2. Metodología

Los autores proponen un Optimizador Levenberg-Marquardt (LM) Factorizado adaptado específicamente para campos de desplazamiento en registro de imágenes.

• Factorización de Rango 1 por Vóxel:

  • En lugar de calcular la Hessiana completa ($J^T J$), que es intratable para campos de desplazamiento densos (con $\sim 10^7$ parámetros), el método aprovecha la estructura local del problema.
  • Aproximan la Hessiana de Gauss-Newton de forma independiente para cada vóxel como un producto externo de rango 1: $\hat{H}_x = g_x g_x^\top$, donde $g_x$ es el gradiente local.
  • Esto permite una actualización en forma cerrada (Ecuación 4) que normaliza el paso del gradiente por su magnitud, sin necesidad de almacenar ni invertir matrices grandes. Solo se requiere calcular la norma escalar del gradiente.

• Adaptación del Parámetro de Amortiguamiento ($\lambda$):

  • Utilizan un enfoque de región de confianza para ajustar dinámicamente $\lambda$.
  • Si el paso reduce la pérdida (paso "bueno"), $\lambda$ disminuye ($\mu^-$), acercándose al método de Newton.
  • Si la pérdida aumenta (paso "malo"), $\lambda$ aumenta ($\mu^+$), acercándose al descenso de gradiente.
  • Se implementa un criterio de rechazo estilo Metropolis-Hastings: si la pérdida aumenta más allá de un umbral tolerable, el paso se descarta, se restaura el campo de deformación y se reintenta con un $\lambda$ mayor.

• Conexión con el Algoritmo Demons:

  • El método se demuestra como una generalización del algoritmo Demons. Mientras que Demons usa un amortiguamiento local dependiente del residuo, el LM propuesto utiliza un amortiguamiento global adaptativo basado en reglas de región de confianza, permitiendo su aplicación a diversas métricas de similitud (no solo MSE) mediante reparametrización de pérdida cuadrática.

3. Contribuciones Clave

1. Eficiencia de Memoria: El optimizador propuesto elimina la necesidad de almacenar buffers de momento, reduciendo el uso de memoria GPU en un 24.6% para volúmenes grandes en comparación con Adam.
2. Configuración Universal: Demuestran que una sola configuración de hiperparámetros, ajustada en un conjunto de datos (MRI cerebral), se transfiere sin modificaciones a otros dominios (CT pulmonar, registro multimodal abdomen) con rendimiento igual o superior.
3. Análisis de Sensibilidad ("El Acantilado de Amortiguamiento"): Identifican un comportamiento crítico en el hiperparámetro de aumento de amortiguamiento ($\mu^+$). Si $\mu^+ > 2.3$, el optimizador entra en un modo de fallo auto-amplificado donde el amortiguamiento se vuelve excesivo, deteniendo el aprendizaje. Esto define un rango de operación seguro estrecho ($\mu^+ \in [1.5, 2.3]$).
4. Actualización Puntual (Tile Size = 1): Confirman que agrupar gradientes en ventanas espaciales (tiles) degrada el rendimiento; la aproximación vóxel a vóxel es óptima en costo y precisión.

4. Resultados Experimentales

Se evaluó el método en cuatro conjuntos de datos públicos:

• MRI Cerebral (LUMIR y OASIS): El LM superó consistentemente a Adam.
  • LUMIR: Dice medio de 0.8715 (LM) vs. 0.8650 (Adam). Mejora en el caso peor (robustez).
  • OASIS: Dice medio de 0.8293 (LM) vs. 0.8028 (Adam).
• CT Pulmonar (NLST): Adam mantuvo una ligera ventaja en el error de registro medio (TRE), pero el LM demostró ser superior en la supresión de valores atípicos catastróficos cuando se activó el criterio de rechazo.
• Registro Abdominal Multimodal (MR a CT): El LM igualó el rendimiento de Adam (Dice ~0.7615), demostrando que no degrada el rendimiento en órganos complejos o modalidades cruzadas.
• Rendimiento y Memoria:
  • El LM nunca fue más lento que Adam.
  • En volúmenes pequeños, fue hasta 1.67x más rápido debido a la menor sobrecarga de actualización de estado.
  • Ahorro de memoria de hasta 3,024 MB en volúmenes de 512³.

5. Significado e Implicaciones

• Habilitación de Escala: Al reducir significativamente los requisitos de memoria, este optimizador permite realizar registro de imágenes difeomórficas en volúmenes de datos que antes eran imposibles de procesar con Adam debido a limitaciones de VRAM.
• Rehabilitación de Métodos de Segundo Orden: El trabajo demuestra que los métodos de segundo orden pueden ser prácticos y eficientes en el contexto de aprendizaje profundo para visión por computadora si se utilizan aproximaciones de Hessiana inteligentes (factorización de rango 1) y estrategias de amortiguamiento robustas.
• Estabilidad y Robustez: La identificación del "acantilado de amortiguamiento" proporciona una guía crucial para la configuración de optimizadores basados en regiones de confianza, evitando configuraciones que lleven a la convergencia prematura o al estancamiento.
• Generalización: La capacidad de transferir hiperparámetros entre modalidades (cerebro, pulmón, abdomen) sugiere que las propiedades del optimizador están intrínsecamente ligadas a la tarea de registro difeomórfico y no a características específicas del conjunto de datos, simplificando el despliegue en entornos clínicos diversos.

En conclusión, el LM Factorizado ofrece una alternativa superior a Adam para el registro de imágenes grandes, combinando la eficiencia de memoria, la velocidad de convergencia de segundo orden y una configuración robusta y universal.