Advanced Training Options

XGBoost Memory Modes

By default, all edge embeddings are collected in CPU RAM before XGBoost training. For large datasets this can exceed available memory. The following modes progressively reduce memory pressure.

Choosing batch_size for batched modes:

batch_size × embedding_dimension × 4 bytes ≤ target GPU memory for data loading
Example:
  embedding_dimension = 256 (2 × 128 hidden + raw features)
  target GPU memory   = 8 GB
  batch_size ≤ 8 GB / (256 × 4 B) = 8,388,608 → use 4,194,304 or lower

A good starting point is batch_size: 65536 (≈ 64 MB for 256-dim embeddings) and double it until you reach the memory limit.

Config example — batched external memory:

xgb:
  num_boost_round: 512
  max_depth: 6
  learning_rate: 0.05
  batched: true
  batch_size: 131072
  extmem: true
  cache_host_ratio: 0.5   # pin 50% of data in host RAM between rounds

How Edge Embeddings Are Generated Before XGBoost Trains

Understanding when to use memmap and extmem requires knowing how the GNN produces the embedding matrix that XGBoost trains on.

memmap: false (default) — in-RAM accumulation:

The GNN processes training edges in mini-batches. After each batch’s forward pass, the resulting edge embedding tensor is appended to a list held in CPU RAM. When all batches complete, the list is concatenated into a single tensor and handed to XGBoost. No disk I/O occurs. This is fastest but requires enough system RAM to hold the full embedding matrix at once (all training edges × embedding dimension × 4 bytes).

memmap: true — incremental disk write (requires batched: true):

Each mini-batch’s embeddings are written immediately to a disk-backed memory-mapped file (edge_emb_train_rank0.mmap) as the GNN produces them. XGBoost then reads from this file in batches — either using QuantileDMatrix (all data in disk file, but processed batch-by-batch) or using ExtMemQuantileDMatrix when extmem: true is also set (XGBoost manages its own disk cache too). Peak RAM during training equals approximately one batch of embeddings rather than the full training set.

The .mmap file is temporary — it is deleted automatically after XGBoost finishes training.

Use memmap: true (requires batched: true) in RAM-constrained environments when system RAM cannot comfortably hold all training edge embeddings at once. extmem: true is a separate flag that enables XGBoost’s own ExtMemQuantileDMatrix disk cache; the two options are independent and can be combined.

Resuming Interrupted Training

Set checkpoint_dir in your config to warm-start GNN weights from a previous run. The GNN saves a single checkpoint (model_final.pt) to output_dir after all epochs complete. If training is interrupted before finishing, no partial checkpoint is available and training will restart from the beginning unless a previous run’s model_final.pt already exists in output_dir or checkpoint_dir.

paths:
  checkpoint_dir: /workspace/checkpoints/run1

To retrain only the XGBoost step using the saved GNN embedding (without re-running GNN training):

paths:
  checkpoint_dir: /workspace/checkpoints/run1
# in models[0]:
hyperparameters:
  gnn:
    skip_gnn_train: true       # skip GNN training; load checkpoint and re-run embedding extraction for XGBoost

Monitoring Training Progress

Real-Time Console Output

All log messages from rank 0 appear on stdout. Key lines to watch:

INFO — Config schema validation passed ✓
INFO — world_size=2  kind=GNN_XGBoost
INFO — Partitioning graph across 2 GPUs ...
INFO — Epoch 01, Loss: 0.4821  lr=0.005000  grad_norm=0.8432  time=12.3s
INFO — Epoch 02, Loss: 0.3409  lr=0.005000  grad_norm=0.7201  time=11.9s
...
INFO — [0]	train-auc:0.91234	train-aucpr:0.88765	train-logloss:0.18234	val-auc:0.87654	val-aucpr:0.84543	val-logloss:0.22876
INFO — [9]	train-auc:0.99123	train-aucpr:0.98765	train-logloss:0.01234	val-auc:0.94321	val-aucpr:0.93210	val-logloss:0.09876
...
INFO — Training complete!

Log Files

Each run writes a timestamped log file: output_dir/training_YYYYMMDD_HHMMSS.log

# Watch live
tail -f /workspace/output/training_*.log
# Search for errors
grep -n "ERROR\|Traceback\|FAILED" /workspace/output/training_*.log

Metrics File

After training completes, output_dir/metrics.json contains the full evaluation report:

{
  "pr_auc": 0.8998,
  "best_threshold": 0.7651,
  "best_threshold_raw": 0.7651,
  "calibrated": false,
  "threshold_0.5": {
    "accuracy": 0.9543,
    "precision": 0.6952,
    "recall": 0.8653,
    "f1": 0.7710
  },
  "threshold_best": {
    "accuracy": 0.9718,
    "precision": 0.9013,
    "recall": 0.7665,
    "f1": 0.8285
  },
  "test": {
    "pr_auc": 0.9243,
    "val_threshold": 0.7651,
    "threshold_val": {
      "accuracy": 0.9757,
      "precision": 0.8310,
      "recall": 0.8788,
      "f1": 0.8542
    },
    "threshold_0.5": {
      "accuracy": 0.9612,
      "precision": 0.7241,
      "recall": 0.8910,
      "f1": 0.7989
    }
  }
}

When prior_test is configured, the file gains six additional top-level fields:

{
  "prior_train": 0.10,
  "prior_test": 0.0003,
  "val_sim_positives": 3,
  "val_sim_total": 10000,
  "threshold_0.5_prod": {
    "accuracy": 0.9992,
    "precision": 0.4231,
    "recall": 0.8412,
    "f1": 0.5631
  },
  "threshold_best_prod": {
    "accuracy": 0.9996,
    "precision": 0.6841,
    "recall": 0.7109,
    "f1": 0.6972
  }
}

threshold_best_prod is the F1-optimal threshold evaluated on a class-ratio-simulated validation set (see Class Imbalance and Calibration). Use this value — not best_threshold — as the production decision threshold when your live fraud rate differs significantly from the training fraud rate.

Metric glossary:

Class Imbalance and Calibration

Production fraud rates (0.01%–0.2%) are typically 50–500× lower than training fraud rates (10–20%). If left unaddressed this gap produces overconfident probability scores and a misaligned decision threshold. Three parameters work together to close it:

scale_pos_weight — Training Balance

Upweights the fraud class during XGBoost training to compensate for imbalance in the training data. The default null computes it automatically as (1 − fraud_rate) / fraud_rate from the training labels. Override only if you want a custom ratio — for example to trade off precision vs. recall:

xgb:
  scale_pos_weight: null   # auto (recommended for most cases)
  # scale_pos_weight: 200  # manual: 200:1 non-fraud:fraud

base_score — XGBoost’s Initial Prediction

Sets XGBoost’s intercept before any trees are built. The default 0.5 is calibrated for balanced datasets. Set it to your training fraud rate so the trees only need to learn the residual signal:

xgb:
  base_score: 0.001   # 0.1% training fraud rate

This reduces the number of trees needed for convergence and improves probability calibration on the training distribution.

prior_test — Production Threshold Calibration

When your evaluation environment (OOT test set, live production) has a much lower fraud rate than training, set prior_test to calibrate scores for that environment. The pipeline downsamples validation-set positives so their rate matches prior_test, then finds the F1-optimal threshold on that simulated distribution — only the threshold is recalibrated. AUC is unaffected.

xgb:
  base_score: 0.10       # training fraud rate
  prior_test: 0.0003     # 0.03% expected production fraud rate

This is essential for meaningful precision/recall at production thresholds — without it, the decision threshold is calibrated on a 10%-fraud validation set and will be misaligned when the model is deployed against 0.03% production traffic, resulting in far too many false positives or false negatives depending on where the threshold falls.

Combined Example — Large Imbalance Gap

models:
  - kind: GNN_XGBoost
    hyperparameters:
      gnn:
        focal_gamma: 1.0    # focus GNN loss on hard fraud examples
        focal_alpha: 0.25   # upweight fraud class in focal loss
      xgb:
        base_score: 0.10    # training data: 10% fraud
        scale_pos_weight: null  # auto = 9x fraud-class weight (for 10% training fraud rate)
        prior_test: 0.0003  # production: 0.03% fraud → recalibrate threshold

Hyperparameter Tuning Guidance

When starting with a new dataset, tune parameters in this order — earlier parameters have larger impact and adjusting them makes later parameters more stable.

Step 1 — Validate basic learning:
Set epochs: 1, num_boost_round: 50. The goal is to confirm the pipeline runs end to end and produces a non-random PR-AUC (> 0.5). If AUC stays at 0.5, check data quality: label distribution, feature normalisation, index correctness.

Step 2 — Tune GNN epochs and learning rate:
Increase epochs until validation AUC plateaus. If AUC drops after plateau, you are overfitting — stop earlier. Learning rates between 0.001 and 0.01 work well for most datasets.

Step 3 — Tune XGBoost num_boost_round and learning_rate:
A smaller learning rate (0.01–0.05) with more rounds (500–1000) usually outperforms a large rate with fewer rounds. Set log_period: 5 to see per-round progress clearly. Important: when using a slow learning rate, increase early_stopping_rounds proportionally (e.g. early_stopping_rounds: 200) — at LR=0.01 the per-round improvement is very small and the default patience of 50 rounds may trigger early stopping before the model converges.

Step 4 — Tune XGBoost regularisation for overfitting:
If training PR-AUC is much higher than validation PR-AUC, increase min_child_weight (try 5, 10, 20), set gamma to 0.1–1.0, and reduce subsample to 0.7–0.8.

Step 5 — Try a deeper GNN or different encoder:
Once the baseline is stable, experiment with num_gnn_layers: 3 (and a corresponding third value in num_neighbors), or switch to encoder: gat.

Saving Node Embeddings for Hybrid Inference (EP Only)

Setting gnn.save_node_embeddings: true enables a third inference mode — hybrid inference — that pre-computes GNN node embeddings at training time and uses fast memmap lookup at inference time for any node that was in the training graph.

What Happens at the End of Training

After the main GNN+XGBoost training loop completes, two additional phases run automatically:

Phase A — Fine-tuning (optional): If final_fit_epochs > 0, the saved model_final.pt is reloaded and trained for final_fit_epochs more epochs on the combined train+val+test edge set, using learning_rate × final_fit_lr_scale as the LR. This improves embedding quality by exposing the model to the full labeled dataset before export.

Phase B — Stage 9: Node embedding export: A full GNN forward pass runs over all nodes of each type (both source and destination node types). Per-rank shard files are written:

output/node_embeddings/
├── {src_type}/
│   ├── rank=0_ids.pt     ← global node IDs for this rank's node partition
│   ├── rank=0_emb.pt     ← GNN output vectors for this rank's node partition
│   ├── rank=0_feat.pt    ← raw node feature vectors (same partition)
│   ├── rank=1_ids.pt
│   ├── rank=1_emb.pt
│   ├── rank=1_feat.pt
│   ├── embeddings.mmap   ← merged flat memmap (created at first client startup)
│   ├── features.mmap     ← merged raw feature memmap (created at first client startup)
│   └── shape.json        ← n_total, hidden_channels, dtype
├── {dst_type}/
│   └── ...
└── manifest.json         ← n_raw_features, world_size, emb_dtype per node type

Inference Routing

At inference time, each edge is routed based on whether its src and dst node IDs fall within the training graph:

The per-path metrics in client_saved_emb.py output tell you what fraction of your inference traffic is hitting each path and the quality of each path independently.

When to Use Saved Node Embeddings

Use save_node_embeddings: true when:

• You need low-latency inference for most edges
• Your inference traffic is dominated by nodes seen during training (the known path)
• You want to decouple the GNN from the serving critical path

Use the standard client.py (no saved embeddings) when:

• You want the simplest possible serving setup
• Most inference nodes are novel (cold-start rate > 50%)
• Latency requirements allow the full GNN inference time