U-Mamba with Heat Conduction Equation

• The paper presents a novel hybrid neural architecture that fuses state-space Mamba modules with heat conduction operators for both selective and global feature propagation.
• It leverages Fourier-based transforms like DCT/IDCT to implement a channelwise isotropic Gaussian filter, ensuring interpretable, efficient diffusion of features.
• Experimental evaluations demonstrate that UMH outperforms conventional U-Net variants, achieving higher segmentation accuracy on CT and MRI datasets.

U-Mamba with Heat Conduction Equation designates a class of hybrid neural architectures integrating state-space deep learning modules—particularly Mamba-based blocks—with mathematical operators derived from the solution theory of the heat conduction partial differential equation (PDE). This approach leverages both efficient, selective global context modeling and isotropic signal diffusion to address semantic segmentation problems, notably in medical imaging. The resulting models demonstrate strong empirical performance, scalable training, and theoretical interpretability rooted in physical analogies to thermal diffusion.

1. Architectural Foundations and Design Principles

The central architecture, referred to as U-Mamba-HCO or "UMH" (Editor's term), follows the encoder–bottleneck–decoder schema typified by U-Net and its variants. Its innovation lies in the fusion of two global reasoning mechanisms at complementary stages:

• U-Mamba Encoder Stages: Each encoder stage employs residual convolutional units followed by Mamba-based state-space modules. The Mamba block achieves linear-time, selective long-range dependency modeling via spatial and channelwise gating and structured recurrence. After each encoding step, spatial subsampling is applied; stage depth is typically 6 (3D) or 7 (2D).
• Heat Conduction Operators (HCOs) in the Bottleneck: The bottleneck replaces conventional attention or convolution with two sequential HCO layers that realize analytic frequency-domain diffusion based on the heat equation.
• Decoder and Output: The decoder is symmetric, merging upsampled features with encoder skips and concluding with a 1×1×1 convolution and softmax.

This design ensures that localized, spatially-detailed information is preserved (skip connections), while both selective and global context are exploited through Mamba blocks and the HCO, respectively (Wu et al., 5 Nov 2025).

2. Heat Conduction Operator: Mathematical Formulation and Implementation

The Heat Conduction Operator models feature propagation as physical heat diffusion over the neural feature manifold. In the 3D case, the abstraction is the classical PDE: $$\frac{\partial u(x,y,z,t)}{\partial t} = k\, (\tfrac{\partial^2 u}{\partial x^2} + \tfrac{\partial^2 u}{\partial y^2} + \tfrac{\partial^2 u}{\partial z^2})$$ where $u$ is a surrogate "temperature" field and $k$ is a learnable or predicted diffusivity.

Application of the spatial Fourier transform (or Discrete Cosine Transform, DCT, for real signals) yields a solution in the frequency domain: $$\widetilde u(\omega, t) = \widetilde f(\omega)\, \exp\left[ -k\, |\omega|^2 t \right]$$ Returning to the spatial domain: $$u(x,t) = \mathcal{F}^{-1}\left\{ \widetilde f(\omega) e^{-k |\omega|^2 t} \right\}$$ In practice, for a feature map block $U^0$, the operation is: $$U^t = \mathrm{IDCT}\left[\, \mathrm{DCT}(U^0)\ \odot\ \exp(-k\,|\Omega|^2 t)\, \right]$$ where $k$ is predicted by a frequency-value embedding subnetwork and $|\Omega|^2$ is the squared spectral norm.

This construction inherently acts as a channelwise isotropic Gaussian filter, promoting low-frequency, global feature propagation in a single differentiable layer, at complexity dominated by the DCT/IDCT transformation (Wu et al., 5 Nov 2025).

3. Empirical Performance and Quantitative Evaluation

Experimental evaluation was conducted on multi-organ abdominal segmentation tasks using CT (3D), MRI (3D), and MRI (2D) data. Key findings include:

Architecture	DSC (CT 3D)	NSD (CT 3D)
nnU-Net	0.8615	0.8972
SegResNet	0.7927	0.8257
UNETR	0.6824	0.7004
SwinUNETR	0.7594	0.7663
U-Mamba_Bot	0.8683	0.9049
U-Mamba_Enc	0.8638	0.8980
UMH (Mamba + HCO)	0.8719	0.9037

Ablation studies demonstrate that the integration of HCO alone gives minimal (sometimes negative) change, Mamba alone yields a modest increase, but the hybridization achieves the largest improvement, evidencing the complementarity of these mechanisms (Wu et al., 5 Nov 2025).

4. Algorithmic Complexity and Scalability

• Mamba Block: $O(N)$ in the number of spatial pixels/voxels due to linear-time structured recurrence.
• HCO Layer: Complexity is dominated by the DCT/IDCT, which is typically $O(N^{1.5})$ in 3D, behaving as $O(N \log N)$ in practice.
• Contrast with Self-Attention: Both Mamba and HCO are substantially less computationally demanding than full self-attention, which is $O(N^2)$.
• Deployment: The model processes high-resolution 3D volumes (e.g., $40 \times 224 \times 192$ patches) on a single NVIDIA A100 GPU with practical batch sizes (2–4), and requires no pretraining.

5. Theoretical Interpretability and Modeling Rationale

The use of the heat equation for feature modulation is justified by its role as the canonical isotropic diffusive operator, promoting analytic smoothing and global context propagation. The learnable diffusivity $k$ controls the degree of abstraction—larger values induce more aggressive global mixing, while smaller values retain detail. This provides an explicit, interpretable degree of control over non-local interactions, distinct from both rigid convolution and anisotropic attention.

Complementary to the selective, directionally-aware global reasoning of Mamba, HCO delivers channelwise isotropic frequency-domain propagation. Their joint application addresses both targeted and blanket context integration, which is particularly valuable for complex, multi-organ segmentation where spatial and semantic ambiguities are significant (Wu et al., 5 Nov 2025).

6. Limitations and Prospects for Extension

Current instantiations of UMH employ isotropic diffusion and restrict HCO layers to the bottleneck. Prospective developments include:

• Extension to anisotropic or learned diffusion kernels to refine boundary localization.
• Adaptive placement and multiplicity of HCO layers.
• Application beyond abdominal datasets (e.g., multi-modal fusion such as PET/CT or other dense prediction tasks).
• Analysis of alternate explicit PDE-based operators for other kinds of context propagation.

A plausible implication is that further spatial adaptation or task-specific learning of the diffusion process could enhance both accuracy and robustness, especially where structures of interest exhibit highly non-uniform contextual dependencies.

7. Connections to Mathematical Modeling and U-Mamba Methodology

The naming U-Mamba references both the UNet-style topology and Mamba state-space modules (Wu et al., 5 Nov 2025), but also relates, semantically, to the "U-Mamba" or Unified Transform Methodology for heat conduction equations in mathematical physics (Farkas et al., 2022, Sheils, 2016). While in the deep learning context, U-Mamba refers to neural network building blocks, several shared mathematical features are noteworthy:

• The explicit spectral representations of thermal diffusion in both PDE analysis and neural HCO layers provide global, frequency-based smoothing operations.
• Discretizations and operator approximations (e.g., by DCT) parallel those used in numerical PDE solvers (Momin et al., 2021).
• Recent advances in the Fokas method for variable and interface-diffusivity problems (Farkas et al., 2022, Sheils, 2016) suggest that more elaborate or non-uniform HCOs could be constructed for learned, spatially inhomogeneous feature diffusion.

In summary, UMH represents an overview of state-space deep learning and physics-inspired operator integration, with measurable benefits on large-scale segmentation benchmarks and a foundation for future methodological advances.