DiGress: Discrete Denoising Diffusion Model for Graph Generation

All experiments are launched via python3 main.py. The project uses Hydra for configuration management, allowing you to override parameters via the CLI.

• Run debug mode (recommended first step): python3 main.py +experiment=debug.yaml
• Run on a few batches (test run): python3 main.py general.name=test
• Run the continuous model: python3 main.py model=continuous
• Run the discrete model (default): python3 main.py
• Specify a dataset: python3 main.py dataset=guacamol (Check configs/dataset for available options).

# Debugging
python3 main.py +experiment=debug.yaml

# Discrete model
python3 main.py

# Continuous model
python3 main.py model=continuous

# Specific dataset
python3 main.py dataset=guacamol

To add a new dataset, create a new file in src/datasets. You can use moses_dataset.py (for molecules) or spectre_datasets.py (for abstract graphs) as templates.

Requirements

1. Implement a Dataset class: Processes the raw data (refer to PyG documentation for details).
2. Implement a DatasetInfos class: Defines the noise model and metrics.

Molecular Dataset Specifics

If working with molecules, DatasetInfos must specify:

• atom_encoder: One-hot encoding of atom types.
• atom_decoder: Inverse mapping of the atom_encoder.
• atomic weight: Weight for each atom type.
• valency: The most common valency for each atom type.

Note: Node counts and distributions of node/edge types can be computed automatically using AbstractDataModule functions. After implementation, add a new config file in configs/dataset and update main.py to handle the new dataset.

The NodeEdgeBlock is the core computational unit that enables interaction between different graph components:

• Attention with Edges: Instead of standard self-attention, it computes unnormalized attention scores $Y$ and modulates them using edge features $E$ via additive and multiplicative FiLM operations: Y = Y * (E1 + 1) + E2.
• Global-to-Edge Modulation: Global features y are used to modulate edge representations: newE = ye1 + (ye2 + 1) * newE.
• Global-to-Node Modulation: Global features y modulate the weighted values from the attention mechanism before they are projected back to node features.
• Edge-to-Global Update: The global feature y is updated by aggregating information from nodes and edges using projection layers (x_y and e_y).

The QM9infos class provides metadata and statistics specific to the QM9 dataset, which vary depending on whether hydrogens are removed (remove_h).

When remove_h=True:

• Atom Encoder: {'C': 0, 'N': 1, 'O': 2, 'F': 3}
• Atom Decoder: ['C', 'N', 'O', 'F']
• Number of Atom Types: 4
• Max Nodes: 9

When remove_h=False:

• Atom Encoder: {'H': 0, 'C': 1, 'N': 2, 'O': 3, 'F': 4}
• Atom Decoder: ['H', 'C', 'N', 'O', 'F']
• Number of Atom Types: 5
• Max Nodes: 29

Standard batch-wise metrics often average within the batch. The SumExceptBatch family of metrics accumulates the total sum and total count across all batches to provide a true global average at the end of an epoch. This is useful for preventing bias in datasets with varying batch sizes or specific distribution properties.

Available implementations:

• SumExceptBatchMetric: A generic accumulator for arbitrary values.
• SumExceptBatchMSE: Computes Mean Squared Error by accumulating squared errors and observation counts across batches.
• SumExceptBatchKL: Computes Kullback-Leibler divergence by accumulating the sum of KL divergence and sample counts.

The XEyTransformerLayer is a single transformer block that simultaneously updates node, edge, and global features. It consists of two main stages:

1. NodeEdgeBlock: A specialized attention mechanism where:
   • Node features (X) generate Queries and Keys.
   • Edge features (E) are incorporated into the attention scores via FiLM (Feature-wise Linear Modulation).
   • Global features (y) modulate both edge features and node features.
   • Global features are updated based on the current state of nodes and edges.
2. Feedforward Networks: Separate MLP blocks for X, E, and y that apply non-linear transformations and residual connections.

To resume training or perform testing on a specific checkpoint, use the general configuration block:

Resuming Training

Set general.resume to the path of the checkpoint. This uses get_resume_adaptive, which allows you to override some parameters from the original run while keeping the core model state.

Testing Only

Set general.test_only to the path of the checkpoint. This uses get_resume, which loads the previous configuration strictly without allowing updates to keys (useful for consistent evaluation).

If general.evaluate_all_checkpoints is true during a test_only run, the script will attempt to evaluate all .ckpt files found in the parent directory of the specified checkpoint.

If you encounter the error: PermissionError: [Errno 13] Permission denied: '.../src/analysis/orca/orca'

This indicates that the orca binary has not been compiled. Navigate to the ./src/analysis/orca directory and compile it using:

g++ -O2 -std=c++11 -o orca orca.cpp

The QM9DataModule is a high-level wrapper used to manage the QM9 dataset splits. It is initialized using a configuration object (cfg).

Expected configuration keys:

• cfg.dataset.datadir: The directory where the dataset is stored.
• cfg.dataset.remove_h: Boolean flag to determine if hydrogens should be removed.
• cfg.general.guidance_target: (Optional) Set to 'mu', 'homo', or 'both' to select specific target properties for regression tasks. If not set or unrecognized, it defaults to removing the target property (RemoveYTransform).

# Conceptual configuration structure
class Config:
    class dataset:
        datadir = 'data/qm9'
        remove_h = True
    class general:
        guidance_target = 'mu'

# Usage
datmodule = QM9DataModule(cfg)

DiGress uses Hydra for configuration. The default configuration is located in ../configs/config.yaml.

Key configuration sections include:

• general: Controls name, gpus, resume, test_only, and evaluate_all_checkpoints.
• model: Defines type ('discrete' or otherwise) and extra_features.
• dataset: Specifies the name (e.g., 'qm9', 'sbm') and dataset-specific parameters.
• train: Controls n_epochs, clip_grad, save_model, and ema_decay.