MolecularDiffusion.modules.models.kgdiff.uni_transformer

KGDiff’s SE(3)-equivariant attention transformer (uni_o2).

Ported from KGDiff models/uni_transformer.py (commit ad893fc), which is itself TargetDiff’s backbone. One block of num_layers alternating x2h / h2x attention updates over a kNN graph rebuilt once per block.

The graph is geometric, not chemical: _connect_edge() is a plain knn_graph(x, k=32) over the joint pocket+ligand point cloud, and the edge_feat_dim: 4 channel is a one-hot of the four ligand/protein incidence combinations (ll / lp / pl / pp), not a bond order. No bond information reaches this network.

Only cutoff_mode='knn' is ported; the released config and the released checkpoint both use it, and the hybrid branch pulled in a batch_hybrid_edge_connection helper nothing else needs.

Module and parameter names are kept byte-for-byte compatible with the released checkpoint (refine_net.base_block.<i>.x2h_layers.0.hk_func.net.* and friends), so conversion is a pure prefix add.

Classes

AttentionLayerO2TwoUpdateNodeGeneral	One layer: num_x2h feature updates then num_h2x coord updates.
BaseH2XAttLayer	Multi-head attention that emits an equivariant coordinate update.
BaseX2HAttLayer	Multi-head attention that updates scalar node features from geometry.
UniTransformerO2TwoUpdateGeneral	The uni_o2 refine net: num_blocks x num_layers attention.

Module Contents

class MolecularDiffusion.modules.models.kgdiff.uni_transformer.AttentionLayerO2TwoUpdateNodeGeneral(hidden_dim: int, n_heads: int, num_r_gaussian: int, edge_feat_dim: int, act_fn: str = 'relu', norm: bool = True, num_x2h: int = 1, num_h2x: int = 1, r_min: float = 0.0, r_max: float = 10.0, num_node_types: int = 8, ew_net_type: str = 'r', x2h_out_fc: bool = True, sync_twoup: bool = False)

Bases: torch.nn.Module

One layer: num_x2h feature updates then num_h2x coord updates.

forward(h, x, edge_attr, edge_index, mask_ligand, e_w=None, fix_x=False)

act_fn = 'relu'

distance_expansion

edge_feat_dim

ew_net_type = 'r'

h2x_layers

hidden_dim

n_heads

norm = True

num_h2x = 1

num_node_types = 8

num_r_gaussian

num_x2h = 1

sync_twoup = False

x2h_layers

x2h_out_fc = True

class MolecularDiffusion.modules.models.kgdiff.uni_transformer.BaseH2XAttLayer(input_dim: int, hidden_dim: int, output_dim: int, n_heads: int, edge_feat_dim: int, r_feat_dim: int, act_fn: str = 'relu', norm: bool = True, ew_net_type: str = 'r')

Bases: torch.nn.Module

Multi-head attention that emits an equivariant coordinate update.

forward(h, rel_x, r_feat, edge_feat, edge_index, e_w=None)

act_fn = 'relu'

edge_feat_dim

ew_net_type = 'r'

hidden_dim

input_dim

n_heads

output_dim

r_feat_dim

xk_func

xq_func

xv_func

class MolecularDiffusion.modules.models.kgdiff.uni_transformer.BaseX2HAttLayer(input_dim: int, hidden_dim: int, output_dim: int, n_heads: int, edge_feat_dim: int, r_feat_dim: int, act_fn: str = 'relu', norm: bool = True, ew_net_type: str = 'r', out_fc: bool = True)

Bases: torch.nn.Module

Multi-head attention that updates scalar node features from geometry.

forward(h, r_feat, edge_feat, edge_index, e_w=None)

act_fn = 'relu'

edge_feat_dim

ew_net_type = 'r'

hidden_dim

hk_func

hq_func

hv_func

input_dim

n_heads

out_fc = True

output_dim

r_feat_dim

class MolecularDiffusion.modules.models.kgdiff.uni_transformer.UniTransformerO2TwoUpdateGeneral(num_blocks: int, num_layers: int, hidden_dim: int, n_heads: int = 1, k: int = 32, num_r_gaussian: int = 50, edge_feat_dim: int = 0, num_node_types: int = 8, act_fn: str = 'relu', norm: bool = True, cutoff_mode: str = 'knn', ew_net_type: str = 'r', num_init_x2h: int = 1, num_init_h2x: int = 0, num_x2h: int = 1, num_h2x: int = 1, r_max: float = 10.0, x2h_out_fc: bool = True, sync_twoup: bool = False)

Bases: torch.nn.Module

The uni_o2 refine net: num_blocks x num_layers attention.

init_h_emb_layer is built (and therefore present in the released checkpoint) but never called – forward() only iterates self.base_block. It is kept so the checkpoint’s key set matches exactly rather than needing a justified drop.

forward(h, x, mask_ligand, batch, return_all=False, fix_x=False)

act_fn = 'relu'

base_block

cutoff_mode = 'knn'

distance_expansion

edge_feat_dim = 0

ew_net_type = 'r'

hidden_dim

init_h_emb_layer

k = 32

n_heads = 1

norm = True

num_blocks

num_h2x = 1

num_init_h2x = 0

num_init_x2h = 1

num_layers

num_node_types = 8

num_r_gaussian = 50

num_x2h = 1

r_max = 10.0

sync_twoup = False

x2h_out_fc = True