MolecularDiffusion.modules.tasks.diffusion_equifm

EquiFM task – equivariant flow matching with hybrid probability transport.

Paper: arXiv:2312.07168 (NeurIPS 2023). Target repo: github.com/AlgoMole/MolFM.

The training objective here is a reconstruction of the paper’s Algorithm 1, not the authors’ released objective. MolFM’s release is sampling-only: it ships no loss, no EOT solver, and no data pipeline. args.pickle shows the released QM9 weights were additionally trained with angle_penalty=True, cat_loss='l2_masked_mean' and ode_regularization=0.001 – none of which appear anywhere in the paper. A model trained with this task will therefore not reproduce the paper’s Table 1 numbers, and the gap must not be reported as a reproduction failure of the paper. The sampler and the converted released checkpoint are faithful ports and are the real test of this integration.

Scope (per the approved INTEGRATION_PLAN.md): unconditional QM9 generation only. No property conditioning / CFG (the release has context_node_nf=0 and no conditional weights), no gradient guidance, no inpainting, no trajectory frames, no sample_chain.

Attributes

ModelTaskFactory

Classes

EquiFMTask	Duck-typed Task (docs/adding_new_models.md §2.1) around Cnflows.
EquiFMTaskFactory	Factory matching cli/train.py's task_factory.build() pattern.

Module Contents

class MolecularDiffusion.modules.tasks.diffusion_equifm.EquiFMTask(cnflows: MolecularDiffusion.modules.models.equifm.cnflows.Cnflows, default_n_timesteps: int, dataset_stats: dict, atom_vocab: List)

Bases: torch.nn.Module

Duck-typed Task (docs/adding_new_models.md §2.1) around Cnflows.

evaluate(pred: torch.Tensor, target: torch.Tensor)

forward(batch: Dict[str, torch.Tensor])

predict_and_target(batch: Dict[str, torch.Tensor])

sample(batch_size: int | None = None, nodesxsample: torch.Tensor | None = None, num_steps: int | None = None, batch: Dict[str, torch.Tensor] | None = None, mode=None, n_frames: int = 0, **kwargs)

Unconditional sampling. Returns (one_hot, charges, coords, node_mask) following the EDM/GeoLDM convention (charges are true nuclear charges), which keeps runmodes/generate/tasks_generate.py on the save_xyz_file branch. mode and n_frames are accepted and ignored (trajectories are out of scope).

T

atom_vocab

cnflows

property device

property model

property n_node_dist

property node_dist_model

property norm_values

prop_dist_model = None

task_type = 'diffusion_equifm'

class MolecularDiffusion.modules.tasks.diffusion_equifm.EquiFMTaskFactory(task_type: str = 'diffusion_equifm', n_dims: int = 3, include_charges: bool = True, normalize_factors=(1.0, 4.0, 10.0), sigma_min: float = 0.0001, beta_min: float = 0.1, beta_max: float = 20.0, discrete_path: str = 'HB_path', default_n_timesteps: int = 250, use_eot: bool = True, eot_max_iters: int = 20, dynamics: dict | None = None, dataset_stats: dict | None = None, atom_vocab: List | None = None, train_set: torch.utils.data.Dataset | None = None, **kwargs)

Factory matching cli/train.py’s task_factory.build() pattern.

train_set is declared as a named parameter so the declarative seam in cli/train.py injects the training dataset (docs/adding_new_models.md §2.5); it is used only to build the atom-count histogram that backs node_dist_model.

build() → EquiFMTask

compute_dataset_stats(dataset) → None

Atom-count histogram from the training set (same shape as FlowMol’s).

atom_vocab

beta_max = 20.0

beta_min = 0.1

dataset_stats

default_n_timesteps = 250

discrete_path = 'HB_path'

dynamics

eot_max_iters = 20

include_charges = True

kwargs

n_dims = 3

normalize_factors = (1.0, 4.0, 10.0)

sigma_min = 0.0001

task_type = 'diffusion_equifm'

train_set = None

use_eot = True

MolecularDiffusion.modules.tasks.diffusion_equifm.ModelTaskFactory