Machine Learning-Augmented Acceleration of Iterative Ptychographic Reconstruction

Résumé technique : Accélération de la reconstruction itérative ptychographique par apprentissage automatique

Énoncé du problème
Les installations modernes de synchrotron et de microscopie électronique connaissent une augmentation dramatique des taux d'acquisition de données en raison des sources de quatrième génération et des détecteurs fonctionnant à des fréquences de cadence de kilohertz. Bien que la ptychographie soit une technique puissante d'imagerie computationnelle permettant de récupérer des informations complexes à haute résolution sur un objet à partir de mesures de diffraction d'intensité uniquement, les algorithmes de reconstruction itérative traditionnels (par exemple, ePIE, carte de différence, vraisemblance maximale) se heurtent à un goulot d'étranglement computationnel critique. Ces méthodes nécessitent souvent des centaines d'itérations et plusieurs minutes à plusieurs heures de temps de traitement par jeu de données, échouant à fournir le retour d'information quasi temps réel nécessaire pour guider la prise de décision expérimentale dans des environnements à haut débit.

Les solutions existantes en apprentissage automatique (ML) tentent souvent d'accélérer la reconstruction en la traitant comme un problème de régression en un seul tir, en mappant directement les motifs de diffraction vers les objets. Cependant, ces approches remplacent généralement l'optimisation basée sur la physique par des prédicteurs boîte noire, sacrifiant les garanties de minimisation basée sur le gradient concernant la cohérence physique, l'interprétabilité et la robustesse face aux imperfections expérimentales (par exemple, erreurs de position de balayage, cohérence partielle). De plus, de nombreuses telles méthodes reposent sur des données d'entraînement synthétiques et des configurations de sonde fixes, limitant leur robustesse dans les lignes de faisceau opérationnelles où les conditions varient.

Méthodologie
Les auteurs proposent un cadre hybride qui complète, plutôt que de remplacer, les solveurs ptychographiques itératifs établis. L'innovation centrale est un opérateur d'avance rapide appris intégré dans la boucle d'optimisation.

• Flux de travail du cadre :

  1. Échauffement : Des mises à jour itératives standard basées sur la physique (en utilisant le cadre CDTools) sont effectuées pour un petit nombre d'itérations initiales (par exemple, 5) afin d'établir une estimation de référence.
  2. Étape d'avance rapide : Un modèle ML entraîné (une architecture U-Net) prend l'estimation intermédiaire de l'objet et prédit un état amélioré plus proche de la convergence. Cette étape est effectuée une seule fois par reconstruction.
  3. Reprise : L'objet prédit est utilisé pour réinitialiser le solveur, et les itérations conventionnelles basées sur le gradient reprennent pour conduire la solution vers la cohérence physique avec les données mesurées.

• Architecture du modèle : L'opérateur est implémenté comme un U-Net avec des convolutions 3×3, un max-pooling et des connexions de saut. Il traite des champs d'objets à valeurs complexes représentés comme deux canaux à valeurs réelles (partie réelle et partie imaginaire). L'architecture est entièrement convolutive, lui permettant de gérer des résolutions spatiales arbitraires sans réentraînement.

• Stratégie d'entraînement : Le modèle est entraîné de manière supervisée en utilisant des jeux de données ptychographiques expérimentaux historiques. Des paires sont construites en prenant l'estimation de l'objet à une itération précoce (par exemple, itération 5) comme entrée et l'estimation à une itération ultérieure, bien convergée (par exemple, itération 100) comme cible. Les données d'entraînement sont constituées d'échantillons diversifiés acquis sur plusieurs années et campagnes expérimentales. Crucialement, le modèle est entraîné et évalué exclusivement sur des données expérimentales, et non sur des simulations synthétiques.

• Intégration : La méthode est intégrée dans le pipeline de reconstruction CDTools, qui utilise la différenciation automatique et l'optimiseur Adam. L'inférence ML ajoute une surcharge computationnelle minimale car elle ne se produit qu'une seule fois.

Résultats clés
La méthode a été évaluée sur des données expérimentales acquises à la Advanced Light Source (ALS), incluant des échantillons de test (catalyseur au nickel et alumine poreuse) collectés une année différente de celle des données d'entraînement pour évaluer la généralisation temporelle.

• Qualité de reconstruction : Les reconstructions augmentées par ML sont visuellement cohérentes avec les références standard, préservant les détails structurels fins à la fois dans l'amplitude et la phase. La méthode récupère également avec succès plusieurs modes de sonde, une capacité souvent absente dans les approches purement basées sur le ML.
• Vitesse de convergence : La vraisemblance négative logarithmique de Poisson (NLL), une métrique physiquement appropriée pour les statistiques de comptage de photons, démontre que la méthode augmentée par ML converge significativement plus rapidement.
  • La méthode réalise une réduction du temps réel de 2,1× à 2,3× par rapport aux solveurs standards pour atteindre le même seuil de convergence.
  • Dans un déploiement sur une ligne de faisceau de production pour la ptycho-tomographie (74 projections), l'approche a réduit le temps par projection d'environ 153 secondes à environ 41 secondes (une réduction d'environ 4× du temps total de résolution pour le jeu de données), permettant un retour d'information à faible latence.
• Robustesse : L'approche montre une robustesse face aux variations des hyperparamètres, y compris l'itération à laquelle l'opérateur ML est inséré ($i_{ML}$), le taux d'apprentissage et la taille du lot. Une augmentation transitoire de la NLL de Poisson est observée immédiatement après l'étape ML, reflétant un mouvement non local dans le paysage de perte qui aide à échapper aux minima sous-optimaux, suivi d'une convergence rapide.
• Évolutivité : La méthode s'étend efficacement sur plusieurs GPU (1 à 4), avec une efficacité parallèle comparable au solveur standard et une surcharge de temps d'exécution modeste (1,9 %–9,4 %).

Signification et revendications
L'article revendique que cette approche hybride comble avec succès le fossé entre la vitesse des méthodes pilotées par les données et le rigueur physique de l'optimisation itérative. En préservant le cadre de minimisation basé sur le gradient, la méthode garantit que la reconstruction finale réside à un minimum local du terme de fidélité des données expérimentales, garantissant la cohérence physique avec les données mesurées.

Les contributions clés mises en avant par les auteurs incluent :

1. Agnosticisme du solveur : L'opérateur d'avance rapide agit sur l'objet reconstruit, le rendant compatible avec diverses tailles de détecteurs, échantillonnages de pixels et solveurs itératifs sous-jacents.
2. Généralisation expérimentale : Le modèle démontre une robustesse et une généralisation temporelle lorsqu'il est entraîné sur des données historiques et testé sur des données d'une année différente, sans dépendre de jeux de données synthétiques.
3. Viabilité opérationnelle : La méthode a été déployée avec succès en production sur une ligne de faisceau de synchrotron opérationnelle, démontrant une utilité pratique pour le pilotage expérimental en temps réel dans des scénarios à haut débit.

Les auteurs concluent que cette approche offre une voie prometteuse pour accélérer l'imagerie computationnelle dans des environnements expérimentaux réalistes tout en maintenant l'interprétabilité et les garanties physiques essentielles à la découverte scientifique.