Logic/mamba_diffusion
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(as_mamba): Remove dt prediction and use fixed dt

Browse Source 
Removes the `dt_head` network and associated configuration parameters
(dt_min, dt_max, lambda_nfe, warmup_epochs). Replaces predicted time
steps with a fixed value derived from sequence length. Eliminates
the warmup phase and NFE loss calculation.
This commit is contained in:
gameloader
2026-01-21 13:07:36 +08:00
parent c58a73ae26
commit 1446f97459
2 changed files with 21 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

80
as_mamba.py
View File

@@ -25,15 +25,11 @@ class TrainConfig:
batch_size: int = 128
steps_per_epoch: int = 50
epochs: int = 60
warmup_epochs: int = 15
seq_len: int = 20
lr: float = 1e-3
weight_decay: float = 1e-2
dt_min: float = 1e-3
dt_max: float = 0.06
lambda_flow: float = 1.0
lambda_pos: float = 1.0
lambda_nfe: float = 0.05
radius_min: float = 0.6
radius_max: float = 1.4
center_min: float = -6.0
@@ -53,7 +49,7 @@ class TrainConfig:
val_every: int = 200
val_samples: int = 256
val_plot_samples: int = 16
val_max_steps: int = 100
val_max_steps: int = 0
class Mamba2Backbone(nn.Module):
@@ -92,8 +88,6 @@ class ASMamba(nn.Module):
def __init__(self, cfg: TrainConfig) -> None:
super().__init__()
self.cfg = cfg
self.dt_min = float(cfg.dt_min)
self.dt_max = float(cfg.dt_max)
args = Mamba2Config(
d_model=cfg.d_model,
@@ -107,27 +101,20 @@ class ASMamba(nn.Module):
self.backbone = Mamba2Backbone(args, use_residual=cfg.use_residual)
self.input_proj = nn.Linear(3, cfg.d_model)
self.delta_head = nn.Linear(cfg.d_model, 3)
self.dt_head = nn.Sequential(
nn.Linear(cfg.d_model, cfg.d_model),
nn.SiLU(),
nn.Linear(cfg.d_model, 1),
)
def forward(
self, x: Tensor, h: Optional[list[InferenceCache]] = None
) -> tuple[Tensor, Tensor, list[InferenceCache]]:
) -> tuple[Tensor, list[InferenceCache]]:
x_proj = self.input_proj(x)
feats, h = self.backbone(x_proj, h)
delta = self.delta_head(feats)
dt_raw = self.dt_head(feats).squeeze(-1)
dt = torch.clamp(F.softplus(dt_raw), min=self.dt_min, max=self.dt_max)
return delta, dt, h
return delta, h
def step(
self, x: Tensor, h: list[InferenceCache]
) -> tuple[Tensor, Tensor, list[InferenceCache]]:
delta, dt, h = self.forward(x.unsqueeze(1), h)
return delta[:, 0, :], dt[:, 0], h
) -> tuple[Tensor, list[InferenceCache]]:
delta, h = self.forward(x.unsqueeze(1), h)
return delta[:, 0, :], h
def init_cache(self, batch_size: int, device: torch.device) -> list[InferenceCache]:
return [
@@ -240,24 +227,23 @@ def sample_batch(
def compute_losses(
delta: Tensor,
dt: Tensor,
x_seq: Tensor,
x0: Tensor,
v_gt: Tensor,
t_seq: Tensor,
cfg: TrainConfig,
) -> tuple[Tensor, Tensor, Tensor]:
target_disp = v_gt[:, None, :] * dt.unsqueeze(-1)
) -> tuple[Tensor, Tensor]:
dt_fixed = 1.0 / cfg.seq_len
target_disp = v_gt[:, None, :] * dt_fixed
flow_loss = F.mse_loss(delta, target_disp)
t_next = t_seq[None, :, None] + dt.unsqueeze(-1)
t_next = t_seq[None, :, None] + dt_fixed
t_next = torch.clamp(t_next, 0.0, 1.0)
x_target = x0[:, None, :] + t_next * v_gt[:, None, :]
x_next_pred = x_seq + delta
pos_loss = F.mse_loss(x_next_pred, x_target)
nfe_loss = (-torch.log(dt)).mean()
return flow_loss, pos_loss, nfe_loss
return flow_loss, pos_loss
def validate(
@@ -271,12 +257,13 @@ def validate(
) -> None:
model.eval()
center_b, radius_b = sphere_b
steps = cfg.seq_len if cfg.val_max_steps <= 0 else cfg.val_max_steps
with torch.no_grad():
x0 = sample_points_in_sphere(
sphere_a[0], float(sphere_a[1].item()), cfg.val_samples, device
)
traj = rollout_trajectory(model, x0, max_steps=cfg.val_max_steps)
traj = rollout_trajectory(model, x0, steps=steps)
x_final = traj[:, -1, :]
center_b_cpu = center_b.detach().cpu()
@@ -291,6 +278,7 @@ def validate(
"val/final_dist_mean": float(dist.mean().item()),
"val/final_dist_min": float(dist.min().item()),
"val/final_dist_max": float(dist.max().item()),
"val/steps": float(steps),
},
step=step,
)
@@ -351,22 +339,15 @@ def train(cfg: TrainConfig) -> tuple[ASMamba, tuple[Tensor, Tensor], tuple[Tenso
global_step = 0
for epoch in range(cfg.epochs):
warmup = epoch < cfg.warmup_epochs
model.train()
for p in model.dt_head.parameters():
p.requires_grad = not warmup
for _ in range(cfg.steps_per_epoch):
x0, x1, x_seq, t_seq = sample_batch(cfg, sphere_a, sphere_b, device)
v_gt = x1 - x0
delta, dt, _ = model(x_seq)
if warmup:
dt = torch.full_like(dt, 1.0 / cfg.seq_len)
flow_loss, pos_loss, nfe_loss = compute_losses(
delta, _ = model(x_seq)
flow_loss, pos_loss = compute_losses(
delta=delta,
dt=dt,
x_seq=x_seq,
x0=x0,
v_gt=v_gt,
@@ -375,8 +356,6 @@ def train(cfg: TrainConfig) -> tuple[ASMamba, tuple[Tensor, Tensor], tuple[Tenso
)
loss = cfg.lambda_flow * flow_loss + cfg.lambda_pos * pos_loss
if not warmup:
loss = loss + cfg.lambda_nfe * nfe_loss
optimizer.zero_grad(set_to_none=True)
loss.backward()
@@ -388,11 +367,6 @@ def train(cfg: TrainConfig) -> tuple[ASMamba, tuple[Tensor, Tensor], tuple[Tenso
"loss/total": float(loss.item()),
"loss/flow": float(flow_loss.item()),
"loss/pos": float(pos_loss.item()),
"loss/nfe": float(nfe_loss.item()),
"dt/mean": float(dt.mean().item()),
"dt/min": float(dt.min().item()),
"dt/max": float(dt.max().item()),
"stage": 0 if warmup else 1,
},
step=global_step,
)
@@ -408,33 +382,20 @@ def train(cfg: TrainConfig) -> tuple[ASMamba, tuple[Tensor, Tensor], tuple[Tenso
def rollout_trajectory(
model: ASMamba,
x0: Tensor,
max_steps: int = 100,
steps: int,
) -> Tensor:
device = x0.device
model.eval()
h = model.init_cache(batch_size=x0.shape[0], device=device)
x = x0
total_time = torch.zeros(x0.shape[0], device=device)
traj = [x0.detach().cpu()]
with torch.no_grad():
for _ in range(max_steps):
delta, dt, h = model.step(x, h)
dt = torch.clamp(dt, min=model.dt_min, max=model.dt_max)
remaining = 1.0 - total_time
overshoot = dt > remaining
if overshoot.any():
scale = (remaining / dt).unsqueeze(-1)
delta = torch.where(overshoot.unsqueeze(-1), delta * scale, delta)
dt = torch.where(overshoot, remaining, dt)
for _ in range(steps):
delta, h = model.step(x, h)
x = x + delta
total_time = total_time + dt
traj.append(x.detach().cpu())
if torch.all(total_time >= 1.0 - 1e-6):
break
return torch.stack(traj, dim=1)
@@ -504,7 +465,8 @@ def run_training_and_plot(cfg: TrainConfig) -> Path:
x0 = sample_points_in_sphere(
sphere_a[0], float(sphere_a[1].item()), plot_samples, device
)
traj = rollout_trajectory(model, x0, max_steps=cfg.val_max_steps)
steps = cfg.seq_len if cfg.val_max_steps <= 0 else cfg.val_max_steps
traj = rollout_trajectory(model, x0, steps=steps)
output_dir = Path(cfg.output_dir)
save_path = output_dir / "as_mamba_trajectory.png"
plot_trajectories(traj, sphere_a, sphere_b, save_path)

1
main.py
View File

@@ -6,7 +6,6 @@ from as_mamba import TrainConfig, run_training_and_plot
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Train AS-Mamba on sphere-to-sphere flow.")
parser.add_argument("--epochs", type=int, default=None)
parser.add_argument("--warmup-epochs", type=int, default=None)
parser.add_argument("--batch-size", type=int, default=None)
parser.add_argument("--steps-per-epoch", type=int, default=None)
parser.add_argument("--seq-len", type=int, default=None)