Contrast transfer functions help quantify neural network out-of-distribution generalization in HRTEM

TL;DR

This work tackles out-of-distribution generalization in high-resolution TEM semantic segmentation by leveraging a data-centric, information-theoretic framework. It introduces two information-transfer metrics, $\epsilon(\chi)$ and $\sigma(\chi, \chi')$, derived from the TEM contrast transfer function $T(\boldsymbol{q})$, to quantify how training and test imaging conditions differ and overlap. Through large-scale synthetic data experiments with thousands of networks trained on multislice TEM simulations, the authors show that OOD performance degrades smoothly under imaging-condition shifts and can be predicted by information-transfer relationships, guiding training-data design. The study provides a principled approach for deploying TEM-based ML workflows while highlighting limitations related to atomic-structure distribution shifts and suggesting avenues for broader applicability and domain-agnostic OOD assessment.

Abstract

Neural networks, while effective for tackling many challenging scientific tasks, are not known to perform well out-of-distribution (OOD), i.e., within domains which differ from their training data. Understanding neural network OOD generalization is paramount to their successful deployment in experimental workflows, especially when ground-truth knowledge about the experiment is hard to establish or experimental conditions significantly vary. With inherent access to ground-truth information and fine-grained control of underlying distributions, simulation-based data curation facilitates precise investigation of OOD generalization behavior. Here, we probe generalization with respect to imaging conditions of neural network segmentation models for high-resolution transmission electron microscopy (HRTEM) imaging of nanoparticles, training and measuring the OOD generalization of over 12,000 neural networks using synthetic data generated via random structure sampling and multislice simulation. Using the HRTEM contrast transfer function, we further develop a framework to compare information content of HRTEM datasets and quantify OOD domain shifts. We demonstrate that neural network segmentation models enjoy significant performance stability, but will smoothly and predictably worsen as imaging conditions shift from the training distribution. Lastly, we consider limitations of our approach in explaining other OOD shifts, such as of the atomic structures, and discuss complementary techniques for understanding generalization in such settings.

Figures (5)

• Figure 1: (a) Example visualization of $\epsilon(\chi)$ for aberrations consisting solely of defocus and spherical aberrations, with magnitudes up to 30nm and 0.1mm, respectively, at 300kV and with a chromatic aberration envelope corresponding to a focal spread of 10Å. Bright bands in regions where defocus and spherical aberration have opposite sign correspond to the formation of passbands in the contrast transfer function. (b) A selection of phase contrast transfer functions corresponding to maxima and minima in (a) are visualized for a lens with a spherical aberration coefficient of 25$\mu$m.
• Figure 2: Comparative visualization of (a) $\sigma(\chi, \chi')$ and (b) $\Delta \epsilon(\chi, \chi')$ for pairs of imaging conditions differing only in defocus. (c) Selected pairs of absolute contrast transfer functions (CTFs) corresponding to defocus conditions indicated by circle, square, and triangle markers. (d) Scatter plot demonstration of how these transfer functions relate to each other via the metrics $\sigma(\chi, \chi')$ and $\Delta \epsilon(\chi, \chi')$. (e) Example image pairs of CdSe nanoparticles on amorphous carbon matching selected CTF pairs. Imaging conditions correspond to a TEM operated at 300kV with a focal spread of 10Å. Scalebar is 1nm.
• Figure 3: Out-of-distribution generalization behavior throughout training of U-Net image segmentation models trained on -10nm (left column), 0nm (middle column) and +10nm (right column) defocus images of CdSe nanoparticles. Defocus is measured from the minimum contrast condition of the CdSe nanoparticles. (a) Example regions from a simulated micrograph at each corresponding defocus for visual reference; scale bar is 1nm. (b--e) Learning curves for series of models trained with learning rates of 2.3e-3, 2.3e-3, 1e-2, 1e-2 and alternating batch sizes of 16, 32, 16, 32; line color corresponds to nominal defocus of validation dataset. (f) Correlation between in- and out-of-distribution losses for models trained with learning rates of 1e-2 and batch size of 16 (d), aggregated across all training times.
• Figure 4: Confusion matrix visualization of neural network loss as trained and evaluated on CdSe image datasets simulated with mean defoci between -25 and +25nm at 300kV and with a defocus spread of 10 Å. From left to right, we visualize model performance at epochs 8, 32, and 128. Visualized loss is taken as mean of twenty models trained with a learning rate of 1e-2 and batch size of 24. Smaller loss is better.
• Figure 5: Neural network loss (a) and induced (excess) loss over the training dataset (b) out of distribution, relative to $\sigma(\chi, \chi')$ and $\Delta \epsilon(\chi, \chi')$. Relative information content increases vertically, and information overlap decreases to the right. More negative induced loss (blue) indicates a model improving OOD, whereas a more positive induced loss (red) indicates the model worsening. OOD evaluations were performed using randomly generated model training datasets and evaluation datasets. Aberrations of training and evaluation datasets were sampled up to 5th order (defocus, axial coma, two- and three-fold astigmatism, and spherical), and models were trained until reaching a loss threshold of 0.05. Each marker corresponds to a unique training dataset--OOD dataset aberration pair, in total encompassing 3089 unique models and 72864 unique dataset pairs.