Energy-Mamba: A Physics-Constrained State-Space Model for Medical Image Classification
TL;DR
Evaluated on four datasets, Energy-Mamba achieves state-of-the-art classification performance with significantly fewer parameters, demonstrating that physics-informed grounding can enhance both efficiency and representational quality in medical vision tasks.
Abstract
State-Space Models (SSMs), particularly Mamba, offer linear-time complexity for long-range dependencies, making them attractive for medical imaging with limited annotated data. However, adapting these sequential models to 2D images through unconstrained state evolution causes representational drift, the dynamic hidden state progressively loses fidelity to local image features. We introduce Energy-Mamba, integrating SSM dynamics with physics-informed constraints via a learnable potential energy function that quantifies compatibility between evolving states and static local features. Our Energy-Mamba Block introduces a gradient-based forcing term, computed dynamically via automatic differentiation, that pulls states toward low-energy configurations maintaining local visual fidelity. This formulation mirrors Hamiltonian dynamics: kinetic energy (SSM scan) plus potential energy (our constraint function) govern state trajectories. This architectural prior enables learning implicit constraints for robust, faithful representations, crucial in medical imaging where fine-grained local detail drives accurate diagnosis. Evaluated on four datasets (retinal OCT, chest X-ray, microscopy, abdominal CT), Energy-Mamba achieves state-of-the-art classification performance with significantly fewer parameters, demonstrating that physics-informed grounding can enhance both efficiency and representational quality in medical vision tasks.