docs/v13d_spatial_beats_design.md

# v13_D 设计文档 — Loss 机制 + 训练细节双重改进

> 日期：2026-05-02
> 起源：v13_B / v13_C 实验 ep3 结果证明 "loss/arch 改动在 cls warmup 结束时贡献微弱"，需要重新从**训练机制**角度切入
> 目标：F20 从 v12 的 0.378 推到 **0.43+**
> 设计原则：延续 v13_B/C 的"增量式、不破坏现有实验"，所有新 flag 默认 False

---

## 0. 为什么 v13_B / v13_C 都不 work

ep3（cls warmup 结束）对比：

|  | v12 ep3 | v13_B ep3 | v13_C ep3 |
|---|---|---|---|
| F20 | 0.3529 | 0.3565 (+0.004) | 0.3849 (+0.032) |
| o_cls | 0.7793 | 0.7757 | 0.7779 |
| aR | 0.120 | 0.194 | 0.090 |
| aP | 0.866 | 0.860 | 0.859 |

**三个实验的 oracle_class_acc 基本一致** → 三组实验的表征学的是同一个东西，新改动都没触达表征。loss/head decision 层面的改动（ASL / gate / soft-F1）在 cls warmup 结束时**基本还没起作用**，因为：

1. **zero-init 的新模块需要 3-5 ep 才能让 layer_scale 爬起来**，而 cls warmup 只有 3 ep
2. **ASL γ-=4 让 loss_activity 绝对值从 0.29 降到 0.07**，实际上是把 activity head 的梯度信号弄弱了
3. **augment 反而降低 activity_precision**（0.95 → 0.86）
4. **v12 自身从 ep3→ep12 只涨了 +0.025**，外推 v13_B/C 的 best 也就 0.38~0.41

## 1. v13_D 的切入点：**训练机制本身**

不碰架构、不碰表征，只改三件事：

- **D-1**: 扩大 cls warmup（3 ep → 8 ep）+ cosine LR + 总 25 ep，让表征学得更稳
- **D-2**: 用 Top-K rank activity loss 替换 BCE，直接针对 activity_recall=0.13 的根本瓶颈
- **D-5**: resume optimizer，从 v12 最后的 Adam momentum 继续，避免前 2-3 ep 梯度方向混乱
- **D-6**: EMA 权重，validate 时用 EMA 模型，ep 间震荡降低 1-2 个点

D-2 是核心，D-1 / D-5 / D-6 是辅助。

## 2. 具体设计

### D-1: cls warmup 拉长 + cosine LR + 25 epochs

#### 诊断
v12 `frame_spatial_loss_warmup_epochs = 3`，ep0-2 只训练 cls + activity，ep3 起放开 spatial loss。但：
- v12 曲线：cls_acc ep0=0.65 → ep1=0.81 → ep2=0.85 → ep3=0.84 → ... → ep12=0.83
- **ep2 已 0.85 但 ep3 反而微降 0.84**，说明 **ep3 放开 spatial loss 的瞬间干扰了 cls**
- cls 没有再涨的机会 —— 之后一直在 0.83 附近波动
- FSD63 的 oracle_cls 卡在 0.78 是 class head 学的，不是 trunk 表征的上限

#### 改动
```python
cfg.frame_spatial_loss_warmup_epochs = 8      # 3 → 8
cfg.frame_spatial_loss_ramp_epochs = 2        # 1 → 2（ramp 更平滑）
cfg.num_epochs = 25                           # 15 → 25

# LR: cosine schedule，峰值在 ep 5-8 之间（warmup 结束附近）
cfg.use_cosine_lr = True                      # 新 flag
cfg.cosine_lr_warmup_epochs = 3               # 前 3 ep linear warmup 到 peak
cfg.cosine_lr_min_ratio = 0.05                # 最后降到 peak * 0.05
```

**训练循环改动**：
```python
# ep 0-2: linear warmup LR 0 → peak_lr
# ep 3-24: cosine decay peak_lr → peak_lr * min_ratio
# ep 0-7: spatial loss weight = warmup_scale (0.0 or 0.1)
# ep 8-9: linear ramp to 1.0
# ep 10-24: full spatial loss
```

### D-2: Top-K rank activity loss（核心）

#### 诊断
当前 BCE（或 ASL）对每个 `(b, k, t)` 位置独立判断"该 slot 是否 active"。问题：
- K=4 slots 在 ov1 数据上永远 3/4 inactive，类别极不平衡
- BCE 的最优解是 `sigmoid(act_logit) ≈ 0.25`（平均 prior），不会敢预测 active
- SELD 评估实际用 **sorted by activity logit, take top-K̂** 的方式决策
- 训练目标和评估目标不对齐

#### 设计

**Top-K rank loss**：在每一帧 `(b, t)`，强制 top-`N_active_gt` 个 slot 的 activity logit **必须比其他 slot 至少高 margin**，而不是独立回归 0/1。

```python
def topk_rank_activity_loss(
    activity_logit: Tensor,     # [B, K, T]
    target_active: Tensor,       # [B, K, T], 0/1
    valid_time: Tensor,          # [B, T]
    margin: float = 2.0,
) -> Tensor:
    """
    Per-frame marginal ranking loss:
      For each active slot i (target=1) and each inactive slot j (target=0),
      enforce logit[i] > logit[j] + margin.
    
    Equivalent to: 
      loss = Σ_{i in A, j in I} max(0, margin + logit[j] - logit[i])
    This gives direct gradient that "ranks" active slots above inactive ones,
    which aligns with the DCASE eval pipeline (take top-K̂ per frame).
    
    Plus a weak binary regularizer to anchor logit magnitude:
      + 0.1 * BCE(activity_logit, target_active)
    """
    # [B, T] active_count per frame
    n_active = target_active.sum(dim=1)  # [B, T]
    # Loop-free formulation using broadcasting:
    # logit_i: [B, K, T] (active side)
    # logit_j: [B, K, T] (inactive side)
    # pairwise diff: [B, K_i, K_j, T]
    # mask: target_active[i] * (1 - target_active[j])  [B, K_i, K_j, T]
    act = target_active.unsqueeze(2)      # [B, K, 1, T]
    ina = (1.0 - target_active).unsqueeze(1)  # [B, 1, K, T]
    pair_mask = act * ina                  # [B, K_i, K_j, T]
    logit_i = activity_logit.unsqueeze(2)  # [B, K, 1, T]
    logit_j = activity_logit.unsqueeze(1)  # [B, 1, K, T]
    diff = logit_j - logit_i + margin      # want this <= 0
    # hinge loss, masked
    hinge = F.relu(diff) * pair_mask       # [B, K, K, T]
    # normalize by valid pairs count
    pair_valid = pair_mask.sum(dim=(1, 2))  # [B, T]
    time_mask = valid_time.float() * (pair_valid > 0).float()
    loss_rank = (hinge.sum(dim=(1, 2)) * time_mask).sum() / time_mask.sum().clamp(min=1.0)
    
    # Anchor term: prevents logits from drifting to ±inf
    loss_bce = F.binary_cross_entropy_with_logits(
        activity_logit, target_active, reduction='none'
    )
    loss_bce = (loss_bce * valid_time.unsqueeze(1)).mean()
    
    return loss_rank + 0.1 * loss_bce
```

#### 为什么比 ASL 好

| 特性 | BCE | ASL | **Top-K rank** |
|---|---|---|---|
| 优化目标 | per-element logprob | per-element logprob (γ-weighted) | **pairwise ranking** |
| 受 K 不平衡影响 | 严重 | 缓解 | 无（只看 rank） |
| 与 DCASE 评估对齐 | ❌ | ❌ | **✓** (top-K̂) |
| 训练稳定性 | 好 | 中 (γ 过大会崩) | **好**（hinge + 小 BCE anchor） |
| 已知效果 | v12 aR=0.13 | v13_B aR=0.19，aP 降 | **未验证**，但机制对 |

#### Config flag

```python
# spatial_loss.py
frame_activity_loss_type: str = "bce"        # + "topk_rank"
topk_rank_margin: float = 2.0
topk_rank_bce_weight: float = 0.1            # anchor 的 BCE 权重
```

### D-5: Resume optimizer（低成本改进）

#### 诊断
v13_B/C 都用 `--no-resume-optimizer`，Adam 的 m/v moment buffer 从零重建。前 2-3 个 epoch 梯度方向不稳，尤其在"换 loss 函数"后更明显。

#### 改动
```bash
# run_ov1_unified_v13d.sh 不加 --no-resume-optimizer
# 但 LR 设为 v12 最后 LR 的 1/3（避免 resume 后太激进）
SPATIAL_LR="${SPATIAL_LR:-7e-6}"   # v12 是 2e-5，这里是 2e-5/3 ≈ 7e-6
```

**注意**：Optimizer state 包含 LR scheduler 状态，如果我们切 cosine schedule 需要 reset schedule 但保留 Adam moments。实现时：
```python
# 加载 optimizer_state_dict
optimizer.load_state_dict(ckpt['optimizer_state_dict'])
# 但把所有 param_group 的 LR 重设为新 LR（cosine scheduler 会从这开始）
for pg in optimizer.param_groups:
    pg['lr'] = new_peak_lr
# 删掉 step count（avoid schedule confusion）
# scheduler 从 epoch=0 重新开始
```

### D-6: EMA 权重

#### 诊断
v12 ep10-14 F20 在 0.367-0.378 震荡，SGD 困在鞍点。EMA = 取最近 N 个权重的平滑平均，能稳定在鞍点中间而非某一端。

#### 改动

**新加 `EMAModel` helper**：
```python
class EMAModel:
    def __init__(self, model: nn.Module, decay: float = 0.9995):
        self.decay = decay
        self.shadow: Dict[str, Tensor] = {
            name: p.detach().clone()
            for name, p in model.named_parameters()
            if p.requires_grad
        }
    
    @torch.no_grad()
    def update(self, model: nn.Module):
        for name, p in model.named_parameters():
            if not p.requires_grad: continue
            self.shadow[name].mul_(self.decay).add_(p.detach(), alpha=1 - self.decay)
    
    def apply_to(self, model: nn.Module) -> Dict[str, Tensor]:
        """Swap model params with EMA shadow, returns backup for restoration."""
        backup = {}
        for name, p in model.named_parameters():
            if name in self.shadow:
                backup[name] = p.data.clone()
                p.data.copy_(self.shadow[name])
        return backup
    
    def restore(self, model: nn.Module, backup: Dict[str, Tensor]):
        for name, p in model.named_parameters():
            if name in backup:
                p.data.copy_(backup[name])
```

**训练循环**：
```python
# 每 step 后：
if ema_model is not None:
    ema_model.update(model)

# validate 前：
if ema_model is not None:
    backup = ema_model.apply_to(model)
val_metrics = evaluate_one_epoch(model, ...)
if ema_model is not None:
    ema_model.restore(model, backup)

# save best.pt 时，保存 EMA 权重而非原始权重
if ema_model is not None:
    backup = ema_model.apply_to(model)
    torch.save({'model_state_dict': model.state_dict(), ...})
    ema_model.restore(model, backup)
```

#### Config flag

```python
# TrainSpatialBEATsConfig
use_ema: bool = False
ema_decay: float = 0.9995
ema_start_epoch: int = 3   # 前 3 ep 不 EMA（避免 warmup 噪声）
```

## 3. v13_D preset

```python
def make_ov1_unified_v13d_config(...):
    cfg = make_ov1_unified_v12_config(...)  # v12 为基础
    
    # D-1: 扩大 cls warmup，cosine schedule
    cfg.frame_spatial_loss_warmup_epochs = 8
    cfg.frame_spatial_loss_ramp_epochs = 2
    cfg.num_epochs = 25
    cfg.use_cosine_lr = True
    cfg.cosine_lr_warmup_epochs = 3
    cfg.cosine_lr_min_ratio = 0.05
    cfg.learning_rate = 1.5e-5   # peak LR
    
    # D-2: Top-K rank activity loss
    cfg.loss.frame_activity_loss_type = "topk_rank"
    cfg.loss.topk_rank_margin = 2.0
    cfg.loss.topk_rank_bce_weight = 0.1
    
    # D-5: resume optimizer (在 run script 里，不写 --no-resume-optimizer)
    
    # D-6: EMA
    cfg.use_ema = True
    cfg.ema_decay = 0.9995
    cfg.ema_start_epoch = 3
    
    cfg.output_dir = "checkpoints/spatial_beats_ov1_unified_v13d_exp/03_ov123_top4"
    return cfg
```

## 4. 实现步骤

| 文件 | 改动 | 对应 D-* |
|---|---|---|
| `spatial_loss.py` | 加 `_topk_rank_activity_loss` + config 字段 + 分支 | D-2 |
| `train_spatial_beats.py` | 加 `EMAModel` class + cosine LR scheduler + 训练循环 hook | D-1, D-6 |
| `train_spatial_beats.py` | 加 `make_ov1_unified_v13d_config` + CLI + choices | - |
| `run_ov1_unified_v13d.sh` | 新脚本，不带 `--no-resume-optimizer` | D-5 |
| `docs/v13d_spatial_beats_design.md` | 本文档 | - |

所有改动都通过 cfg flag 控制，默认 False → v12/v13_B/v13_C 不受影响。

## 5. 预期结果

| 指标 | v12 best | v13_D 预期 | 机制 |
|---|---|---|---|
| F20 | 0.378 | **0.42 ~ 0.46** | Top-K rank + EMA + 长 warmup |
| aR | 0.126 | **0.25 ~ 0.40** | Top-K 强制拉高活跃 slot |
| aP | 0.855 | **0.80 ~ 0.85** | 可能小降（recall ↑ 的代价），但 hinge 保留 rank 信号 |
| class_acc | 0.834 | **0.86 ~ 0.88** | 长 warmup 让 cls 真的学完 |
| azi_mae | 19.7° | **18~20°** | 不变，不是目标 |

## 6. 风险和兜底

| 风险 | 概率 | 兜底 |
|---|---|---|
| Top-K rank 的 hinge 梯度饱和（margin 太大） | 中 | 降 margin 到 1.0 |
| margin=2.0 导致 logit 分布爆炸（两端拉开） | 低 | anchor BCE 权重从 0.1 升到 0.3 |
| EMA 反而降 F（热启动 EMA 初始化问题） | 低 | ema_start_epoch 提到 5 |
| Cosine LR 峰值太高毁掉 v12 表征 | 中 | peak_lr 降到 1e-5（v12 也是这个） |
| resume optimizer 把 v12 的 Adam moment 固化在错方向 | 低 | 如果 ep0-2 loss 爆炸，退回 --no-resume-optimizer |
| 总体不涨（所有改动都没用） | 中 | 写 ablation，跑 v13_D_noema / v13_D_nocos 诊断 |

## 7. 和 v13_B/C 的关系

- **v13_D 不依赖 v13_B/C 的改动**，直接从 v12 best.pt 热启动
- v13_B/C 可以看作"改模块结构"的尝试，v13_D 是"改训练机制"的尝试
- 如果 v13_D 成功（F ≥ 0.42），**可以在它之上加回 v13_C 的 refinement 2-layer**，那才是真正的 v14

## 8. 验证步骤

1. `python -c "import ast; ast.parse(open('spatial_loss.py').read())"` 语法
2. Top-K rank loss 单测：给已知 activity_logit + target 手算验证
3. 模型构造 + hot-start v12 best.pt：确认 missing=0, unexpected=0
4. 单 batch 前向 + backward：确认 loss 是 scalar、梯度非 NaN
5. EMA 单测：update 后 shadow 权重正确
6. 1 epoch dry-run：看 cosine LR 曲线 + EMA shadow 随 step 变化

---

## 附：为什么不加更多改动（D-7 per-class expert 等）

诊断：v13_B/C 失败的主因不是"改动不够多"，而是"改动不对靶"。v13_D 只碰 loss 机制 + 训练 schedule，属于**最小必要改动**：

- D-2 Top-K 直接针对 activity_recall 瓶颈
- D-1 扩大 warmup 给表征更多学习时间
- D-6 EMA 降低末期震荡
- D-5 resume optimizer 让 LR 轨迹连续

加更多改动（per-class expert、class-conditional gate v2）会重复 v13_B 的错误——改了 head 但没触达瓶颈，而且同时改太多东西无法 ablation。