Cryo-SWAN: the Multi-Scale Wavelet-decomposition-inspired Autoencoder Network for molecular density representation of molecular volumes

TL;DR

Cryo-SWAN is a novel voxel-based variational autoencoder inspired by wavelet decomposition, designed to accurately represent molecular density volumes by capturing both global and high-frequency details, advancing 3D shape analysis in structural biology.

Contribution

The paper introduces Cryo-SWAN, a multi-scale wavelet-inspired autoencoder that improves volumetric density representation for molecular structures, especially in cryo-EM data.

Findings

Outperforms existing 3D autoencoders in reconstruction quality.

Organizes molecular densities in latent space based on geometric features.

Enables denoising and conditional shape generation with diffusion models.

Abstract

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.