MolecularDiffusion.modules.models.transformer_dynamics

Dense self-attention adapter for EnVariationalDiffusion.

Binds a plain (non-equivariant) Transformer to the dynamics._forward(t, xh, node_mask, edge_mask, context) contract that EnVariationalDiffusion.phi (modules/models/en_diffusion.py:234) calls.

Novel-model ablation (docs/model_novel/diffusion_transformer/ INTEGRATION_PLAN.md): swaps EGCL’s equivariant pairwise-difference message passing (modules/models/egcl.py::EGNN_dynamics) for plain multi-head self-attention over absolute coordinates, with no distance/edge features and no built-in rotation-equivariance – the hypothesis is that the platform’s existing rotation-augmentation flag (GeomMolecularGenerative.data_augmentation) substitutes for the architectural guarantee.

Revision 2026-09-09: a real cluster run (300k steps, hyperparameters matched to the EGCL baseline) still landed far below EGCL’s valid_posebuster. Two concrete asymmetries against both EGCL and this platform’s other non-equivariant 3D transformer (ADiT’s DiT, modules/models/ldm/denoisers/dit.py) were identified and fixed here:

  1. No near-zero output-head init – EGCL’s coordinate-update head is Xavier-init’d with gain=0.001 (modules/layers/conv.py); DiT zero-inits its final projection and AdaLN gates. This file’s output head previously had neither.

  2. Time was injected once, additively, at the input – every one of the 9 attention blocks was otherwise unconditioned, unlike DiT’s per-block AdaLN-zero conditioning.

Both are fixed by reusing DiT’s own blocks (DiTBlock, FinalLayer) by import rather than re-deriving the same recipe – this platform’s other non-equivariant 3D transformer already solved this exact problem, on the same platform, so there’s no reason to reinvent it. Transformer/ AttentionBlock (modules/layers/tabasco/) are no longer used here as a result; TimeFourierEncoding still is, feeding the same time embedding into every block’s AdaLN conditioning instead of only the input sum.

Classes

TransformerDynamics

Denoising network: plain multi-head self-attention, EDM interface.

Module Contents

class MolecularDiffusion.modules.models.transformer_dynamics.TransformerDynamics(in_node_nf: int, context_node_nf: int = 0, n_dims: int = 3, hidden_dim: int = 192, num_layers: int = 9, num_heads: int = 8, mlp_dim: int | None = None, dropout: float = 0.0, activation_type: str = 'gelu', add_sinusoid_posenc: bool = False)

Bases: torch.nn.Module

Denoising network: plain multi-head self-attention, EDM interface.

Parameters:

  • in_node_nf – Node feature channels the diffusion model expects back (atom-type one-hot + atomic number [+ extra values]), excluding time and context, which are added internally.

  • context_node_nf – Conditioning channels, concatenated to the node features before embedding.

  • n_dims – Spatial dimensions (3).

  • hidden_dim – Transformer token width.

  • num_layers – Number of ``DiTBlock``s.

  • num_heads – Multi-head self-attention heads.

  • mlp_dim – Feed-forward hidden width; None defaults to 4 * hidden_dim (DiTBlock’s own default mlp_ratio).

  • dropout – Currently a no-op – DiTBlock’s attention/MLP hardcode zero dropout internally and silently ignore extra kwargs. Kept as a constructor arg (this platform’s other dynamics wrappers all expose one) rather than removed, since the config default is already 0.0; flagged here so a nonzero value doesn’t look like it did something.

  • activation_type – Currently a no-op for the same reason – DiTBlock’s MLP hardcodes GELU(approximate="tanh"). Kept for config-shape compatibility with the pre-revision file; the bundled config’s default was already "gelu".

  • add_sinusoid_posenc – Ablation-only knob, off by default. Atoms are an unordered set here, so this has no principled reason to help; it exists only for later ablation curiosity and is never load-bearing.

add_sinusoid_posenc = False
blocks
context_node_nf = 0
feat_embed
final_layer
in_node_nf
n_dims = 3
pos_embed
time_encoding