CryoGeo: Latent Pose Equivariance for Amortized Ab Initio Cryo-EM Reconstruction

Cryo-electron microscopy (cryo-EM) reconstructs 3D macromolecular structures from noisy 2D particle images by jointly estimating the density and projection pose of each particle. Amortized inference can greatly accelerate this process, but remains vulnerable to pose collapse because known image-space symmetries are often learned only implicitly through reconstruction losses. We ask whether this instability can be mitigated by enforcing cryo-EM transformation laws directly in latent projection-pose space. We propose CryoGeo, a geometry-constrained framework for amortized ab initio cryo-EM reconstruction. CryoGeo enforces latent pose equivariance by requiring in-plane transformations to induce valid group actions on the predicted projection pose: rotations act as right multiplications on the predicted $SO(3)$ orientation, while shifts transform frame-consistently. At the core of CryoGeo is an output-level group-action objective that forces predicted cryo-EM projection poses to obey physical transformation laws induced by in-plane image rotations and shifts, rather than relying on equivariant feature extraction alone. Across synthetic and real cryo-EM datasets, CryoGeo improves pose consistency, mitigates collapse-prone behavior, and yields more reliable amortized ab initio reconstructions than unconstrained variants and neural baselines. These results show that latent pose equivariance turns in-plane symmetry from an implicitly learned nuisance factor into a self-supervised geometric constraint for stable, efficient amortized structural reconstruction. Code will be released.