schrodinger.application.matsci.flywheel.ldmol.dit_model module

Diffusion Transformer (DiT) model for LDMol molecular generation.

Adapted from Meta’s DiT for 1-D latent molecular sequences.

schrodinger.application.matsci.flywheel.ldmol.dit_model.approx_gelu()

Return GELU with tanh approximation.

Matches LDMol training activation; Hendrycks & Gimpel (2016).

Return type:

nn.GELU

Returns:

Approximate GELU activation module

schrodinger.application.matsci.flywheel.ldmol.dit_model.modulate(hidden, shift, scale)

Apply adaptive modulation (adaLN).

Computes: hidden * (1 + scale) + shift, broadcasting over sequence dim.

Parameters:

  • hidden (torch.Tensor) – Input (batch, seq_len, dim)

  • shift (torch.Tensor) – Shift parameter (batch, dim)

  • scale (torch.Tensor) – Scale parameter (batch, dim)

Return type:

torch.Tensor

Returns:

Modulated tensor (batch, seq_len, dim)

class schrodinger.application.matsci.flywheel.ldmol.dit_model.Attention(dim, num_heads=8, qkv_bias=False, qk_norm=False, attn_drop=0.0, proj_drop=0.0, norm_layer=<class 'torch.nn.modules.normalization.LayerNorm'>)

Bases: Module

Multi-head self-attention with optional fused kernel.

__init__(dim, num_heads=8, qkv_bias=False, qk_norm=False, attn_drop=0.0, proj_drop=0.0, norm_layer=<class 'torch.nn.modules.normalization.LayerNorm'>)
Parameters:

  • dim (int) – Input dimension

  • num_heads (int) – Number of attention heads

  • qkv_bias (bool) – Use bias in QKV projection

  • qk_norm (bool) – Normalize Q and K

  • attn_drop (float) – Attention dropout rate

  • proj_drop (float) – Projection dropout rate

  • norm_layer (nn.Module) – Normalization layer

forward(hidden)
Parameters:

hidden (torch.Tensor) – Input (batch, seq_len, dim)

Return type:

torch.Tensor

Returns:

Output (batch, seq_len, dim)

class schrodinger.application.matsci.flywheel.ldmol.dit_model.CrossAttention(dim, dim_y, num_heads=8, qkv_bias=False, qk_norm=False, attn_drop=0.0, proj_drop=0.0, norm_layer=<class 'torch.nn.modules.normalization.LayerNorm'>)

Bases: Module

Cross-attention: queries from latent, keys/values from text.

__init__(dim, dim_y, num_heads=8, qkv_bias=False, qk_norm=False, attn_drop=0.0, proj_drop=0.0, norm_layer=<class 'torch.nn.modules.normalization.LayerNorm'>)
Parameters:

  • dim (int) – Query dimension

  • dim_y (int) – Key/Value dimension

  • num_heads (int) – Number of attention heads

  • qkv_bias (bool) – Use bias in QKV projection

  • qk_norm (bool) – Normalize Q and K

  • attn_drop (float) – Attention dropout rate

  • proj_drop (float) – Projection dropout rate

  • norm_layer (nn.Module) – Normalization layer

forward(hidden, text, pad_mask=None)
Parameters:

  • hidden (torch.Tensor) – Query (batch, seq_q, dim)

  • text (torch.Tensor) – Key/Value (batch, seq_k, dim_y)

  • pad_mask (torch.Tensor) – Text attention mask (batch, seq_k)

Return type:

torch.Tensor

Returns:

Output (batch, seq_q, dim)

class schrodinger.application.matsci.flywheel.ldmol.dit_model.Mlp(in_features, hidden_features=None, out_features=None, act_layer=<class 'torch.nn.modules.activation.GELU'>, drop=0.0)

Bases: Module

Two-layer MLP with GELU activation.

__init__(in_features, hidden_features=None, out_features=None, act_layer=<class 'torch.nn.modules.activation.GELU'>, drop=0.0)
Parameters:

  • in_features (int) – Input dimension

  • hidden_features (int) – Hidden dimension

  • out_features (int) – Output dimension

  • act_layer (nn.Module) – Activation function

  • drop (float) – Dropout rate

forward(hidden)
Parameters:

hidden (torch.Tensor) – Input

Return type:

torch.Tensor

Returns:

Output

class schrodinger.application.matsci.flywheel.ldmol.dit_model.TimestepEmbedder(hidden_size, frequency_embedding_size=256)

Bases: Module

Embed diffusion timesteps via sinusoidal frequencies + 2-layer MLP.

__init__(hidden_size, frequency_embedding_size=256)
Parameters:

  • hidden_size (int) – Output embedding dimension

  • frequency_embedding_size (int) – Sinusoidal frequency dimension

static timestepEmbedding(timesteps, dim, max_period=10000)

Create sinusoidal timestep embeddings (OpenAI GLIDE style).

Frequency: cos/sin of t * 10000^(-2i/dim) for i in [0, dim/2).

Parameters:

  • timesteps (torch.Tensor) – 1-D batch of timestep indices

  • dim (int) – Embedding dimension

  • max_period (int) – Controls minimum embedding frequency

Return type:

torch.Tensor

Returns:

Embeddings (batch, dim)

forward(timesteps)
Parameters:

timesteps (torch.Tensor) – Diffusion timestep indices (batch,)

Return type:

torch.Tensor

Returns:

Timestep embeddings (batch, hidden_size)

class schrodinger.application.matsci.flywheel.ldmol.dit_model.DiTBlock(hidden_size, num_heads, mlp_ratio=4.0, cross_attn=0, **block_kwargs)

Bases: Module

DiT block with adaLN-Zero conditioning.

Self-attention + optional cross-attention to text + MLP, each gated by learned scalars from the timestep embedding.

__init__(hidden_size, num_heads, mlp_ratio=4.0, cross_attn=0, **block_kwargs)
Parameters:

  • hidden_size (int) – Hidden dimension

  • num_heads (int) – Number of attention heads

  • mlp_ratio (float) – MLP hidden size multiplier

  • cross_attn (int) – Cross-attention key/value dim (0 = disabled)

  • block_kwargs – Extra kwargs forwarded to attention modules

forward(hidden, cond, text=None, pad_mask=None)
Parameters:

  • hidden (torch.Tensor) – Input (batch, seq_len, hidden_size)

  • cond (torch.Tensor) – Timestep conditioning (batch, hidden_size)

  • text (torch.Tensor) – Text cross-attention input

  • pad_mask (torch.Tensor) – Text attention mask

Return type:

torch.Tensor

Returns:

Output (batch, seq_len, hidden_size)

class schrodinger.application.matsci.flywheel.ldmol.dit_model.FinalLayer(hidden_size, patch_size, out_channels)

Bases: Module

adaLN normalization + linear output projection.

__init__(hidden_size, patch_size, out_channels)
Parameters:

  • hidden_size (int) – Hidden dimension

  • patch_size (int) – Patch size (always 1 for LDMol)

  • out_channels (int) – Output channels

forward(hidden, cond)
Parameters:

  • hidden (torch.Tensor) – Input (batch, seq_len, hidden_size)

  • cond (torch.Tensor) – Timestep conditioning (batch, hidden_size)

Return type:

torch.Tensor

Returns:

Output (batch, seq_len, out_channels)

schrodinger.application.matsci.flywheel.ldmol.dit_model.get_1d_sincos_pos_embed_from_grid(embed_dim, pos)

Build 1D sinusoidal positional embeddings from a position grid.

Frequencies: sin/cos of pos * 10000^(-2i/embed_dim) for i in [0, D/2).

Parameters:

  • embed_dim (int) – Embedding dimension (must be even)

  • pos (numpy.ndarray) – Position indices (M,)

Return type:

numpy.ndarray

Returns:

Positional embeddings (M, embed_dim)

schrodinger.application.matsci.flywheel.ldmol.dit_model.get_2d_sincos_pos_embed(embed_dim, grid_size)

Build 1D positional embeddings for a sequence of length grid_size.

Named get_2d_sincos_pos_embed to match the original DiT checkpoint interface; LDMol only uses the 1D variant (no spatial patching).

Parameters:

  • embed_dim (int) – Embedding dimension

  • grid_size (int) – Sequence length

Return type:

numpy.ndarray

Returns:

Positional embeddings (grid_size, embed_dim)

class schrodinger.application.matsci.flywheel.ldmol.dit_model.DiT(input_size=127, patch_size=1, in_channels=64, hidden_size=768, depth=12, num_heads=16, mlp_ratio=4.0, cross_attn=768, condition_dim=1024, learn_sigma=True)

Bases: Module

Diffusion Transformer for LDMol molecular generation.

Takes noisy latents (batch, channels, seq_len, 1) and a timestep, conditioned on MolT5 text embeddings via cross-attention, and predicts noise + variance (batch, 2*channels, seq_len, 1).

__init__(input_size=127, patch_size=1, in_channels=64, hidden_size=768, depth=12, num_heads=16, mlp_ratio=4.0, cross_attn=768, condition_dim=1024, learn_sigma=True)
Parameters:

  • input_size (int) – Latent sequence length

  • patch_size (int) – Patch size (always 1 for LDMol)

  • in_channels (int) – Input channel dimension (VAE latent dim)

  • hidden_size (int) – Transformer hidden dimension

  • depth (int) – Number of DiTBlocks

  • num_heads (int) – Attention heads per block

  • mlp_ratio (float) – MLP hidden size multiplier

  • cross_attn (int) – Cross-attention key/value dimension

  • condition_dim (int) – Text encoder output width

  • learn_sigma (bool) – Predict variance (doubles out_channels)

initializeWeights()

Xavier-initialize all linear layers; zero adaLN and output projections.

Zero-init of adaLN-Zero gates ensures identity residuals at the start of training (Peebles & Xie, 2023, Section 3).

forward(noisy, timesteps, text=None, pad_mask=None)

Predict noise and variance from noisy latents.

Parameters:

  • noisy (torch.Tensor) – Noisy latents (batch, channels, seq_len, 1)

  • timesteps (torch.Tensor) – Diffusion timestep indices (batch,)

  • text (torch.Tensor) – MolT5 embeddings (batch, text_len, cond_dim)

  • pad_mask (torch.Tensor) – Text attention mask (batch, text_len)

Return type:

torch.Tensor

Returns:

Noise + variance predictions (batch, out_channels, seq_len, 1)

forwardWithCfg(noisy, timesteps, text, pad_mask, cfg_scale)

Forward pass with classifier-free guidance (CFG).

Caller must pass doubled batch: noisy = cat([z, z]), text = cat([text_cond, text_null]), pad_mask likewise.

Parameters:

  • noisy (torch.Tensor) – (2*batch, channels, seq_len, 1)

  • timesteps (torch.Tensor) – (2*batch,)

  • text (torch.Tensor) – (2*batch, text_len, condition_dim)

  • pad_mask (torch.Tensor) – (2*batch, text_len)

  • cfg_scale (float) – Guidance strength (>= 1.0)

Return type:

torch.Tensor

Returns:

Guided predictions (2*batch, out_channels, seq_len, 1)

schrodinger.application.matsci.flywheel.ldmol.dit_model.create_ldmol_dit(device='cpu')

Build the LDMol DiT model with its pre-trained architecture.

All hyperparameters are fixed by the pre-trained checkpoint; only device placement is a runtime choice.

Parameters:

device (str) – Compute device (‘cpu’, ‘cuda’, ‘mps’)

Return type:

DiT

Returns:

DiT model in eval mode on the specified device