159 lines
4.9 KiB
Python
159 lines
4.9 KiB
Python
"""Conditional Flow Matching model with Transformer backbone.
|
|
|
|
Predicts velocity field v_θ(x_t, t, cond) that transports noise to data.
|
|
- x_t: (B, T, 3) noisy trajectory at time t
|
|
- t: (B,) interpolation time ∈ [0,1]
|
|
- cond: (B, 3) condition [dist_norm, log_dist, log_dur]
|
|
- Output: (B, T, 3) velocity field
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Sinusoidal time embedding
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class SinusoidalTimeEmbedding(nn.Module):
|
|
"""Map scalar timestep t ∈ [0,1] to a d_model-dimensional vector."""
|
|
|
|
def __init__(self, d_model: int):
|
|
super().__init__()
|
|
self.d_model = d_model
|
|
|
|
def forward(self, t: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Args:
|
|
t: (B,) scalar timesteps
|
|
|
|
Returns:
|
|
(B, d_model) embedding vectors
|
|
"""
|
|
half = self.d_model // 2
|
|
freqs = torch.exp(
|
|
-math.log(10000.0) * torch.arange(half, device=t.device, dtype=t.dtype) / half
|
|
)
|
|
args = t.unsqueeze(-1) * freqs.unsqueeze(0) # (B, half)
|
|
return torch.cat([torch.sin(args), torch.cos(args)], dim=-1) # (B, d_model)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TrajectoryFlowModel
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TrajectoryFlowModel(nn.Module):
|
|
"""Conditional Flow Matching model for mouse trajectory generation.
|
|
|
|
Architecture:
|
|
- SinusoidalTimeEmbedding: t scalar → d_model vector
|
|
- Input projection: 3 → d_model
|
|
- Learned positional embedding: (1, seq_len, d_model)
|
|
- Time embed + Condition embed added as bias to all tokens
|
|
- 4-layer TransformerEncoder (pre-norm, GELU, batch_first=True)
|
|
- Output projection: d_model → 3
|
|
|
|
Args:
|
|
seq_len: Number of time steps (default 64).
|
|
d_model: Transformer hidden dimension (default 128).
|
|
nhead: Number of attention heads (default 4).
|
|
num_layers: Number of transformer layers (default 4).
|
|
dim_feedforward: Feedforward hidden size (default 256).
|
|
dropout: Dropout rate (default 0.1).
|
|
cond_dim: Condition vector size (default 3).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
seq_len: int = 64,
|
|
d_model: int = 128,
|
|
nhead: int = 4,
|
|
num_layers: int = 4,
|
|
dim_feedforward: int = 256,
|
|
dropout: float = 0.1,
|
|
cond_dim: int = 3,
|
|
):
|
|
super().__init__()
|
|
self.seq_len = seq_len
|
|
self.d_model = d_model
|
|
self.cond_dim = cond_dim
|
|
|
|
# Input projection: (forward, lateral, log_dt) → d_model
|
|
self.input_proj = nn.Linear(3, d_model)
|
|
|
|
# Learned positional embedding
|
|
self.pos_embed = nn.Parameter(torch.randn(1, seq_len, d_model) * 0.02)
|
|
|
|
# Time embedding
|
|
self.time_embed = SinusoidalTimeEmbedding(d_model)
|
|
self.time_mlp = nn.Sequential(
|
|
nn.Linear(d_model, d_model),
|
|
nn.GELU(),
|
|
nn.Linear(d_model, d_model),
|
|
)
|
|
|
|
# Condition embedding
|
|
self.cond_mlp = nn.Sequential(
|
|
nn.Linear(cond_dim, d_model),
|
|
nn.GELU(),
|
|
nn.Linear(d_model, d_model),
|
|
)
|
|
|
|
# Transformer encoder (pre-norm via norm_first=True)
|
|
encoder_layer = nn.TransformerEncoderLayer(
|
|
d_model=d_model,
|
|
nhead=nhead,
|
|
dim_feedforward=dim_feedforward,
|
|
dropout=dropout,
|
|
activation="gelu",
|
|
batch_first=True,
|
|
norm_first=True,
|
|
)
|
|
self.transformer = nn.TransformerEncoder(
|
|
encoder_layer, num_layers=num_layers, enable_nested_tensor=False
|
|
)
|
|
|
|
# Output projection: d_model → 3
|
|
self.output_proj = nn.Linear(d_model, 3)
|
|
|
|
def forward(
|
|
self,
|
|
x_t: torch.Tensor,
|
|
t: torch.Tensor,
|
|
cond: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Predict velocity field.
|
|
|
|
Args:
|
|
x_t: (B, T, 3) noisy trajectory at interpolation time t
|
|
t: (B,) interpolation timestep ∈ [0,1]
|
|
cond: (B, 3) condition vector [dist_norm, log_dist, log_dur]
|
|
|
|
Returns:
|
|
(B, T, 3) predicted velocity field
|
|
"""
|
|
# Project input tokens
|
|
h = self.input_proj(x_t) # (B, T, d_model)
|
|
|
|
# Add positional embedding
|
|
h = h + self.pos_embed # (B, T, d_model)
|
|
|
|
# Time embedding → bias added to all tokens
|
|
t_emb = self.time_mlp(self.time_embed(t)) # (B, d_model)
|
|
h = h + t_emb.unsqueeze(1) # broadcast over T
|
|
|
|
# Condition embedding → bias added to all tokens
|
|
c_emb = self.cond_mlp(cond) # (B, d_model)
|
|
h = h + c_emb.unsqueeze(1) # broadcast over T
|
|
|
|
# Transformer
|
|
h = self.transformer(h) # (B, T, d_model)
|
|
|
|
# Output projection
|
|
return self.output_proj(h) # (B, T, 3)
|