MolecularDiffusion.modules.models.syncogen.diffusion.sampling.discrete_strategies.p2¶
P2 (Path Planning) discrete sampling strategy.
Path planning uses confidence-based or random scoring to decide which tokens to unmask at each step, allowing more control over the denoising trajectory.
Classes¶
Path Planning sampling strategy for discrete graph features. |
Module Contents¶
- class MolecularDiffusion.modules.models.syncogen.diffusion.sampling.discrete_strategies.p2.PathPlanning(discrete_noise=None, constrain_edge_sampling=True, score_type: Literal['confidence', 'random'] = 'confidence', temperature: float = 1.0, eta: float = 1.0)¶
-
Path Planning sampling strategy for discrete graph features.
Instead of using the MDLM probabilistic update, path planning: 1. Samples predictions from logits/probabilities 2. Scores each token by confidence (or randomly) 3. Uses top-k masking to decide which tokens to remask 4. Reveals tokens that were masked but shouldn’t be anymore
This provides more control over the generation trajectory.
- 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 = None) Tuple[torch.Tensor, torch.Tensor]¶
Perform one path planning 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 (unused by PathPlanning, but accepted for API consistency)
- Returns:
Updated features
- Return type:
X_next, E_next
- eta = 1.0¶
- score_type = 'confidence'¶
- temperature = 1.0¶