AtomDiffuser: Time-Aware STEM Model

• AtomDiffuser is a degradation modeling framework that separates geometric drift from radiometric decay in sequential STEM and cryo-STEM images.
• It uses a dual-stream U-Net with temporal conditioning to predict an affine drift transform and a spatially varying attenuation map for each frame.
• Experimental results show superior performance over CNN baselines, offering interpretable, high-resolution maps for both drift correction and damage assessment.

AtomDiffuser is a time-aware degradation modeling framework for sequential atomic-resolution STEM and cryo-STEM images that separates two effects that are usually entangled in time-resolved electron microscopy: sample drift and beam-induced damage. In the reported setting, frames are acquired about $1.6$–$2.0$ seconds apart, so atomic-scale structures can both move geometrically because of stage or specimen instability and fade or disappear radiometrically because of cumulative electron-beam exposure. AtomDiffuser models the transition between two frames as a physically motivated degradation process composed of a global affine transformation for drift and a spatially varying attenuation map for beam damage, rather than as generic denoising, restoration, or registration (Wang et al., 14 Aug 2025).

1. Imaging context and problem definition

Scanning transmission electron microscopy is used to observe atomic columns and lattice structures as they evolve over time under external conditions. In sparse time-lapse STEM and cryo-STEM, however, two degradation sources corrupt the sequence simultaneously. Sample drift arises from mechanical instability, thermal expansion, or charging, and changes geometry but ideally not signal strength. Beam-induced damage or signal attenuation arises from electron irradiation and appears as intensity loss, blurring, or disappearance of atomic features. In beam-sensitive materials such as LLZO, distinguishing these mechanisms is essential because a disappearing or displaced feature may reflect either true structural instability or microscope and sample motion (Wang et al., 14 Aug 2025).

The reported difficulty is not merely low image quality. Once beam damage has attenuated or erased a feature, registration becomes ill-posed because the feature needed for matching may no longer exist. Conversely, if frames are misaligned, intensity changes caused by motion can be mistaken for damage. The paper therefore treats sparse STEM frame transitions as a structured physical degradation process. This is especially relevant in the regime described in the paper, where motion may exceed tens or hundreds of pixels and there may be no stable reference frame or consistent pixel correspondence (Wang et al., 14 Aug 2025).

The paper explicitly argues that several existing method classes are insufficient for this setting. Classical drift-correction methods such as RevSTEM, orthogonal scan correction, or scan-path redesign focus on scan distortion or hardware correction rather than explicit modeling of inter-frame atomic evolution. Single-frame restoration methods improve image quality but do not separate geometric motion from radiometric decay across time. Registration-only methods assume a stable scene and usable correspondence structure. Optical flow and generic motion estimation are likewise not physically tailored to sparse STEM sequences with large accumulated drift and non-conservative intensity changes. This positioning is central to AtomDiffuser’s identity: it is designed to decompose frame-to-frame change into interpretable physical factors rather than to collapse all differences into a generic image-to-image mapping (Wang et al., 14 Aug 2025).

2. Degradation model and conceptual framing

AtomDiffuser is presented as a diffusion-variant or diffusion-inspired degradation model, but it is not a standard denoising diffusion probabilistic model. The paper contrasts the stochastic DDPM forward process,

$$x_t = \sqrt{\bar{\alpha}_t}\, x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon, \qquad \epsilon \sim \mathcal{N}(0,\mathbf{I}),$$

with a Cold-Diffusion-style deterministic degradation,

$x_t = \mathcal{D}(x_0,t).$

AtomDiffuser follows the second viewpoint and defines an interpretable deterministic degradation operator for STEM evolution (Wang et al., 14 Aug 2025).

The core formulation writes the STEM frame at time $t$ as

$$x_t = \mathcal{D}(x_0,t) = \Phi(\lambda_t \odot x_0;\,\mathcal{T}_t),$$

where $x_0$ is the reference image, $\lambda_t \in [0,1]^{H\times W}$ is a pixel-wise attenuation or survival map, $\mathcal{T}_t \in \mathbb{R}^{2\times 3}$ is an affine transform representing drift, and $\Phi(\cdot;\mathcal{T}_t)$ is a spatial warping operator. Operationally, the image is first attenuated spatially to account for beam-induced signal loss and then warped globally to model drift. The paper emphasizes this decay-before-drift ordering as physically motivated: damage accumulates during exposure, and slow stage or sample drift affects the observed frame afterward (Wang et al., 14 Aug 2025).

The model takes a reference frame $2.0$0, a target or terminal degraded frame $2.0$1, and a time-step signal $2.0$2, and predicts

$2.0$3

The output therefore consists of a spatially varying decay map and a restricted affine transform parameterized by in-plane rotation plus translations. The explicit matrix parameterization shown in the paper includes rotation and translation but no scale or shear, even though the text refers to the transformation as affine. Reconstruction is then factored into attenuation and warping rather than direct pixel regression: $2.0$4 This factorization is the basis of the method’s interpretability, because the predicted latent variables correspond directly to radiometric decay and global motion (Wang et al., 14 Aug 2025).

A key feature is time awareness. The same frame pair $2.0$5 can be queried at different intermediate times $2.0$6, and the model predicts the corresponding degradation state. The paper frames this as continuous inference between frames rather than true multi-frame sequence modeling. Conceptually, $2.0$7 acts as a structural prompt that conditions the latent path from $2.0$8 toward the observed terminal frame (Wang et al., 14 Aug 2025).

3. Architecture, synthetic data generation, and supervision

AtomDiffuser uses a lightweight dual-stream U-Net. Two inputs, $2.0$9 and $$x_t = \sqrt{\bar{\alpha}_t}\, x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon, \qquad \epsilon \sim \mathcal{N}(0,\mathbf{I}),$$0, are processed by shared encoders to obtain comparable multi-scale features. These are fused at the bottleneck, where a time embedding $$x_t = \sqrt{\bar{\alpha}_t}\, x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon, \qquad \epsilon \sim \mathcal{N}(0,\mathbf{I}),$$1 is injected for temporal conditioning. A fully connected head predicts the drift parameters, while a decoder with dual skip connections reconstructs the full-resolution decay map. The paper does not specify the exact mathematical form of $$x_t = \sqrt{\bar{\alpha}_t}\, x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon, \qquad \epsilon \sim \mathcal{N}(0,\mathbf{I}),$$2, so it is only described as a time embedding (Wang et al., 14 Aug 2025).

Training uses synthetic degradations because real paired STEM data with ground-truth drift and damage decompositions are not available. The base images are the TEMImageNet raw atom maps in a noise-free and no-background version, with $$x_t = \sqrt{\bar{\alpha}_t}\, x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon, \qquad \epsilon \sim \mathcal{N}(0,\mathbf{I}),$$3 samples. For each training iteration, a clean image $$x_t = \sqrt{\bar{\alpha}_t}\, x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon, \qquad \epsilon \sim \mathcal{N}(0,\mathbf{I}),$$4 is sampled and a degraded endpoint $$x_t = \sqrt{\bar{\alpha}_t}\, x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon, \qquad \epsilon \sim \mathcal{N}(0,\mathbf{I}),$$5 is synthesized on the fly. Final beam-damage fields are generated using nonlinear Perlin noise, and intermediate decay is modeled by

$$x_t = \sqrt{\bar{\alpha}_t}\, x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon, \qquad \epsilon \sim \mathcal{N}(0,\mathbf{I}),$$6

Final drift is sampled as $$x_t = \sqrt{\bar{\alpha}_t}\, x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon, \qquad \epsilon \sim \mathcal{N}(0,\mathbf{I}),$$7, and intermediate drift is interpolated through

$$x_t = \sqrt{\bar{\alpha}_t}\, x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon, \qquad \epsilon \sim \mathcal{N}(0,\mathbf{I}),$$8

The paper presents these as practical interpolation schemes rather than a more elaborate temporal dynamics model (Wang et al., 14 Aug 2025).

After decay and drift are applied, synthetic acquisition noise is added. The reported noise components are Poisson noise, scan-line jitter, and readout noise, but exact parametric definitions are not given. Supervision is described through intermediate-frame reconstruction,

$$x_t = \sqrt{\bar{\alpha}_t}\, x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon, \qquad \epsilon \sim \mathcal{N}(0,\mathbf{I}),$$9

Although the synthetic generator contains the underlying $x_t = \mathcal{D}(x_0,t).$0 and $x_t = \mathcal{D}(x_0,t).$1, the paper reports only this explicit $x_t = \mathcal{D}(x_0,t).$2 reconstruction loss and does not report separate direct supervision losses on decay maps or affine parameters. This suggests that the principal stated supervision is factorized reconstruction through the predicted physical variables rather than parameter-wise regression (Wang et al., 14 Aug 2025).

Implementation details are concrete. Training uses PyTorch on a single NVIDIA RTX 3090, with batch size $x_t = \mathcal{D}(x_0,t).$3, input resolution $x_t = \mathcal{D}(x_0,t).$4, AdamW, initial learning rate $x_t = \mathcal{D}(x_0,t).$5, cosine decay, $x_t = \mathcal{D}(x_0,t).$6 epochs, and $x_t = \mathcal{D}(x_0,t).$7 hours of training time. The paper also claims flexible input size at inference, with support up to $x_t = \mathcal{D}(x_0,t).$8 (Wang et al., 14 Aug 2025).

4. Inference behavior, interpretability, and practical role

At inference time, AtomDiffuser takes a pair $x_t = \mathcal{D}(x_0,t).$9 and a chosen time $t$0, and returns an inferred drift transform $t$1, an inferred attenuation map $t$2, and an inferred intermediate frame $t$3 obtained by composing the two. The method is therefore fundamentally pairwise between two frames, but continuous in the queried time variable. The paper does not describe recurrent, transformer-based, or long-horizon aggregation over more than two frames at once (Wang et al., 14 Aug 2025).

This pairwise formulation supports two main practical operations. The first is drift alignment: the predicted global transform can be used to align frames and compensate for sample drift. Because damage is modeled separately, the intended claim is that alignment is more robust than conventional registration when image regions have decayed. The second is high-resolution degradation inference: the decay map is produced at input resolution, rather than as a scalar damage score or coarse mask. The paper emphasizes this as an advantage for localizing damage boundaries and interpreting spatially heterogeneous beam damage (Wang et al., 14 Aug 2025).

Interpretability is a central stated property. The attenuation map can be read as a spatial map of signal survival under irradiation, while the affine parameters provide an explicit global motion estimate. This permits visualization and quantification of degradation patterns that correlate with radiation-induced structural instability. In the context emphasized by the paper—sparse time-resolved STEM or cryo-STEM image pairs of beam-sensitive materials—this decomposition is intended to answer a specific analytical question: how much of the inter-frame difference is motion, and how much is dose-induced degradation (Wang et al., 14 Aug 2025).

The paper also states that AtomDiffuser generalizes well from synthetic degradations to real-world cryo-STEM data and applies the model to consecutive real cryo-STEM frames from LLZO. Because this real-data decomposition lacks ground-truth factorization into drift and damage, the claim remains qualitative rather than fully benchmarked (Wang et al., 14 Aug 2025).

5. Experimental evaluation and reported performance

The experiments are organized around atomic damage assessment, sample drift correction, and full degradation inference on real cryo-STEM data. In the synthetic damage-assessment experiment, attenuation is isolated without drift. A fixed synthetic atomic image is degraded over $t$4 frames with damage intensity increasing from $t$5 to $t$6. Three degradation styles are tested: Gaussian noise with black-hole artifacts, low-frequency Perlin masks, and random noise mimicking scan jitter and long-range distortions. A scalar damage intensity is computed from the predicted decay map as

$t$7

Against the CNN baseline from Zheng et al. (2024), AtomDiffuser is reported to perform much better across MAE, MSE, RMSE, and $t$8 in all three settings (Wang et al., 14 Aug 2025).

Setting	CNN baseline	AtomDiffuser
Gaussian	MAE 0.0994; MSE 0.0113; RMSE 0.1065; $t$9	MAE 0.0147; MSE 0.0004; RMSE 0.0198; $$x_t = \mathcal{D}(x_0,t) = \Phi(\lambda_t \odot x_0;\,\mathcal{T}_t),$$0
Perlin	MAE 0.0573; MSE 0.0059; RMSE 0.0767; $$x_t = \mathcal{D}(x_0,t) = \Phi(\lambda_t \odot x_0;\,\mathcal{T}_t),$$1	MAE 0.0193; MSE 0.0006; RMSE 0.0239; $$x_t = \mathcal{D}(x_0,t) = \Phi(\lambda_t \odot x_0;\,\mathcal{T}_t),$$2
Random	MAE 0.0635; MSE 0.0074; RMSE 0.0859; $$x_t = \mathcal{D}(x_0,t) = \Phi(\lambda_t \odot x_0;\,\mathcal{T}_t),$$3	MAE 0.0236; MSE 0.0009; RMSE 0.0306; $$x_t = \mathcal{D}(x_0,t) = \Phi(\lambda_t \odot x_0;\,\mathcal{T}_t),$$4

In the drift-correction experiment, attenuation is removed by setting $$x_t = \mathcal{D}(x_0,t) = \Phi(\lambda_t \odot x_0;\,\mathcal{T}_t),$$5 everywhere. Real STEM images are cropped and then synthetically warped. The predicted transform is compared to the known transform with drift error

$$x_t = \mathcal{D}(x_0,t) = \Phi(\lambda_t \odot x_0;\,\mathcal{T}_t),$$6

and rotation error

$$x_t = \mathcal{D}(x_0,t) = \Phi(\lambda_t \odot x_0;\,\mathcal{T}_t),$$7

The reported results are: for $$x_t = \mathcal{D}(x_0,t) = \Phi(\lambda_t \odot x_0;\,\mathcal{T}_t),$$8 and $$x_t = \mathcal{D}(x_0,t) = \Phi(\lambda_t \odot x_0;\,\mathcal{T}_t),$$9 px drift, drift error $x_0$0 and rotation error $x_0$1; for $x_0$2 and $x_0$3 px, $x_0$4 and $x_0$5; for $x_0$6 and $x_0$7 px, $x_0$8 and $x_0$9; for $\lambda_t \in [0,1]^{H\times W}$0 and $\lambda_t \in [0,1]^{H\times W}$1 px, $\lambda_t \in [0,1]^{H\times W}$2 and $\lambda_t \in [0,1]^{H\times W}$3; and for $\lambda_t \in [0,1]^{H\times W}$4 and $\lambda_t \in [0,1]^{H\times W}$5 px, $\lambda_t \in [0,1]^{H\times W}$6 and $\lambda_t \in [0,1]^{H\times W}$7. The authors interpret these results as indicating strong robustness under moderate motion and reasonable generalization to larger or asymmetric motions (Wang et al., 14 Aug 2025).

For full degradation inference, the paper applies AtomDiffuser to two consecutive real cryo-STEM frames from LLZO in which both drift and beam damage are present. By varying $\lambda_t \in [0,1]^{H\times W}$8, the model generates intermediate states between $\lambda_t \in [0,1]^{H\times W}$9 and $\mathcal{T}_t \in \mathbb{R}^{2\times 3}$0. The reported qualitative behavior is smooth progressive decay, gradual alignment of the frames, and interpretable separation of geometric and radiometric effects. No quantitative ground-truth metrics are reported for this real-data decomposition (Wang et al., 14 Aug 2025).

The paper also reports a model-size and inference comparison with the CNN baseline: the baseline has $\mathcal{T}_t \in \mathbb{R}^{2\times 3}$1M parameters, $\mathcal{T}_t \in \mathbb{R}^{2\times 3}$2 MB, $\mathcal{T}_t \in \mathbb{R}^{2\times 3}$3 input, and $\mathcal{T}_t \in \mathbb{R}^{2\times 3}$4 s/sample inference time; AtomDiffuser has $\mathcal{T}_t \in \mathbb{R}^{2\times 3}$5M parameters, $\mathcal{T}_t \in \mathbb{R}^{2\times 3}$6 MB, support up to $\mathcal{T}_t \in \mathbb{R}^{2\times 3}$7, and $\mathcal{T}_t \in \mathbb{R}^{2\times 3}$8 s/sample inference time. The paper does not discuss in depth whether timings were measured under identical resolution settings, so that comparison is best read cautiously even though the reported figures are substantially different (Wang et al., 14 Aug 2025).

6. Scope, limitations, and significance

A recurrent source of confusion is the model’s name. Despite the term “diffuser,” AtomDiffuser is best understood as a diffusion-inspired degradation model rather than as a conventional stochastic diffusion generative model. Its primary object is not Gaussian noising and denoising in the DDPM sense, but a deterministic, interpretable transition operator for STEM degradation. This distinction matters because the model’s scientific contribution lies less in generic image synthesis than in factorizing temporally correlated atomic-resolution image change into radiometric decay and global motion (Wang et al., 14 Aug 2025).

Within the paper’s framing, AtomDiffuser contributes a physically structured alternative to pipelines that treat all inter-frame changes as either pure geometry or generic corruption. Its central formulation,

$\mathcal{T}_t \in \mathbb{R}^{2\times 3}$9

makes explicit the idea that apparent atomic evolution in sparse STEM sequences may combine reversible geometric displacement and irreversible signal loss. This suggests a broader methodological significance for microscopy analysis: modeling degradation itself can be more informative than attempting to remove it blindly (Wang et al., 14 Aug 2025).

Several limitations can be inferred from the reported setup. Training relies on synthetic degradation, and real STEM beam damage may not follow the assumed Perlin-based decay or linear temporal interpolation. The drift model is global affine only, which suggests limited capacity for nonrigid, spatially varying, or scan-dependent distortions. Temporal modeling is pairwise rather than multi-frame. Beam damage is represented as attenuation, which captures signal loss but may not fully describe defect creation, appearance of new features, or more complex structural transformation. Real-data validation for the full drift-damage decomposition remains mostly qualitative because ground-truth factorization is unavailable. The interpolation assumptions,

$\Phi(\cdot;\mathcal{T}_t)$0

are also explicitly simple (Wang et al., 14 Aug 2025).

The paper’s stated future-facing uses include studying time-dependent beam effects, virtual time-lapse reconstruction from sparse acquisitions, and improved analysis of dynamic material behavior. A plausible implication is that interpretable, high-resolution damage maps could also inform dose management and acquisition planning, since the method is designed to visualize where radiation-induced instability is localized. In that sense, AtomDiffuser occupies a specific niche within atomic-resolution microscopy: not general restoration, not classical registration, and not standard generative diffusion, but time-aware degradation modeling for disentangling drift and beam damage in sparse STEM and cryo-STEM sequences (Wang et al., 14 Aug 2025).