> This page is for Kumo Engine.

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/sdgm/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/sdgm/_mcp/server.

# Temporal Hetero-GNN

> The architecture behind Kumo Engine's built-in model, the tasks it supports, and the research it builds on

Kumo Engine's built-in model is a temporal heterogeneous graph neural network (hetero-GNN) for relational data. It follows the relational deep learning formulation of [Fey et al. (2024)](https://arxiv.org/abs/2312.04615): each table row is a node, each primary key/foreign key link is a typed edge, and each prediction is computed from a subgraph sampled around the target entity, using only rows from before the prediction time. The same architecture handles classification, regression, forecasting, and link prediction.

It combines ideas from several papers, including GraphSAGE's neighbor sampling, R-GCN's relation-specific weights, PNA's multiple aggregators, and the ID-GNN and ContextGNN models for link prediction. The combined architecture has been tested on hundreds of datasets and runs in production at many organizations.

## Supported Task Types

Kumo reads the task from the predictive query, including its target, labels, and time horizon, and attaches the matching prediction head and loss.

| Task                                 | Output                                        | Example                                   |
| ------------------------------------ | --------------------------------------------- | ----------------------------------------- |
| Binary and multiclass classification | Class probabilities for each entity           | Whether a transaction is fraudulent       |
| Multilabel classification            | A probability for each label                  | Which categories a customer will buy from |
| Multilabel ranking                   | Labels in score order                         | A customer's interests, ranked            |
| Regression                           | One number for each entity                    | A customer's spend over the next month    |
| Forecasting                          | One number for each future time window        | Daily demand over the next week           |
| Static link prediction               | A score for each candidate link               | User-product links missing from the data  |
| Temporal link prediction             | A score for each candidate future interaction | The products a customer will buy next     |

Classification, regression, and forecasting are node tasks, where the output belongs to the entity at the root of the sampled subgraph. Link prediction scores pairs of a source entity and a candidate target. Both kinds can be temporal, in which case each training example has a prediction time, with inputs from rows before it and labels from the horizon after it. Static link prediction instead holds out a random subset of existing links for evaluation.

See [Predictive Queries](/sdgm/kumo-engine/user-code-guide/train-a-model/predictive-queries).

## Feature Encoders

Before message passing, each row is encoded into a fixed-width vector that becomes its node's initial representation. Encoders are defined per table and, within a table, per column. Kumo picks each column's encoder from its semantic type (stype) and statistics, then combines the column outputs. This per-stype design comes from PyTorch Frame ([Hu et al., 2024](https://arxiv.org/abs/2404.00776)).

| Column stype     | Encoders                                                                                                                                  |
| ---------------- | ----------------------------------------------------------------------------------------------------------------------------------------- |
| Numerical        | `Numerical` rescales the value. `MaxLogNumerical` and `MinLogNumerical` compress heavy tails with a log transform.                        |
| Categorical, ID  | `Index` learns one vector per category. `Hash` maps categories into a fixed number of buckets, bounding memory at the cost of collisions. |
| Multicategorical | `MultiCategorical` combines the vectors of every category in a cell.                                                                      |
| Text             | `TextEmbedding` averages pretrained word vectors.                                                                                         |
| Timestamp        | `Datetime` encodes calendar components: minute, hour, day of week, day of month, day of year, and year.                                   |
| Numerical vector | `NumericalList` takes fixed-length vectors, including embeddings computed outside Kumo.                                                   |

`Null` excludes a column from the features. A key column excluded this way still defines its edges in the graph. See [Encoder Selection](/sdgm/kumo-engine/user-code-guide/set-up-a-graph/encoder-selection).

Each sampled row is also encoded with the time between its timestamp and the prediction time, which lets the model tell a purchase made ten minutes before a transaction from one made a year earlier. `use_seq_id` optionally adds the order of the sampled events, and `prediction_time_encodings` adds the absolute prediction time.

## Model Architectures

### Heterogeneous GNN

The GNN runs on a subgraph sampled around the target entity with a fixed fanout per hop, as in GraphSAGE ([Hamilton et al., 2017](https://arxiv.org/abs/1706.02216)). Fanouts can be overridden per edge type, and [adaptive sampling](/sdgm/research/engineering-guides/adaptive-sampling-guide) shifts budget to later hops when earlier hops come up short.

Each table is a node type, and each foreign key defines an edge type in each direction. The GNN learns separate weights for every edge type, the construction introduced by R-GCN ([Schlichtkrull et al., 2018](https://arxiv.org/abs/1703.06103)) and provided in PyG by `HeteroConv` and `to_hetero` ([Fey and Lenssen, 2019](https://arxiv.org/abs/1903.02428)).

On the one side of a one-to-many relationship, a node aggregates many neighbors, as an account does with its transactions. Following PNA ([Corso et al., 2020](https://arxiv.org/abs/2004.05718)), Kumo combines several aggregators, such as sum, mean, and max, because each loses different information: mean discards neighborhood size, which sum preserves ([Xu et al., 2019](https://arxiv.org/abs/1810.00826)), and max keeps only the extreme. On the many side, each row receives the representation of the single row it references. A learned skip connection carries each node's own representation into the next layer.

Schematically, for a prediction at time $t$, one layer updates a node $v$ of type $\tau(v)$ as

$$
h_v^{(\ell+1)} = \sigma\Big( W_{\tau(v)}^{(\ell)}\, h_v^{(\ell)} \;+ \sum_{r \in \mathcal{R}(v)} \operatorname{AGG}_r \big\{ W_r^{(\ell)}\, h_u^{(\ell)} \;:\; u \in \mathcal{N}_r^{t}(v) \big\} \Big)
$$

where $\mathcal{R}(v)$ is the set of edge types that end at $v$'s table, $\mathcal{N}_r^{t}(v)$ is the set of $v$'s sampled neighbors under edge type $r$ that are visible at time $t$, and $\operatorname{AGG}_r$ combines the configured aggregators. Normalization, dropout, and the MLP layers before and after message passing are omitted.

Each layer extends the receptive field by one hop, and the number of layers equals the sampling depth. For a transaction $v$ scored for fraud, the first hop reaches its account $a$ and its merchant $m$, and the second reaches the account's earlier transactions and the merchant's other transactions.

![Left: transaction v at time t links to card account a and merchant m; the account's earlier transactions and the merchant's other transactions are sampled, and a transaction after t is excluded. Right: layer 1 aggregates each hop-2 group with sum, mean, and max under edge-type weights and adds a skip connection, and also updates v from a and m; layer 2 combines the layer-1 states of a, m, and v, and a head produces the prediction.](/sdgm/_fern-files/nvidia-sdgm.docs.buildwithfern.com/459f30e4e44655cdaebe6c18a9c2ba235282494931e9c370ef25d97c5e7b6f82/_dot_dot_/images/diagrams/hetero-gnn-aggregation.svg)

*A two-hop temporal subgraph for a transaction scored at time t, and two layers of message passing over it.*

The remaining layer choices follow the GNN design space of GraphGym ([You et al., 2020](https://arxiv.org/abs/2011.08843)): MLP layers before and after message passing, activation, normalization, and dropout. The model plan can search over these along with hidden width and the aggregator set.

### Graph Transformer

`GraphTransformerModelPlan` replaces hop-by-hop message passing with self-attention over every node in the sampled subgraph, letting any two sampled rows exchange information in one layer rather than through the nodes between them. Since attention ignores graph structure, each node also gets positional encodings. `hop` embeds the node's distance from the target entity. `random_gnn` adds the output of a small GNN run over random initial features, which captures the node's local structure. These correspond to the hop-distance and subgraph-GNN encodings in the Relational Graph Transformer ([Dwivedi et al., 2026](https://arxiv.org/abs/2505.10960)), and the use of random features to break symmetry between otherwise identical nodes goes back to [Sato et al. (2021)](https://arxiv.org/abs/2002.03155).

Unlike the GNN's, the transformer's depth is not tied to sampling depth; you set the number of layers and attention heads separately. It supports node tasks, including forecasting, but not link prediction. See [Relational Graph Transformers](/sdgm/research/concepts-models/relational-graph-transformers).

### Identity-Aware Message Passing

Without distinguishing node features, a standard GNN cannot tell a node in a triangle from a node in a square, because their computation trees are identical. ID-GNN ([You et al., 2021](https://arxiv.org/abs/2101.10320)) colors the node being embedded wherever it recurs in its own tree and gives messages from colored nodes separate weights, which makes the trees distinguishable. For link prediction, the source node is colored instead, so a candidate's representation depends on how it connects to that source. The `"pairwise"` and `"context"` link prediction modules are built on this.

![Node v1 sits in a triangle and node v2 in a square. A standard GNN produces identical three-layer computation trees for both. ID-GNN colors the root wherever it recurs: in v1's tree it recurs at depths 2 and 3, in v2's tree only at depth 2, so the trees and the representations differ.](/sdgm/_fern-files/nvidia-sdgm.docs.buildwithfern.com/814a313a320e03da7c0d57a963ad14d361ed529593ab217c1e0ecaeff257592a/_dot_dot_/images/diagrams/id-gnn.svg)

*After You et al. (2021), Figure 1. Nodes have no distinguishing features.*

Kumo applies the same idea when a classification or regression query involves several entities: it marks each entity's role and combines their representations in a joint prediction head.

## Link Prediction Variants

The main design choice in link prediction is how to represent candidate targets. A two-tower model embeds sources and targets independently and ranks pairs by similarity. Target embeddings can then be computed once, which scales to large catalogs, but a target looks the same to every source. A pair-wise model such as ID-GNN conditions each target on the source at the cost of one representation per pair, which in practice limits it to candidates inside the source's sampled subgraph. ContextGNN ([Yuan et al., 2025](https://arxiv.org/abs/2411.19513)) fuses the two, using pair-wise scores for candidates near the source and two-tower scores for the rest.

`plan.model_architecture.module` selects among these:

| Module          | Architecture                                                                                                                                                                                                               | Task support        |
| --------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------- |
| `"shallow_rhs"` | A GNN representation of the source is scored against a shallow embedding of each target, which is learned per target, built from its encoded columns, or both. This has the same form as ContextGNN's two-tower component. | Static and temporal |
| `"two_tower"`   | The GNN encodes the source's and the target's neighborhoods separately and ranks pairs by embedding similarity.                                                                                                            | Static and temporal |
| `"pairwise"`    | ID-GNN. Message passing conditioned on the source scores each candidate in the source's sampled subgraph.                                                                                                                  | Temporal            |
| `"context"`     | ContextGNN. `"pairwise"` scores for candidates in the sampled subgraph are fused with `"shallow_rhs"` scores over the full candidate set.                                                                                  | Temporal            |

For `"shallow_rhs"` and `"context"`, `target_embedding_mode` sets how target embeddings are built, and with it whether the model can score targets that were not in the training data. `"lookup"` learns a free vector per training target and cannot score new ones. `"feature"` and `"shallow_feature"` derive the vector from the target's encoded columns, `"feature"` through a deeper learned projection and `"shallow_feature"` through normalization alone. `"fusion"` adds the lookup and feature vectors. To score new targets, use a feature mode and set `handle_new_target_entities=True`.

![Entity tower: a heterogeneous GNN aggregates the source customer's transactions and their products into a vector h\_c. Target tower: each product's vector e\_p is the sum of an MLP over its features and its row in a lookup table, with no message passing. A dot product of h\_c and e\_p feeds a softmax over products.](/sdgm/_fern-files/nvidia-sdgm.docs.buildwithfern.com/ffccef8767126604c072b9d811ab35f95bf55dec0e84f602e0a21545c7dcacea/_dot_dot_/images/diagrams/shallow-rhs.svg)

*`"shallow_rhs"` with `target_embedding_mode="fusion"`.*

For learned embeddings, we recommend `"shallow_rhs"` with hybrid lookup and feature embeddings. `"two_tower"` is fully inductive: it computes both source and target embeddings from their features and graph neighborhoods.

See [Link Prediction and Embeddings](/sdgm/kumo-engine/user-code-guide/train-a-model/model-plan-guide#link-prediction-and-embeddings), [Product Recommendations](/sdgm/kumo-engine/user-code-guide/code-examples/linux-deployment/product-recommendations), and [Hybrid GNNs for Recommendations](/sdgm/research/concepts-models/hybrid-graph-neural-networks).

## Forecasting Variants

The forecasting head combines the entity's graph representation with the target's own recent history, taken from the lag windows generated with the training table. `past_encoder` selects how that history is encoded: `"decomposed"` (the suggested default), `"normalized"`, `"mlp"`, or `"transformer"`. A separate encoder handles previous-year values when the training table includes them. The head is either shared across entities, which lets it forecast for entities absent from training (`handle_new_entities=True`), or extended with learned parameters for each training entity (`False`).

See [Forecasting Architecture and Historical Inputs](/sdgm/kumo-engine/user-code-guide/train-a-model/model-plan-guide#forecasting-architecture-and-historical-inputs) and [Time Series Forecasting](/sdgm/research/applications/time-series-forecasting).

## Putting It Together

The model plan defines the architecture search space. Kumo Engine automatically runs hyperparameter search using [Optuna's Tree-structured Parzen Estimator (TPE)](https://optuna.readthedocs.io/en/stable/reference/samplers/generated/optuna.samplers.TPESampler.html). To define your search space, list the options you want to try for each parameter.

Here's an example using the `query`, `training_table`, and `snapshot` from [Getting Started](/sdgm/kumo-engine/get-started-with-kumo-engine/overview#getting-started):

```python
from kumo.engine import Trainer
from kumoapi.model_plan import GNNModelPlan

plan = query.suggest_model_plan(snapshot)
plan.model_architecture.model = [GNNModelPlan(
    channels=[64, 128],
    aggregation=[["sum", "mean", "max"], ["mean"]],
    dropout=[0.0, 0.1],
)]
plan.neighbor_sampling.num_neighbors = [[20, 10], [20, 10, 5]]
```

The sampling configuration also sets the GNN depth: `[20, 10]` gives two message-passing layers, while `[20, 10, 5]` gives three.

Next, set a budget of eight trials, each with up to 20 training epochs, and train:

```python
plan.training_job.num_experiments = 8
plan.optimization.max_epochs = 20

trainer = Trainer(plan)
result = trainer.fit(training_table, snapshot=snapshot)
print(result.metrics)
```

The fitted `trainer` holds the model selected by validation performance. Use it to save predictions and customer embeddings as Parquet. Choose a new output directory for each run:

```python
from kumo.engine import FileConnector, PredictionOutputConfig

prediction_table = query.generate_prediction_table(snapshot)
predictions = trainer.predict(
    prediction_table,
    snapshot=snapshot,
    binary_classification_threshold=0.5,
    output_config=PredictionOutputConfig(
        connector=FileConnector(root_dir="./results/run-001"),
        output_paths={
            "predictions": "predictions/",
            "embeddings": "embeddings/",
        },
    ),
)
print(predictions.frame.head())
```

See the [Model Plan Guide](/sdgm/kumo-engine/user-code-guide/train-a-model/model-plan-guide), [Making Predictions](/sdgm/kumo-engine/user-code-guide/making-predictions), and [Generate and Save Embeddings](/sdgm/kumo-engine/user-code-guide/advanced-prediction-configurations#generate-and-save-embeddings).

## References

* Corso, G., et al. "Principal Neighbourhood Aggregation for Graph Nets." NeurIPS 2020. [arXiv:2004.05718](https://arxiv.org/abs/2004.05718)
* Dwivedi, V. P., et al. "Relational Graph Transformer." ICLR 2026. [arXiv:2505.10960](https://arxiv.org/abs/2505.10960)
* Fey, M., and Lenssen, J. E. "Fast Graph Representation Learning with PyTorch Geometric." ICLR Workshop on Representation Learning on Graphs and Manifolds, 2019. [arXiv:1903.02428](https://arxiv.org/abs/1903.02428)
* Fey, M., et al. "Relational Deep Learning: Graph Representation Learning on Relational Databases." ICML 2024. [arXiv:2312.04615](https://arxiv.org/abs/2312.04615)
* Hamilton, W. L., Ying, R., and Leskovec, J. "Inductive Representation Learning on Large Graphs." NeurIPS 2017. [arXiv:1706.02216](https://arxiv.org/abs/1706.02216)
* Hu, W., et al. "PyTorch Frame: A Modular Framework for Multi-Modal Tabular Learning." 2024. [arXiv:2404.00776](https://arxiv.org/abs/2404.00776)
* Sato, R., Yamada, M., and Kashima, H. "Random Features Strengthen Graph Neural Networks." SDM 2021. [arXiv:2002.03155](https://arxiv.org/abs/2002.03155)
* Schlichtkrull, M., et al. "Modeling Relational Data with Graph Convolutional Networks." ESWC 2018. [arXiv:1703.06103](https://arxiv.org/abs/1703.06103)
* Xu, K., Hu, W., Leskovec, J., and Jegelka, S. "How Powerful are Graph Neural Networks?" ICLR 2019. [arXiv:1810.00826](https://arxiv.org/abs/1810.00826)
* You, J., Gomes-Selman, J., Ying, R., and Leskovec, J. "Identity-aware Graph Neural Networks." AAAI 2021. [arXiv:2101.10320](https://arxiv.org/abs/2101.10320)
* You, J., Ying, R., and Leskovec, J. "Design Space for Graph Neural Networks." NeurIPS 2020. [arXiv:2011.08843](https://arxiv.org/abs/2011.08843)
* Yuan, Y., et al. "ContextGNN: Beyond Two-Tower Recommendation Systems." ICLR 2025. [arXiv:2411.19513](https://arxiv.org/abs/2411.19513)

See the research pages on [Relational Deep Learning](/sdgm/research/concepts-models/relational-deep-learning), [RelBench](/sdgm/research/data-benchmarks/relbench), [Relational Graph Transformers](/sdgm/research/concepts-models/relational-graph-transformers), and [Hybrid GNNs for Recommendations](/sdgm/research/concepts-models/hybrid-graph-neural-networks), and the [PyG Practitioner's Guide](/sdgm/research/engineering-guides/pyg-practitioners-guide).