Spatial-BEATs / docs /v13_spatial_beats_design.md
dieKarotte's picture
Add files using upload-large-folder tool
bf04039 verified
|
Raw
History Blame Contribute Delete
21.9 kB

v13_B + v13_C 设计文档

日期:2026-05-01 作者:Claude + user 目标:在 v12(F20=0.378)基础上,分两个正交实验把 F20 推到 0.45+,合并实验 v14 目标 0.55~0.62 设计原则:所有改动增量式,通过 cfg flag 控制,默认全关 → 不破坏任何现有实验(v7/v11/v12 全可正常复现)


0. 背景与瓶颈诊断(摘自 v12 per-subset 分析)

v12 best.pt (F20=0.3779 聚合) 按子集拆分:

子集           N      F20    ER20   LE_CD  LR_CD  o_cls  o_azi   aP    aR
ov1_sim       4800   0.386  0.686   26.8° 0.640  0.796  25.5°  0.950 0.046
ov2_sim       1718   0.299  0.926   29.6° 0.546  0.653  30.0°  0.916 0.151
ov3_sim       1612   0.270  0.917   30.5° 0.499  0.599  31.8°  0.888 0.165
ov1_real      3374   0.140  0.924  117.7° 0.232  0.809  27.4°  0.767 0.144
ov2_real      2230   0.098  0.866  121.1° 0.198  0.747  38.3°  0.655 0.227
ov3_real       740   0.052  0.766  146.8° 0.125  0.624  51.5°  0.612 0.232
dcase_starss  4560   0.071  1.171  130.3° 0.185  0.698  36.0°  0.625 0.176
unified       35021  (~0.40, 推算)

三大诊断

  1. 瓶颈是 activity_recall (0.13),不是 cls 也不是 spatial。 oracle_class_acc 在 real 上仍有 0.80,oracle_azi_mae_deg 27-51°,说明表征是健康的,被 activity gate 挡住了。
  2. Real/Sim gap 极大:同 ov 级别下 real F20 比 sim 低 64~81%。原因是 train 里 dcase_real 只占 6%。
  3. Overlap 惩罚显著:sim ov3 F20 比 ov1 低 30%(0.27 vs 0.39),K=4 track head 在 overlap 下 slot 分配混乱。

1. 实验切分:B 和 C 是正交维度

维度 v13_B v13_C
改 loss / head 输出接口 ❌(沿用 v12 loss)
改数据比例 / augment 仅 augment ✅ replication
改主干架构 / 容量 ✅ refinement + adapter V3
热启动 v12 best.pt (strict=False) v12 best.pt (strict=False)
能解的子集 sim + real 的 activity 瓶颈 ov2/ov3 overlap + real gap

理由:让 B 的收益和 C 的收益互不污染,便于 ablation。v14 = B + C 合并。


2. 共同约束(关键:不破坏现有实验)

所有 B/C 改动都要遵守:

2.1 Cfg flag 默认全关

每条新改动对应一个 SpatialBEATsConfig / SpatialLossConfig / TrainSpatialBEATsConfig 字段,**默认值就是"这条改动不启用"**。

# spatial_modules 侧
self.use_class_activity_bias: bool = False        # [B-1]
self.use_class_conditional_gate: bool = False     # [B-3]
self.use_track_refinement: bool = False           # [C-2]
self.track_refinement_layers: int = 2
self.patch_adapter_version: str = "v1"            # "v1"/"v2"/"v3"  [C-3] 加 v3
self.use_log_distance_head: bool = False          # [C-4]

# spatial_loss 侧
self.activity_loss_type: str = "bce"              # "bce"/"asymmetric"  [B-2]
self.asl_gamma_neg: float = 4.0
self.asl_gamma_pos: float = 0.0
self.asl_probability_margin: float = 0.05
self.soft_f1_weight: float = 0.0                  # 0 = disabled  [B-4]
self.distance_loss_type: str = "l1"               # "l1"/"laplace_nll"  [C-4]

# dataset 侧
self.use_spec_augment: bool = False               # [B-5]
self.spec_augment_time_mask_ratio: float = 0.0
self.spec_augment_freq_mask_ratio: float = 0.0
self.random_gain_db: float = 0.0
self.channel_dropout_prob: float = 0.0
self.lowpass_sim_real_prob: float = 0.0

ov1_unified_v12 preset 和之前所有 preset 都不动,因为它们没设置这些 flag → 新代码分支不走。

2.2 热启动安全

新加的 nn.Module / nn.Parameter 全部零初始化identity 等价初始化

  • class_activity_biastorch.zeros(num_classes) → logit 不变
  • GatingMLP 最后一层 bias/weight zero-init → gate_logit = 0
  • TrackRefinementDecoderlayer_scale = zeros(num_layers) → 残差 x + 0 = x
  • log_distance_head:bias 初始化为 log(mean_distance_v12) ≈ log(1.5)

从 v12 best.pt strict=False 加载时:

  • 不存在的 key → 走零初始化
  • 存在的 key(所有 v12 组件)→ 正常加载
  • ep0 forward 输出应与 v12 best.pt 完全一致(或数值上差 < 1e-5)

2.3 完全向后兼容

任何旧脚本跑起来(比如 run_ov1_unified_v12.sh不需要改一行,因为所有新字段都有默认值且默认 disabled。

2.4 Preset 命名

  • ov1_unified_v13b:启用所有 B 相关 flag
  • ov1_unified_v13c:启用所有 C 相关 flag
  • ov1_unified_v14:B + C 全开,热启动 max(v13b.best, v13c.best)

3. v13_B 详细设计:Loss + Decision 全面重写

[B-1] Per-class learnable logit bias

为什么:全局阈值 0.5 对所有类别一视同仁不合理。稀有类(jackhammer)和常见类(singing)的 activity 先验差异巨大,应当让模型自己学各类的 logit bias(等价于 per-class threshold)。

实现点

  • 文件:spatial_modules.py
  • 类:FrameTrackPredictionHeads(定位:class_logitsactivity_logits 的出口处)
  • 新增 parameter:self.class_activity_bias = nn.Parameter(torch.zeros(num_classes))
  • 新增 buffer:self.use_class_activity_bias: bool
  • Forward 改动:
# 原:
activity_logit = self.activity_head(token)  # [B, K, T, 1]

# 新:
activity_logit_raw = self.activity_head(token)  # [B, K, T, 1]
if self.use_class_activity_bias:
    class_probs = F.softmax(class_logits, dim=-1)  # [B, K, T, C]
    # 用 class_probs 作为加权软分配(避免 argmax 阻断梯度)
    expected_bias = torch.einsum('bktc,c->bkt', class_probs, self.class_activity_bias)
    activity_logit = activity_logit_raw + expected_bias.unsqueeze(-1)
else:
    activity_logit = activity_logit_raw

训练/推理一致性:bias 在训练的 BCE loss 和推理的 sigmoid 里都是同一个量,因此不需要 threshold sweep。推理时 threshold 始终 = 0.0(logit 空间)或 0.5(prob 空间),完全等价。

参数量:63 个标量,忽略不计。

[B-2] Asymmetric Loss 替换 BCE

为什么:BCE 把 FN 和 FP 等权重。当前 activity_recall=0.13 说明 FN 惩罚严重不够。ASL 对 easy negatives 用 (1-p)^γ- 下压,对 positives 用弱 γ+=0,正负不均衡下表现显著好于 BCE。

实现点

  • 文件:spatial_loss.py
  • 新增函数:asymmetric_loss_with_logits(logits, targets, gamma_neg=4, gamma_pos=0, margin=0.05)
  • compute_frame_track_losses(或同名函数)里根据 config.activity_loss_type 分支:
if config.activity_loss_type == "asymmetric":
    loss_act = asymmetric_loss_with_logits(
        activity_logit, target_active,
        gamma_neg=config.asl_gamma_neg,
        gamma_pos=config.asl_gamma_pos,
        margin=config.asl_probability_margin,
    )
else:  # "bce"
    loss_act = F.binary_cross_entropy_with_logits(activity_logit, target_active, ...)

数学

p = sigmoid(logit)
positive: -( (1-p)**γ+ ) * log(p)
negative: p_shifted = max(p - m, 0)
          -( p_shifted**γ- ) * log(1 - p_shifted)

参数γ+ = 0, γ- = 4, margin = 0.05(ASL paper 推荐起点)

[B-3] Class-conditional gating MLP

为什么:activity 当前只看 token embedding。应该让 activity 也依赖 class/DOA 的确信度 —— class softmax 尖锐、DOA 稳定时更大胆判 active。

实现点

  • 文件:spatial_modules.py
  • 新增类:ClassConditionalGate(embed_dim, num_classes, hidden_dim=128)
    • 输入:fused_token [B, K, T, D], class_logits [B, K, T, C], pred_dir [B, K, T, 3]
    • 融合:gate_input = concat(token, class_emb_avg, dir_vec) → MLP → gate_logit [B, K, T, 1]
    • class_emb 用 class_logits.softmax() 加权的 class embedding(新增 nn.Embedding(C, 32)
  • 在 FrameTrackPredictionHeads 里:
if self.use_class_conditional_gate:
    gate_logit = self.class_conditional_gate(token, class_logits, pred_dir)
    activity_logit = activity_logit + self.gate_scale * gate_logit

初始化:MLP 最后一层 weight=zero, bias=zero → gate_logit = 0 → ep0 等价 v12。

参数量:~80K。

[B-4] Soft-F1 auxiliary loss

为什么:BCE/ASL 仍是 per-sample 损失,优化目标和 macro-F20 评测有 gap。Soft-F1 直接按类聚合,和 DCASE 评估同构。

实现点

  • 文件:spatial_loss.py
  • 新增函数:soft_macro_f1_loss(activity_logits, class_logits, target_active, target_class)
    • 对每个类 c
      • p_c = sigmoid(act_logit) * softmax(class)[c] (class-c 的软 activity)
      • y_c = (target_active and target_class==c)
      • tp_c = sum(p_c * y_c), fp_c = sum(p_c * (1-y_c)), fn_c = sum((1-p_c) * y_c)
      • f1_c = 2 tp_c / (2 tp_c + fp_c + fn_c + eps)
    • loss = 1 - mean(f1_c)
  • 在总 loss 里:
if config.soft_f1_weight > 0:
    total_loss = total_loss + config.soft_f1_weight * soft_macro_f1_loss(...)

warmup(已确认采用):前 3 ep soft_f1_weight=0.1,第 3 ep 起硬切到 0.3

实现方式:在 train_spatial_beats.py 的 epoch 循环里根据 epoch >= soft_f1_warmup_epochs 动态设置 train_cfg.loss.soft_f1_weight,新增 config 字段:

cfg.loss.soft_f1_weight_warmup: float = 0.1        # ep < warmup_epochs 时使用
cfg.loss.soft_f1_weight: float = 0.3               # ep >= warmup_epochs
cfg.loss.soft_f1_warmup_epochs: int = 3

[B-5] Real-distribution augment

为什么:sim_static 混响干净,模型学到的 activity 判据在低 SNR 下崩溃。augment 让模型见到各种"污染"的 spec,对 real 数据更鲁棒。

实现点

  • 文件:spatial_dataset.py
  • SpatialDataset.__getitem__ 或 collate 里加 augment pipeline
  • 只在训练集(split='train')启用,valid/test 不启用
  • 新增 config flag:
    • use_spec_augment(默认 False)
    • spec_augment_time_mask_ratio(0.2 = 20% time 长度)
    • spec_augment_freq_mask_ratio(0.15)
    • random_gain_db(±8)
    • channel_dropout_prob(0.1)
    • lowpass_sim_real_prob(0.1,cutoff ∈ U[4000, 8000] Hz)

顺序:waveform-level augment(gain, lowpass, channel_dropout)→ feature-level augment(SpecAugment)。

重要:augment 只作用在 FOA 4 通道 waveform / delta feature 上,target labels 不变

B 实验 preset: make_ov1_unified_v13b_config

热启动:v12_best.pt (strict=False) 开关:

cfg.model.use_class_activity_bias = True          # [B-1]
cfg.model.use_class_conditional_gate = True       # [B-3]
cfg.loss.activity_loss_type = "asymmetric"        # [B-2]
cfg.loss.asl_gamma_neg = 4.0
cfg.loss.asl_probability_margin = 0.05
cfg.loss.soft_f1_weight = 0.3                     # [B-4]
cfg.dataset.use_spec_augment = True               # [B-5]
cfg.dataset.spec_augment_time_mask_ratio = 0.2
cfg.dataset.spec_augment_freq_mask_ratio = 0.15
cfg.dataset.random_gain_db = 8.0
cfg.dataset.channel_dropout_prob = 0.1
cfg.dataset.lowpass_sim_real_prob = 0.1
cfg.learning_rate = 1e-5
cfg.num_epochs = 15
cfg.output_dir = "checkpoints/spatial_beats_ov1_unified_v13b_exp/03_ov123_top4"

数据 manifest 完全复用 v12:unified train/valid + old ov1/2/3 sim/real + dcase_starss 作为 val。


4. v13_C 详细设计:Data + Architecture 全面重写

[C-1] Real data upsampling (replication)

为什么:real (dcase_real) 在 train 里占 6%,梯度感受不到。DCASE 社区标准做法是 20-30% real。

实现点

  • 预处理脚本:scripts/split_unified_train_by_source.py
    • unified_spatial_foa_fsd63_all/train.jsonl
    • data_source 字段拆成三份:
      • train_sim_static.jsonl
      • train_qa_sim.jsonl
      • train_dcase_real.jsonl
    • 写到 unified_spatial_foa_fsd63_all/ 同目录下
  • Preset:
cfg.train_manifest_paths = (
    unified_root / "train_sim_static.jsonl",
    unified_root / "train_qa_sim.jsonl",
    unified_root / "train_dcase_real.jsonl",
)
cfg.train_manifest_replication = (1, 1, 6)

影响估算(基于 v12 已知分布):

  • sim_static 304K × 1 = 304K
  • qa_sim ~? × 1 = ~?
  • dcase_real 20K × 6 = 120K
  • real 占比从 6% → ~25% (取决于 qa_sim 规模)

兼容性train_manifest_replication 机制在 train_spatial_beats.py 已经存在(v7j 用过),不需要新加框架代码。只改 preset。

[C-2] Track-wise Refinement Transformer(2 layers)

为什么:K=4 track slots 之间互相不知道对方在干嘛,overlap 时同一源被多个 slot 抢,或同一 slot 被多个源抢。引入 self-attention 让 slot 互相"排斥"。

实现点

  • 文件:spatial_modules.py
  • 新增类:
class TrackRefinementDecoder(nn.Module):
    def __init__(self, num_tracks=4, embed_dim=768, num_layers=2,
                 num_heads=8, dim_feedforward=2048, dropout=0.0):
        self.track_queries = nn.Parameter(torch.randn(num_tracks, embed_dim) * 0.02)
        self.layers = nn.ModuleList([
            nn.TransformerDecoderLayer(
                d_model=embed_dim, nhead=num_heads,
                dim_feedforward=dim_feedforward, dropout=dropout,
                activation='gelu', norm_first=True, batch_first=True,
            ) for _ in range(num_layers)
        ])
        # Zero-init layer scale: ep0 refinement = identity
        self.layer_scale = nn.Parameter(torch.zeros(num_layers))
    
    def forward(self, memory):
        # memory: fused_spatial_embeddings [B, T, D]
        # 输出:refined track tokens [B, K, T, D]
        B, T, D = memory.shape
        K = self.track_queries.size(0)
        # 复制 K queries 到时间维度:[B, K, T, D]
        q = self.track_queries[None, :, None, :].expand(B, K, T, D)
        # 每个时间步独立做 decoder
        # 为简化,把 T 维 flatten 进 batch:[B*T, K, D] cross-attn with [B*T, 1, D]
        q_flat = q.permute(0, 2, 1, 3).reshape(B * T, K, D)
        mem_flat = memory.reshape(B * T, 1, D)
        for i, layer in enumerate(self.layers):
            out = layer(q_flat, mem_flat)
            q_flat = q_flat + self.layer_scale[i] * (out - q_flat)
        # reshape 回 [B, K, T, D]
        refined = q_flat.reshape(B, T, K, D).permute(0, 2, 1, 3).contiguous()
        return refined
  • SpatialBEATs 里:
if cfg.use_track_refinement:
    self.track_refinement = TrackRefinementDecoder(
        num_tracks=cfg.num_tracks,
        embed_dim=cfg.encoder_embed_dim,
        num_layers=cfg.track_refinement_layers,
    )

# encode_patches 之后、送入 head 之前:
if self.track_refinement is not None:
    track_tokens = self.track_refinement(encoder_memory)  # [B, K, T, D]
    # 传给 FrameTrackPredictionHeads 的输入从 [B, T, D] 改成 [B, K, T, D]
else:
    track_tokens = None  # head 沿用旧 expand 逻辑
  • FrameTrackPredictionHeads 的 forward 加个 track_tokens: Optional[Tensor] 参数:
    • 传入 None → 沿用现有的"[B,T,D] 复制到 K slots"
    • 传入 [B,K,T,D] → 用 refined tokens 走 head

参数量:2 layer × (self_attn + cross_attn + FFN) ≈ 2 × 2M = 4M。

Zero-init 校验layer_scale = zeros(2) + 残差公式 q + scale * (out - q) → ep0 输出 = track_queries(静态,和 memory 无关)。但这会丢掉时间信息 —— 修正:改用 q + scale * layer_out,并且把 track_queries 初始化成 memory 投影:

实际更安全的等价初始化:

# Zero-init 方案:layer 不改 query,query 本身先吸收 memory 信息
# 思路:在 refine 前先做一次 "identity fallback":如果 scale=0,输出 = memory 广播到 K
def forward(self, memory):
    B, T, D = memory.shape
    K = ...
    # 初始 track_tokens = memory 广播到 K(+ 一个很小的 query 偏移)
    track_tokens = memory[:, None, :, :].expand(B, K, T, D).contiguous()
    track_tokens = track_tokens + 0.02 * self.track_queries[None, :, None, :]
    # refine
    for i, layer in enumerate(self.layers):
        ...
        track_tokens = track_tokens + self.layer_scale[i] * delta
    return track_tokens

这样 layer_scale=0 时 refinement 输出 ≈ memory 广播到 K,和 v12 "把 [B,T,D] 复制到 [B,K,T,D]" 等价。热启动安全。

[C-3] Multi-scale Patch Adapter V3

为什么:v12 用的 V2 adapter 只看 3 个时间 bin(30 ms),抓不到房间冲激响应的 early reflection (50-150ms)。V3 加多尺度 + dilated conv。

实现点

  • 文件:spatial_modules.py
  • 新增类:SpatialDeltaPatchAdapterV3
    • 三路 branch:
      • branch_3x3: Conv2d(C, H, kernel=3, padding=1) (同 V2)
      • branch_5x5: Conv2d(C, H, kernel=5, padding=2) (中尺度)
      • branch_dilated: Conv2d(C, H, kernel=3, padding=2, dilation=2) (长时)
    • fuse: torch.cat along channel → Conv2d(3H, H, kernel=1)
    • 接现有 V2 的 SE block + residual + patchify
  • cfg:patch_adapter_version: str = "v1" 增加选项 "v3"

参数量:比 V2 多 ~1M。

[C-4] Log-distance head + Laplace NLL loss

为什么:dist_mae=0.57 很差。距离分布长尾,log 后近似高斯。加 uncertainty 头允许模型对不确信的距离给大 variance,减少高 bias 样本的损失。

实现点

  • 文件:spatial_modules.py, 类 FrameTrackPredictionHeads

  • 把现有 distance_head: Linear(D, 1) 升级为 distance_head: Linear(D, 2) 输出 [log_dist, log_var]

  • 初始化:bias[0] = log(1.5)(v12 平均距离附近),bias[1] = log(0.2^2)(var=0.04 起点)

  • cfg:use_log_distance_head: bool = False, distance_loss_type: str = "l1" / "laplace_nll"

  • 文件:spatial_loss.py, Laplace NLL:

def laplace_nll_loss(pred_log_dist, pred_log_var, target_dist, mask):
    # target_dist > 0 的位置才算 loss
    pred_dist = pred_log_dist.exp()
    pred_b = (0.5 * pred_log_var).exp()  # Laplace scale
    nll = (target_dist - pred_dist).abs() / pred_b + pred_log_var * 0.5
    return (nll * mask).sum() / mask.sum().clamp(min=1)

推理时 pred_distance = exp(pred_log_dist)

初期稳定性(已确认):v13c 从第 0 ep 就切 Laplace NLL(不做 L1 warmup)。如果训练前期 loss_dist 不稳,再回来调。

C 实验 preset: make_ov1_unified_v13c_config

热启动:v12_best.pt (strict=False) 开关:

cfg.model.use_track_refinement = True             # [C-2]
cfg.model.track_refinement_layers = 2
cfg.model.patch_adapter_version = "v3"            # [C-3]
cfg.model.use_log_distance_head = True            # [C-4]
cfg.loss.distance_loss_type = "laplace_nll"       # [C-4]

cfg.train_manifest_paths = (sim_static, qa_sim, dcase_real)   # [C-1]
cfg.train_manifest_replication = (1, 1, 6)

cfg.learning_rate = 1e-5
cfg.num_epochs = 20
cfg.output_dir = "checkpoints/spatial_beats_ov1_unified_v13c_exp/03_ov123_top4"

loss 完全沿用 v12 默认(BCE + L1 → Laplace),soft_f1_weight=0, activity_loss_type="bce"


5. v14 合并实验(后续)

在 B 和 C 都验证有效(F20 > 0.42)后启动:

  • 热启动:max(v13b.best, v13c.best).pt
  • 所有 B 和 C 的 flag 全开
  • LR = 5e-6(更保守,防止双改动发散)
  • epochs = 20

预期 F20 = 0.55~0.62。


6. 预期结果 & 风险矩阵

B 预期

  • 聚合 F20: 0.378 → 0.45~0.52
  • sim ov1: 0.386 → 0.50~0.55
  • real ov1: 0.140 → 0.22~0.28
  • dcase: 0.071 → 0.15~0.20
  • activity_recall: 0.13 → 0.40~0.55

C 预期

  • 聚合 F20: 0.378 → 0.44~0.50
  • ov3_sim: 0.270 → 0.38~0.42
  • real ov1: 0.140 → 0.20~0.26
  • dcase: 0.071 → 0.14~0.19
  • dist_mae: 0.566 → 0.38~0.42

风险

风险 发生概率 兜底
B-2 ASL γ- 太大发散 先 γ-=2 跑 1 ep 验证,再拉到 4
B-3 gate 挡掉好样本 gate_scale 从 0.5 改 0.2 重跑
B-4 soft-F1 和 ASL 冲突 soft_f1_weight 从 0.3 降到 0.1
B-5 augment 太强 sim 掉点 减半 augment 比例重跑
C-1 real 6× 导致 sim 掉点 降到 4×
C-2 refinement 不学 手动设 layer_scale warmup
C-3 多尺度显存爆 去掉 dilated branch
C-4 log-dist 初期不稳 前 3 ep 用 L1 再切
v14 合并发散 降 LR 到 3e-6,freeze trunk 前半段

7. 落地文件清单

文件 改动类型 B/C
spatial_modules.py 新增 3 个类 + 现有类加 forward 分支 B+C
spatial_loss.py 新增 2 个 loss 函数 + config 分支 B+C
spatial_dataset.py 新增 augment 逻辑 + config 字段 B
spatial_beats.py config 字段 + 可选模块构造 + forward 分支 B+C
train_spatial_beats.py 新增 2 个 preset 工厂 + CLI dispatch + choices B+C
scripts/split_unified_train_by_source.py 新文件,预处理脚本 C
run_ov1_unified_v13b.sh 新文件 B
run_ov1_unified_v13c.sh 新文件 C
docs/v13_spatial_beats_design.md 本文档

所有现有文件的改动都是新增分支,不删除/修改任何现有逻辑。


8. 验证步骤

每完成一步,按顺序验证:

  1. 语法检查python -c "import ast; ast.parse(open(path).read())"
  2. 旧 preset 回归python train_spatial_beats.py --preset ov1_unified_v12 --dry-run(或者 ep=1 跑到第一个 batch),确认 F20 和 v12 一致
  3. 新模块零初始化等价:在 v13b_config / v13c_config 下跑 ep=0 valid,确认和 v12 best.pt 的 valid 指标差异 < 1%
  4. B/C 训练:完整跑 15/20 ep,观察 F20 曲线
  5. per-subset eval:用 eval_v12_per_subset.py --preset ov1_unified_v13b --checkpoint ... 看每个子集涨点
  6. test eval:用同脚本加 --split test 跑最终测试