tutorials.examples.train_bayesian_structure

The goal of this script is to reproduce the results of DAG-GFlowNet for Bayesian structure learning (Deleu et al., 2022) using the GraphEnv.

Specifically, we consider a randomly generated (under the Erdős-Rényi model) linear-Gaussian Bayesian network over num_nodes nodes. We generate 100 datapoints from it, and use them to calculate the BGe score. The GFlowNet is learned to generate directed acyclic graphs (DAGs) proportionally to their BGe score, using the modified DB loss.

Some expected results on the Erdős-Rényi model with 5 nodes and 5 edges:

(On-policy training) python train_bayesian_structure.py –seed 0 –no_buffer –max_epsilon 0.0 –min_epsilon 0.0 >> Expected SHD: 6.99 >> Expected edges: 10.00 >> ROC-AUC: 0.68 >> Jensen-Shannon divergence: 0.36

(Off-policy training with epsilon-noisy exploration, which is the default) python train_bayesian_structure.py –seed 0 >> Expected SHD: 6.31 >> Expected edges: 8.80 >> ROC-AUC: 0.78 >> Jensen-Shannon divergence: 0.02

Attributes

parser

Classes

DAGEdgeActionGNN

Simple GNN-based edge action module

DAGEdgeActionGNNv2

GNN-based edge action module, adapted from the implementation of

DAGEdgeActionMLP

Network that processes flattened adjacency matrices to predict graph actions.

Functions

main(args)

Module Contents

class tutorials.examples.train_bayesian_structure.DAGEdgeActionGNN(n_nodes, num_node_classes, num_edge_classes, num_conv_layers=1, embedding_dim=128, is_backward=False)

Bases: gfn.utils.modules.GraphActionGNN

Simple GNN-based edge action module

Parameters:

  • n_nodes (int) – The number of nodes in the graph.

  • num_edge_classes (int) – The number of edge classes.

  • num_conv_layers (int) – The number of GNN layers.

  • embedding_dim (int) – The dimension of embeddings.

  • is_backward (bool) – Whether the module is used for backward action prediction.

  • num_node_classes (int)

property edges_dim: int
Return type:

int

forward(states_tensor)
Parameters:

states_tensor (torch_geometric.data.Batch)

Return type:

tensordict.TensorDict

n_nodes
class tutorials.examples.train_bayesian_structure.DAGEdgeActionGNNv2(n_nodes, num_node_classes, num_edge_classes, num_conv_layers=2, embedding_dim=128, num_heads=4, is_backward=False)

Bases: torch.nn.Module

GNN-based edge action module, adapted from the implementation of https://github.com/GFNOrg/GFN_vs_HVI/blob/master/dags/dag_gflownet/nets/gnn/gflownet.py

Parameters:

  • n_nodes (int) – The number of nodes in the graph.

  • num_edge_classes (int) – The number of edge classes.

  • num_conv_layers (int) – The number of GNN layers.

  • embedding_dim (int) – The dimension of embeddings.

  • num_heads (int) – The number of attention heads.

  • is_backward (bool) – Whether the module is used for backward action prediction.

  • num_node_classes (int)

_input_dim = 1
_n_nodes
_output_dim
attention
edge_embedding
property edges_dim: int
Return type:

int

embedding_dim = 128
forward(states_tensor)
Parameters:

states_tensor (torch_geometric.data.Batch)

Return type:

tensordict.TensorDict

graph_network
property input_dim
is_backward = False
property n_nodes
node_embedding
num_edge_classes
num_node_classes
property output_dim: int
Return type:

int

projection
receivers_mlp
senders_mlp
stop_mlp
temperature
class tutorials.examples.train_bayesian_structure.DAGEdgeActionMLP(n_nodes, num_node_classes, num_edge_classes, n_hidden_layers=3, n_hidden_layers_exit=2, embedding_dim=128, is_backward=False)

Bases: gfn.utils.modules.GraphEdgeActionMLP

Network that processes flattened adjacency matrices to predict graph actions.

Unlike the GNN-based GraphActionGNN, this module uses standard MLPs to process the entire adjacency matrix as a flattened vector. This approach:

  1. Can directly process global graph structure without message passing.

  2. May be more effective for small graphs where global patterns are important.

  3. Does not require complex graph neural network operations.

The module architecture consists of: - An MLP to process the flattened adjacency matrix into an embedding. - An edge MLP that predicts logits for each possible edge action. - An exit MLP that predicts a logit for the exit action.

Parameters:

  • n_nodes (int) – Number of nodes in the graph.

  • directed – Whether the graph is directed or undirected.

  • n_hidden_layers (int) – Number of hidden layers in the MLP for the edge actions.

  • n_hidden_layers_exit (int) – Number of hidden layers in the MLP for the exit action.

  • embedding_dim (int) – Dimension of internal embeddings.

  • is_backward (bool) – Whether this is a backward policy.

  • num_node_classes (int)

  • num_edge_classes (int)

property edges_dim: int
Return type:

int

tutorials.examples.train_bayesian_structure.main(args)
Parameters:

args (argparse.Namespace)

tutorials.examples.train_bayesian_structure.parser