Most fraud teams we meet have a very good tabular model and a nagging feeling that it is missing something. It is. A gradient-boosted tree scores one transaction at a time from that transaction's own features, plus whatever aggregates someone remembered to engineer: card velocity in the last hour, device count in the last day. What it cannot see is structure - that this brand-new account shares a device fingerprint with four accounts charged back last month, or that a cluster of merchants settle into the same bank account through three hops of intermediaries.
That structure is a graph, and graph neural networks are how you learn from it. On the TensorFlow stack, the library is TF-GNN (tensorflow_gnn), which builds on Keras and tf.data and slots into the same training and serving pipeline you already run. This tutorial builds a fraud-detection GNN end to end: schema, sampling, model, training, evaluation, and the parts that break in production.
When a GNN is worth it
Be honest about this before you spend a quarter on it. A GNN earns its complexity when:
- The label depends on neighbours. Fraud rings, money-laundering layering, collusive reviews, seller-buyer rings in marketplaces.
- The graph is dense enough to carry signal. If the average node has one edge, you have a table with extra steps.
- Relationships are hard to flatten. You can hand-engineer "number of chargebacks among accounts sharing this device", but not the two- and three-hop versions of it, and not for every relation type.
A GNN is usually not the answer when your signal is dominated by per-transaction attributes (amount, MCC, hour of day), or when you have fewer than a few thousand labelled positives. Our honest advice to several clients has been: fix the feature store first, revisit graphs next year. Where we do recommend it, we almost always ship the GNN as a second score alongside the existing model rather than a replacement - see the ensembling note below.
Install
pip install tensorflow tensorflow_gnn
TF-GNN tracks TensorFlow releases fairly closely; pin both in your lockfile, and check the TF-GNN documentation for the compatible pairing before upgrading. If you are mid-migration to Keras 3, note that TF-GNN's modelling API is built on Keras layers and you may need TF_USE_LEGACY_KERAS=1 depending on your versions - the same class of pinning problem we covered in Upgrading to TensorFlow 2.21.
Step 1: define the graph schema
TF-GNN is explicitly heterogeneous: multiple node sets and edge sets, each with its own features. For card fraud, a workable starting schema is accounts, devices, and merchants, with transaction edges.
import tensorflow as tf
import tensorflow_gnn as tfgnn
schema = tfgnn.parse_schema("""
node_sets {
key: "account"
value {
features {
key: "dense"
value { dtype: DT_FLOAT shape { dim { size: 16 } } }
}
features {
key: "label"
value { dtype: DT_INT32 }
}
}
}
node_sets {
key: "device"
value {
features {
key: "dense"
value { dtype: DT_FLOAT shape { dim { size: 8 } } }
}
}
}
node_sets {
key: "merchant"
value {
features {
key: "dense"
value { dtype: DT_FLOAT shape { dim { size: 12 } } }
}
}
}
edge_sets {
key: "uses_device"
value {
source: "account"
target: "device"
features {
key: "recency_days"
value { dtype: DT_FLOAT }
}
}
}
edge_sets {
key: "paid_merchant"
value {
source: "account"
target: "merchant"
features {
key: "amount_log"
value { dtype: DT_FLOAT }
}
}
}
""")
graph_spec = tfgnn.create_graph_spec_from_schema_pb(schema)
Two rules that save weeks of rework:
- Every raw feature gets normalised before it enters the graph. Log-scale amounts, clip recency, hash high-cardinality IDs to embeddings. A GNN will not rescue an un-normalised feature.
- Do not put the label anywhere except the node set you are classifying, and never leak an aggregate that was computed using the label.
Step 2: sample subgraphs
You do not train on the whole graph. You train on subgraph samples: for each labelled seed node, pull its neighbourhood out to two or three hops, cap the fan-out per hop, and serialise it as a tf.train.Example. TF-GNN ships a sampler (tensorflow_gnn.experimental.sampler, with a Beam-based pipeline for large graphs) that does exactly this and writes TFRecords.
Fan-out is the knob that decides both cost and quality:
# conceptual: two hops, 20 neighbours at hop 1, 10 at hop 2
sampling_spec = {
"uses_device": [20, 10],
"paid_merchant": [20, 10],
}
In fraud graphs there are always hub nodes - a shared corporate proxy IP, a payment processor merchant ID - with millions of edges. Without a fan-out cap they dominate every sampled subgraph and the model learns nothing. Cap fan-out, and consider dropping degree outliers above a threshold entirely; they are infrastructure, not behaviour.
Read the TFRecords back with the graph spec:
def make_dataset(file_pattern, batch_size, shuffle):
ds = tf.data.TFRecordDataset(tf.io.gfile.glob(file_pattern))
if shuffle:
ds = ds.shuffle(10_000)
ds = ds.batch(batch_size)
ds = ds.map(lambda s: tfgnn.parse_example(graph_spec, s),
num_parallel_calls=tf.data.AUTOTUNE)
return ds.prefetch(tf.data.AUTOTUNE)
train_ds = make_dataset("gs://bucket/fraud/train-*.tfrecord", 128, True)
val_ds = make_dataset("gs://bucket/fraud/val-*.tfrecord", 128, False)
The tf.data discipline from tf.data Pipelines That Keep Your GPU Busy applies unchanged here - sampling is I/O heavy and it is easy to leave the accelerator idle.
Step 3: build the model
A TF-GNN model is three stages: map raw features into hidden states, run several rounds of message passing, then read out the seed node and classify it.
from tensorflow_gnn.models import mt_albis
HIDDEN = 128
def set_initial_states(node_set, *, node_set_name):
return tf.keras.layers.Dense(HIDDEN, activation="relu")(node_set["dense"])
def build_model(graph_spec, num_layers=2):
inputs = tf.keras.layers.Input(type_spec=graph_spec)
graph = inputs.merge_batch_to_components()
graph = tfgnn.keras.layers.MapFeatures(
node_sets_fn=set_initial_states
)(graph)
for i in range(num_layers):
graph = mt_albis.MtAlbisGraphUpdate(
units=HIDDEN,
message_dim=HIDDEN,
attention_type="none", # try "multi_head" once the baseline works
simple_conv_reduce_type="mean|sum",
normalization_type="layer",
next_state_type="residual",
edge_dropout_rate=0.1,
state_dropout_rate=0.1,
)(graph)
seed = tfgnn.keras.layers.ReadoutFirstNode(
node_set_name="account", feature_name=tfgnn.HIDDEN_STATE
)(graph)
logits = tf.keras.layers.Dense(1)(seed)
return tf.keras.Model(inputs, logits)
model = build_model(graph_spec)
MtAlbisGraphUpdate is TF-GNN's batteries-included convolution: it covers the common design space (sum/mean aggregation, optional attention, residual or gated state updates) with one configuration object, which is why we start there instead of hand-rolling a GraphSAGE or GAT layer.
Number of layers = number of hops the model can see. Two is the sweet spot for fraud; three occasionally helps; beyond that you hit over-smoothing, where every node's embedding converges to the graph average and AUC quietly drops.
Step 4: train with the class imbalance in mind
Fraud is rare - often well under 1% positive. Train on the imbalance rather than pretending it away:
model.compile(
optimizer=tf.keras.optimizers.Adam(1e-3),
loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
metrics=[
tf.keras.metrics.AUC(curve="PR", from_logits=True, name="auprc"),
tf.keras.metrics.Precision(thresholds=0.5, name="precision"),
tf.keras.metrics.Recall(thresholds=0.5, name="recall"),
],
)
def extract_label(graph):
y = tfgnn.keras.layers.Readout(
node_set_name="account", feature_name="label"
)(graph)
graph = graph.remove_features(node_sets={"account": ["label"]})
return graph, tf.cast(y, tf.float32)
model.fit(
train_ds.map(extract_label),
validation_data=val_ds.map(extract_label),
epochs=20,
callbacks=[
tf.keras.callbacks.EarlyStopping(monitor="val_auprc", mode="max",
patience=3, restore_best_weights=True),
tf.keras.callbacks.TensorBoard("./logs"),
],
)
Note remove_features: the label lives on the node set, so you must strip it from the input or the model trains on the answer. This is the single most common self-inflicted wound in GNN projects, and it produces a validation AUC of 0.999 that collapses on the first day of production traffic.
Other imbalance tactics that work: downsample negatives to roughly 10:1 at sampling time and correct the threshold afterwards; or use focal loss. Track AUPRC, not AUC-ROC - with 0.3% positives, ROC flatters everything.
Step 5: split by time, not at random
Random splits leak in graphs twice over: a fraud ring's members end up on both sides of the split, and future information flows backwards through edges. Do both of these:
- Split by time. Train on seeds labelled before date T, validate on T..T+2 weeks, test after that.
- Respect the as-of timestamp when sampling. Every edge and feature in a training subgraph must have existed at the seed's decision time. If your sampler cannot filter edges by timestamp, build that before you build the model.
A GNN that scores 0.85 AUPRC on a random split and 0.55 on a time split does not have a modelling problem. It has a leakage problem, and the time split is the true number.
Step 6: serve it
Two viable patterns:
Batch scoring (most fraud teams start here). Re-sample subgraphs nightly for all active accounts, score them in a batch job, write scores into the feature store, and let the real-time rules engine read a score that is at most a few hours old. Cheap, simple, and enough for account-level risk.
Real-time scoring. The model's export is a normal SavedModel and serves fine under TensorFlow Serving - our serving guide covers versioning and batching. The hard part is not inference; it is doing the k-hop sample inside your latency budget, which means a graph store that answers neighbourhood queries in single-digit milliseconds and a fan-out small enough to stay there. Budget most of the engineering for the sampler, not the model.
model.export("./serving/fraud_gnn/1")
A useful middle path: precompute and cache node embeddings for stable entities (devices, merchants) nightly, and at request time only sample the one-hop neighbourhood of the account. You get most of the structural signal at a fraction of the latency.
Step 7: explain the score
Fraud decisions get contested - by customers, by risk committees, and increasingly by regulators. "The graph model said so" does not survive that conversation. Two things we insist on:
- Subgraph evidence. When the model fires, return the sampled neighbourhood with edge weights or attention scores, so an analyst sees this account, this device, these three charged-back accounts. That narrative is what an investigator acts on.
- A model card and a monitoring plan - the practices in our explainability tutorial apply, plus graph-specific drift checks: node and edge counts per day, degree distribution, and share of seeds with an empty neighbourhood. A sampler outage shows up as "suddenly everything looks safe", which is the most dangerous failure mode there is.
Ensembling with what you already have
The realistic production shape is not GNN-instead-of-GBDT. It is: keep the tabular model, add the GNN's output (score plus the seed embedding) as extra features, and let the existing model arbitrate. This preserves your calibration and your incident playbooks, and it makes the GNN's contribution measurable - you can compute exactly how much recall it bought at a fixed alert budget. That number, not AUPRC, is what gets the project renewed.
What to expect
On the engagements we have run, a well-built two-hop GNN on a genuinely relational fraud problem adds somewhere in the range of a few points to a couple of dozen points of recall at a fixed false-positive budget, over a strong tabular baseline. The variance is enormous and it is driven by graph quality, not model architecture. Teams with clean entity resolution - one device fingerprint that actually means one device - get the big numbers. Teams whose identity graph is glued together from three legacy systems get the small ones, and their first project is entity resolution.
Plan for roughly: two to four weeks on schema and entity resolution, one week to get a training loop running, then most of the remaining time on the sampler, the time-correct splits, and the serving path.
Need a second opinion on whether graph ML is right for your problem? DeepAdvisors' senior TensorFlow consultants have built fraud, AML, and recommendation systems on the TensorFlow stack across a range of industries. We are happy to run a short feasibility review of your graph before you commit a team to it - get in touch, or see our TensorFlow ML consulting services.