Flow Matching & Diffusion Transformers (DiT): Continuous-Time Generative Physics

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    # x: [B, N_tokens, D], shift: [B, D], scale: [B, D]
    return x * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)

class DiTBlockAdaLNZero(nn.Module):
    """Diffusion Transformer Block with adaLN-Zero conditioning (Peebles & Xie, 2023)."""
    def __init__(self, d_model: int = 128, n_heads: int = 4, mlp_ratio: int = 4):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model, elementwise_affine=False, eps=1e-6)
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.norm2 = nn.LayerNorm(d_model, elementwise_affine=False, eps=1e-6)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, d_model * mlp_ratio),
            nn.GELU(approximate="tanh"),
            nn.Linear(d_model * mlp_ratio, d_model)
        )
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(d_model, 6 * d_model, bias=True)
        )
        # Zero-initialize the modulation linear layer so alpha_1 = alpha_2 = 0 at step 0!
        nn.init.zeros_(self.adaLN_modulation[1].weight)
        nn.init.zeros_(self.adaLN_modulation[1].bias)

    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        # x: [B, N_tokens, D], cond: [B, D] (sum of timestep and class/text embeddings)
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
            self.adaLN_modulation(cond).chunk(6, dim=-1)
        )
        h = modulate(self.norm1(x), shift_msa, scale_msa)
        attn_out, _ = self.attn(h, h, h, need_weights=False)  # [B, N_tokens, D]
        x = x + gate_msa.unsqueeze(1) * attn_out
        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
        return x

def sample_logit_normal_timesteps(
    batch_size: int,
    mu: float = 0.0,
    sigma: float = 1.0,
    res_shift: float = 1.0,
    device: str = "cpu"
) -> torch.Tensor:
    u = torch.randn(batch_size, device=device) * sigma + mu
    t = torch.sigmoid(u)
    if res_shift != 1.0:
        t = (res_shift * t) / (1.0 + (res_shift - 1.0) * t)
    return t  # [B]

def rectified_flow_loss(model_fn, x_1: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
    """Optimal Transport Conditional Flow Matching loss: E[||v_theta(x_t, t, c) - (x_1 - x_0)||^2]."""
    B = x_1.size(0)
    x_0 = torch.randn_like(x_1)  # [B, ...]
    t = sample_logit_normal_timesteps(B, device=x_1.device)
    t_broadcast = t.view(B, *([1] * (x_1.ndim - 1)))  # [B, 1, ...]

    x_t = (1.0 - t_broadcast) * x_0 + t_broadcast * x_1
    target_velocity = x_1 - x_0
    pred_velocity = model_fn(x_t, t, cond)
    return F.mse_loss(pred_velocity, target_velocity)

import torch

@torch.no_grad()
def solve_flow_ode(
    velocity_fn,             # Callable(x_t, t_tensor, cond) -> velocity tensor
    x_0: torch.Tensor,       # Initial Gaussian noise [B, ...]
    cond: torch.Tensor,      # Conditioning tensor [B, ...]
    uncond: torch.Tensor,    # Null conditioning tensor [B, ...]
    num_steps: int = 20,
    cfg_scale: float = 4.5,
    solver: str = "heun"     # 'euler' (1st order) or 'heun' (2nd order predictor-corrector)
) -> torch.Tensor:
    B = x_0.size(0)
    ts = torch.linspace(0.0, 1.0, num_steps + 1, device=x_0.device)
    x_t = x_0.clone()

    def guided_v(x: torch.Tensor, t_scalar: float) -> torch.Tensor:
        t_vec = torch.full((B,), t_scalar, device=x.device, dtype=x.dtype)
        if cfg_scale == 1.0:
            return velocity_fn(x, t_vec, cond)
        v_u = velocity_fn(x, t_vec, uncond)
        v_c = velocity_fn(x, t_vec, cond)
        return v_u + cfg_scale * (v_c - v_u)

    for i in range(num_steps):
        t_cur = ts[i].item()
        t_next = ts[i + 1].item()
        dt = t_next - t_cur

        v_cur = guided_v(x_t, t_cur)
        x_euler = x_t + dt * v_cur

        if solver == "heun":
            v_next = guided_v(x_euler, t_next)
            x_t = x_t + 0.5 * dt * (v_cur + v_next)
        else:
            x_t = x_euler

    return x_t

import torch
import torch.nn as nn

def verify_flow_matching_dit():
    torch.manual_seed(42)

    # 1. Verify DiTBlockAdaLNZero exact identity at step 0
    B, N, D = 4, 16, 64
    x_lat = torch.randn(B, N, D)
    cond = torch.randn(B, D)
    dit_block = DiTBlockAdaLNZero(d_model=D, n_heads=4)
    out_init = dit_block(x_lat, cond)
    identity_diff = (out_init - x_lat).abs().max().item()
    assert identity_diff == 0.0, f"adaLN-Zero failed identity init: {identity_diff}"

    # 2. Train a tiny Conditional Rectified Flow velocity network on 2D conditional target modes
    # Class c=+1 -> target mode at (+3.0, +3.0); Class c=-1 -> target mode at (-3.0, -3.0)
    net = nn.Sequential(
        nn.Linear(2 + 1 + 1, 64),
        nn.SiLU(),
        nn.Linear(64, 64),
        nn.SiLU(),
        nn.Linear(64, 2)
    )
    opt = torch.optim.AdamW(net.parameters(), lr=1e-2)

    def model_fn(x_t: torch.Tensor, t: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
        inp = torch.cat([x_t, t.unsqueeze(-1), c.unsqueeze(-1)], dim=-1)
        return net(inp)

    for _ in range(200):
        c_batch = torch.randint(0, 2, (128,)).float() * 2.0 - 1.0  # {-1, +1}
        x1_batch = c_batch.unsqueeze(-1) * 3.0 + 0.1 * torch.randn(128, 2)
        loss = rectified_flow_loss(model_fn, x1_batch, c_batch)
        opt.zero_grad()
        loss.backward()
        opt.step()

    # 3. Integrate Probability Flow ODE using 2nd-order Heun solver for class c=+1
    x0_test = torch.randn(64, 2)
    c_pos = torch.ones(64)
    c_null = torch.zeros(64)
    x1_gen = solve_flow_ode(model_fn, x0_test, c_pos, c_null, num_steps=16, cfg_scale=1.0, solver="heun")
    gen_mean = x1_gen.mean(dim=0)
    dist_to_target = (gen_mean - torch.tensor([3.0, 3.0])).norm().item()
    assert dist_to_target < 0.35, f"Flow ODE transport error too high: {dist_to_target}"

    print(f"[PASS] DiT adaLN-Zero Step-0 Identity Diff : {identity_diff:.1e} (Exact Identity!)")
    print(f"[PASS] Final Rectified Flow Training Loss  : {loss.item():.4f}")
    print(f"[PASS] Heun ODE Generated Mode Mean        : ({gen_mean[0]:.2f}, {gen_mean[1]:.2f}) vs Target (3.00, 3.00)")

if __name__ == "__main__":
    verify_flow_matching_dit()