FrostByte is a continuous score-based generative diffusion framework for 3D macromolecular electron density reconstruction from noisy 2D Cryo-EM projections (SNR < -5 dB). The project evolves from geometric point-cloud message passing (Phase 1–3) → continuous 3D volumetric electron density fields with Diffusion Posterior Sampling (Phase 5) → scalable continuous Tri-Plane neural fields for 128³+ resolution (Phase 6–7).
Cryo-EM single-particle reconstruction and Cryo-Electron Tomography (Cryo-ET) are severely ill-posed inverse problems due to extreme radiation damage dose limits (SNR < -5 dB), Contrast Transfer Function (CTF) phase flips, and missing wedge geometries. Traditional regularizers (RELION, cryoSPARC) rely on empirical low-pass filtering and solvent masks.
Generative Diffusion Priors provide physical, learnable structural regularization:
- Differentiable Physics Modeling: Forward CTF modulation and differentiable 3D Radon line-integral projections.
- Diffusion Posterior Sampling (DPS): Reverse SDE trajectories guided by measurement loss gradients
∇_{x_t} ||y - R_R(x̂_0)||². - Workstation-Accessible Execution: Optimized for single-GPU mobile workstations (e.g. NVIDIA RTX A2000 Laptop GPU) via PyTorch FP16 Automatic Mixed Precision (AMP).
- Scale Calibration & Stability: Elimination of prior-induced volume collapse via coordinate scaling (λ = 1.59).
graph TD;
A[Noisy 2D Projections y via CTF] --> B(DPS Inverse Solver);
C[3D Score Network Prior] --> B;
D[Coordinate Calibration λ=1.59] -->|prevents collapse| B;
B --> E[Reconstructed 3D Density Map];
E --> F{Evaluation Metrics};
F --> G[Pearson CC / FSC 0.143 / Rg];
subgraph Phase 1-3: Point Cloud Prior
H[SE3-Equivariant GNN Score Model]
end
subgraph Phase 5: Volumetric Prior
J[3D U-Net + Differentiable Radon Projector]
end
subgraph Phase 6-7: Scalable Tri-Plane INR
K[Tri-Plane Feature Encoder + MLP Decoder]
end
H --> C
J --> C
K --> C
diffusion-cryoem-prior/
├── data/
│ ├── volume_dataset.py # 3D Voxelization via Gaussian density splatting
│ └── triplane_dataset.py # Continuous coordinate sampling for Tri-Plane INR
├── models/
│ ├── diffusion.py # Continuous DDPM & DPS posterior sampling
│ ├── unet_3d.py # 3D Volumetric U-Net score network
│ ├── unet_2d.py # 2D Tri-Plane U-Net prior
│ ├── triplane.py # Implicit Neural Representation (INR) MLP decoder
│ └── triplane_encoder.py # Continuous 3D feature encoder
├── projection/
│ ├── radon.py # Differentiable 3D Radon Transform operator
│ └── neural_radon.py # Differentiable ray-marching projection module
├── utils/
│ └── metrics.py # 3D Pearson CC and Fourier Shell Correlation (FSC 0.143)
├── scripts/
│ ├── benchmark_a2000_workstation.py # Hardware latency, VRAM scaling & throughput benchmark
│ ├── prolonged_stress_test.py # Continuous GPU saturation and FP16 endurance suite
│ ├── remote_runner.py # Paramiko SSH runner for remote workstation execution
│ ├── train_volume_prior.py # Volumetric 3D DDPM training pipeline
│ ├── verify_volume_reconstruction.py # Single-protein overfitting validation
│ └── visualize_volume_hd.py # High-definition Z-slice visualizer
└── assets/ # Architecture diagrams, figures, and animations
Evaluated live on a mobile workstation equipped with an NVIDIA RTX A2000 Laptop GPU (4 GB physical VRAM, 3.68 GiB usable), CUDA 13.0:
| Spatial Grid | Batch Size (B) | Precision | Latency / Volume | Peak Memory | GPU Compute Utilization |
|---|---|---|---|---|---|
| 32³ (32×32×32) | 1 | FP32 | 19.17 ms | 132 MB | 12% |
| 32³ (32×32×32) | 4 | FP16 AMP | 4.46 ms | 132 MB | 95–100% |
| 64³ (64×64×64) | 1 | FP32 | 30.65 ms | 342 MB | 15% |
| 64³ (64×64×64) | 4 | FP16 AMP | 16.16 ms | 343 MB | 90–98% |
| 128³ (128×128×128) | 1 | FP32 | 228.53 ms | 2.05 GB | 85–90% |
Key Finding: Batched FP16 Automatic Mixed Precision (AMP) delivers a 4.30x speedup at 32³ and 1.90x speedup at 64³, eliminating GPU dispatch starvation while remaining safely within the 3.68 GiB VRAM envelope.
- Equivariance Verification: SE(3) equivariance error validated at 1.0 × 10⁻⁶.
- Scale Mismatch Discovery & Fix: Identified that normalized latent sampling contracts physical protein densities. Applying coordinate scale factor λ = 1.59 restored true Radius of Gyration (Rg) bounds (< 0.8 Å RMSD).
Figure: Simulated Contrast Transfer Function applied to a 2D projection with visible phase reversals.
Figure: Calibration sweep over guidance strength α. α=1.0 with λ=1.59 achieves <0.8 Å aligned RMSD.
Transitioned to continuous 3D spatial grids (64³) compatible with experimental Cryo-EM MRC densities:
VolumeDataset: Voxelization of macromolecular PDB coordinates via 3D Gaussian kernels.UNet3D: 3D volumetric convolutional score network.RadonTransform: Differentiable line-integral projection operator.
Figure: Left — input 2D projection. Centre — ground truth central slice. Right — reconstructed density slice.
Figure: High-resolution central slice comparison showing recovered tertiary contour density.
Animation: 6-Panel Continuous Z-Axis Density Sweep Video Simulation across novel protein structures (1A3N Hemoglobin Alpha & 1CQY Flavodoxin Fold). Displays Ground Truth 3D Density (left), Noisy Observation at -5dB SNR (centre), and Score-Matching Diffusion Prior Reconstruction (right).
Figure: 3D Volumetric Electron Density Recovery Benchmark across novel unseen PDB protein folds (1A3N Hemoglobin Alpha, 1CQY Flavodoxin Fold, 1TFG Transcription Factor) under severe -5dB phase noise.
To overcome O(N³) volumetric memory scaling for 128³+ grids:
- Tri-Plane Representation: Three orthogonal 2D feature planes (XY, XZ, YZ) decoded by a shared continuous MLP.
- Latent 2D Diffusion: Denoising prior trained over compressed Tri-Plane feature maps.
Figure: High-capacity Tri-Plane reconstruction gallery at 128³ spatial resolution across benchmark structures.
# Clone the repository
git clone https://github.com/QntmSeer/FrostByte.git
cd FrostByte
# Install dependencies
pip install -r requirements.txt
# Run workstation hardware latency & VRAM saturation benchmark
python scripts/benchmark_a2000_workstation.py
# Run prolonged multi-volume stress test
python scripts/prolonged_stress_test.py
# Verify volumetric reconstruction pipeline
python scripts/verify_volume_reconstruction.pyTo maintain scientific rigor and transparency, the current implementation operates under the following explicit boundary conditions:
- Known Pose Orientations: The DPS likelihood guidance assumes projection viewing angles R_i ∈ SO(3) are known or pre-estimated. Joint blind pose estimation and volume refinement (as in RELION) is an active area of future development.
- Synthetic Noise vs Real Micrographs: Current benchmarks use simulated additive Gaussian noise (SNR ∈ [-10 dB, 0 dB]) with Contrast Transfer Function (CTF) modulation. Experimental Cryo-EM micrographs exhibit non-Gaussian shot noise, beam-induced motion blur, and ice gradient artifacts.
- Dataset Scale & Generalization: Training on limited structural subsets provides strong fold-specific regularization; generalized zero-shot foundation priors require training across 10,000+ diverse structures from RCSB PDB and EMDB.
- Volumetric Memory Scaling: Direct O(N³) voxel diffusion requires ~2.05 GB VRAM for 128³ volumes. Sub-Ångström full-micrograph reconstructions (512³+) require Tri-Plane neural representations (Phase 6–7) or spatial patch decomposition.
- Iterative Sampling Latency: Continuous reverse-SDE sampling requires multiple denoising steps (e.g. 50–1,000 steps), which is computationally more demanding than single-pass feed-forward inversion networks.
- DPS: Chung et al., "Diffusion Posterior Sampling for General Noisy Inverse Problems," ICLR, 2023.
- DDPM / Score SDE: Song et al., "Score-Based Generative Modeling Through Stochastic Differential Equations," ICLR, 2021; Ho et al., "Denoising Diffusion Probabilistic Models," NeurIPS, 2020.
- Tri-Planes / EG3D: Chan et al., "Efficient Geometry-Aware 3D Generative Adversarial Networks," CVPR, 2022.
- Cryo-EM Bayesian Foundations: Scheres, "RELION: Implementation of a Bayesian approach to cryo-EM structure determination," JSB, 2012; Punjani et al., "cryoSPARC: algorithms for rapid unsupervised cryo-EM structure determination," Nature Methods, 2017.
- CryoDRGN: Zhong et al., "CryoDRGN: Reconstruction of Heterogeneous Cryo-EM Structures Using Neural Networks," Nature Methods, 2021.
- SE(3)-EGNN: Satorras et al., "E(n) Equivariant Graph Neural Networks," ICML, 2021.