Module diffusion 

AF3-style EDM diffusion module for all-atom coordinate generation.

Architecture:

1. Single repr projected to token dim: [B, N, d_single] → [B, N, c_token=768]
2. Fourier noise embedding encodes σ_t → [fourier_dim=256]
3. 12 token-level transformer blocks (c_token=768, 16 heads, with pair bias)
4. 3 atom-level transformer blocks (c_atom=128, 4 heads) — TODO
5. Output projection: [B, N, c_token] → [B, N*n_atoms_per_token, 3]

EDM noise schedule (inference):

sigma_t = s_max * (s_min / s_max)^((t / T)^p)
  s_max = 160.0,  s_min = 0.0004,  T = num_steps,  p = 7.0

Stochastic correction: gamma_t = min(gamma_0, sqrt(sigma_next/sigma_t) - 1)
sigma_hat = sigma_t * (1 + gamma_t)
x_hat = x + sqrt(sigma_hat^2 - sigma_t^2) * noise * noise_scale
D = score_network(x_hat, sigma_hat)
x = x_hat + step_scale * (D - x_hat) * (sigma_next / sigma_hat - 1)

Weight layout (rooted at structure_head):

token_proj.*                            — d_single → c_token (no bias)
noise_embedding.*                       — Fourier frequencies (learnable)
noise_proj.*                            — fourier_dim → c_token (no bias)
token_transformer.blocks.{0..11}.*     — 12 token transformer blocks
token_transformer.blocks.{i}.norm.*
token_transformer.blocks.{i}.attn.*    — MHA with pair bias
token_transformer.blocks.{i}.ffn.*     — SwiGLU FFN
out_proj.*                             — c_token → 3 (Cα coords, no bias)

Structs

DiffusionModule

AF3-style EDM diffusion module.