MolecularDiffusion.modules.models.syncogen.diffusion.sampling.discrete_strategies.mdlm

MDLM (Masked Diffusion Language Model) sampling strategy.

Classes

MDLM	MDLM sampling strategy for discrete graph features.

Module Contents

class MolecularDiffusion.modules.models.syncogen.diffusion.sampling.discrete_strategies.mdlm.MDLM(discrete_noise=None, constrain_edge_sampling=True)

Bases: MolecularDiffusion.modules.models.syncogen.diffusion.sampling.discrete_strategies.base.DiscreteStrategyBase

MDLM sampling strategy for discrete graph features.

Uses the MDLM update rule: interpolate between staying masked and transitioning to the predicted distribution.

Parameters:

• discrete_noise – Noise schedule function for discrete features.

• constrain_edge_sampling – Whether to constrain edge sampling step.

step(graph: MolecularDiffusion.modules.models.syncogen.api.graph.graph.BBRxnGraph, p_x0: torch.Tensor, p_e0: torch.Tensor, t: torch.Tensor, dt: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor]

Perform one MDLM denoising step.

Parameters:

• graph – BBRxnGraph with current noisy state

• p_x0 – Predicted node probabilities (B, N, D_node)

• p_e0 – Predicted edge probabilities (B, N, N, D_edge)

• t – Current timestep (B, 1)

• dt – Time step size

Returns:

Updated features

Return type:

X_next, E_next