Energy-Mamba: A Physics-Constrained State-Space Model for Medical Image Classification

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.

• The paper introduces Energy-Mamba, a state-space model that combines SS2D long-range processing with a learned potential-energy force anchored to initial patch features to reduce representational drift.
• The model achieves leading or near-leading results across four MedMNIST benchmarks, including 87.60% OCTMNIST accuracy, 99.91% BloodMNIST AUC, and 97.04% PneumoniaMNIST AUC, while using only 2.0M parameters and 2.18G FLOPs.
• The experiments show stronger intermediate feature similarity to the input than a pure SSM baseline, but missing component ablations, formal stability analysis, and 3D or segmentation tests leave the mechanism’s causal benefits unresolved.

Overview and motivation

"Energy-Mamba: A Physics-Constrained State-Space Model for Medical Image Classification" (2608.19813) addresses a specific failure mode of vision state-space models (SSMs): representational drift. In Mamba-style architectures adapted to 2D images via multi-directional selective scanning (SS2D), the hidden state evolves under purely "kinetic" dynamics—each token modulates the SSM parameters, but nothing anchors the evolving state to the local image content it was initialized from. Earlier spatial information progressively migrates into different representational subspaces, and scan order and small perturbations accumulate across layers. The authors argue this is especially damaging in medical imaging, where fine-grained local cues (lesion boundaries, nodule texture) are directly diagnostic.

The paper's diagnosis is grounded in prior structural analyses by the same group (2608.19813), which observed instability in hidden states under sequential scanning but did not supply a corrective mechanism. Existing hybrid CNN–SSM designs mitigate drift symptoms through architectural diversity; the authors' claim is that none prevents the underlying cause.

Method

Energy-Mamba augments each Mamba block with a learnable potential energy function inspired by Hamiltonian mechanics. Each block receives the dynamic feature map $\mathbf{Z}$ and a static reference $\mathbf{\Phi}$—the initial patch embedding, held fixed throughout the network. A three-layer MLP computes a scalar energy $$E_{i,j} = f_{\text{energy}}([\mathbf{h}_{i,j}; \boldsymbol{\phi}_{i,j}])$$ at every spatial location, where low energy indicates compatibility between the current state and the original local content. A corrective forcing term $$\mathbf{F}_{i,j} = -\eta \nabla_{\mathbf{h}}\mathcal{E}$$ is computed via automatic differentiation; because the Jacobian $\partial\mathbf{z}/\partial\mathbf{h}$ selects only the first $C$ components of the full gradient, this is computationally cheap. The learnable scalar $\eta$ controls force magnitude.

Architecturally, the block has two parallel branches: a kinetic branch applying SS2D to the normalized state for long-range dependencies, and the potential branch producing the spatial force map. Their outputs are fused residually with stochastic depth, followed by a channel-wise MLP with DropPath. The authors frame this as a discretized Hamiltonian system with total energy $T + V$: SS2D supplies the kinetic term governing information propagation, while the potential defines a low-energy manifold of content-consistent configurations that the gradient force continuously enforces. Unlike standard Mamba, perturbations are not free to propagate uncorrected.

Notably, Energy-Mamba uses a uniform non-hierarchical design: after patch embedding to $28\times28$, resolution and channel width remain constant across all blocks, in contrast to pyramidal designs like MedMamba (2608.19813).

Experimental results

Evaluation uses four MedMNIST v2 benchmarks spanning retinal OCT (OCTMNIST), chest X-ray (PneumoniaMNIST), blood cell microscopy (BloodMNIST), and abdominal CT (OrganCMNIST), at $224\times224$ resolution with standard augmentation and official splits. Seven baselines span CNNs (ResNet-18/50), Transformers (ViT-B, Swin-T, MedViT), and SSM/KAN hybrids (MedMamba, MedKAFormer-T). Training used Adam with cosine annealing on a single RTX 3090 for up to 50 epochs.

The headline result is efficiency: Energy-Mamba uses 2.0M parameters and 2.18G FLOPs, roughly 6× fewer than the next smallest competitor (MedKAFormer-T, 12.47M) and over 43× fewer than ViT-B (86.57M), yet achieves state-of-the-art or near-best results everywhere:

Dataset	Metric	Energy-Mamba	Best competitor
OCTMNIST	ACC / F1 (%)	87.60 / 87.50	81.03 / 80.00
BloodMNIST	AUC / ACC / F1 (%)	99.91 / 97.79 / 97.68	99.77 / 96.52 / 96.68
OrganCMNIST	AUC / F1 (%)	99.21 / 90.18	99.10 / 89.51

On PneumoniaMNIST, Energy-Mamba ranks first in AUC (97.04%) and accuracy (90.74%), with MedKAFormer-T slightly ahead in F1 (89.99% vs. 87.44%). The margin over MedKAFormer-T on OCTMNIST—+6.57% accuracy and +7.50% F1—is large relative to typical benchmark variance. ViT-B consistently underperforms despite its capacity, supporting the paper's claim that raw scale is not the primary driver here.

Direct evidence for the drift hypothesis comes from cosine-similarity tracking: mean similarity between the static input representation and intermediate features is 0.65 at block 1 versus 0.31 for a pure SSM baseline. Both converge to ~0.18–0.19 by block 8, consistent with necessary abstraction at depth, but Energy-Mamba maintains higher fidelity throughout intermediate layers. This supports—but does not fully prove—the causal claim: the paper provides no ablation isolating the contribution of the forcing term itself, no comparison against alternative anchoring mechanisms (e.g., simple skip connections to $\mathbf{\Phi}$0), and no sensitivity analysis over $\mathbf{\Phi}$1. The reported gains therefore conflate the physics-informed mechanism with other design choices such as the uniform-resolution backbone and stochastic depth.

Limitations and open questions

The evaluation is confined to 2D classification on relatively small, standardized benchmark datasets; extension to 3D volumetric data and dense tasks such as segmentation remains untested, as the authors acknowledge. The Hamiltonian framing is interpretive rather than rigorous—no formal stability or convergence guarantees are provided for the combined discrete update, and the learned energy function's structure is not analyzed. A systematic ablation of the kinetic/potential branch interplay across modalities and dataset scales is explicitly left open, as is exploration of alternative energy function parameterizations. Whether the restorative constraint helps most precisely in data-scarce regimes, as the motivation implies, is asserted rather than demonstrated experimentally.

Conclusion

Energy-Mamba introduces an explicit restorative constraint on SSM hidden-state evolution via a learnable potential energy function and gradient-based forcing term, framed within a Hamiltonian interpretation. It delivers state-of-the-art or competitive results on four MedMNIST classification benchmarks with only 2.0M parameters and shows empirically reduced representational drift relative to a pure SSM baseline. The principal caveat is the absence of component-level ablations tying the performance gains specifically to the potential-energy mechanism; establishing that attribution, and extending the approach beyond 2D classification, constitute the paper's main open questions.