Benchmarking Machine Learning Approaches for Polarization Mapping in Ferroelectrics Using 4D-STEM

This study benchmarks various machine learning models for automating polarization mapping in ferroelectric potassium sodium niobate using 4D-STEM data, demonstrating that while synthetic training faces a simulation-to-experiment domain gap, specific training regimes and PCA-based methods can bridge this divide while also revealing that model errors correlate with crystal defects.

The Gist

Imagine you are trying to figure out the direction of a tiny, invisible wind blowing inside a microscopic crystal. This crystal is made of a special material called a ferroelectric, which acts like a tiny magnet for electricity. Knowing which way this "electric wind" (polarization) is blowing is crucial because it tells us how the material will behave in things like sensors, memory chips, or medical devices.

The scientists in this paper used a super-powerful microscope called 4D-STEM. Think of this microscope not just as a camera, but as a flashlight that scans a sample point-by-point. At every single point, it doesn't just take a photo; it captures a complex "shadow" or diffraction pattern. These patterns are like unique fingerprints that change depending on the direction of the electric wind.

The Challenge: The "Simulation vs. Reality" Gap

The problem is that reading these fingerprints by hand is incredibly slow and difficult. So, the team tried to teach computers (using Machine Learning) to read them automatically.

Here's the catch: They didn't have enough real-world examples to teach the computer. Instead, they created a virtual training school using computer simulations. They taught the AI thousands of times on perfect, clean, digital fingerprints.

But when they sent the AI out into the real world to look at actual crystals, it struggled. It was like teaching a student to drive only on a perfect, empty video game track, and then expecting them to drive immediately in a rainy, pothole-filled city street. The "domain gap" between the clean simulation and the messy reality was too wide.

The Experiment: Testing Different "Students"

The researchers tested four different types of AI "students" to see which one could learn best:

1. ResNet & VGG: These are like students who went to a famous, general-purpose school (trained on millions of regular photos of cats, dogs, and cars) and tried to apply that knowledge to crystals.
2. Custom CNN: A student built from scratch specifically for this job.
3. PCA + k-NN: A student using a classic, old-school math approach (finding patterns by simplifying the data first).

They also tried three different ways of teaching them:

• Classification: "Is this wind blowing North, South, East, or West?"
• Regression: "What is the exact angle of the wind?"
• Prototype Representation: "Here is a 'perfect example' of a North wind. Find the one that looks most like this."

The Results: Who Passed the Test?

• The Video Game Students (ResNet/VGG): They were great at the simulation (the video game track) but failed miserably in the real world. They couldn't adapt their "cat and dog" knowledge to the strange crystal patterns.
• The Custom Student: Did okay, but still struggled with the messy real-world data.
• The Old-School Math Student (PCA): Surprisingly, this simple, non-deep-learning approach turned out to be the most robust. It was like a seasoned mechanic who didn't need a fancy computer to fix the car; they just knew the basics of how the engine worked.
• The Secret Sauce (Filtering & Augmentation): The researchers realized that the simulation data was too perfect. They started "messing up" the training data on purpose (adding digital noise, blurring it slightly) and throwing away the "boring" examples where the wind was too weak to see. This forced the AI to learn the real features, not just the perfect ones. When they did this, the Prototype Representation method (the one that learns "perfect examples") became the best at bridging the gap between simulation and reality.

The Bonus Discovery: Finding "Scars" in the Crystal

The team also asked a clever question: Can these AI models find defects in the crystal, like a missing atom or a crack?

They simulated a crystal with a tiny "scar" (an oxygen vacancy).

• The Custom AI (the deep learning one) smoothed over the scar because it was trained to ignore noise.
• The Old-School Math Student (PCA) actually spotted the scar! It noticed that the "wind" pattern changed strangely right where the defect was.

It turns out that when the AI gets confused or makes a mistake in a specific pattern, that confusion often points directly to a physical defect in the crystal. It's like a doctor noticing that a patient's heartbeat is slightly irregular in a specific spot, which reveals a hidden heart issue.

The Bottom Line

This paper is a roadmap for teaching computers to understand the invisible forces inside materials.

1. Don't just trust the fancy AI: Sometimes simple, old-school math works better when the data is messy.
2. Training matters: You can't just feed an AI perfect data; you have to teach it with "noisy" data that looks like the real world.
3. Mistakes are useful: When the AI gets confused, it might actually be pointing out a hidden flaw in the material, which is a huge bonus for finding defects.

While the AI isn't perfect yet (it still struggles with very thick crystals), this study gives scientists a much better toolkit for mapping the invisible electric winds inside the materials that power our future technology.

Technical Summary

1. Problem Statement

Context: Four-dimensional scanning transmission electron microscopy (4D-STEM) provides atomic-scale insights into material structures. A critical application is mapping the spontaneous polarization direction in ferroelectric materials, such as Potassium Sodium Niobate (KNN), which is essential for understanding their functional properties.
Challenge: Traditional analysis relies on manual inspection or rigid algorithms (e.g., Center of Mass calculations) that are time-consuming and struggle with noise, subtle shifts, or anomalous features. While Machine Learning (ML) offers automation, there is a significant domain gap between idealized synthetic training data and noisy experimental 4D-STEM data. Furthermore, it is unclear which ML architectures and training paradigms best bridge this gap for polarization classification.

2. Methodology

Data Generation and Preprocessing

• Synthetic Data: The authors generated 4D-STEM diffraction patterns using QSTEM simulations for a lead-free KNN crystal (orthorhombic symmetry).
  • Model: 8 distinct polarization directions were simulated based on Nb atom displacements.
  • Configurations: Datasets included varying thicknesses (20 nm and 50 nm) and lateral sizes (2×2, 4×4, 6×6 unit cells) to test robustness.
  • Defect Simulation: A synthetic defect dataset was created by introducing an oxygen vacancy and relaxing the structure using DFT (Quantum ESPRESSO) to test anomaly detection.
• Experimental Data: Real 4D-STEM data was acquired from a 20 nm KNN crystal using a Cs-corrected ARM200CF microscope with a Merlin pixelated detector. Ground truth polarization labels were derived manually from HAADF image analysis.
• Preprocessing:
  • Cropping: Images were cropped to 64×64 pixels centered on the diffraction disk using a global Center of Mass (CoM) correction to align the optical axis.
  • Filtering: To prevent the model from learning on ambiguous, near-centrosymmetric patterns, the training set was filtered to include only samples with above-average electron beam shift magnitudes.
  • Augmentation: Training data underwent randomized augmentation (Adaptive Gaussian Smoothing, Brightness/Contrast adjustment, and Noise Injection) to simulate experimental variations.

Machine Learning Models and Paradigms

The study benchmarked four model architectures across three training paradigms:

1. Architectures:
   • ResNet and VGG: Pretrained on ImageNet (modified for 1-channel input).
   • Custom CNN: A shallow, scratch-trained network (3 convolutional layers) designed to prevent overfitting.
   • PCA + k-NN: A traditional pipeline using Principal Component Analysis (retaining 20 components) followed by a k-Nearest Neighbors classifier ($k=5$).
2. Training Paradigms:
   • Classification: Standard cross-entropy loss to predict one of 8 discrete polarization classes.
   • Regression: Predicting a 2D unit vector (polarization angle) using MSE loss, leveraging the geometric continuity of polarization directions.
   • Prototype Representation: A novel approach learning an embedding space where class prototypes are defined. It uses a circular distance metric and Kullback-Leibler (KL) divergence loss to account for the cyclic nature of polarization angles (e.g., 0° is close to 360°).

Evaluation Strategy

Models were tested in three phases:

1. Synthetic-to-Synthetic: Testing on unseen synthetic geometries (different thicknesses/sizes).
2. Synthetic-to-Experimental: Testing on real experimental data.
3. Anomaly Detection: Evaluating if models trained on pristine structures could identify regions with simulated defects via misclassification patterns.

3. Key Contributions

• Systematic Benchmark: A comprehensive comparison of deep learning (ResNet, VGG, Custom CNN) and traditional ML (PCA+k-NN) for 4D-STEM polarization mapping.
• Paradigm Evaluation: Explicit comparison of classification, regression, and prototype-based learning objectives, highlighting the importance of geometric constraints in polarization data.
• Domain Gap Analysis: Demonstrated that while standard deep learning models excel on synthetic data, they fail on experimental data without specific interventions. The study identifies prototype representation learning combined with data filtering and augmentation as the most effective strategy to bridge the synthetic-experimental gap.
• Defect Detection Insight: Showed that misclassification patterns in supervised models can correlate with structural defects, though preprocessing trade-offs exist.

4. Results

Synthetic-to-Synthetic Performance

• Models achieved near-perfect accuracy (>99%) on synthetic test sets identical to training data.
• Thickness Limitation: All models failed significantly on 50 nm thick samples. The increased dynamical scattering in thicker samples alters image contrast non-linearly, breaking the assumptions of the thin-sample simulations used for training.
• Lateral Size: Models performed well on larger lateral dimensions (4×4, 6×6), indicating scalability for larger fields of view.

Synthetic-to-Experimental Performance

• Baseline Failure: Without augmentation or filtering, all models failed on experimental data (accuracy <20%), confirming a massive domain gap.
• Regression Failure: Regression approaches (predicting angles) tended to collapse, predicting only two classes (RU and RM) regardless of the input, failing to learn the multi-class distribution.
• Pretrained Models: ImageNet-pretrained models (ResNet/VGG) did not transfer well, often producing trivial predictions or wrong majority classes.
• Best Performers:
  • PCA + k-NN: Showed surprising robustness, correctly predicting the majority polarization class for all experimental datasets, though with high variance.
  • Custom CNN (Prototype): When combined with filtering (removing low-signal samples) and augmentation, this approach correctly predicted the majority class for all experimental datasets. It was the only neural approach to partially bridge the domain gap.

Error Analysis & Anomaly Detection

• Misclassification Patterns: Errors were not random; they clustered in inter-columnar regions where the electron beam shift magnitude was low (indicating low polarization information).
• Defect Detection:
  • PCA successfully identified the location of a simulated oxygen vacancy by detecting polarization changes caused by the defect.
  • Conv (Proto) failed to visualize the defect in the misclassification map.
  • Trade-off: While preprocessing (augmentation/filtering) improved polarization classification accuracy, it obscured the defect signal, suggesting a conflict between robust classification and sensitive defect detection.

5. Significance and Conclusion

This study establishes that no single "off-the-shelf" deep learning model is currently sufficient for robust polarization mapping in 4D-STEM due to the synthetic-experimental domain gap.

• Key Takeaway: Specialized approaches, specifically prototype representation learning combined with rigorous data filtering (selecting high-signal training examples) and augmentation, are superior to standard classification or regression for this task.
• Limitation: The current methodology is limited to thin samples (<20 nm). Thicker samples require new physical models or training data that account for dynamical scattering.
• Future Outlook: The work suggests that while ML can automate polarization mapping, the "black box" nature of deep learning may obscure subtle structural defects. Hybrid approaches or specific feature engineering (like PCA) may be necessary for simultaneous polarization mapping and defect detection.

The authors provide open access to their code and datasets to facilitate further research in automated electron microscopy analysis.