# Policy-Embedded Graph Expansion: Networked HIV Testing with Diffusion-Driven Network Samples

arXiv: https://arxiv.org/abs/2601.16233 · HTML: https://mlsanity.com/paper/2601.16233/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 13 of its 44 references that are in the map.
- First posted on arXiv 20 Jan 2026.

## Builds on (most cited first, 13 of 13)

- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [End to end learning and optimization on graphs](https://mlsanity.com/paper/1905.13732/) (2019)
- [Clinical trial of an AI-augmented intervention for HIV prevention in youth experiencing homelessness](https://mlsanity.com/paper/2009.09559/) (2020)
- [Robustification of Online Graph Exploration Methods](https://mlsanity.com/paper/2112.05422/) (2021)
- [No Time to Observe: Adaptive Influence Maximization with Partial Feedback](https://mlsanity.com/paper/1609.00427/) (2016)
- [IM-META: Influence Maximization Using Node Metadata in Networks With Unknown Topology](https://mlsanity.com/paper/2106.02926/) (2021)
- [Adaptive Influence Maximization with Myopic Feedback](https://mlsanity.com/paper/1905.11663/) (2019)
- [DeepTrace: Learning to Optimize Contact Tracing in Epidemic Networks with Graph Neural Networks](https://mlsanity.com/paper/2211.00880/) (2022)
- [Selective Harvesting over Networks](https://mlsanity.com/paper/1703.05082/) (2017)
- [Neural Online Graph Exploration](https://mlsanity.com/paper/2012.03345/) (2020)
- [Dynamic group testing to control and monitor disease progression in a population](https://mlsanity.com/paper/2106.10765/) (2021)
- [Adaptive Frontier Exploration on Graphs with Applications to Network-Based Disease Testing](https://mlsanity.com/paper/2505.21671/) (2025)

## Built on by (most cited first, 3 of 3)

- [Latent Spherical Flow Policy for Reinforcement Learning with Combinatorial Actions](https://mlsanity.com/paper/2601.22211/) (2026)
- [Network-Based Interventions for HIV Prevention via Cascade-Aware Suppression of Transmission](https://mlsanity.com/paper/2605.20218/) (2026)
- [Generative Frontier Planning for Adaptive Peer-Referral Recruitment under Covariate-Dependent Arrivals](https://mlsanity.com/paper/2606.08360/) (2026)

## Abstract (arXiv)

HIV is a retrovirus that attacks the human immune system and can lead to death without proper treatment. In collaboration with the WHO and the University of Witwatersrand, we study how to improve the efficiency of HIV testing with the goal of eventual deployment, directly supporting progress toward UN Sustainable Development Goal 3.3. While prior work has demonstrated the promise of intelligent algorithms for sequential, network-based HIV testing, existing approaches rely on assumptions that are impractical in our real-world implementations. Here, we study sequential testing on incrementally revealed disease networks and introduce Policy-Embedded Graph Expansion (PEGE), a novel framework that directly embeds a generative distribution over graph expansions into the decision-making policy rather than attempting explicit topological reconstruction. We further propose Dynamics-Driven Branching (DDB), a diffusion-based graph expansion model that supports decision making in PEGE and is designed for data-limited settings where forest structures arise naturally, as in our real-world referral process. Experiments on real HIV transmission networks show that the combined approach (PEGE + DDB) consistently outperforms baselines (e.g., 17.3% improvement in discounted reward and 15.4% more HIV detections with 25% of the population tested) and explore key tradeoffs that drive solution quality.