Hi everyone!

I have just submitted my new journal paper on using Generative Flow Networks (GFlowNets) to speed up radio propagation modeling.

• Preprint on arXiv
• Tutorial notebook
• GitHub repository

The problem and our solution

Traditional point-to-point ray tracing suffers from exponential computational complexity, scaling with the number of objects raised to the interaction order. To fix this bottleneck, we define path finding as a sequential decision process and trained a generative model to intelligently sample valid ray paths instead of relying on an exhaustive search.

This work extends previous work I presented at ICMLCN 2025, but with much better results and details. Specifically, the proposed model achieves speedups of up to 10x on GPU and 1000x on CPU while maintaining high coverage accuracy!

Comparison of the coverage map between the ground truth (upper left) and the prediction (upper right) using 20 samples. Lower left and right figures show the relative and log-relative differences (in dB) between the two coverage maps, as defined in the paper.

Improvements from previous model

While working on this project, I researched a lot about reinforcement learning and GFlowNets. Applying GFlowNets here meant traversing a tree rather than a generic directed graph, which led to a number of standard solutions not being applicable. However, a few of them led to positive outcomes:

• Sparse Rewards: Finding valid geometric paths is rare, leading to a massive sparse reward issue and model collapse. After exploring goal-oriented RL with no success, I solved this by introducing a successful experience replay buffer to capture and store rare valid paths.
• Exploration: Using a uniform exploratory policy (ε-greedy) turned out to slightly improve performance on higher-order paths (i.e., deeper trees).
• Action Masking: I applied a physics-based action masking strategy to filter out physically impossible paths before the model even considers them, drastically pruning the search space.
• Muon Optimizer: Finally, I recently tried the Muon optimizer instead of the traditional Adam I was always using, and noticed much better training performance and convergence speed.

ML framework and hardware

Everything was built using the JAX ecosystem (Equinox, Optax, and my own library DiffeRT). Sadly, sharing code isn’t super common in my specific research community, but I strongly believe open-sourcing research data can only benefit everyone. As a result, I put a lot of effort into making the code clean and well-documented.

I’m not an ML expert but a telecom researcher, and I performed these experiments entirely on my own using a single NVIDIA RTX 3070. FYI, training the three models (as shown in the tutorial) takes about 3 hours on my computer. It might not be ready to completely replace exhaustive ray tracing just yet, but the results are really promising.

I’m very happy to receive questions, comments, or criticisms about this work. I hope you like it! 🙂