MolecularDiffusion.runmodes.train.tasks_transformer

Factory wiring TransformerDynamics – a plain (non-equivariant) multi-head self-attention backbone – into the existing, unmodified diffusion task, structured exactly like runmodes/train/tasks_painn.py::ModelTaskFactory (a new dynamics network dropped into EnVariationalDiffusion + GeomMolecularGenerative, zero edits to either).

Per the approved ledger (docs/model_novel/diffusion_transformer/INTEGRATION_PLAN.md) this is a novel-model ablation against the EGCL parent (tasks_egcl.py): only the backbone (and the platform’s own data_augmentation flag) differ – every diffusion-math hyperparameter is copied verbatim from configs/tasks/diffusion.yaml so the ablation isolates the backbone swap and nothing else. Follows the parent’s simpler constructor convention (a plain condition_names parameter) rather than tasks_painn.ModelTaskFactory’s train_set/task_names legacy-alias pair, which this backbone does not need.

Attributes

logger

Classes

ModelTaskFactory

Build the Transformer-backbone diffusion model + task.

Module Contents

class MolecularDiffusion.runmodes.train.tasks_transformer.ModelTaskFactory(task_type: str, atom_vocab: list = None, condition_names: list = [], hidden_dim: int = 192, num_layers: int = 9, num_heads: int = 8, mlp_dim: int = None, dropout: float = 0.0, activation_type: str = 'gelu', add_sinusoid_posenc: bool = False, chkpt_path: str = None, **kwargs)

Build the Transformer-backbone diffusion model + task.

Parameters:

  • task_type (str) – must be "diffusion".

  • atom_vocab (list) – atom vocabulary used for encoding.

  • condition_names (list) – condition names for conditional generation.

  • hidden_dim (int) – transformer token width.

  • num_layers (int) – transformer block depth (Transformer’s depth).

  • num_heads (int) – multi-head self-attention heads.

  • mlp_dim (int) – feed-forward hidden width; None defaults to 4 * hidden_dim inside Transformer.

  • dropout (float) – dropout probability.

  • activation_type (str) – Transformer’s feed-forward activation string knob (e.g. "gelu").

  • add_sinusoid_posenc (bool) – ablation-only knob, off by default.

  • chkpt_path (str) – optional path to a model checkpoint.

  • **kwargs – diffusion keyword arguments, see configs/tasks/diffusion_transformer.yaml.

build()

Build and return the GeomMolecularGenerative task.

activation_type = 'gelu'
add_sinusoid_posenc = False
atom_vocab = None
chkpt_path = None
condition_names = []
context_node_nf
dropout = 0.0
hidden_dim = 192
in_node_nf
kwargs
mlp_dim = None
num_heads = 8
num_layers = 9
task_type
MolecularDiffusion.runmodes.train.tasks_transformer.logger