# Graph machine learning Understand four graph-learning tasks through their data, labels, splits, model architectures, training objectives, evaluation protocols, saved models, and inference workflows. ## Tasks and model compatibility Use graph machine learning through train, infer, and visualize_model_architecture. Node classification, link prediction, and pretraining read the session's active single graph; graph classification reads a separate collection of graphs. These tasks predict node classes, score node pairs, predict graph-level labels, and produce reusable node embeddings, respectively. | task / method | Model and input | Objective and result | | --- | --- | --- | | node_classification | Required model: APPNP / GCN / GAT / GATV2 / GEN / GIN / SAGE / TRANSFORMER; single-graph node features and labels | Cross-entropy; node classes and softmax probabilities | | link_prediction | Required model: APPNP / GCN / GAT / GATV2 / GEN / GIN / SAGE / TRANSFORMER; single-graph structure | Binary cross-entropy on positive and negative edges; pair scores and connection probabilities | | graph_classification | Required model: APPNP / GCN / GAT / GATV2 / GEN / GIN / SAGE / TRANSFORMER; HIV / PCBA graph collections | Binary cross-entropy excluding missing labels; single-task or multi-task graph predictions | | pretrain / node2vec | No model required; single-graph structure, trained on CPU | Positive/negative random-walk context learning; saved node embeddings | | pretrain / contrastive_learning, graph_contrastive_learning, mask_autoencoding, graph_mask_autoencoding | Required model: APPNP / GCN / GAT / GATV2 / GEN / GIN / SAGE / TRANSFORMER; single-graph features and structure | These names currently share the same GNN edge-reconstruction training branch | Registered backbone names are APPNP / GCN / GAT / GATV2 / GEN / GIN / SAGE / TRANSFORMER; SAGE means GraphSAGE. pretrain requires method, while all three supervised tasks and GNN pretraining require model. The GNN pretraining names do not establish separate graph-augmentation, feature-masking or masked-autoencoder implementations; use the actual loss described below. ## Single-graph tensors, features, and labels Let N be the node count, E the stored edge count, F the input feature dimension, and D the embedding dimension. Models use GraphData's node order and edge_index directly, without converting the graph to text. Message passing uses a [2, E] edge index. Semantically undirected graphs receive a bidirectional copy for aggregation without modifying the active graph. Edge weights and edge attributes are currently not passed to the current GNN backbones. | Input | Shape and processing | | --- | --- | | node_feat | Converted to float32. [N] becomes [N, 1]; tensors with more than two dimensions are flattened per node to [N, F]. Existing features are not automatically standardized. | | Missing node_feat | Generate [in-degree, out-degree, total degree, constant 1] with shape [N, 4]; divide each column by the larger of its maximum and 1. | | node_label: [N] | Scalar class values are mapped to contiguous class indices. Missing labels are excluded from supervision. Training needs enough labeled nodes and at least two classes. | | node_label: [N, C] | Take argmax across classes; zero-sum rows are treated as unlabeled. This is a single-label classification path, not multi-label node classification. | | data.feature_source = embedding | Use embedding.npy from upstream_model_id, aligned by names in node_names.json; unmatched nodes in the new graph receive zero vectors. | data.feature_source defaults to graph. Setting it to embedding also requires an upstream_model_id in the same session scope, and that model must contain saved node embeddings. This reuses input features; it neither resumes the upstream GNN's weights nor restores optimizer state. Node-name changes affect alignment, so check naming consistency and missing matches before reuse. ## Node classification: labels, splits, and training Node classification maps [N, F] features through a GNN to [N, D] embeddings, then applies a linear head to produce [N, C] class logits. Each epoch performs a full-graph training forward pass and computes cross-entropy only on training nodes. After optimization, evaluation mode recomputes training and validation accuracy on the same graph with dropout disabled. - data.target_nodes can restrict supervision to named nodes; unrecognized or unlabeled nodes are not usable supervised examples. - If train_mask exists, training uses its intersection with labeled nodes. Validation prefers val_mask, falls back to test_mask, and otherwise uses the remaining labeled nodes. - Without train_mask, labeled nodes are shuffled using seed and split by train / (train + val), defaulting to 0.8 / 0.2. The compatibility field train_config.split.test is not used by this flow. - Random splits are local to the training run and are not written into graph masks. infer query.split reads the original graph's masks and cannot retrieve an unpersisted random split. Model selection uses val_acc. A checkpoint improves only when validation accuracy is strictly greater than best_val_acc + early_stop_min_delta; equal accuracy does not replace it for a lower training loss. Training restores the best backbone and head and reports train_loss, train_acc, val_acc, best_epoch, trained_epochs, and stopped_early. If test_mask served as validation, the run has no independent test evaluation; report that protocol explicitly. ## Link prediction: sampling, decoding, and metrics Link prediction takes positive examples from logical graph edges or named pairs in data.target_edges; at least two positive edges are required. Positives are shuffled with seed and split into training and validation sets. Validation negatives are unique, non-self pairs absent from the graph, sampled with seed + 11 and undirected equivalence where applicable. Attempts are bounded, so dense graphs can yield fewer negatives than requested. The backbone produces node embeddings z. A pair's raw score is the dot product of z[u] and z[v], converted to prob with sigmoid. Each epoch uses PyG negative_sampling, and the loss sums binary cross-entropy for positive and negative edges. This decoder is symmetric: u→v and v→u have the same score even on a directed graph. There is no separate relation-type or directional decoder. Each epoch computes val_auc on fixed validation positives and negatives and searches a threshold for validation accuracy. With many distinct scores, candidates are reduced to quantiles. Checkpoint selection primarily uses val_auc, breaking comparisons within min_delta using val_acc. The selected threshold is saved. AUC uses this module's rank-based implementation, which can differ from other libraries in how tied scores are handled. This flow splits the positives used for loss and metrics, but GNN message passing still uses the active graph's complete adjacency, including validation positives. Its metrics are therefore not equivalent to a benchmark that removes validation edges from training adjacency. Strict unseen-edge generalization experiments require a separate adjacency-isolation and negative-sampling protocol. ## Pretraining and embedding reuse Node2Vec samples random walks and learns node vectors using positive and negative context losses. It uses sparse parameters and SparseAdam on CPU; selecting a GPU does not move this branch's walk training to the GPU. walk_length, context_size, walks_per_node, num_negative_samples, p, and q control sampling. Its batch size comes from algo_config.batch_size, default 128, rather than graph classification's train_config.batch_size. Node2Vec defaults to walk_length=20, context_size=10, walks_per_node=10, num_negative_samples=1, and p=q=1; its trainer uses embedding_dim=128 when omitted. Examples set the dimension explicitly to avoid confusion between task defaults. This branch runs all epochs, ignores weight_decay and early-stopping options, and saves final embedding.npy without a resumable Node2Vec checkpoint. GNN pretraining reads graph features and the complete edge index, reconstructing positive edges with embedding dot products and applying binary cross-entropy to sampled negatives. contrastive_learning, graph_contrastive_learning, mask_autoencoding, and graph_mask_autoencoding all follow this flow; they have no separate view augmentation, masking ratio, or feature-reconstruction loss. It uses Adam and saves the final backbone and embeddings, without validation selection or early stopping. Pretraining infer looks up saved embeddings: query.nodes selects nodes present during training, or infer_config.limit returns an initial subset, default 200. Even when a GNN checkpoint exists, this inference path does not run the GNN on a new graph. Rerun pretraining to generate embeddings for an updated graph; to use saved vectors for node classification or link prediction, set data.feature_source=embedding and provide the upstream model_id. ## Molecular graph-classification data structure The unit of graph classification is an entire graph. HIV has one binary task; PCBA has 128. load_graph_classification_dataset reads a .pt file from the session upload directory with dataset_profile set to auto, hiv, or pcba, validates it, binds it to the session, and returns dataset_id. This collection is separate from the active single graph. load_graph_classification_graph can extract one sample for inspection without converting graph classification into node classification. | Field | Required structure | | --- | --- | | graphs | A nonempty list of PyG Data objects; each graph contains x, xe, edge_index, and y tensors and at least one node. | | x | An integer [N] vector of nonnegative node-category IDs within the node vocabulary. These are not global node identities or a two-dimensional continuous feature matrix. | | xe / edge_index | xe is an integer [E] vector of edge-category IDs; edge_index is an integer [2, E] tensor with valid endpoints. Categories must be covered by the edge vocabulary. | | y | Per-graph [T] or [1, T], normalized to float32 [1, T]; only 0, 1, and NaN are allowed, with infinities and other values rejected. T must be consistent across graphs. | | texts | At least three string lists for node, edge, and label text; these provide vocabularies and descriptions, not automatic language-model embeddings. | | split | split[0] contains train, valid (or val), and test graph-index lists with no duplicates or overlap and complete graph coverage. split[1] is prompt/few-shot metadata whose summary is retained without changing training sampling. | Graph-classification training uses the file's fixed splits and rejects train_config.split. It does not accept graph_version, method, upstream_model_id, data.target_nodes, data.target_edges, or embedding feature input. Supply data.dataset_id explicitly or use the bound dataset; other tasks cannot use that field. The loader computes a full-file SHA-256 fingerprint that inference checks for dataset identity; an identical filename alone is insufficient. ## Architecture, training loop, and evaluation Registered backbones include multilayer graph convolutions, attention aggregation and APPNP propagation. Common defaults are hidden_dim=64, num_layers=2 and dropout=0.2. GNN embedding_dim follows hidden_dim when omitted, so it is 64 when hidden_dim is unchanged. GAT, GATV2 and TRANSFORMER use gat_heads=2 with concat=False. APPNP applies multilayer linear prediction followed by personalized propagation, defaulting to appnp_k=10 and appnp_alpha=0.1. use_residual_connection and use_normalization default to false; normalization supports rms_norm, layer_norm or batch_norm, with pre or post placement. Graph classification begins with a trainable nn.Embedding of node categories, followed by the GNN, global_mean_pool, and a linear graph head. Batching concatenates nodes and edges and records graph membership in a batch vector: node embeddings are [ΣN, D], pooled features [B, D], and final logits [B, T]. train_config.batch_size defaults to 64 and num_workers to 0. Training batches shuffle; validation and test preserve order, and incomplete final batches are retained. Graph classification uses masked binary cross-entropy with logits, averaging only non-NaN labels. Batches with no observed labels skip optimization; an epoch with no valid training labels fails. HIV selects models by ROC-AUC. PCBA computes AP for each evaluable task and macro-averages those tasks, recording valid and skipped counts. Tasks with only one class or entirely missing labels do not contribute, so the result is not an unconditional mean over all 128 tasks. All three supervised tasks use Adam. early_stop_patience=0 disables early stopping; a positive integer sets consecutive non-improving epochs before stopping. early_stop_min_delta defaults to 0; neither value may be negative. Graph classification evaluates validation each epoch and updates the checkpoint only on strict improvement beyond min_delta. It restores the best weights and evaluates test once afterward. Neither pretraining branch implements this validation-selection and early-stopping flow. | Setting | Behavior and scope | | --- | --- | | epochs / lr / weight_decay | Defaults are 200 / 0.05 / 0.0. Examples explicitly set the learning rate without claiming optimal tuning. Node2Vec uses SparseAdam and ignores weight_decay. | | seed | Defaults to 42 and seeds Python, NumPy, PyTorch, and CUDA. Multi-graph DataLoaders also use seeded generators and worker initialization. | | algo_config.device | cpu / cuda / auto; requesting unavailable CUDA fails. auto uses CPU without GPUs and chooses among GPUs using free memory and available utilization data. Node2Vec itself remains on CPU. | | Metrics and progress | Training reports per-epoch loss and task metrics. Supervised tasks include best_epoch, trained_epochs, and stopped_early to distinguish scheduled epochs, actual epochs, and selected weights. | ## Tool calls: submit training and retrieve results The JSON below contains tool argument objects for an authenticated, session-bound MCP or agent tool environment. These examples demonstrate request structure, not measured performance or optimal hyperparameters. Load a suitable labeled or structured single graph before single-graph training; for molecular classification, upload and load the dataset first. Replace filenames, dataset_id, model_id, and node names with actual results and data names. Node classification: call train with active-graph features and validation-accuracy model selection. split controls the run's random partition when train_mask is absent. ```json { "task": "node_classification", "model": "GCN", "data": { "feature_source": "graph" }, "train_config": { "split": { "train": 0.8, "val": 0.2 }, "epochs": 100, "lr": 0.001, "weight_decay": 0.0005, "seed": 42, "early_stop_patience": 10, "early_stop_min_delta": 0.0001 }, "algo_config": { "hidden_dim": 64, "embedding_dim": 64, "num_layers": 2, "dropout": 0.2, "device": "auto" } } ``` Link prediction: call train with logical graph edges as positives. To restrict supervised positives, provide data.target_edges as arrays of two node names. ```json { "task": "link_prediction", "model": "SAGE", "data": { "feature_source": "graph" }, "train_config": { "epochs": 100, "lr": 0.001, "seed": 42, "early_stop_patience": 10 }, "algo_config": { "hidden_dim": 64, "embedding_dim": 64, "num_layers": 2, "device": "auto" } } ``` Node2Vec: call train with an explicit embedding dimension, random-walk settings, and this branch's batch size. ```json { "task": "pretrain", "method": "node2vec", "train_config": { "epochs": 30, "lr": 0.01, "seed": 42 }, "algo_config": { "embedding_dim": 64, "walk_length": 20, "context_size": 10, "walks_per_node": 10, "num_negative_samples": 1, "batch_size": 128, "p": 1.0, "q": 1.0, "device": "cpu" } } ``` Molecular graph classification: first call load_graph_classification_dataset. After loading completes, insert its dataset_id into the second object and call train. ```json { "filename": "hiv.pt", "dataset_profile": "hiv" } ``` ```json { "task": "graph_classification", "model": "GCN", "data": { "dataset_id": "DATASET_ID_FROM_LOAD" }, "train_config": { "epochs": 100, "lr": 0.001, "weight_decay": 0.0, "seed": 42, "early_stop_patience": 10, "early_stop_min_delta": 0.0001, "batch_size": 64, "num_workers": 0 }, "algo_config": { "node_embedding_dim": 64, "hidden_dim": 64, "embedding_dim": 64, "num_layers": 2, "dropout": 0.2, "device": "auto" } } ``` Direct MCP calls may initially return job_id for long-running work. Poll get_job_status; start with after_seq=0, then pass the returned progress_cursor as after_seq for incremental progress, continuing while progress_has_more=true. Completion is job.status=done, with the training result under job.result. Confirm result success=true, read model_id, then call infer. The agent's GraphMCPTools wrapper automatically waits for and resolves these jobs, normally delivering the final result. ```json { "job_id": "JOB_ID_FROM_TRAIN", "after_seq": 0, "progress_limit": 100 } ``` get_job_result(job_id=...) also retrieves final output. Successful lookup is not training completion: first confirm status=done, then check the inner training success and model_id. Large-graph algorithm output consists of summaries, previews and artifact references, not a graph-learning model_id. ## Inference queries and response structure infer dispatches using saved task metadata. Node classification and link prediction rebuild the backbone, load its checkpoint, and run on the current graph. Graph classification resolves and verifies its dataset. Pretraining returns stored node vectors. The returned model_id, task, and graph-version or dataset identifier identify the model and input behind the result. For node-name queries, predictions contain class information such as pred_class_idx and pred_label, plus true labels and correctness when available. Enable return_prob for class probabilities. result_limit and result_offset paginate details while evaluation covers all requested nodes. An empty query returns split evaluation summaries and sampled predictions; query.split requires an actual graph mask. ```json { "model_id": "NODE_MODEL_ID_FROM_TRAIN", "query": { "nodes": [ "node_a", "node_b" ] }, "infer_config": { "return_prob": true, "result_limit": 20, "result_offset": 0 } } ``` Link prediction accepts query.edges or source_node with candidate_nodes for ranking. Each prediction contains score, prob, pred, and exists_in_graph, using the saved validation threshold. Candidates may include existing edges, which callers can filter by exists_in_graph. An empty query searches nonedges: max_candidate_pairs bounds retained/scored candidates, but generation still iterates over node pairs. Prefer explicit candidates for large graphs. ```json { "model_id": "LINK_MODEL_ID_FROM_TRAIN", "query": { "source_node": "node_a", "candidate_nodes": [ "node_b", "node_c" ], "topk": 2 } } ``` Graph classification accepts query.split=train, val, or test, defaulting to test for an empty query, or explicit graph_indices; these options are mutually exclusive. HIV details include probability, predicted_label, and true_label. PCBA defaults to the five highest-probability tasks, configurable with top_k_tasks; return_all_task_probs=true includes complete task probabilities. evaluation covers the full query, count is the requested graph count, returned_count is the page size, and truncated indicates omitted details. ```json { "model_id": "GRAPH_MODEL_ID_FROM_TRAIN", "query": { "split": "test" }, "infer_config": { "batch_size": 128, "device": "auto", "result_limit": 20, "result_offset": 0, "threshold": 0.5 } } ``` Graph-classification inference accepts infer_config.device. Node-classification and link-prediction inference currently select a device automatically, so the same field should not be assumed to control them. For a pretrained model, the following call only reads saved node vectors, returning an embeddings array of node and embedding entries. ```json { "model_id": "PRETRAIN_MODEL_ID_FROM_TRAIN", "query": { "nodes": [ "node_a", "node_b" ] }, "infer_config": { "limit": 200 } } ``` ## Model files, version binding, and reproducibility Training outputs live in Graph Engine's model directory, defaulting to artifacts/graph_ml and configurable with GRAPH_ML_ARTIFACT_ROOT. Directories are isolated by runtime user/session scope, and model_id identifies one training output. Graph ML initializes lazily. Controlled workers write staging outputs, which Graph Engine publishes after validating that the job remains eligible. Receiving job_id does not mean the model is already available. | Task | Main files in a model directory | | --- | --- | | Node classification | checkpoint.pt: backbone and head; meta.json: architecture, feature source, label mapping, and graph_version; metrics.json and node_names.json. | | Link prediction | checkpoint.pt: backbone; meta.json also records threshold; metrics.json and node_names.json. | | Node2Vec pretraining | embedding.npy, node_names.json, meta.json, and metrics.json; no checkpoint.pt. | | GNN pretraining | checkpoint.pt: backbone; plus embedding.npy, node_names.json, meta.json, and metrics.json. | | Graph classification | checkpoint.pt: model weights; meta.json: dataset fingerprint, vocabulary sizes, task count, model configuration, and split sizes; metrics.json and per_task_metrics.json. | Single-graph train treats graph_version as an assertion about the active version, not a request to load history. Call load_graph_version first to train on an older version. Node-classification and link-prediction inference require the active version to match training by default. infer_config.allow_version_mismatch=true skips this version check but cannot resolve incompatible input dimensions, label semantics, or upstream embeddings. Pretrained-vector lookup checks saved model version without relying on the active graph; graph classification binds dataset_id and the full dataset fingerprint. For reproducible experiments, separately preserve the complete train request, source data and fingerprint, graph version, upstream model_id, dependency versions, and actual device. Metadata records architecture and results but does not preserve every training hyperparameter, RNG state, or optimizer state, and there is no public resume-training endpoint. seed controls major random sources but deterministic algorithms are not enforced; GPU kernels and environment changes can still affect results, so matching seed alone does not guarantee bitwise reproducibility. visualize_model_architecture accepts model_id and reads saved model configuration for architecture inspection. The visualization explains layers and computation, not predictive quality. When comparing runs, inspect best epoch, actual epochs, split protocol, validation/test independence, and PCBA's evaluable-task count alongside metrics. ## Capability boundaries and experimental interpretation - Graph classification uses node-category embeddings, a registered backbone (APPNP / GCN / GAT / GATV2 / GEN / GIN / SAGE / TRANSFORMER), fixed global mean pooling and a graph-level linear head. xe is loaded, validated and retained but does not enter the forward pass; metadata records uses_edge_features=false. This implementation does not establish use of molecular bond types. - GIN is registered; GINE and selectable graph pooling are not provided. The HIV/PCBA loader does not imply support for arbitrary OGB datasets. Multi-graph tasks require the specified .pt structure and fixed splits; train does not accept arbitrary new metrics or resplit the test set. - Single-graph supervised tasks and GNN pretraining use full-graph forward passes without neighbor-sampled minibatches. train_config.batch_size does not reduce their full-graph feature or activation memory; it applies to graph classification, while Node2Vec has its own batch-size field. - Two-dimensional node labels enter single-label classification through argmax. Molecular graph classification instead uses [B, T] independent binary tasks. Similar tensor shapes do not imply the same learning objective. - Link-prediction validation edges remain in message passing and its dot-product decoder is direction-invariant. Named GNN pretraining methods share edge reconstruction. Describe these boundaries alongside metrics for meaningful comparison with research baselines. - The current interface provides no automatic hyperparameter search, class-weighted loss, learning-rate scheduler, optimizer-state restoration, or cross-environment determinism guarantee. The examples explain request configurations and do not claim unperformed performance measurements.