MolecularDiffusion.modules.models.midi.transformer_model

MiDi’s relational graph transformer with an EGNN-style coordinate update.

Ported verbatim in behaviour from midi/models/transformer_model.py; the only changes are import paths, the Dims dims container and docstrings.

Two upstream oddities are deliberately preserved because the released checkpoints depend on them:

  • NodeEdgeBlock.pre_softmax is never used in forward but is still registered (upstream comment: “Unused, but needed to load old checkpoints”).

  • every layer is built with last_layer=False, including the last one, so the y-branch modules exist in all 12 layers (upstream line 339).

Removing either would change the parameter set and silently break the weight mapping, since cli/generate.py loads with strict=False.

Classes

GraphTransformer

MiDi's denoiser: PlaceHolder -> PlaceHolder.

NodeEdgeBlock

Self-attention over nodes that also updates edges, y and positions.

XEyTransformerLayer

One transformer layer updating nodes, edges, positions and y.

Module Contents

class MolecularDiffusion.modules.models.midi.transformer_model.GraphTransformer(input_dims: MolecularDiffusion.modules.models.midi.placeholder.Dims, n_layers: int, hidden_mlp_dims: dict, hidden_dims: dict, output_dims: MolecularDiffusion.modules.models.midi.placeholder.Dims, adapter_indices: list | None = None, concat_indices: list | None = None)

Bases: torch.nn.Module

MiDi’s denoiser: PlaceHolder -> PlaceHolder.

Parameters:

  • input_dims – channel counts of the noised input (y includes the timestep column).

  • n_layers – number of XEyTransformerLayer blocks.

  • hidden_mlp_dims – widths of the input/output MLPs, keys X/E/y/pos.

  • hidden_dims – transformer widths, keys dx/de/dy/n_head/ dim_ffX/dim_ffE/dim_ffy.

  • output_dims – channel counts of the prediction (y is 0 – MiDi never asks the network to predict a property value).

  • adapter_indices – indices into an external per-graph condition tensor routed through a per-layer adapter MLP, additively injected into the evolving y after every transformer layer.

  • concat_indices – indices into condition concatenated onto y once, before the input embedding – reaches every layer for free via y’s existing per-layer FiLM injection into X/E.

forward(data: MolecularDiffusion.modules.models.midi.placeholder.PlaceHolder, condition: torch.Tensor | None = None) → MolecularDiffusion.modules.models.midi.placeholder.PlaceHolder

Denoise one batch; returns logits for the categorical modalities.

condition, when given, is a per-graph (B, D) property tensor: columns in self.concat_indices are appended to y before the input embedding, columns in self.adapter_indices are routed through a dedicated per-layer MLP added to y after each layer.

adapter_indices = []
concat_indices = []
condition_adapters
mlp_in_E
mlp_in_X
mlp_in_pos
mlp_in_y
mlp_out_E
mlp_out_X
mlp_out_pos
n_layers
out_dim_E
out_dim_X
out_dim_charges
out_dim_y
tf_layers
class MolecularDiffusion.modules.models.midi.transformer_model.NodeEdgeBlock(dx: int, de: int, dy: int, n_head: int, *, last_layer: bool = False)

Bases: torch.nn.Module

Self-attention over nodes that also updates edges, y and positions.

forward(X: torch.Tensor, E: torch.Tensor, y: torch.Tensor, pos: torch.Tensor, node_mask: torch.Tensor) → tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor]

Return (Xout, Eout, y_out, vel), all with the input shapes.

a
de
df
dist_add_e
dist_mul_e
dx
dy
e_att_mul
e_out
e_pos1
e_pos2
e_x_mul
in_E
k
last_layer = False
lin_dist1
lin_norm_pos1
lin_norm_pos2
n_head
out
pos_att_mul
pos_x_mul
pre_softmax
q
v
x_e_mul1
x_e_mul2
x_out
y_e_add
y_e_mul
y_x_add
y_x_mul
class MolecularDiffusion.modules.models.midi.transformer_model.XEyTransformerLayer(dx: int, de: int, dy: int, n_head: int, dim_ffX: int = 2048, dim_ffE: int = 128, dim_ffy: int = 2048, dropout: float = 0.1, layer_norm_eps: float = 1e-05, device: torch.device | None = None, dtype: torch.dtype | None = None, *, last_layer: bool = False)

Bases: torch.nn.Module

One transformer layer updating nodes, edges, positions and y.

forward(features: MolecularDiffusion.modules.models.midi.placeholder.PlaceHolder) → MolecularDiffusion.modules.models.midi.placeholder.PlaceHolder

Residual update of X, E, y and pos.

activation
dropoutE1
dropoutE2
dropoutE3
dropoutX1
dropoutX2
dropoutX3
last_layer = False
linE1
linE2
linX1
linX2
normE1
normE2
normX1
normX2
norm_pos1
self_attn