Cryo-SWAN: the Multi-Scale Wavelet-decomposition-inspired Autoencoder Network for molecular density representation of molecular volumes
Authors: Rui Li, Artsemi Yushkevich, Mikhail Kudryashev, Artur Yakimovich
Organizations: Center for Advanced Systems Understanding (CASUS), Görlitz, Germany. · Helmholtz-Zentrum Dresden-Rossendorf e. V. (HZDR), Dresden, Germany. · In situ Structural Biology, Max Delbrück Center for Molecular Medicine in the Helmholtz Association, Berlin, Germany. · Department of Physics, Humboldt University of Berlin, Berlin, Germany. · Institute of Medical Physics and Biophysics, Charite-Universitätsmedizin, Berlin, Germany. · Institute of Computer Science, University of Wrocaw, Wrocaw, Poland. · Cluster of Excellence Physics of Life, TU Dresden, Dresden, Germany.
Learning robust representations of 3D shapes from voxelized data is essential for advancing AI methods in biomedical imaging. However, most contemporary 3D computer vision approaches operate on point clouds, meshes, or octrees, while volumetric density maps, the native format of structural biology and cryo-EM, remain comparatively underexplored. We present Cryo-SWAN, a voxel-based variational autoencoder inspired by multi-scale wavelet decomposition. The model performs conditional coarse-to-fine latent encoding and recursive residual quantization across perception scales, enabling accurate capture of both global geometry and high-frequency structural detail in molecular density volumes. Evaluated on ModelNet40, BuildingNet, and a newly curated dataset of cryo-EM volumes, ProteinNet3D, Cryo-SWAN consistently improves reconstruction quality over state-of-the-art 3D autoencoders. We demonstrate that the molecular densities organize in learned latent space according to shared geometric features, while integration with diffusion models enables denoising and conditional shape generation. Together, Cryo-SWAN is a practical framework for data-driven structural biology and volumetric imaging.
Figures & tables
Figure 1: Principle architecture of Multi-Scale Wavelet-decomposition-inspired Autoencoder Network . (a) Multi-scale decomposition of the input signal. (b) The residual quantization process. A residual loss is computed during codebook lookup at each scale. (c) Downstream applications of the voxel-based 3D shape representations.
Figure 2: cryo-SWAN performance on common 3D shape benchmarks . (a) Performance on three representative objects from BuildingNet dataset. (b) Performance on three representative objects from ModelNet dataset. Here, low, mix, and high denote the frequency components, where mix captures a combination of low- and high-frequency details.
Figure 3: Representation learning on the ProteinNet3D and explorations in the latent space. (a) The representation learning on the ProteinNet3D. (b) The FSC evaluation of all the representations. (c) Dimensionality reduction comparison: Cryo-SWAN latent vectors vs. original EMDB volumes. (d) UMAP projection of latent vectors reveals distinct hubs reflecting structural similarity. Molecules with similar 3D structures at 8 Å resolution (siblings).
Figure 4: Ablation study for codebook collapsing and multi-scale representations. (a) The calls for all the codebook entries. (b) The latent code representation for both scales (coarse and fine). (c) The molecular density volume: original (GT), at a certain representation scale (RQ1) and at a full representation (RQ1+RQ2). (d) Evaluation of representation performance at a certain representation scale (RQ1) compared to the full representation (RQ1+RQ2).
Figure 5: Downstream applications of cryo-SWAN including shape denoising and conditional molecular shape generation. (a) The conditional generation from diffusion models based on cryo-SWAN representations. (b) Examples of conditional generations. The variants are similar in the geometric perspective to the real protein densities in EMDB (anchor point). (c) The unsupervised denoising based on the representation learning for 3D densities from ModelNet.
Figure S1 : Optimal choice of the pixel size and the volume size. (top) Volume size distribution at various resampled pixel size. (bottom) Volumes distribution respective to the originally reported resolution and the pixel size: ProteinNet3D volumes, preserving original resolution after resampling (blue) and degraded to 8A˚ after resampling due to the Nyquist limit (green), as well as filtered out volumes due to the cutoffs applied (white).
Figure S3 : ProteinNet3D complexity: membrane-related signal and experimental noise presence. (top) Normal protein densities in ProteinNet. (middle) ProteinNet example volumes containing membrane-related signal. (bottom) ProteinNet example volumes containing experimental noise
Figure S4 : Molecular shapes similarity within hubs identified with Cryo-SWAN.
Figure S5 : The latent diffusion generation based on cryo-SWAN representation.
Advanced Light Source, Lawrence Berkeley National Laboratory, Berkeley, CA, USA · Center for Advanced Mathematics for Energy Research Applications, Lawrence Berkeley National Laboratory, Berkeley, CA, USA · Molecular Biophysics & Integrated Bioimaging Division, Lawrence Berkeley National Laboratory, Berkeley, CA, USA +9
Institute of AI for Health & Helmholtz AI, Computational Health Center, Helmholtz Munich – German Research Center for Environmental Health, Neuherberg, Germany · Munich Center for Machine Learning (MCML), Munich, Germany · Department of Medicine III, Ludwig-Maximilian-University Hospital, Munich, Germany +2