Skip to navigation
Models

Kumo Relational

View as Markdown

Kumo Relational extends in-context learning from one table to related tables. It predicts a target for task rows using the rows linked to them through primary and foreign keys, without flattening the data into one table.

At a Glance

PropertyValue
Classsdm.models.KumoRelational
TasksClassification and regression (999 quantiles)
InputsTask rows (x_context, y_context, x_query) plus RelatedTables for context and query
Parameters29.93M to 30.94M
Weightsnvidia/Kumo-Relational, OpenMDW-1.1 with BSD-3-Clause notice
Available inSDM 0.2.1 and later
ReferenceKumoRFM-2: Scaling Foundation Models for Relational Learning

When to Use It

  • Your prediction target is about an entity (a user, an item, a patient) and the evidence is in other tables, such as orders or visits.
  • You want to avoid hand-built aggregate features.
  • Choose Kumo Tabular when all features already sit in one row.

How It Works

Kumo Relational in three panels: (a) task rows linked to related Users, Orders, and Items tables; (b) row and column attention within each table; (c) inter-table message passing and attention across samples.

Kumo Relational in three panels: (a) task rows linked to related Users, Orders, and Items tables; (b) row and column attention within each table; (c) inter-table message passing and attention across samples.

  1. Intra-table row embeddings. Each related table is embedded with a TabICLv2-style row embedding stack. Context targets are distributed over the relational graph.
  2. Inter-table message passing. A schema-agnostic GNN exchanges row embeddings along the relationships for num_hops rounds, then reads out the rows linked to the task table.
  3. In-context learning. Query task rows attend to labeled context task rows to produce class logits or quantiles.

Describe Your Tables

A RelatedTables object has three parts:

  • tables: a name for each related table and its TableTensor.
  • relationships: key joins between tables. Single or composite keys.
  • task_links: exactly one entry point from task rows into a related table.

Related tables and their keys: task rows with a churn target link by user_id to a users table through a dashed task link. Views and orders reference users by user_id and items by item_id through solid relationship links.

Related tables and their keys: task rows with a churn target link by user_id to a users table through a dashed task link. Views and orders reference users by user_id and items by item_id through solid relationship links.

Columns with the id semantic type define relationships and are not used as features.

Example

This toy example predicts churn for two users from two labeled users and their orders.

import torch
import sdm
task_table = sdm.TableTensor.from_columns(
    {"user_id": [0, 1, 2, 3], "churn": [True, False, True, False]},
    stypes={"user_id": "id", "churn": "categorical"},
)
related = sdm.RelatedTables(
    tables={
        "users": sdm.TableTensor.from_columns(
            {"user_id": [0, 1, 2, 3], "age": [42, 23, 31, 26]},
            stypes={"user_id": "id", "age": "numerical"},
        ),
        "orders": sdm.TableTensor.from_columns(
            {"user_id": [0, 0, 1, 3, 3, 3],
             "amount": [9.99, 4.99, 12.99, 7.99, 3.99, 5.99]},
            stypes={"user_id": "id", "amount": "numerical"},
        ),
    },
    relationships=[{"left_table": "orders", "left_columns": "user_id",
                    "right_table": "users", "right_columns": "user_id"}],
    task_links=[{"task_columns": "user_id",
                 "table": "users", "table_columns": "user_id"}],
)
# Context rows (users 0, 1) carry labels; query rows (users 2, 3) do not.
related_context = related.replace_tables({
    "users": related.tables["users"][:2], "orders": related.tables["orders"][:3]})
related_query = related.replace_tables({
    "users": related.tables["users"][2:], "orders": related.tables["orders"][3:]})
model = sdm.models.KumoRelational(task="classification")
probs = model(
    x_context=task_table[:2].drop_columns("churn"),
    y_context=task_table[:2, "churn"],
    x_query=task_table[2:].drop_columns("churn"),
    related_context_tables=related_context,
    related_query_tables=related_query,
    num_hops=1,
    generator=torch.Generator().manual_seed(0),
)
print(probs.to_pandas())
Expected output
      False      True
0  0.321723  0.678277
1  0.203182  0.796818

With two labeled users, the probabilities only show that the pipeline runs. Real tasks need many more context rows.

Work with Real Data

For real datasets, use SDM’s subgraph samplers to build RelatedTables for each batch of task rows, with CPU (pyg-lib) or CUDA (cugraph) backends. When two task rows refer to the same entity at different times, add a task ID to composite keys so each row only sees records available at its prediction time.

The RelBench example in the SDM repository shows this end to end. From that example’s directory, run:

python main.py --dataset=rel-f1 --task=driver-dnf

Inputs and Limits

  • Pass related_context_tables and related_query_tables together; the model requires related tables.
  • Exactly one task link to an entity table.
  • Tables must be two-dimensional.
  • One target column per call.

Run as a Service

The Kumo Relational NIM accepts predictive queries through the Structured Data Models Client. Use SDM when you want the model in your own Python process and full control over sampling.

Learn More