Skip to content

Pretrained Foundation Models

K3 Node provides native implementations of modern Graph Foundation Models with built-in checkpoint downloaders and weight loaders for official pre-trained models.

1. GraphMAE2 (Masked Autoencoder for Graphs)

GraphMAE2 is an advanced masked autoencoder for self-supervised graph representation learning featuring multi-task reconstruction (scaled cosine error) and projection heads.

Architecture

k3_node.models.graphmae2.GraphMAE2

Bases: Layer

The GraphMAE2 model from "GraphMAE2: A Decoding-Enhanced Masked Self-Supervised Learning Framework for Graphs" <https://arxiv.org/abs/2304.04779>_.

Parameters:

Name Type Description Default
in_dim int

Dimensionality of input node features.

 required
num_hidden int

Dimensionality of hidden node representations.

 required
num_layers int

Number of encoder layers. (default: 4)

 4
num_dec_layers int

Number of decoder layers. (default: 1)

 1
num_remasking int

Number of remasking views in decoder. (default: 3)

 3
nhead int

Number of attention heads in encoder. (default: 8)

 8
nhead_out int

Number of attention heads in decoder output. (default: 1)

 1
activation str

Activation function. (default: "prelu")

 'prelu'
feat_drop float

Node feature dropout rate. (default: 0.2)

 0.2
attn_drop float

Attention dropout rate. (default: 0.1)

 0.1
negative_slope float

LeakyReLU negative slope. (default: 0.2)

 0.2
residual bool

Whether to use residual connections. (default: True)

 True
norm str

Normalization type ("layernorm", "batchnorm", or None). (default: "layernorm")

 'layernorm'
mask_rate float

Fraction of input nodes to mask. (default: 0.5)

 0.5
remask_rate float

Fraction of latent nodes to remask. (default: 0.5)

 0.5
remask_method str

Remasking method ("random" or "fixed"). (default: "random")

 'random'
loss_fn str

Reconstruction loss type ("sce" or "mse"). (default: "sce")

 'sce'
alpha_l float

Power exponent in Scaled Cosine Error loss. (default: 2.0)

 2.0
lam float

Weight of the latent prediction loss term. (default: 1.0)

 1.0
momentum float

Teacher EMA update momentum. (default: 0.996)

 0.996
delayed_ema_epoch int

Epoch to begin EMA teacher updates. (default: 0)

 0

call(x, edge_index, training=False)

Forward pass: returns node embeddings by default.

ema_update(momentum=None)

Updates teacher EMA parameters.

embed(x, edge_index)

Generates node embeddings with the encoder.

encoding_mask_noise(x, mask_rate=None, mask_nodes=None)

Masks node features for encoder input.

from_pretrained(dataset='ogbn-arxiv', folder='checkpoints', download=True, **kwargs) classmethod

Instantiates a GraphMAE2 model with pre-trained weights downloaded from Google Drive.

Parameters:

Name Type Description Default
dataset str

Dataset name ("ogbn-arxiv", "ogbn-products", "mag-scholar-f", or "ogbn-papers100M").

 'ogbn-arxiv'
folder str

Directory to store/find checkpoints. (default: "checkpoints")

 'checkpoints'
download bool

Whether to download checkpoint if missing locally. (default: True)

 True
**kwargs

Overrides for model hyperparameters.

 {}

Returns:

Name Type Description
GraphMAE2 GraphMAE2

Model instance loaded with pre-trained weights.

load_weights_from_checkpoint(checkpoint_path=None, dataset=None, folder='checkpoints', download=True)

Loads weights from a PyTorch state dict checkpoint or Google Drive.

loss(x, edge_index, mask_nodes=None, targets=None, epoch=0, training=True)

Computes GraphMAE2 loss: attribute reconstruction loss + latent prediction loss.

random_remask(rep, remask_rate=None, remask_nodes=None)

Remasks latent representation for decoder input.

Checkpoint Loading

k3_node.models.graphmae2.load_graphmae2_weights(model, checkpoint_path=None, dataset=None, folder='checkpoints', download=True)

Loads pre-trained weights from a GraphMAE2 PyTorch checkpoint (.pt).

If the checkpoint does not exist locally and download=True, it will be automatically downloaded from the official Google Drive folder using download_google_url.

Parameters:

Name Type Description Default
model GraphMAE2

The target GraphMAE2 model instance.

 required
checkpoint_path str

Local path to .pt file or dataset name.

 None
dataset str

Dataset name if downloading from Google Drive.

 None
folder str

Directory to store downloaded checkpoints. (default: "checkpoints")

 'checkpoints'
download bool

Whether to download checkpoint if missing locally. (default: True)

 True

Checkpoint Downloader

k3_node.models.graphmae2.download_graphmae2_checkpoint(dataset, folder='checkpoints', log=True)

Downloads a pre-trained GraphMAE2 checkpoint from Google Drive using download_google_url.

Google Drive folder: https://drive.google.com/drive/folders/1GiuP0PtIZaYlJWIrjvu73ZQCJGr6kGkh

Parameters:

Name Type Description Default
dataset str

Name of dataset ("ogbn-arxiv", "ogbn-products", "mag-scholar-f", or "ogbn-papers100M").

 required
folder str

Target directory to save the checkpoint. (default: "checkpoints")

 'checkpoints'
log bool

Whether to print download progress. (default: True)

 True

Returns:

Name Type Description
str str

Absolute path to the downloaded checkpoint file.

Usage Example

from k3_node.models.graphmae2 import GraphMAE2, load_graphmae2_weights, download_graphmae2_checkpoint

# Download checkpoint
ckpt_path = download_graphmae2_checkpoint("cora")

# Initialize and load model
model = GraphMAE2(in_dim=1433, num_hidden=512, out_dim=1433, num_layers=2)
load_graphmae2_weights(model, ckpt_path)

2. Graphormer & Graphormer3D

Graphormer is a Graph Transformer with spatial, degree, and edge attention biases. Graphormer3D extends this architecture with 3D Gaussian RBF distance biases for molecular conformation and quantum chemistry.

Graphormer (2D)

k3_node.models.graphormer.Graphormer

Bases: Model

Graphormer model for molecular graph representation and property prediction from "Do Transformers Really Perform Badly for Graph Representation?" <https://arxiv.org/abs/2106.05234>_.

Parameters:

Name Type Description Default
num_atoms int

Maximum atom vocabulary size. (default: 512)

 512
num_in_degree int

Maximum in-degree. (default: 512)

 512
num_out_degree int

Maximum out-degree. (default: 512)

 512
num_edges int

Maximum edge vocabulary size. (default: 512)

 512
num_spatial int

Maximum spatial distance. (default: 512)

 512
num_edge_dis int

Maximum edge distance multiplier. (default: 128)

 128
edge_type str

Type of edge encoding ("multi_hop" or "single_hop"). (default: "multi_hop")

 'multi_hop'
multi_hop_max_dist int

Max distance for multi-hop paths. (default: 20)

 20
num_encoder_layers int

Number of Transformer layers. (default: 12)

 12
embedding_dim int

Hidden embedding dimension. (default: 768)

 768
ffn_embedding_dim int

FFN intermediate dimension. (default: 768)

 768
num_attention_heads int

Number of attention heads. (default: 32)

 32
dropout float

Dropout rate. (default: 0.0)

 0.0
attention_dropout float

Attention dropout rate. (default: 0.1)

 0.1
activation_dropout float

Activation dropout rate. (default: 0.1)

 0.1
encoder_normalize_before bool

Whether to apply LayerNorm before encoder blocks. (default: True)

 True
pre_layernorm bool

Whether to use Pre-LN. (default: False)

 False
num_classes int

Output dimension for prediction head. (default: 1)

 1
activation_fn str

Activation function. (default: "gelu")

 'gelu'
**kwargs

Additional model arguments.

 {}

call(batched_data=None, x=None, in_degree=None, out_degree=None, attn_bias=None, spatial_pos=None, edge_input=None, attn_edge_type=None, perturb=None, return_all=False, training=False)

Forward pass for Graphormer.

Accepts either a single dictionary batched_data containing graph tensors, or individual tensor arguments.

Returns:

Type Description

Tensor or Tuple[Tensor, Tensor]: Graph prediction tensor of shape [batch_size, num_classes],

or if return_all=True, a tuple of (graph_pred, all_node_features).

embed(batched_data=None, x=None, in_degree=None, out_degree=None, attn_bias=None, spatial_pos=None, edge_input=None, attn_edge_type=None)

Computes graph and node embeddings without applying the prediction head.

from_pretrained(pretrained_name='pcqm4mv1_graphormer_base', folder='checkpoints', download=True, **kwargs) classmethod

Instantiates a Graphormer model with pre-trained weights.

Graphormer3D (3D)

k3_node.models.graphormer_3d.Graphormer3D

Bases: Model

Graphormer-3D model for 3D molecular structure modeling, energy, and force prediction from "Benchmarking Graphormer on Large-Scale Molecular Modeling Datasets" <https://arxiv.org/abs/2203.04810>_.

Parameters:

Name Type Description Default
layers int

Number of encoder layers per block. (default: 12)

 12
blocks int

Number of repeated encoder blocks. (default: 4)

 4
embed_dim int

Hidden embedding dimension. (default: 768)

 768
ffn_embed_dim int

FFN intermediate dimension. (default: 768)

 768
attention_heads int

Number of attention heads. (default: 48)

 48
num_kernel int

Number of Gaussian basis kernels. (default: 128)

 128
atom_types int

Number of atom types. (default: 64)

 64
dropout float

Dropout probability. (default: 0.1)

 0.1
attention_dropout float

Attention dropout. (default: 0.1)

 0.1
activation_dropout float

FFN activation dropout. (default: 0.0)

 0.0
input_dropout float

Input features dropout. (default: 0.0)

 0.0
**kwargs

Additional model arguments.

 {}

call(atoms, tags, pos, real_mask=None, training=False)

Forward pass for Graphormer-3D predicting total energy and atomic forces.

Parameters:

Name Type Description Default
atoms Tensor

Atom indices of shape [batch_size, num_nodes].

 required
tags Tensor

Tag indices of shape [batch_size, num_nodes] (0: fixed, 1: sub-surface, 2: surface).

 required
pos Tensor

3D atomic coordinates of shape [batch_size, num_nodes, 3].

 required
real_mask Tensor

Valid non-padding mask of shape [batch_size, num_nodes]. If None, non-zero atom indices are considered valid.

 None
training bool

Training flag. (default: False)

 False

Returns:

Type Description

Tuple[Tensor, Tensor]: Tuple of predicted energy [batch_size] and atomic forces

[batch_size, num_nodes, 3].

from_pretrained(pretrained_name='oc20is2re_graphormer3d_base', folder='checkpoints', download=True, **kwargs) classmethod

Instantiates a Graphormer3D model with pre-trained weights.

Checkpoint Utilities

k3_node.models.graphormer.load_graphormer_weights(model, checkpoint_path=None, pretrained_name=None, folder='checkpoints', download=True)

Loads weights from a PyTorch (.pt or .bin) checkpoint into a Keras Graphormer model.

Parameters:

Name Type Description Default
model Graphormer

Target model instance.

 required
checkpoint_path str

Path to .pt or .bin file.

 None
pretrained_name str

Name of pre-trained model to load or download.

 None
folder str

Checkpoints directory. (default: "checkpoints")

 'checkpoints'
download bool

Whether to download if missing. (default: True)

 True

Returns:

Name Type Description
Graphormer Graphormer

The model with loaded weights.

k3_node.models.graphormer.download_graphormer_checkpoint(name, folder='checkpoints', log=True)

Downloads a pre-trained Graphormer checkpoint to the specified folder.

Parameters:

Name Type Description Default
name str

Pretrained checkpoint name (e.g., "pcqm4mv1_graphormer_base", "pcqm4mv2_graphormer_base").

 required
folder str

Destination folder. (default: "checkpoints")

 'checkpoints'
log bool

Whether to print download progress. (default: True)

 True

Returns:

Name Type Description
str str

Absolute path to the downloaded file.

3. GraphGPS (General Powerful Scalable Graph Transformer)

GraphGPS is a hybrid graph transformer architecture combining local message-passing neural networks (CustomGatedGCN) and global linear/full attention with Random Walk Structural Encodings (RWSE).

GPSModel

k3_node.models.gps_model.GPSModel

Bases: Model

GraphGPS: General Powerful Scalable Graph Transformer from the "Recipe for a General, Powerful, Scalable Graph Transformer" <https://arxiv.org/abs/2205.12454>_ paper (NeurIPS 2022).

Parameters:

Name Type Description Default
dim_in int

Initial input feature dimension. (default: 256)

 256
dim_out int

Target output dimension. (default: 1)

 1
num_layers int

Number of GPS layers. (default: 16)

 16
dim_hidden int

Hidden embedding dimension. (default: 256)

 256
num_heads int

Number of attention heads. (default: 8)

 8
local_gnn_type str

Local MPNN layer type. (default: "CustomGatedGCN")

 'CustomGatedGCN'
act str

Activation function. (default: "gelu")

 'gelu'
dropout float

Dropout probability. (default: 0.1)

 0.1
attn_dropout float

Attention dropout probability. (default: 0.1)

 0.1
batch_norm bool

Whether to use batch normalization. (default: True)

 True
layer_norm bool

Whether to use layer normalization. (default: False)

 False
node_encoder_type str

Node encoder type ("Atom+RWSE", "Atom", "Linear", or None). (default: "Atom+RWSE")

 'Atom+RWSE'
edge_encoder_type str

Edge encoder type ("Bond", "Linear", or None). (default: "Bond")

 'Bond'
atom_feature_dims List[int]

Categorical feature vocabulary sizes for atom features.

 None
bond_feature_dims List[int]

Categorical feature vocabulary sizes for bond features.

 None
rwse_num_steps int

Number of RWSE steps. (default: 16)

 16
rwse_dim_pe int

RWSE embedding dimension. (default: 20)

 20
graph_pooling str

Graph pooling type ('mean', 'add', 'max'). (default: "mean")

 'mean'
head_layers int

Number of hidden layers in prediction head. (default: 2)

 2
**kwargs

Additional model arguments.

 {}

Checkpoint Utilities

k3_node.models.gps_model.load_gps_weights(model, checkpoint_path)

Loads trained PyTorch GraphGPS checkpoint weights into a Keras 3 GPSModel.

Parameters:

Name Type Description Default
model GPSModel

Target GPSModel instance.

 required
checkpoint_path str

Path to PyTorch .ckpt or .pt checkpoint file.

 required

Returns:

Name Type Description
GPSModel

The model with loaded weights.

k3_node.models.gps_model.download_gps_checkpoint(checkpoint_name='pcqm4m-GPS+RWSE.deep', cache_dir=None)

Downloads and extracts a pretrained GraphGPS checkpoint.

Parameters:

Name Type Description Default
checkpoint_name str

Name of the checkpoint. Currently supported: "pcqm4m-GPS+RWSE.deep".

 'pcqm4m-GPS+RWSE.deep'
cache_dir str

Cache directory to store downloaded checkpoint.

 None

Returns:

Name Type Description
str str

Absolute path to the extracted .ckpt checkpoint file.

Usage Example

from k3_node.models.gps_model import GPSModel, load_gps_weights, download_gps_checkpoint

# Download official pcqm4m-GPS+RWSE.deep checkpoint
ckpt_path = download_gps_checkpoint("pcqm4m-GPS+RWSE.deep")

# Initialize and load weights
model = GPSModel(dim_in=256, dim_out=1, num_layers=16, dim_hidden=256, num_heads=8)
load_gps_weights(model, ckpt_path)

4. GROVER (Graph Representation from Self-Supervised Message Passing Transformer)

GROVER incorporates dual-track message passing across directed bonds and atoms with transformer self-attention and scope-based molecular readouts.

GROVER

k3_node.models.grover.GROVER

Bases: Model

Complete GROVER Model.

call(inputs, training=False)

Inputs can be a tuple/list: (f_atoms, f_bonds, a2b, b2a, b2revb, a_scope, b_scope, a2a).

get_fingerprint(inputs, fingerprint_source='both', features_batch=None, training=False)

Generate molecule-level fingerprints using Readout.

Checkpoint Utilities

k3_node.models.grover.load_grover_weights(model, checkpoint_path)

Loads PyTorch GROVER checkpoint state dict into Keras 3 GROVER model.

k3_node.models.grover.download_grover_checkpoint(checkpoint_name='grover_base', cache_dir=None)

Download official GROVER pre-trained model checkpoint from Google Drive.

Usage Example

from k3_node.models.grover import GROVER, load_grover_weights, download_grover_checkpoint

ckpt_path = download_grover_checkpoint("grover_base")

model = GROVER(hidden_size=800, num_attn_head=4, depth=6, num_mt_block=1)
load_grover_weights(model, ckpt_path)

5. Mole-BERT (Self-Supervised Molecular GNN)

Mole-BERT pre-trains a 5-layer GIN backbone with categorical atom and bond embeddings, batch normalization, and Jumping Knowledge for downstream molecular property prediction.

MoleBERT

k3_node.models.mole_bert.MoleBERT

Bases: Model

Complete Mole-BERT Model with graph-level pooling and property prediction head.

call(inputs, training=False)

Call MoleBERT model.

inputs can be a tuple: (x, edge_index, edge_attr) or (x, edge_index, edge_attr, batch).

Checkpoint Utilities

k3_node.models.mole_bert.load_mole_bert_weights(model, checkpoint_path)

Loads official PyTorch Mole-BERT.pth checkpoint state dict into Keras 3 MoleBERT model.

k3_node.models.mole_bert.download_mole_bert_checkpoint(cache_dir=None)

Downloads official Mole-BERT.pth checkpoint from GitHub.

Usage Example

from k3_node.models.mole_bert import MoleBERT, load_mole_bert_weights, download_mole_bert_checkpoint

ckpt_path = download_mole_bert_checkpoint()

model = MoleBERT(num_layer=5, emb_dim=300, num_tasks=1)
load_mole_bert_weights(model, ckpt_path)