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

TL;DR

The paper tackles efficient HIV testing under partial observability in networked referrals by formulating Sequential Acting on Partially Observed Graphs (SAPOG) and introducing Policy-Embedded Graph Expansion (PEGE) coupled with Dynamics-Driven Branching (DDB). PEGE embeds a generative graph expansion model into the decision loop, sampling multiple expansions from the frontier to inform action selection via an oracle (Gittins index-based) under a forest-structured network; DDB provides data-scarce, diffusion-GPR-based generation of neighbor dynamics while preserving forest structure. The authors formalize SAPOG, analyze why single-completion strategies fail, and demonstrate that the PEGE framework with DDB yields consistent improvements over multiple baselines on real HIV networks, achieving substantial gains in discounted reward and additional HIV detections at practical testing budgets. The work integrates interdisciplinary collaboration with WHO and Wits to enable deployment in South Africa, advancing SDG 3.3 by accelerating diagnosis, treatment initiation, and prevention efforts through intelligent testing sequences in resource-constrained settings. The approach also offers a generalizable framework for other problems with incrementally revealed graphs and limited data, where robust decision-making must contend with uncertain topology near the exploration frontier.

Abstract

HIV is a retrovirus that attacks the human immune system and can lead to death without proper treatment. In collaboration with the WHO and Wits University, 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 existing baselines (e.g., 13% improvement in discounted reward and 9% more HIV detections with 25% of the population tested) and explore key tradeoffs that drive decision quality.

Figures (6)

• Figure 1: Photographs of our HIV self-test kit distribution in South Africa.
• Figure 2: An illustration of a SAPOG instance at time $t$. Solid nodes and edges represent the induced subgraph $\mathcal{G}^{(t)} = \mathcal{G}[\mathbf{V}^{(t)}]$. Grey nodes represent nodes that have already been acted on, i.e., $\mathbf{V}^{(t)}_\mathbf{Y}$. Green nodes represent frontier nodes, which are nodes that have been discovered but not acted on, and thus have an unrevealed status. These frontier nodes make up the action space $\mathbf{V}^{(t)} \setminus\mathbf{V}^{(t)}_\mathbf{Y}$. Dashed nodes and edges represent unrevealed nodes, i.e., $\mathbf{V} \setminus \mathbf{V}^{(t)}$ and their unrevealed incident edges. (1) Agent evaluates the state $s^{(t)}$. (2) Agent acts on the red node $a \in \mathbf{V}^{(t)} \setminus\mathbf{V}_\mathbf{Y}^{(t)}$. (3) Agent observes the label of node $a$ and reveals its previously undiscovered neighbors in the next state.
• Figure 3: Workflow Diagram for PEGE. Grey nodes represent $\mathbf{V}^{(t)}_{\mathbf{Y}}$ while green nodes represent frontier nodes, i.e., $\mathbf{V}^{(t)} \setminus\mathbf{V}^{(t)}_\mathbf{Y}$. The red node represents an action. Blue nodes and edges represent graph expansions generated by $\mathcal{M}$ that are treated as true observed nodes and covariates by the oracle. 1) The state $s^{(t)}$ is observed. 2) We sample $k$ graph expansions up to depth $d$ from frontier nodes. In our case, we use a diffusion-driven model for $\mathcal{M}$, as described in Section \ref{['application']}. In this example, $k = 3$ and $d = 1$. 3) The oracle solves each of the $k$ instances separately and the resulting action evaluations are aggregated. 4) The agent takes an action $v$. 5) Reward, the label of $v$, and new nodes adjacent to $v$ and their covariates are observed as we transition to the next state.
• Figure 4: Overview of the ICPSR HIV sexual contact network dataset after forest projection (root nodes shown in red). Nodes in the graph represent people and edges represent sexual contacts.
• Figure 5: Performance of our method (PEGE + DDB) compared to several baselines and the FOG upper bound, as measured by cumulative discounted reward.

Theorems & Definitions (3)

• Definition 1: Sequential Acting on Partially Observed Graphs (SAPOG)
• Proposition 1: Non-Optimality of Maximum-Likelihood Graph Completion
• proof