4D-PreNet: Unified 4D-STEM Preprocessing

• The paper introduces 4D-PreNet, a unified framework that replaces multiple manual corrections with a deep-learning pipeline for denoising, center calibration, and distortion correction in 4D-STEM data.
• It utilizes a multi-stage architecture combining an attention-enhanced U-Net for noise suppression and beam-center localization with a ResNet-50 for estimating elliptical distortion parameters.
• Experimental validation demonstrates significant reductions in mean error and improved sub-pixel accuracy, facilitating real-time, high-throughput analysis and robust downstream quantitative analysis.

Searching arXiv for the exact 4D-PreNet paper and closely related 4D-STEM preprocessing work. 4D-PreNet is a unified preprocessing framework for 4D scanning transmission electron microscopy (4D-STEM) data analysis that addresses three acquisition artifacts within one deep-learning pipeline: diffraction noise, beam-center drift, and elliptical distortion (Liu et al., 5 Aug 2025). It is designed for high-throughput and real-time 4D-STEM workflows, where conventional correction algorithms are often material-specific and where preprocessing has become a critical bottleneck for automated experimentation. The framework integrates attention-enhanced U-Net and ResNet architectures to perform denoising, center correction, and elliptical distortion calibration, and it is trained on large simulated datasets spanning varied noise levels, drift magnitudes, and distortion types to support generalization to experimental data acquired under varying conditions (Liu et al., 5 Aug 2025).

1. Problem setting and motivation

4D-STEM couples high-throughput diffraction acquisition with downstream quantitative analysis, but raw diffraction patterns are often corrupted by Poisson and Gaussian noise, beam-center drift, and elliptical distortions (Liu et al., 5 Aug 2025). The paper emphasizes that these are not merely cosmetic defects. Beam-center errors can bias strain and orientation estimates, and uncorrected ellipticity can introduce systematic errors into diffraction analysis. A central motivation of 4D-PreNet is therefore to replace a sequence of manual or dataset-specific corrections with a single automated workflow suitable for real-time or closed-loop use (Liu et al., 5 Aug 2025).

The framework is positioned against two classes of prior practice. One consists of classical correction procedures that are manual and material-specific. The other consists of deep-learning approaches that address only one preprocessing task at a time (Liu et al., 5 Aug 2025). In this setting, 4D-PreNet is presented as a unified solution at the pipeline level, with dedicated modules for each subproblem.

A plausible implication is that the method is intended not only to improve image quality, but also to stabilize any downstream analysis that depends on reciprocal-space geometry, including strain mapping, symmetry analysis, and phase classification. The paper directly connects preprocessing quality to those later tasks, especially through its discussion of sub-pixel center localization (Liu et al., 5 Aug 2025).

2. Pipeline architecture and task decomposition

The framework is organized as a three-stage pipeline (Liu et al., 5 Aug 2025). First, a denoising module removes diffraction noise. Second, a center calibration module localizes the beam center and aligns each diffraction frame. Third, an ellipse calibration module estimates and corrects elliptical distortion parameters.

Stage	Architecture	Function
Denoising module	U-Net with CBAM attention	Suppress noise while preserving fine diffraction structure
Center calibration module	U-Net	Heatmap regression for beam-center localization
Ellipse calibration module	ResNet-50 regression model	Predict geometric parameters of elliptical distortion

The architectural choices are task-specific. U-Net is used for dense image-to-image tasks where local structure is important, while ResNet-50 is used for global geometric regression where the model must infer overall distortion parameters from pattern-wide context (Liu et al., 5 Aug 2025). This decomposition is explicitly described as specialization rather than a fully shared multi-task backbone.

The practical inference sequence is also specified. Raw 4D-STEM data are first passed to the denoising network, then to the center heatmap network. Weighted center-of-mass is applied to the predicted heatmap to recover sub-pixel center coordinates, followed by affine realignment to the geometric image center. The averaged diffraction pattern is then passed to the ResNet-50 ellipse regression model, whose predicted parameters drive the final affine ellipse correction (Liu et al., 5 Aug 2025).

This staged design suggests a pipeline in which early corrections improve the operating conditions of later modules. The paper does not present a single unified multi-task loss in the provided material, but it does describe a fully specified sequence of supervised subtasks (Liu et al., 5 Aug 2025).

3. Denoising and beam-center calibration

The denoising module is a U-Net augmented with a Convolutional Block Attention Module (CBAM) (Liu et al., 5 Aug 2025). CBAM is included to improve spatial feature representation and to help the network focus on informative diffraction structures while suppressing noise. The output layer uses a sigmoid activation so that predicted intensities remain non-negative, consistent with physical diffraction measurements (Liu et al., 5 Aug 2025).

The module is trained to remove both Poisson noise and Gaussian noise (Liu et al., 5 Aug 2025). The former is associated with electron counting, while the latter models additive detector or background noise. The denoiser is intended to preserve high-frequency diffraction features such as disks and ring structure, since these are required for subsequent center localization and ellipse calibration. The paper qualitatively states that denoised outputs “closely match the clean ground truth,” while residual maps indicate suppression of both uniformly distributed Gaussian noise and signal-dependent Poisson noise (Liu et al., 5 Aug 2025).

Beam-center calibration is implemented by a second U-Net, but its output is not a reconstructed image. Instead, the network predicts a 2D heatmap representing the beam-center probability distribution (Liu et al., 5 Aug 2025). A weighted center-of-mass calculation is then used to extract sub-pixel coordinates from the heatmap. The resulting center is used to apply an affine transformation that aligns the diffraction frame to a fixed image center, for example $(127.5, 127.5)$ in a $256 \times 256$ image (Liu et al., 5 Aug 2025).

A notable architectural decision is that CBAM is excluded from the center calibration module. The reason given is to keep inference efficient and avoid overfitting; the authors report that the global context required for center localization can be captured without attention in this branch (Liu et al., 5 Aug 2025). This makes the two U-Net stages structurally similar but not identical.

The emphasis on sub-pixel localization is central. The paper states that even tiny center errors propagate into strain and orientation measurement errors, and it treats frame-by-frame center correction as necessary for real-time workflows that cannot rely on a global average center (Liu et al., 5 Aug 2025).

4. Elliptical distortion calibration

Elliptical distortion correction is formulated as a global geometric regression problem and is handled by a ResNet-50 model (Liu et al., 5 Aug 2025). Rather than predicting a pixelwise correction field, the network estimates three parameters describing the affine correction: $\cos(\theta)$, $\sin(\theta)$, and a scale factor representing the aspect ratio between principal axes (Liu et al., 5 Aug 2025).

The use of $\cos(\theta)$ and $\sin(\theta)$ instead of directly regressing $\theta$ is explicitly motivated by optimization considerations. According to the paper, this avoids discontinuities at periodic wrap-around boundaries and improves optimization stability; the authors note that the same parameterization is common in pose or orientation estimation because it is continuous and differentiable (Liu et al., 5 Aug 2025).

These predicted parameters are then used to correct elliptical distortion and restore circular symmetry in the diffraction pattern, including amorphous patterns where diffuse rings are broad and less visually obvious to calibrate (Liu et al., 5 Aug 2025). The evaluation of this module is primarily qualitative and geometric rather than summarized by a single scalar score. The paper reports restoration of circular symmetry in both crystal and amorphous patterns, narrowing of radial intensity peaks in crystalline data, and reduced radial fluctuations measured by radial standard deviation in polar coordinates (Liu et al., 5 Aug 2025).

The paper further notes that uncorrected patterns exhibit sinusoidal distortions in polar coordinates, whereas corrected patterns become straightened and symmetric (Liu et al., 5 Aug 2025). This is used to argue that the learned model captures global geometric distortion rather than merely smoothing local appearance.

A plausible implication is that the ellipse module depends on sufficiently informative global structure in averaged diffraction patterns. The paper makes this explicit in its limitations, noting that the method works best when averaged diffraction patterns contain clean, discernible rings (Liu et al., 5 Aug 2025).

5. Simulated training data and optimization protocol

A major component of 4D-PreNet is its simulated training corpus (Liu et al., 5 Aug 2025). The paper states that 1,024 structures were selected from public CIF databases such as the Materials Project, and that both crystalline and amorphous materials were included. Diffraction patterns were simulated using abTEM based on the multislice STEM simulation algorithm (Liu et al., 5 Aug 2025).

Two separate datasets were constructed because the learning objectives differ (Liu et al., 5 Aug 2025). The first dataset supports denoising and center detection. For each structure, crystal orientations were randomized by rotation, convergence semi-angle and other experimental parameters were randomly varied, artificial beam-center shifts were introduced via random x/y translations, elliptical distortions were simulated by random rotations and scaling, and noise was injected at the input stage only, while the noiseless image served as the denoising target (Liu et al., 5 Aug 2025).

The second dataset supports ellipse calibration and is described as simulating more realistic stacked or averaged diffraction scenarios (Liu et al., 5 Aug 2025). Each CIF was simulated using a fixed experimental parameter set, with two parameter sets per CIF. Each set generated 128 patterns, which were augmented by 30 random rotations to yield 3840 patterns per CIF per parameter set. Then 3000 patterns were randomly selected and combined with mixed noise and random affine transformations, producing 100 composite diffraction patterns for each CIF-parameter combination. Across all CIFs, this yielded roughly 102,400 simulated patterns (Liu et al., 5 Aug 2025). The associated labels are $\cos(\theta)$, $\sin(\theta)$, and the scale or aspect factor.

All datasets were partitioned into training, validation, and test sets (Liu et al., 5 Aug 2025). Training used Adam optimizer, learning-rate scheduling, and early stopping based on validation loss (Liu et al., 5 Aug 2025). The paper does not give a single explicit multi-task loss in the provided material, but it specifies the supervision regime for each task: denoising is an image reconstruction problem evaluated with MSE and PSNR; center calibration is trained through heatmap regression; ellipse calibration is a three-parameter regression problem (Liu et al., 5 Aug 2025).

This simulated-data strategy is crucial to the framework’s claim of generalization. The paper interprets successful transfer to experiments as evidence that the simulated distortions were sufficiently realistic and diverse (Liu et al., 5 Aug 2025).

6. Quantitative performance and experimental generalization

The denoising module improves both MSE and PSNR across noise levels and across both crystal and amorphous test data (Liu et al., 5 Aug 2025). The paper reports an MSE reduction of more than 50% on average and PSNR improvement of approximately 5–10 dB, including on test data not seen during training (Liu et al., 5 Aug 2025).

For beam-center localization, the framework is benchmarked against conventional center-of-mass (CoM), Friedel symmetry, and a DPC segmentation network on a test set of 4,096 simulated patterns (Liu et al., 5 Aug 2025). The reported results are:

Metric	Value
MAE-R	0.034
MAE-X	0.022
MAE-Y	0.021

The paper states that the average error is reduced by more than 60% relative to the second-best method, Friedel symmetry, and by more than 90% relative to the DPC segmentation network (Liu et al., 5 Aug 2025). It also highlights that localization is sub-pixel, with average errors below 0.04 pixels (Liu et al., 5 Aug 2025).

For ellipse calibration, the evidence is mainly qualitative and based on radial statistics rather than a single headline scalar (Liu et al., 5 Aug 2025). The reported outcomes include restoration of circular symmetry, narrowing of crystalline radial peaks, and substantial reduction in radial standard deviation after calibration (Liu et al., 5 Aug 2025).

A key claim is that the framework generalizes from synthetic training data to real experimental 4D-STEM data without manual parameter tuning (Liu et al., 5 Aug 2025). The experimental data are from a silicon sample containing crystal, amorphous, and mixed-phase regions. According to the paper, denoising suppresses both Poisson and Gaussian noise while preserving structure, center calibration aligns patterns to the image center with sub-pixel precision, and ellipse correction restores radial symmetry in averaged and polar-transformed patterns (Liu et al., 5 Aug 2025).

The runtime reported for the full pipeline is about 7 minutes for a full-size $256 \times 256 \times 256 \times 256$ data cube on an NVIDIA RTX 4090 with batch size 128 (Liu et al., 5 Aug 2025). This is presented as evidence that the method is compatible with high-throughput workflows.

7. Significance, limitations, and relation to broader 4D methods

The practical significance of 4D-PreNet lies in moving 4D-STEM preprocessing from a sequence of manual, dataset-specific operations into a single automated pipeline (Liu et al., 5 Aug 2025). The paper explicitly links this to high-throughput screening, automated acquisition, closed-loop experiments, and real-time beam-sensitive studies where manual correction is infeasible. It also states that the framework improves consistency for downstream strain mapping, symmetry analysis, and phase classification (Liu et al., 5 Aug 2025).

The paper is also explicit about limitations (Liu et al., 5 Aug 2025). First, ellipse calibration assumes well-defined diffraction rings and may struggle on structurally complex datasets with overlapping or irregular features. Second, performance may degrade under extreme noise if the signal is too corrupted for reliable denoising. Third, generalization depends on the realism of the synthetic simulation pipeline; unmodeled experimental artifacts could reduce robustness. Fourth, although the framework is unified at the pipeline level, the subtasks are handled by dedicated networks rather than a jointly optimized shared representation (Liu et al., 5 Aug 2025).

The name “4D-PreNet” may invite confusion with unrelated uses of “4D” in other areas, such as 4D human reconstruction from point-cloud sequences or 4D point-cloud video pretraining (Tang et al., 2021, Liu et al., 1 Dec 2025). In the present case, however, the term refers specifically to 4D-STEM preprocessing rather than dynamic scene reconstruction or robotic perception. This distinction is essential because the input modality, supervision, objective, and evaluation criteria are entirely different.

Taken together, the method is best understood as a microscopy-specific preprocessing system rather than a generic 4D representation learner. Its core contribution is the integration of CBAM-enhanced U-Net denoising, U-Net heatmap-based sub-pixel center localization, and ResNet-50 ellipse-parameter regression into a single workflow that reduces denoising error, achieves sub-pixel center localization, restores diffraction symmetry, and operates quickly enough for high-throughput 4D-STEM analysis (Liu et al., 5 Aug 2025).