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

1. Problématique

L'article aborde le problème de l'enregistrement d'images diféomorphiques (diffeomorphic image registration) formulé comme un problème d'optimisation au moment du test (test-time optimization) sur GPU, tel que proposé par le framework FireANTs.

• Contexte : L'optimisation vise à trouver un champ de déplacement $u$ qui minimise une fonction de perte combinant une similarité d'images et une régularisation.
• Limitation actuelle : L'optimiseur par défaut, Adam, bien que robuste et rapide, est une méthode quasi-du second ordre qui nécessite le stockage de variables d'état (moments et estimations de moments au carré). Pour des images volumineuses (ex: volumes 3D de haute résolution), ces buffers de mémoire consomment une quantité significative de RAM GPU, limitant l'applicabilité de la méthode sur de grands ensembles de données.
• Défi : Trouver une méthode d'optimisation du second ordre qui soit aussi efficace en mémoire qu'Adam, tout en conservant ou améliorant la précision de l'enregistrement, sans nécessiter le calcul explicite et coûteux de la matrice Hessienne complète (qui serait de l'ordre de $10^{13}$ paramètres pour un champ de déplacement standard).

2. Méthodologie : Optimiseur Levenberg–Marquardt Factorisé

Les auteurs proposent une variante modifiée de l'algorithme Levenberg–Marquardt (LM) adaptée aux champs de déformation denses.

• Approximation Hessienne Rank-1 par Voxel :
  Au lieu de calculer la Hessienne complète $J^T J$, l'article exploite la structure locale de la perte d'enregistrement. L'Hessienne est approximée voxel par voxel par un produit extérieur de rang 1 ($\hat{H}_x = g_x g_x^T$), où $g_x$ est le gradient local.

  • Cela permet d'obtenir une mise à jour en forme fermée (closed-form) pour chaque voxel :
    $$\Delta u_x = -\frac{r g_x}{\|g_x\|^2 + \lambda}$$
  • Cette approche ne nécessite aucun stockage de matrice ni d'inversion coûteuse, seulement le calcul d'une norme scalaire.

• Mise à jour Composite Diffeomorphique :
  L'optimisation se fait dans un cadre eulérien. Le champ de déformation est mis à jour de manière composite pour garantir que le résultat reste un diféomorphisme :
  $$u^{(t+1)} = \epsilon v^{(t)} + u^{(t)} \circ (Id + \epsilon v^{(t)})$$
  où $v$ est le champ de vitesse produit par l'optimiseur LM.

• Adaptation du Paramètre d'Amortissement ($\lambda$) :
  Un paramètre d'amortissement global $\lambda$ est ajusté dynamiquement selon une approche de région de confiance (trust-region) :

  • Si la perte augmente (étape "mauvaise"), $\lambda$ est augmenté par un facteur $\mu_+$.
  • Si la perte diminue (étape "bonne"), $\lambda$ est diminué par un facteur $\mu_-$.
  • Cela permet de basculer entre une approche de type Newton (près de l'optimum) et une descente de gradient (loin de l'optimum).

• Critère de Rejet (Optionnel) :
  Une variante implémente un critère de rejet de style Metropolis-Hastings. Si la dégradation de la perte dépasse un seuil par rapport à l'amélioration précédente, l'étape est rejetée, le champ est restauré, et $\lambda$ est augmenté avant de réessayer.

3. Contributions Clés

1. Réduction de la mémoire : L'optimiseur LM factorisé élimine le besoin de buffers de moments, réduisant la consommation de mémoire GPU de 11 % à 24,6 % selon la taille du volume, permettant ainsi le traitement de volumes plus grands.
2. Performance équivalente ou supérieure : Malgré sa simplicité, l'approche LM égale ou dépasse Adam sur trois des quatre benchmarks utilisés, avec une seule configuration d'hyperparamètres transférable entre différents types d'images (IRM cérébrale, CT pulmonaire, abdomen multimodal).
3. Analyse de sensibilité critique (La "Falaise d'Amortissement") : L'étude révèle une sensibilité critique au paramètre d'augmentation $\mu_+$. Si $\mu_+ > 2.3$, l'algorithme entre dans un mode de défaillance auto-amplifié où l'amortissement devient trop fort, figeant la déformation et dégradant drastiquement les résultats.
4. Configuration Universelle : Identification d'une configuration d'hyperparamètres robuste ($\lambda_0=0.006, \mu_+=1.5, \mu_-=0.975$) qui fonctionne sans réajustement sur des modalités et des organes très différents.

4. Résultats Expérimentaux

Les expériences ont été menées sur quatre benchmarks publics :

• IRM Cérébrale (LUMIR, OASIS) : LM surpasse Adam (Dice moyen de 0.8715 vs 0.8650 pour LUMIR et 0.8293 vs 0.8028 pour OASIS). La robustesse est améliorée, avec une meilleure performance dans les pires cas.
• CT Pulmonaire (NLST) : Adam reste légèrement supérieur en moyenne (TRE 0.85 mm vs 1.68 mm), mais cette moyenne est biaisée par une seule paire aberrante. En excluant cette paire, LM est compétitif (0.95 mm vs 0.83 mm). Le rejet d'étapes permet de supprimer l'aberration.
• Abdomen Multimodal (MR vers CT) : LM égale Adam (Dice ~0.7615) sur tous les organes (foie, reins, rate).
• Efficacité Mémoire et Temps :
  • Pour un volume de $512^3$, LM économise environ 3 Go de mémoire GPU par rapport à Adam.
  • LM n'est jamais plus lent qu'Adam ; il est même 1.67 fois plus rapide sur de petits volumes (où la mise à jour des moments d'Adam domine le temps de calcul).

5. Signification et Conclusion

Cet article démontre que les méthodes du second ordre peuvent être rendues pratiques pour l'enregistrement d'images à grande échelle grâce à des approximations locales intelligentes (factorisation Rank-1).

• Impact pratique : La réduction de la mémoire ouvre la voie à l'enregistrement de volumes médicaux de très haute résolution qui étaient auparavant impossibles à traiter avec Adam.
• Insight théorique : La découverte de la "falaise d'amortissement" (damping cliff) pour le paramètre $\mu_+$ est cruciale pour la stabilité des algorithmes LM adaptatifs. Elle souligne que la robustesse de ces méthodes dépend d'un réglage fin des paramètres de mise à jour du taux d'apprentissage.
• Futur : L'approche encourage l'utilisation d'optimiseurs basés sur la région de confiance pour des tâches d'optimisation complexes en imagerie médicale, offrant un compromis optimal entre vitesse, précision et consommation de ressources.

En résumé, les auteurs ont réussi à "réhabiliter" la méthode Levenberg–Marquardt pour l'enregistrement diféomorphique, la rendant plus efficace en mémoire qu'Adam tout en maintenant, voire en améliorant, la qualité de l'enregistrement sur des données réelles.