Invertible Diffusion for Low-Memory Channel Gain Map Construction in Wireless Communication Networks

Abstract: Channel gain maps (CGMs) enable propagation-aware services in edge-intelligent wireless communication networks, while diffusion-based CGM construction is memory intensive for on-device training or adaptation. This letter proposes InvDiff-CGM, an invertible diffusion framework that constructs CGMs from sparse measurements and environmental priors. By adopting invertible architectures in both the diffusion process and the U-Net noise estimator, InvDiff-CGM achieves near-constant training memory consumption. A prior-informed multi-scale injector further integrates environmental priors with sparse measurements to improve physical consistency and detail preservation. Experiments on RadioMap3DSeer show about an 85\% reduction in peak training memory and a PSNR of 38.02~dB, outperforming representative recent baselines. This validates the practicality of InvDiff-CGM for high-fidelity CGM construction under edge resource constraints.

• The paper presents InvDiff-CGM, which leverages invertible design principles to achieve nearly constant memory usage, cutting GPU memory by 85.59%.
• It integrates a multi-scale prior injection that effectively fuses environmental layouts with sparse measurements for sharp and accurate map recovery.
• Empirical results demonstrate state-of-the-art performance with PSNR of 38.02 dB and SSIM of 0.9634, enabling online adaptation on consumer-grade hardware.

Invertible Diffusion for Resource-Efficient Channel Gain Map Construction in Wireless Edge Networks

Introduction

Channel gain maps (CGMs) offer high-resolution spatial representations of large-scale path loss and shadowing, enabling environment-aware resource allocation, coverage prediction, and adaptive communications in 6G edge-intelligent wireless networks. Deep learning-based radio map construction techniques, particularly recent diffusion models, have demonstrated state-of-the-art fidelity in reconstructing CGMs from sparse measurement data. However, existing diffusion-based solutions suffer from prohibitive training-time memory requirements due to the necessity of storing intermediate activations for unrolled iterative networks, limiting their practicality for deployment on memory-constrained edge nodes.

This paper introduces InvDiff-CGM, an invertible diffusion-based generative framework for CGM construction that explicitly addresses the memory bottleneck. InvDiff-CGM employs invertible design principles at both the iterative solver and the neural backbone (U-Net) levels, yielding near-constant peak memory usage during training. This innovation is paired with a prior-informed multi-scale injector, enabling the efficient integration of environmental layout information and sparse measurements, thereby ensuring physical consistency and improved preservation of propagation details.

Problem Formulation and Theoretical Contributions

The CGM construction task is formalized as a spatial inverse problem: given a sparse set of path-loss measurements and environment layout priors within a grid, estimate a high-fidelity map of the propagation landscape. CGMs $\mathbf{X}\in\mathbb{R}^{H\times W}$ are reconstructed from sparse observations $\mathbf{Y} = \mathcal{A}(\mathbf{X})$ and priors $\mathbf{C}$ (e.g., building geometry, BS location).

InvDiff-CGM recasts this as an unrolled, deterministic denoising diffusion process parameterized by a neural network $\mu_{\theta}$, with $T$ iterative diffusion steps. Unlike conventional approaches, each denoising step and backbone block is engineered to be invertible, permitting on-the-fly computation of all intermediate states during training, thus obviating the need for activation caching. The theoretical implication is a transition from $\mathcal{O}(T)$ memory complexity to $\mathcal{O}(1)$, a requirement for online adaptation at the edge.

Physically consistent map construction is enforced via a deterministic update comprising (i) noise prediction informed by environmental priors through an invertible U-Net, (ii) orthogonal projection onto the measurement-consistent affine space, and (iii) a DDIM-based state update. The multi-scale prior injection further regularizes the generative process, ensuring sharp transitions and preservation of building-induced effects without excessive spatial smoothing.

Architectural Innovations

Invertible Iterative Sampler

The core iterative solver adopts a dual-channel structure, propagating states via invertible coupling transformations. At each step $t$, the update produces

$$(\mathbf{\hat{X}_{t-1}}, \mathbf{H}_{t-1}) = \text{Invert}(\mathbf{\hat{X}_{t}}, \mathbf{H}_{t})$$

where the design preserves strict mathematical invertibility. This enables all intermediate variables required for gradient computation to be deterministically reconstructed, minimizing memory footprint irrespective of the number of diffusion steps.

Invertible U-Net Noise Predictor

The noise estimation module implements a U-Net constructed from reversible blocks: all down-sampling and up-sampling stages are composed of modules with invertible coupling layers. Non-invertible operations, such as skip-connections and resolution changes, are confined to minimal segments of the architecture. This reduces memory growth with network depth while maintaining expressiveness.

Multi-Scale Prior-Informed Injector

To bridge the gap between learned features and physical priors, the architecture integrates fused features at multiple scales within the U-Net. Sparse measurement back-projections and environmental priors are concatenated and injected at each resolution level, both in the encoder and decoder pathways. This ensures that physical constraints and geometric context are reinforced throughout the network hierarchy, substantially improving high-frequency detail recovery and the delineation of occlusion boundaries.

Empirical Results

Experiments are conducted on the public RadioMap3DSeer dataset, employing $256\times256$ grid maps (1 m$\mathbf{Y} = \mathcal{A}(\mathbf{X})$0 resolution), with carrier operating at 3.5 GHz and diverse urban building layouts. InvDiff-CGM utilizes a Stable Diffusion v1.5 U-Net backbone, Adam optimization, and mixed-precision training, and is evaluated against six contemporary baselines, including RadioUNet, RME-GAN, and four diffusion-based models (e.g., RadioDiff, RadioDiff-$\mathbf{Y} = \mathcal{A}(\mathbf{X})$1, RadioDiff-Flux, RadioDiff-Inverse).

Quantitative Performance

InvDiff-CGM achieves a PSNR of 38.02 dB, SSIM of 0.9634, NMSE of 0.00196, and RMSE of 0.0123—establishing new state-of-the-art values across all metrics for the RadioMap3DSeer benchmark. Notably, InvDiff-CGM surpasses RadioDiff-$\mathbf{Y} = \mathcal{A}(\mathbf{X})$2 by 1.83 dB PSNR, demonstrating the efficacy of explicit prior fusion over implicit physics-informed regularization. Additionally, compared to all baselines, InvDiff-CGM reconstructs CGMs with sharper delineation of building boundaries and more accurate renderings of shadowed and occluded regions.

Memory Efficiency

Ablation studies reveal that, for three diffusion steps ($\mathbf{Y} = \mathcal{A}(\mathbf{X})$3), peak GPU memory is reduced by 85.59% (from 49 GB to 6.9 GB) relative to a non-invertible baseline. This substantial reduction enables end-to-end training and online adaptation on consumer-grade hardware, a configuration unattainable for existing diffusion-based methods, which require multi-GPU data parallelism. During online adaptation experiments, conventional baselines suffer from out-of-memory failures, whereas InvDiff-CGM completes adaptation runs successfully and efficiently.

Injector and Ablation Analysis

Removal of the multi-scale inverter module leads to significantly increased blurring and degradation at building-induced attenuation boundaries. This verifies that continuous, hierarchical fusion of physical priors and measurements is essential for detailed, physically plausible CGM recovery.

Runtime Trade-offs

While invertibility confers immense memory savings, it incurs a per-epoch runtime increase of approximately 38% during training, due to recomputation overhead. However, inference latency remains unaffected (0.15 s per map). The net effect is a favorable trade-off for edge-deployed adaptation scenarios, where memory, not compute, is the critical bottleneck.

Implications, Limitations, and Future Perspectives

InvDiff-CGM advances the design of resource-aware generative models for wireless map inference, yielding a practical framework for on-device CGM construction and rapid online adaptation even under severe hardware constraints. The strict invertibility paradigm is shown to be a superior alternative to checkpointing or layer-wise caching for deep, iterative generative networks in the wireless domain.

Two limitations are articulated: the added training latency due to recomputation may be restrictive for time-critical adaptation; and the underlying spatial model assumes flat terrain, limiting generalizability to mountainous or non-planar environments. Incorporating efficient attention mechanisms and terrain-aware propagation priors represents a promising avenue for further efficiency gains and modeling capacity.

Potential future directions include the extension of InvDiff-CGM toward dynamic, spatio-temporal radio map construction under non-stationary channel conditions (e.g., vehicular shadowing), leveraging real-time traffic states as additional priors and exploiting temporal invertibility for memory-efficient video-like diffusion processes.

Conclusion

InvDiff-CGM introduces a mathematically invertible, diffusion-based generative architecture for CGM construction, achieving both state-of-the-art accuracy and unprecedented training-time memory efficiency. By tightly integrating sparse measurement data, physical priors, and neural generative modeling through an invertible, multi-scale architecture, the framework creates a strong foundation for practical, high-fidelity propagation-aware services at the wireless network edge.