Initial commit: code, paper, small artifacts

Mixed_CFM/model.py

@@ -0,0 +1,244 @@
+from __future__ import annotations
+import math
+from dataclasses import dataclass, field
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import importlib.util as _ilu
+import sys as _sys
+from pathlib import Path as _Path
+_UNIFIED_NAME = 'unified_cfm_model'
+if _UNIFIED_NAME not in _sys.modules:
+    _unified_spec = _ilu.spec_from_file_location(_UNIFIED_NAME, _Path(__file__).resolve().parents[1] / 'Unified_CFM' / 'model.py')
+    _unified = _ilu.module_from_spec(_unified_spec)
+    _sys.modules[_UNIFIED_NAME] = _unified
+    _unified_spec.loader.exec_module(_unified)
+else:
+    _unified = _sys.modules[_UNIFIED_NAME]
+AdaLNBlock = _unified.AdaLNBlock
+SinusoidalTimeEmb = _unified.SinusoidalTimeEmb
+_sinkhorn_coupling = _unified._sinkhorn_coupling
+
+@dataclass
+class MixedCFMConfig:
+    T: int = 64
+    flow_dim: int = 20
+    n_cont_pkt: int = 3
+    n_disc_pkt: int = 6
+    cont_pkt_idx: tuple[int, ...] = (0, 1, 8)
+    disc_pkt_idx: tuple[int, ...] = (2, 3, 4, 5, 6, 7)
+    n_disc_classes: int = 2
+    token_dim: int | None = None
+    d_model: int = 128
+    n_layers: int = 4
+    n_heads: int = 4
+    mlp_ratio: float = 4.0
+    time_dim: int = 64
+    sigma: float = 0.1
+    use_ot: bool = False
+    reference_mode: str | None = None
+    lambda_disc: float = 1.0
+    disc_path: str = 'uniform'
+    disc_embed_scale: float = 1.0
+
+    def __post_init__(self) -> None:
+        if len(self.cont_pkt_idx) != self.n_cont_pkt:
+            raise ValueError('cont_pkt_idx length mismatch n_cont_pkt')
+        if len(self.disc_pkt_idx) != self.n_disc_pkt:
+            raise ValueError('disc_pkt_idx length mismatch n_disc_pkt')
+        if self.disc_path != 'uniform':
+            raise NotImplementedError(f'disc_path={self.disc_path}')
+
+class MixedVelocity(nn.Module):
+
+    def __init__(self, token_dim: int, seq_len: int, n_disc: int, n_classes: int, d_model: int=128, n_layers: int=4, n_heads: int=4, mlp_ratio: float=4.0, time_dim: int=64, reference_mode: str | None=None) -> None:
+        super().__init__()
+        if reference_mode not in (None, 'causal_packets', 'causal_all'):
+            raise ValueError(f'reference_mode={reference_mode!r}')
+        self.token_dim = token_dim
+        self.seq_len = seq_len
+        self.n_disc = n_disc
+        self.n_classes = n_classes
+        self.reference_mode = reference_mode
+        self.input_proj = nn.Linear(token_dim, d_model)
+        self.pos_emb = nn.Parameter(torch.zeros(1, seq_len, d_model))
+        self.type_emb = nn.Embedding(2, d_model)
+        nn.init.trunc_normal_(self.pos_emb, std=0.02)
+        nn.init.normal_(self.type_emb.weight, std=0.02)
+        self.time_emb = SinusoidalTimeEmb(time_dim)
+        self.cond_mlp = nn.Sequential(nn.Linear(time_dim, d_model), nn.SiLU(), nn.Linear(d_model, d_model))
+        self.blocks = nn.ModuleList([AdaLNBlock(d_model, n_heads, mlp_ratio, cond_dim=d_model) for _ in range(n_layers)])
+        self.out_norm = nn.LayerNorm(d_model, elementwise_affine=False)
+        self.head_v = nn.Linear(d_model, token_dim)
+        self.head_disc = nn.Linear(d_model, n_disc * n_classes)
+        for layer in (self.head_v, self.head_disc):
+            nn.init.zeros_(layer.weight)
+            nn.init.zeros_(layer.bias)
+        type_ids = torch.ones(seq_len, dtype=torch.long)
+        type_ids[0] = 0
+        self.register_buffer('type_ids', type_ids, persistent=False)
+
+    def _attn_mask(self, L: int, device: torch.device) -> torch.Tensor | None:
+        if self.reference_mode is None:
+            return None
+        if self.reference_mode == 'causal_packets':
+            mask = torch.zeros((L, L), dtype=torch.bool, device=device)
+            if L > 1:
+                mask[1:, 1:] = torch.triu(torch.ones(L - 1, L - 1, dtype=torch.bool, device=device), diagonal=1)
+            return mask
+        return torch.triu(torch.ones(L, L, dtype=torch.bool, device=device), diagonal=1)
+
+    def forward(self, x: torch.Tensor, t: torch.Tensor, key_padding_mask: torch.Tensor | None=None) -> tuple[torch.Tensor, torch.Tensor]:
+        (B, L, _) = x.shape
+        if t.dim() == 0:
+            t = t.expand(B)
+        h = self.input_proj(x)
+        h = h + self.pos_emb[:, :L, :] + self.type_emb(self.type_ids[:L])[None, :, :]
+        cond = self.cond_mlp(self.time_emb(t))
+        attn_mask = self._attn_mask(L, x.device)
+        for block in self.blocks:
+            h = block(h, cond, key_padding_mask, attn_mask=attn_mask)
+        h = self.out_norm(h)
+        v = self.head_v(h)
+        d = self.head_disc(h).view(B, L, self.n_disc, self.n_classes)
+        return (v, d)
+
+class MixedTokenCFM(nn.Module):
+
+    def __init__(self, cfg: MixedCFMConfig) -> None:
+        super().__init__()
+        self.cfg = cfg
+        cont_size = cfg.n_cont_pkt + cfg.n_disc_pkt
+        self.token_dim = cfg.token_dim or 1 + max(cfg.flow_dim, cont_size)
+        if self.token_dim < 1 + max(cfg.flow_dim, cont_size):
+            raise ValueError('token_dim too small')
+        self.seq_len = cfg.T + 1
+        self.velocity = MixedVelocity(token_dim=self.token_dim, seq_len=self.seq_len, n_disc=cfg.n_disc_pkt, n_classes=cfg.n_disc_classes, d_model=cfg.d_model, n_layers=cfg.n_layers, n_heads=cfg.n_heads, mlp_ratio=cfg.mlp_ratio, time_dim=cfg.time_dim, reference_mode=cfg.reference_mode)
+
+    def _embed_disc(self, x_disc_int: torch.Tensor) -> torch.Tensor:
+        s = self.cfg.disc_embed_scale
+        return (x_disc_int.float() - 0.5) * s
+
+    def build_tokens(self, flow: torch.Tensor, packets_cont: torch.Tensor, x_disc_t_int: torch.Tensor) -> torch.Tensor:
+        (B, T, Cp) = packets_cont.shape
+        assert T == self.cfg.T and Cp == self.cfg.n_cont_pkt
+        z = packets_cont.new_zeros((B, T + 1, self.token_dim))
+        z[:, 0, 0] = -1.0
+        z[:, 0, 1:1 + self.cfg.flow_dim] = flow
+        z[:, 1:, 0] = 1.0
+        z[:, 1:, 1:1 + self.cfg.n_cont_pkt] = packets_cont
+        z[:, 1:, 1 + self.cfg.n_cont_pkt:1 + self.cfg.n_cont_pkt + self.cfg.n_disc_pkt] = self._embed_disc(x_disc_t_int)
+        return z
+
+    def key_padding_mask(self, lens: torch.Tensor) -> torch.Tensor:
+        B = lens.shape[0]
+        idx = torch.arange(self.cfg.T, device=lens.device)[None, :]
+        packet_real = idx < lens[:, None]
+        real = torch.cat([torch.ones(B, 1, dtype=torch.bool, device=lens.device), packet_real], dim=1)
+        return ~real
+
+    def _loss_mask(self, lens: torch.Tensor) -> torch.Tensor:
+        return (~self.key_padding_mask(lens)).float()
+
+    def compute_loss(self, flow: torch.Tensor, packets_cont: torch.Tensor, packets_disc: torch.Tensor, lens: torch.Tensor, *, return_components: bool=False) -> torch.Tensor | dict[str, torch.Tensor]:
+        (B, T, _) = packets_cont.shape
+        device = packets_cont.device
+        mask = self._loss_mask(lens)
+        kpm = mask == 0
+        x_1_cont = self.build_tokens(flow, packets_cont, torch.zeros_like(packets_disc))
+        x_0_cont = torch.randn_like(x_1_cont)
+        if self.cfg.use_ot:
+            flat0 = (x_0_cont * mask[:, :, None]).reshape(B, -1)
+            flat1 = (x_1_cont * mask[:, :, None]).reshape(B, -1)
+            col = _sinkhorn_coupling(torch.cdist(flat0.float(), flat1.float()))
+            x_1_cont = x_1_cont[col]
+            packets_cont = packets_cont[col]
+            packets_disc = packets_disc[col]
+            flow = flow[col]
+            lens = lens[col]
+            mask = self._loss_mask(lens)
+            kpm = mask == 0
+        t = torch.rand(B, device=device)
+        x_t_cont = (1.0 - t[:, None, None]) * x_0_cont + t[:, None, None] * x_1_cont
+        if self.cfg.sigma > 0:
+            std = self.cfg.sigma * torch.sqrt(t * (1.0 - t))[:, None, None]
+            x_t_cont = x_t_cont + std * torch.randn_like(x_t_cont)
+        target_cont = x_1_cont - x_0_cont
+        u = torch.rand(B, T, self.cfg.n_disc_pkt, device=device)
+        keep = u < t[:, None, None]
+        rand_disc = torch.randint(0, self.cfg.n_disc_classes, packets_disc.shape, device=device)
+        x_disc_t = torch.where(keep, packets_disc, rand_disc)
+        disc_start = 1 + self.cfg.n_cont_pkt
+        x_t_full = x_t_cont.clone()
+        x_t_full[:, 1:, disc_start:disc_start + self.cfg.n_disc_pkt] = self._embed_disc(x_disc_t)
+        (v_pred, d_logits) = self.velocity(x_t_full, t, key_padding_mask=kpm)
+        v_err = (v_pred - target_cont).square()
+        v_err[:, :, disc_start:disc_start + self.cfg.n_disc_pkt] = 0.0
+        v_per_token = v_err.mean(dim=-1)
+        per_sample = (v_per_token * mask).sum(dim=-1) / mask.sum(dim=-1).clamp_min(1.0)
+        L_cont = per_sample.mean()
+        pkt_logits = d_logits[:, 1:]
+        pkt_real = mask[:, 1:].bool()
+        corrupt = ~keep & pkt_real[:, :, None]
+        flat_logits = pkt_logits.reshape(-1, self.cfg.n_disc_classes)
+        flat_targets = packets_disc.reshape(-1).long()
+        flat_ce = F.cross_entropy(flat_logits, flat_targets, reduction='none')
+        flat_ce = flat_ce.view(B, T, self.cfg.n_disc_pkt)
+        flat_ce = flat_ce * corrupt.float()
+        denom = corrupt.float().sum().clamp_min(1.0)
+        L_disc = flat_ce.sum() / denom
+        total = L_cont + self.cfg.lambda_disc * L_disc
+        if return_components:
+            return {'total': total, 'main': L_cont.detach(), 'aux_disc': L_disc.detach(), 'aux_flow': L_cont.new_zeros(()), 'aux_packet': L_cont.new_zeros(())}
+        return total
+
+    @torch.no_grad()
+    def trajectory_metrics(self, flow: torch.Tensor, packets_cont: torch.Tensor, packets_disc: torch.Tensor, lens: torch.Tensor, n_steps: int=16) -> dict[str, torch.Tensor]:
+        z = self.build_tokens(flow, packets_cont, packets_disc)
+        mask = self._loss_mask(lens)
+        kpm = mask == 0
+        B = z.shape[0]
+        dt = 1.0 / n_steps
+        disc_start = 1 + self.cfg.n_cont_pkt
+        disc_end = disc_start + self.cfg.n_disc_pkt
+        disc_embed = z[:, 1:, disc_start:disc_end].clone()
+        for k in range(n_steps):
+            t_val = 1.0 - k * dt
+            t = torch.full((B,), t_val, device=z.device)
+            (v, _) = self.velocity(z, t, key_padding_mask=kpm)
+            v[:, :, disc_start:disc_end] = 0.0
+            z = z - v * dt
+            z[:, 1:, disc_start:disc_end] = disc_embed
+        z_real = z * mask[:, :, None]
+        z_cont = z_real.clone()
+        z_cont[:, 1:, disc_start:disc_end] = 0.0
+        packet_count = mask[:, 1:].sum(dim=-1).clamp_min(1.0)
+        terminal = z_cont.reshape(B, -1).norm(dim=-1) / (mask.sum(dim=-1) * self.token_dim).clamp_min(1.0).sqrt()
+        terminal_flow = z_cont[:, 0].norm(dim=-1) / math.sqrt(self.token_dim)
+        terminal_packet = (z_cont[:, 1:] * mask[:, 1:, None]).reshape(B, -1).norm(dim=-1) / (packet_count * self.token_dim).sqrt()
+        return {'terminal_norm': terminal, 'terminal_flow': terminal_flow, 'terminal_packet': terminal_packet}
+
+    @torch.no_grad()
+    def disc_nll_score(self, flow: torch.Tensor, packets_cont: torch.Tensor, packets_disc: torch.Tensor, lens: torch.Tensor, t_eval: float=0.5) -> dict[str, torch.Tensor]:
+        (B, T, _) = packets_cont.shape
+        device = packets_cont.device
+        mask = self._loss_mask(lens)
+        kpm = mask == 0
+        z = self.build_tokens(flow, packets_cont, packets_disc)
+        t = torch.full((B,), float(t_eval), device=device)
+        (_, d_logits) = self.velocity(z, t, key_padding_mask=kpm)
+        pkt_logits = d_logits[:, 1:]
+        flat_logits = pkt_logits.reshape(-1, self.cfg.n_disc_classes)
+        flat_targets = packets_disc.reshape(-1).long()
+        ce = F.cross_entropy(flat_logits, flat_targets, reduction='none')
+        ce = ce.view(B, T, self.cfg.n_disc_pkt)
+        pkt_real = mask[:, 1:].bool().float()
+        per_sample = (ce.sum(dim=-1) * pkt_real).sum(dim=-1) / pkt_real.sum(dim=-1).clamp_min(1.0)
+        per_ch = (ce * pkt_real[:, :, None]).sum(dim=1) / pkt_real.sum(dim=1).clamp_min(1.0)[:, None]
+        out = {'disc_nll_total': per_sample}
+        for (c, idx) in enumerate(self.cfg.disc_pkt_idx):
+            out[f'disc_nll_ch{idx}'] = per_ch[:, c]
+        return out
+
+    def param_count(self) -> int:
+        return sum((p.numel() for p in self.parameters()))