Sparse Mixture of Mambas for Domain Generalized Atomic Electron Tomography Augmentation.

Abstract

Atomic electron tomography (AET) is essential for characterizing the atomic structure of functional materials. However, raw 3-D tomograms often exhibit severe artifacts caused by geometric constraints and low radiation doses. Although point-attention-based ensemble augmentation methods effectively remove artifacts in simulated datasets with varying structure factors, they struggle with the complexity of real tomograms that demand multidomain feature learning. Moreover, existing models degrade in multidomain scenarios and incur high parameter counts introduced by point-attention mechanisms, which increase hardware demands. To address these challenges, we propose a sparse mixture of Mambas (MoMambas), a novel 3-D augmentation method that enhances domain generalization. MoMambas decouple domain-specific parameters by integrating a sparse mixture-of-experts (MoE) framework with Mamba-based experts, resolve positional ambiguity in sparse input sequences through positional information enhancement, and boost MoE accuracy via a multihead routing algorithm. Our approach achieves a 22% accuracy improvement over state-of-the-art AET augmentation methods in multidomain learning, reduces the parameter count to just 2.9% of the original, and lowers computational cost by 6%. Codes and data are publicly available at https://github.com/yuy38457/MoMambas.