ReStack-Mamba: Stage-wise Residual Refinement for Volume Electron Microscopy Registration

Volume electron microscopy (vEM) assembles three-dimensional tissue volumes from physically cut sections, but cutting, collection, and imaging introduce rigid motion and nonuniform elastic deformation. The central challenge is to correct these section-specific distortions while preserving long-range axial continuity without propagating errors through the stack. We present ReStack-Mamba, a self-supervised framework for this coupled local-and-global registration problem. A Local Elastic Aligner first estimates in-plane correspondence between neighboring sections. Two independently trained, z-preserving StackMamba stages then use XY, XZ, and YZ context to correct successive residual errors without changing the section order. Training the second stage on errors remaining after the first allows it to learn a distinct correction; robust axial post-processing suppresses isolated outliers. Across six simulated volumes spanning five tissues and one real-world dataset, ReStack-Mamba improves XZ/YZ SSIM by approximately 10% over existing methods and achieves the best aggregate PSNR, 3D NCC, and 3D MI. Per-section and deformation-severity analyses show that the advantage persists through long stacks and becomes more pronounced under stronger distortions. The real-volume result further shows improved axial continuity, supporting robust vEM registration across diverse tissue structures.