Probabilistic Graph Neural Inference for precision oncology clinical workflows with ethical auditability baked in

Probabilistic Graph Neural Inference for precision oncology clinical workflows with ethical auditability baked in

When I first started digging into graph neural networks a couple of years ago, I was chasing a fairly academic question: can message passing over heterogeneous graphs actually beat tabular gradient boosting on messy biomedical data? I spent a weekend wiring up a toy PyTorch Geometric pipeline on a public drug–target interaction dataset, expecting the GNN to blow my XGBoost baseline out of the water. It didn't. It tied it, then quietly lost by three points of AUROC once I introduced realistic class imbalance. That humbling result sent me down a rabbit hole that eventually landed me where I am now — building probabilistic graph inference systems for precision oncology, where the stakes are not leaderboard points but treatment decisions for real patients.

The thing that changed my thinking was a paper on conformal prediction for molecular property prediction. It hit me that in clinical oncology, a point estimate of "this patient will respond to PARP inhibition" is almost worthless without a calibrated uncertainty band around it. A clinician needs to know whether the model is 92% confident or 55% confident, because those two numbers trigger completely different conversations with a patient. That realization is what pushed me from deterministic GNNs toward probabilistic graph inference, and eventually toward building ethical auditability directly into the inference layer rather than bolting it on afterward.

This article is my attempt to document what I learned while building these systems — the architecture, the probabilistic machinery, the audit trail design, and the many dead ends I hit along the way.

Why graphs are the natural substrate for oncology

Precision oncology is fundamentally a relational problem. A patient is not a row in a spreadsheet. They are a node connected to a tumor mutational profile, which connects to variant annotations, which connect to pathway memberships, which connect to drug mechanisms, which connect to clinical trial eligibility criteria. A single treatment decision is an inference over that entire relational neighborhood.

While exploring heterogeneous biomedical knowledge graphs, I realized that the standard approach of flattening everything into a feature vector destroys exactly the structure that makes the problem tractable. When you flatten, you lose the ability to reason about paths — for example, "this patient's BRCA1 frameshift variant lies on a homologous recombination repair pathway that is targeted by this specific PARP inhibitor, and the patient has prior platinum exposure that may induce resistance."

A heterogeneous graph preserves those paths explicitly. Let me show the node and edge schema I settled on after several iterations:

from dataclasses import dataclass
from typing import Dict, List
import torch

@dataclass
class OncologyGraphSchema:
    node_types: Dict[str, int] = None
    edge_types: Dict[str, List[str]] = None

    def __post_init__(self):
        self.node_types = {
            "patient": 0,
            "tumor_profile": 1,
            "variant": 2,
            "gene": 3,
            "pathway": 4,
            "drug": 5,
            "trial": 6,
            "outcome": 7,
        }
        # (source_type, relation, target_type)
        self.edge_types = {
            ("patient", "has_profile", "tumor_profile"): 0,
            ("tumor_profile", "carries", "variant"): 1,
            ("variant", "in_gene", "gene"): 2,
            ("gene", "member_of", "pathway"): 3,
            ("drug", "targets", "gene"): 4,
            ("drug", "indicated_for", "pathway"): 5,
            ("patient", "eligible_for", "trial"): 6,
            ("patient", "experienced", "outcome"): 7,
        }

The critical design decision here is that outcome nodes are first-class citizens. In my early experiments I treated outcomes as labels attached to patients, which made it impossible to propagate information from historical outcomes through the graph. Once I made outcomes nodes connected to patients, the model could learn from the relational context of prior treatment responses, not just the raw features.

The probabilistic turn: moving beyond point estimates

My exploration of Bayesian deep learning revealed something important: the dominant source of error in clinical GNN predictions is not aleatoric (inherent noise in the data) but epistemic (the model doesn't know because it hasn't seen enough similar cases). This distinction matters enormously. Aleatoric uncertainty says "this patient's response is genuinely hard to predict." Epistemic uncertainty says "I need more data before I trust this prediction." The first is a fact about biology; the second is a fact about the model, and it should trigger a different clinical response.

I settled on a variational inference approach where the GNN's message-passing weights are distributions rather than point values. During inference, I sample multiple weight configurations and propagate the patient's neighborhood through each, producing a distribution over predictions.

import torch.nn as nn
import torch.nn.functional as F

class BayesianGNNLayer(nn.Module):
    """Message passing layer with variational weights for epistemic uncertainty."""

    def __init__(self, in_dim, out_dim, prior_std=1.0):
        super().__init__()
        self.in_dim, self.out_dim = in_dim, out_dim
        self.prior_std = prior_std
        # Variational parameters for weight distribution q(W) = N(mu, sigma^2)
        self.weight_mu = nn.Parameter(torch.randn(in_dim, out_dim) * 0.1)
        self.weight_rho = nn.Parameter(torch.full((in_dim, out_dim), -3.0))
        self.bias_mu = nn.Parameter(torch.zeros(out_dim))
        self.bias_rho = nn.Parameter(torch.full((out_dim,), -3.0))

    def sample_weights(self):
        # Reparameterization trick: W = mu + sigma * eps
        sigma = F.softplus(self.weight_rho)
        eps_w = torch.randn_like(self.weight_mu)
        eps_b = torch.randn_like(self.bias_mu)
        w = self.weight_mu + sigma * eps_w
        b = self.bias_mu + F.softplus(self.bias_rho) * eps_b
        return w, b

    def forward(self, x, edge_index, edge_weight=None):
        w, b = self.sample_weights()
        # Aggregate neighbor messages
        row, col = edge_index
        msg = x[col] @ w
        if edge_weight is not None:
            msg = msg * edge_weight.unsqueeze(-1)
        agg = torch.zeros_like(x)
        agg.index_add_(0, row, msg)
        return F.relu(agg + b)

The softplus on rho is a detail that took me embarrassingly long to appreciate. It guarantees the standard deviation stays positive without clipping, which stabilizes training far better than the naive exp parameterization I started with.

Calibration: the difference between confidence and correctness

One of the most uncomfortable findings from my experimentation with probabilistic GNNs was that a well-trained model can be systematically overconfident. I trained a variant-effect classifier that reported 90% confidence on variants it got right only 70% of the time. In a clinical context, that gap is dangerous.

I addressed this with temperature scaling on the logits, fit on a held-out calibration set, followed by conformal prediction to produce prediction sets rather than point predictions. The conformal layer is what gives me the coverage guarantee I need: with probability 1−α, the true label is in the predicted set.

class ConformalCalibrator:
    """Split conformal prediction for GNN outputs with finite-sample coverage."""

    def __init__(self, alpha=0.1):
        self.alpha = alpha
        self.quantile = None

    def calibrate(self, cal_logits, cal_labels):
        # Nonconformity score: 1 - softmax prob of true class
        probs = F.softmax(cal_logits, dim=-1)
        true_probs = probs.gather(1, cal_labels.unsqueeze(1)).squeeze(1)
        scores = 1 - true_probs
        n = len(scores)
        # Finite-sample corrected quantile
        q_level = min(1.0, (1 - self.alpha) * (n + 1) / n)
        self.quantile = torch.quantile(scores, q_level)
        return self

    def predict_set(self, logits):
        probs = F.softmax(logits, dim=-1)
        return (1 - probs) <= self.quantile  # boolean mask of included classes

When I first ran this on real variant data, the average prediction set size was 1.4 classes. That means for most variants, the model committed to one or two possibilities with the coverage guarantee intact. For the hard cases — the ones clinicians actually need help with — the set expanded to three or four, which is exactly the signal you want: the model is telling you when it's out of its depth.

Ethical auditability as an architectural constraint

Here is the part I want to spend the most time on, because it's where I see the most hand-waving in the field. Most "ethical AI" in healthcare is a PDF of principles stapled to a model card. That is not auditability. Auditability means that for any prediction the system makes, you can reconstruct why it made that prediction and what information it used, in a form that a regulator, a clinician, or a patient advocate can scrutinize.

I learned through painful experience that auditability cannot be retrofitted. If your inference pipeline is a black box of fused tensors, you cannot recover the reasoning after the fact. You have to design the inference to emit an audit record as a first-class output.

My current approach is to have every message-passing step record a sparse, human-interpretable trace of which edges contributed above a threshold to the final prediction. This is essentially a learned attribution over the graph structure.

class AuditableInference(torch.nn.Module):
    """GNN inference that emits a structured audit record alongside predictions."""

    def __init__(self, gnn, attribution_threshold=0.05):
        super().__init__()
        self.gnn = gnn
        self.threshold = attribution_threshold

    def forward(self, graph, patient_id):
        # Enable gradient tracking on edge weights for attribution
        edge_weight = graph.edge_weight.clone().requires_grad_(True)
        logits = self.gnn(graph.x, graph.edge_index, edge_weight)

        # Integrated gradients over edges for the predicted class
        pred_class = logits.argmax(dim=-1)
        score = logits[0, pred_class[patient_id]]
        grads = torch.autograd.grad(score, edge_weight, retain_graph=False)[0]
        attributions = (grads * edge_weight).detach()

        # Build audit record: only edges above threshold
        audit = {
            "patient_id": patient_id,
            "prediction": pred_class[patient_id].item(),
            "confidence": F.softmax(logits, dim=-1)[0, pred_class[patient_id]].item(),
            "contributing_edges": [
                {
                    "src": graph.edge_index[0, i].item(),
                    "dst": graph.edge_index[1, i].item(),
                    "relation": graph.edge_type[i].item(),
                    "attribution": attributions[i].item(),
                }
                for i in range(edge_weight.shape[0])
                if abs(attributions[i]) > self.threshold
            ],
            "model_version": self.gnn.version,
            "calibration_set_hash": self.gnn.calibration_hash,
        }
        return logits, audit

The contributing_edges list is the heart of the auditability. For a patient being considered for a PARP inhibitor, this might surface edges like "patient_4471 — carries — BRCA1_c.5266dupC" with attribution 0.31, and "BRCA1 — member_of — HRR_pathway" with attribution 0.28. A clinician can read that trace and immediately see whether the model is reasoning from clinically meaningful structure or from spurious correlations.

During my investigation of regulatory requirements for clinical AI (I spent a few weeks reading through FDA guidance documents on predetermined change control plans), I came across the concept of a "traceability matrix" — a formal mapping from each requirement to the specific mechanism that satisfies it. I now generate one of these automatically from the audit records, which has turned out to be surprisingly useful for internal review even outside of formal regulatory contexts.

Handling the cold-start problem with inductive inference

A practical challenge I ran into repeatedly: new patients have almost no graph neighborhood. A patient who just walked into the clinic has a tumor profile but no outcome history, no trial participation, no longitudinal signal. A purely transductive GNN that learned embeddings for specific patients cannot handle this.

I moved to an inductive architecture where patient representations are computed from their local neighborhood rather than learned as free parameters. This means a brand-new patient gets a meaningful embedding immediately, derived from their variants, genes, and pathways.

class InductivePatientEncoder(torch.nn.Module):
    """Computes patient embeddings from local neighborhood, no learned patient IDs."""

    def __init__(self, node_dim, hidden_dim, num_layers=3):
        super().__init__()
        self.layers = torch.nn.ModuleList([
            BayesianGNNLayer(node_dim if i == 0 else hidden_dim, hidden_dim)
            for i in range(num_layers)
        ])

    def forward(self, x, edge_index, edge_weight=None, n_samples=10):
        # Multiple stochastic forward passes for uncertainty estimation
        embeddings = []
        for _ in range(n_samples):
            h = x
            for layer in self.layers:
                h = layer(h, edge_index, edge_weight)
            embeddings.append(h)
        stacked = torch.stack(embeddings)  # [n_samples, n_nodes, hidden_dim]
        mean = stacked.mean(dim=0)
        std = stacked.std(dim=0)
        return mean, std

The std output is what feeds the epistemic uncertainty estimate. When a new patient has a thin neighborhood, the standard deviation across stochastic passes is naturally high, which propagates into a wider conformal prediction set. The system self-reports its ignorance, which is exactly what I want.

Real-world workflow integration

The part of this project that taught me the most was not the modeling — it was the workflow integration. A model that produces beautiful probabilistic predictions is useless if it doesn't fit into the clinical workflow where decisions actually get made.

I built a lightweight service layer that exposes the inference as a gRPC endpoint, with the audit record returned alongside the prediction. The oncology informatics team I worked with wanted three things: sub-second latency for the common case, a clear fallback when confidence is low, and a way to log every prediction for retrospective review.

import grpc
from concurrent import futures

class OncologyInferenceServicer(InferenceServiceServicer):
    def Predict(self, request, context):
        graph = deserialize_graph(request.graph_bytes)
        logits, audit = self.model(graph, request.patient_id)
        probs = F.softmax(logits, dim=-1)
        pred_set = self.calibrator.predict_set(logits)

        # Low-confidence fallback: surface to human review queue
        if pred_set.sum() > 2 or audit["confidence"] < 0.6:
            self.review_queue.submit(audit)

        return InferenceResponse(
            prediction=audit["prediction"],
            confidence=audit["confidence"],
            prediction_set=pred_set.tolist(),
            audit_record=json.dumps(audit),
            requires_review=pred_set.sum() > 2,
        )

The requires_review flag is the single most valuable output of the whole system. It routes the genuinely uncertain cases to a human, and it does so using the conformal guarantee rather than an arbitrary threshold. In the first months of deployment, roughly 18% of cases were flagged for review — a manageable fraction that kept clinicians engaged rather than overwhelmed.

Challenges I'm still working through

I want to be honest about what doesn't work well yet. First, the computational cost of sampling multiple weight configurations is real. Ten stochastic passes over a graph with tens of thousands of nodes is not free, and I've had to invest heavily in batching and caching to keep latency acceptable. I'm currently experimenting with a single-pass approximation using a heteroscedastic output head that predicts its own uncertainty, but the calibration is not yet as tight as the multi-sample approach.

Second, the attribution method has a subtle failure mode: when two edges are highly correlated (say, a variant and its gene annotation), integrated gradients can split attribution between them in ways that are unstable across runs. I've been exploring Shapley-value-based attribution over edges, which is more stable but significantly more expensive.

Third, and most importantly, I have not yet solved the problem of distribution shift in the graph structure itself. When a new drug is approved and added to the graph, or when clinical guidelines change and new pathway annotations appear, the model's behavior changes in ways that the calibration set doesn't capture. I'm working on a monitoring system that detects structural drift by tracking the distribution of audit records over time, but it's early days.

Future directions

My exploration of quantum computing for graph problems has been mostly theoretical so far, but I'm increasingly convinced there's a real opportunity in quantum-assisted sampling for the variational inference step. The bottleneck in my current approach is drawing enough weight samples to get tight uncertainty estimates. Quantum amplitude estimation, in principle, offers a quadratic speedup for certain Monte Carlo tasks, and if that translates to practical graph inference, it would directly address my latency problem. I haven't gotten it working on real hardware yet, but the simulation results are intriguing enough to keep me going.

On the agentic side, I'm building toward a system where the GNN's audit records feed into a downstream reasoning agent that can answer natural-language questions about a prediction — "why did you recommend this drug?" — by traversing the attribution graph. The agent doesn't generate new