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

1. Problema e Motivação

O trabalho aborda o problema de registro de imagens difeomórficas (alinhamento de imagens médicas preservando a topologia), formulado como um problema de otimização no tempo de teste (test-time optimization) acelerado por GPU, utilizando o framework FireANTs.

• Limitação Atual: O otimizador padrão do FireANTs é o Adam. Embora eficiente e robusto, o Adam é um método de segunda ordem aproximado que requer o armazenamento de variáveis de estado (momento e estimativas de momento quadrado) para cada voxel da imagem. Para volumes de imagem grandes (ex: resoluções 3D altas), isso consome uma quantidade significativa de memória da GPU, limitando a aplicação em datasets voluminosos.
• Desafio: Métodos de segunda ordem clássicos, como Levenberg–Marquardt (LM), geralmente são inviáveis para campos de deformação densos devido à necessidade de calcular e inverter a matriz Hessiana completa ($J^T J$), que teria dimensões da ordem de $10^{13}$ parâmetros para uma imagem 3D típica.
• Objetivo: Desenvolver um otimizador de segunda ordem que seja computacionalmente eficiente, consuma menos memória que o Adam e mantenha (ou supere) a precisão do registro, sem exigir o armazenamento de grandes buffers de estado.

2. Metodologia

Os autores propõem um otimizador Levenberg–Marquardt (LM) Fatorado, adaptado especificamente para campos de deslocamento em registros difeomórficos.

A. Atualização Compositiva e Descida Euleriana

O método mantém a difeomorfia através de uma regra de atualização compositiva. Em vez de otimizar diretamente o campo de deslocamento $u$, o otimizador gera um campo de velocidade $v(t)$, e o warp é atualizado iterativamente:
$$u(t+1) = \epsilon v(t) + u(t) \circ (Id + \epsilon v(t))$$
Isso permite o uso de otimizadores definidos para espaços euclidianos (como Adam ou LM) sobre a variedade Riemanniana dos difeomorfismos, evitando operações caras de transporte paralelo.

B. Fatoração de Hessiana Rank-1 por Voxel

A inovação central é a aproximação da Hessiana. Em vez de formar a matriz Hessiana global intransponível, o método assume que o gradiente do registro em cada voxel depende principalmente do conteúdo local da imagem.

• Aproximação: A Hessiana de Gauss-Newton é aproximada independentemente para cada voxel como um produto externo de posto 1 ($\hat{H}_x = g_x g_x^T$).
• Atualização de Forma Fechada: A inversa amortecida desta matriz $3 \times 3$ (ou $d \times d$) tem uma forma fechada simples:
  $$\Delta u_x = -\frac{r g_x}{\|g_x\|^2 + \lambda}$$
  Onde $r$ é o resíduo global e $\lambda$ é o parâmetro de amortecimento. Isso elimina a necessidade de armazenar ou inverter matrizes grandes, exigindo apenas o cálculo de uma norma escalar.
• Contexto Espacial: Embora seja possível agrupar aproximações em janelas (tiles), os experimentos mostram que a atualização ponto a ponto ($tile\_size=1$) é a mais eficiente e precisa.

C. Amortecimento Adaptativo e Critério de Rejeição

• Amortecimento (Trust Region): O parâmetro $\lambda$ é ajustado dinamicamente após cada passo. Se a perda (loss) aumentar (passo ruim), $\lambda$ é multiplicado por um fator $\mu_+ > 1$; se diminuir (passo bom), é multiplicado por $\mu_- < 1$.
• Rejeição de Passo (Metropolis-Hastings style): Para evitar atualizações que pioram a função de perda, implementa-se um critério de rejeição. Se a degradação da perda exceder um limiar em relação à melhoria anterior, o passo é rejeitado, o campo de deformação é restaurado e o passo é refeito com $\lambda$ aumentado. Um limite máximo ($\lambda_{max}$) é essencial para evitar divergência catastrófica.

3. Principais Contribuições

1. Otimizador LM Fatorado: Introdução de um método de segunda ordem viável para registros de imagem densos, reduzindo o custo de memória ao eliminar buffers de momento, exigindo apenas um parâmetro escalar de estado ($\lambda$).
2. Descoberta da "Cliff" de Amortecimento (Damping Cliff): Identificação de uma sensibilidade crítica no hiperparâmetro de aumento de amortecimento ($\mu_+$). Valores acima de ~2.3 causam uma queda catastrófica no desempenho (o algoritmo para de atualizar), enquanto valores entre 1.5 e 2.3 são ótimos.
3. Universalidade de Hiperparâmetros: Demonstração de que uma única configuração de hiperparâmetros, ajustada em um dataset (MRI cerebral), transfere-se sem modificações para outros datasets (CT pulmonar, registro multimodal abdômen), mantendo ou superando o Adam.
4. Eficiência de Memória: Redução de até 24,6% no uso de memória de pico da GPU para volumes grandes, permitindo o processamento de imagens maiores que seriam impossíveis com o Adam.

4. Resultados Experimentais

Os autores avaliaram o método em quatro benchmarks públicos:

• MRI Cerebral (LUMIR e OASIS): O LM superou o Adam consistentemente.
  • LUMIR: Dice médio de 0.8715 (LM) vs. 0.8650 (Adam).
  • OASIS: Dice médio de 0.8293 (LM) vs. 0.8028 (Adam).
• Registro Multimodal Abdômen (Abdomen-L2R): O LM igualou o desempenho do Adam (Dice ~0.7615).
• CT Pulmonar (NLST): O Adam teve desempenho ligeiramente superior em média, mas o LM mostrou-se mais robusto a outliers extremos quando o mecanismo de rejeição de passos foi ativado.
• Memória e Tempo:
  • Para volumes grandes (ex: 512³), o LM economiza ~3 GB de memória.
  • O tempo de execução é comparável ou ligeiramente superior ao Adam (até 1.67x mais rápido em volumes pequenos devido à menor sobrecarga de atualização de momento).

5. Significado e Conclusão

O trabalho demonstra que métodos de segunda ordem podem ser viáveis e superiores para registro de imagens médicas se forem projetados com restrições de memória específicas (fatoração de Hessiana).

• Impacto Prático: A redução de memória permite aplicar registro difeomórfico de alta precisão em datasets clínicos maiores e de maior resolução, que antes eram proibitivos para otimizadores como Adam.
• Robustez: A configuração proposta é robusta e universal, eliminando a necessidade de ajuste fino de hiperparâmetros para cada novo tipo de imagem ou órgão.
• Insight Teórico: A descoberta da "cliff" de amortecimento alerta a comunidade sobre a instabilidade potencial de esquemas de amortecimento adaptativo mal calibrados em métodos de segunda ordem, fornecendo diretrizes claras para configurações seguras.

Em resumo, o LM Fatorado oferece uma alternativa eficiente, econômica em memória e de alta precisão ao Adam para o registro de imagens difeomórficas, tornando-se uma nova configuração padrão recomendada para o framework FireANTs.