docs/spatial_beats_file_level_checklist.md

Spatial-BEATs 文件级实现清单

1. 文档目的

这份文档把当前确定的 Spatial-BEATs 简化版方案，收紧成可以直接指导编码的文件级实现清单。

目标是回答四件事：

• 哪些现有文件直接复用
• 哪些新文件需要实现
• 每个模块的职责是什么
• 从 FOA 到 LLM spatial tokens 的张量形状如何变化

这份清单默认遵循当前主线：

• 重点训练 FOA -> BEATs trunk 的空间表征
• 后面的模块只承担 readout / supervision
• 不以复杂 slot decoder 为主线
• 最终给 LLM 的 token 是固定 2.5 Hz
• 10 s 输入默认输出 25 个 spatial tokens

2. 默认配置

• 输入采样率：16 kHz
• 输入通道：4 个 FOA 通道，分别为 W / X / Y / Z
• 预处理输出通道：7
• 空间特征：W_logmel / X_logmel / Y_logmel / Z_logmel / IVx / IVy / IVz
• mel bins：128
• Qwen-like mel 参数：
  • sample_rate=16000
  • num_mel_bins=128
  • n_fft=400
  • win_length=400
  • hop_length=160
  • dither=0.0
• source vocabulary：/apdcephfs_cq12/share_302080740/user/schmittzhu/data/fsd50k/FSD50K.ground_truth/final_vocabulary.csv
• source vocabulary fields：label_id / final_label
• source class count：65
• 主干初始化：BEATs_iter3+ AS2M pre-trained
• patch embed dim：512
• encoder dim：768
• encoder layers：12
• 最终时间 token rate：2.5 Hz
• 最大同时源数：4
• 当前监督：弱时间窗口 + 固定位置
• batch 时间组织：每个样本自己的 T_s_i = round(duration_i * 2.5)，batch 内 pad 到 T_s_max

3. 推荐文件布局

建议保留原始 BEATs 文件不动，新实现集中放在新增文件里：

• 复用：BEATs.py
• 复用：backbone.py
• 新增：spatial_beats.py
• 新增：spatial_modules.py
• 新增：spatial_loss.py
• 新增：spatial_dataset.py
• 新增：train_spatial_beats.py

推荐原则：

• 原始 BEATs.py 作为参考和权重来源，不直接改
• 原始 backbone.py 直接复用，不改 Transformer 内部
• 新的 Spatial-BEATs 在新文件里组装

4. 现有文件如何复用

4.1 BEATs.py

当前 BEATs 里最重要的可复用部分是：

• BEATsConfig
• post_extract_proj
• layer_norm
• TransformerEncoder 的使用方式
• 预训练 checkpoint 的参数命名

当前不直接复用的部分是：

• preprocess()
• patch_embedding = Conv2d(1, embed_dim, ...)
• predictor

原因：

• preprocess() 只支持单通道 fbank
• patch_embedding 只支持 1 通道输入
• predictor 是 AudioSet 分类头，不适合当前 spatial encoder

4.2 backbone.py

这里建议直接复用：

• TransformerEncoder
• TransformerSentenceEncoderLayer
• attention / FFN / pos_conv / layer_norm

当前 Spatial-BEATs 不需要改这些内部结构。

5. 新文件实现清单

5.1 spatial_beats.py

这个文件放主模型类，负责把所有模块串起来。

建议包含：

• SpatialBEATsConfig
• SpatialBEATs
• load_beats_pretrained()

5.1.1 SpatialBEATsConfig

职责：

• 在 BEATsConfig 基础上补充 spatial encoder 需要的超参

建议新增字段：

• foa_feature_channels = 7
• sample_rate = 16000
• num_mel_bins = 128
• target_token_rate = 2.5
• source_vocab_path
• source_label_id_field = "label_id"
• source_label_name_field = "final_label"
• source_num_classes = 65
• max_sources = 4
• readout_layers = 1 or 2
• slot_hidden_dim = 768
• distance_head_type = "regression"
• llm_hidden_dim

5.1.2 SpatialBEATs.__init__

职责：

• 初始化所有子模块
• 建立 trunk 与 heads
• 提供 BEATs checkpoint 兼容结构

建议持有的子模块：

• preprocessor: SpatialBEATsPreprocessor
• patch_embedding: SpatialPatchEmbedding
• layer_norm
• post_extract_proj
• dropout_input
• encoder: TransformerEncoder
• frequency_pool: FrequencyPool
• temporal_resampler: TemporalResampler
• temporal_readout: ShallowTemporalReadout
• slot_readout: FixedSlotReadout
• prediction_heads: SpatialPredictionHeads
• projector: SpatialTokenProjector

5.1.3 SpatialBEATs.forward()

建议拆成几个清晰阶段：

1. preprocessor
2. patch embedding
3. BEATs trunk
4. frequency pooling
5. temporal resampling
6. temporal readout
7. slot readout + heads
8. projector

建议输出一个结构化对象或 dict，至少包含：

{
    "foa_feat": foa_feat,
    "patch_tokens": patch_tokens,
    "encoder_memory": encoder_memory,
    "temporal_tokens": temporal_tokens,
    "spatial_embeddings": spatial_embeddings,
    "slot_latents": slot_latents,
    "pred_activity": pred_activity,
    "pred_azi_logits": pred_azi_logits,
    "pred_ele_logits": pred_ele_logits,
    "pred_dist": pred_dist,
    "pred_class_logits": pred_class_logits,
    "llm_spatial_tokens": llm_spatial_tokens,
    "temporal_padding_mask": temporal_padding_mask,
    "target_num_steps": target_num_steps,
}

这样训练和调试都更方便。

5.2 spatial_modules.py

这个文件放所有新增模块。

建议包含以下类。

5.2.1 SpatialBEATsPreprocessor

职责：

• 把 FOA waveform 变成 7 通道空间特征图

输入：

• waveform: [B, 4, T]

默认 10 秒示例：

• waveform: [B, 4, 160000]

内部步骤：

1. 对 W/X/Y/Z 做 STFT
2. 计算四个通道的 log-mel
3. 计算 IVx / IVy / IVz
4. 将 IV 映射到 mel 维
5. 按通道拼接

输出：

• foa_feat: [B, 7, T_f, 128]

默认示例：

• 如果 frame_shift = 10 ms
• 对 10 s，T_f 约为 1000
• 因此输出近似：
  • foa_feat: [B, 7, 1000, 128]

作用：

• 用显式 FOA 物理特征替代原始 BEATs 的单通道 fbank

5.2.2 SpatialPatchEmbedding

职责：

• 把 7 通道空间特征图变成 patch tokens

建议实现：

• Conv2d(7, 512, kernel_size=patch, stride=patch, bias=cfg.conv_bias)

输入：

• foa_feat: [B, 7, T_f, 128]

输出中间张量：

• patch_grid: [B, 512, T_p, F_p]

flatten 后输出：

• patch_tokens: [B, N_p, 512]

其中：

• T_p = floor((T_f - p_t) / s_t) + 1
• F_p = floor((128 - p_f) / s_f) + 1
• N_p = T_p * F_p

如果第一版保持和当前 BEATs 一样的等步长 patch，且 patch=(16,16)：

• T_f ≈ 1000 -> T_p ≈ 62
• 128 -> F_p = 8
• N_p ≈ 496

则一个典型输出是：

• patch_tokens: [B, 496, 512]

作用：

• 用多通道 spatial features 替换原始单通道 patch stem

5.2.3 FrequencyPool

职责：

• 把 trunk 的二维 patch grid 沿频率维压缩成时间序列

输入：

• encoder_memory: [B, N_p, 768]
• grid_size = (T_p, F_p)

先 reshape：

• grid_memory: [B, T_p, F_p, 768]

再沿 F_p 做 pooling：

• temporal_patch_tokens: [B, T_p, 768]

推荐第一版：

• mean pooling over F_p

作用：

• 把 patch 级二维结构转成时间序列

5.2.4 TemporalResampler

职责：

• 把 patch 级时间步压到固定 2.5 Hz

输入：

• temporal_patch_tokens: [B, T_p, 768]

输出：

• temporal_tokens: [B, T_s_max, 768]
• temporal_padding_mask: [B, T_s_max]

其中：

• 每个样本：
  • T_s_i = round(duration_i * 2.5)
• batch 内：
  • T_s_max = max_i T_s_i

对 10 s：

• T_s_i = 25

因此单个 10 s 样本的有效长度是：

• 25

推荐实现：

• 第一版：线性插值
• 备选：轻量 Conv1d 下采样

作用：

• 明确固定最终时间基准
• 保证最终给 LLM 的 token rate 就是 2.5 Hz

5.2.5 ShallowTemporalReadout

职责：

• 把 trunk 输出的时间序列再整理一遍
• 让它更适合作为空间 embedding

建议实现：

• 1~2 层轻量 TransformerEncoder
• 或最小版 LayerNorm + MLP

输入：

• temporal_tokens: [B, T_s_max, 768]

输出：

• spatial_embeddings: [B, T_s_max, 768]

作用：

• 这就是当前主线里最重要的 LLM 前空间 embedding

5.2.6 FixedSlotReadout

职责：

• 给 encoder-only 训练提供多源监督出口
• 不引入复杂 decoder

输入：

• spatial_embeddings: [B, T_s_max, 768]

建议最小实现：

• Linear(768, 4 * 768)
• reshape 成：
  • slot_latents: [B, T_s_max, 4, 768]

备选更稳实现：

• Linear(768, 768) -> GELU -> Linear(768, 4 * 768)
• reshape 后：
  • slot_latents: [B, T_s_max, 4, 768]

作用：

• 每个时间步提供 4 个固定 source 槽位
• 用于承载多源监督

注意：

• 这里的 4 只是 supervision 容量
• 不影响最终输出给 LLM 的 token 数

5.2.7 SpatialPredictionHeads

职责：

• 从 slot_latents 上读出显式空间监督

输入：

• slot_latents: [B, T_s_max, 4, 768]

输出：

• pred_activity: [B, T_s_max, 4]
• pred_azi_logits: [B, T_s_max, 4, 360]
• pred_ele_logits: [B, T_s_max, 4, 180]
• pred_dist: [B, T_s_max, 4, 1]
• pred_class_logits: [B, T_s_max, 4, C_cls]

作用：

• pred_activity：弱时间窗口下的源存在性
• pred_azi_logits：方位角分类
• pred_ele_logits：仰角分类
• pred_dist：连续距离回归
• pred_class_logits：辅助类别监督

当前默认建议：

• C_cls = 65
• labels 来自 final_vocabulary.csv

5.2.8 SpatialTokenProjector

职责：

• 把 spatial_embeddings 投影到 LLM hidden size

输入：

• spatial_embeddings: [B, T_s_max, 768]

输出：

• llm_spatial_tokens: [B, T_s_max, d_llm]

推荐实现：

• Linear(768, D_mid)
• GELU
• LayerNorm
• Linear(D_mid, d_llm)

注意：

• encoder-only 第一阶段可以不训练这个模块
• 但接口最好先保留

5.3 spatial_loss.py

这个文件放 target 组织、matching 和 loss 计算。

建议包含：

• build_time_window_mask()
• match_fixed_slots()
• compute_spatial_losses()

5.3.1 build_time_window_mask()

输入：

• source start / end
• clip duration
• 每个样本的 T_s_i
• batch 的 T_s_max

输出：

• window_mask: [B, num_gt, T_s_max]

作用：

• 把弱时间标注映射到 2.5 Hz 时间轴

5.3.2 match_fixed_slots()

输入：

• pred_activity: [B, T_s_max, 4]
• pred_azi_logits: [B, T_s_max, 4, 360]
• pred_ele_logits: [B, T_s_max, 4, 180]
• pred_dist: [B, T_s_max, 4, 1]
• pred_class_logits: [B, T_s_max, 4, C_cls]
• gt targets

输出：

• 每个时间步的 slot 分配结果

推荐第一版：

• per-step fixed-slot matching

作用：

• 不依赖复杂 decoder，也能对 4 个固定槽位施加多源监督

5.3.3 compute_spatial_losses()

建议输出各项 loss：

• L_activity
• L_azi
• L_ele
• L_dist
• L_cls_aux
• L_temp

总损失：

L_total =
  lambda_act * L_activity
  + lambda_azi * L_azi
  + lambda_ele * L_ele
  + lambda_dist * L_dist
  + lambda_cls * L_cls_aux
  + lambda_temp * L_temp

各项作用：

• L_activity：窗口外负样本，窗口内弱正样本
• L_azi / L_ele：方向监督
• L_dist：距离连续回归
• L_cls_aux：帮助保持 source-level 语义区分
• L_temp：固定位置场景下的时间平滑

5.4 spatial_dataset.py

这个文件负责把数据组织成模型和 loss 可直接使用的格式。

建议一个 batch sample 至少包含：

• waveform: [4, T]
• source_labels
• source_azimuth
• source_elevation
• source_distance
• source_start_time
• source_end_time
• clip_duration
• source label id / name should map through final_vocabulary.csv

建议 collate 后的 batch 字段：

• waveform: [B, 4, T]
• padding_mask: [B, T] 或 None
• gt_sources: 变长列表，或 pad 成固定 max_gt
• clip_duration_seconds: [B]

建议在 dataset 或 collate 阶段完成：

• 时间窗口离散化到每个样本自己的 T_s_i，再 pad 到 T_s_max
• 距离归一化
• 角度离散化成 bin index
• label id/name 映射到统一 vocabulary index

5.5 train_spatial_beats.py

这个文件负责训练阶段组织。

建议职责：

• 构建 model
• 加载 BEATs checkpoint
• 配置参数组和学习率
• 调用 dataset / loss
• 执行 encoder-only 训练

建议阶段：

Stage 1

• 训练：
  • preprocessor
  • patch_embedding
  • temporal_resampler
  • temporal_readout
  • slot_readout
  • prediction_heads
• projector 默认不训练
• trunk 可以：
  • 先只解冻高层
  • 或全量解冻但用更小 lr

Stage 2

• 接上 projector
• 再和 LLM 接口联合训练

6. 预训练权重加载清单

建议在 spatial_beats.py 里实现 load_beats_pretrained()

6.1 直接加载的参数

从 BEATs checkpoint 里直接复用：

• layer_norm.*
• post_extract_proj.*
• encoder.pos_conv.*
• encoder.layers.*
• encoder.layer_norm.*

6.2 部分继承或重新初始化的参数

patch_embedding

原始 BEATs：

• Conv2d(1, 512, ...)

新的 Spatial-BEATs：

• Conv2d(7, 512, ...)

推荐初始化：

• W_logmel 通道拷贝原始卷积权重
• 其他 6 个通道初始化为：
  • 零
  • 或小随机值

这样可以最大程度保留原始单通道初始化稳定性。

新模块随机初始化

以下模块默认新初始化：

• SpatialBEATsPreprocessor
• FrequencyPool
• TemporalResampler
• ShallowTemporalReadout
• FixedSlotReadout
• SpatialPredictionHeads
• SpatialTokenProjector

7. 10 秒输入的完整张量流

下面给一个默认示例。

假设：

• batch size 为 B
• 输入时长 10 s
• 采样率 16 kHz
• T = 160000
• mel bins = 128
• patch 为 16 x 16

7.1 从波形到 trunk

1. 输入波形

   • [B, 4, 160000]

2. 预处理后

   • [B, 7, 1000, 128]

3. patch conv 后

   • [B, 512, 62, 8]

4. flatten + transpose 后

   • [B, 496, 512]

5. layer_norm

   • [B, 496, 512]

6. post_extract_proj

   • [B, 496, 768]

7. TransformerEncoder

   • [B, 496, 768]

7.2 从 trunk 到 spatial embeddings

8. reshape 回 patch grid

   • [B, 62, 8, 768]

9. frequency pooling

   • [B, 62, 768]

10. temporal resample 到 2.5 Hz

    • 单个 10 s 样本有效长度为 25
    • mixed-length batch 中张量形状为 [B, T_s_max, 768]

11. shallow temporal readout

    • [B, T_s_max, 768]

这一步得到：

• spatial_embeddings: [B, T_s_max, 768]

7.3 从 spatial embeddings 到训练 heads

12. fixed-slot expand

    • [B, T_s_max, 4, 768]

13. activity head

    • [B, T_s_max, 4]

14. azimuth head

    • [B, T_s_max, 4, 360]

15. elevation head

    • [B, T_s_max, 4, 180]

16. distance head

    • [B, T_s_max, 4, 1]

17. class auxiliary head

    • [B, T_s_max, 4, C_cls]

7.4 从 spatial embeddings 到 LLM 接口

18. spatial projector
    • [B, T_s_max, d_llm]

这就是最终给 LLM 的 spatial tokens。

关键点：

• 单个 10 s 样本：10 s -> 25 个有效 token
• mixed-length batch：张量按 T_s_max padding
• 不是 25 * 4
• 4 只存在于监督头里

8. 推荐实现顺序

建议按下面顺序写代码。

1. spatial_modules.py

   • 先实现 preprocessor / patch embedding / frequency pool / temporal resampler / readout / heads / projector

2. spatial_beats.py

   • 把主链路串起来
   • 先跑通 forward 和 shape

3. spatial_dataset.py

   • 先把数据组织成 2.5 Hz 时间轴上的 supervision，并支持 mixed-length batch padding

4. spatial_loss.py

   • 先实现 L_activity / L_azi / L_ele / L_dist / L_cls_aux
   • 最后再补 L_temp

5. train_spatial_beats.py

   • 最后接训练循环

9. 第一版验收标准

在不接 LLM 的前提下，第一版至少应满足：

• forward 全链路 shape 正确
• 能成功加载 BEATs trunk 预训练参数
• 10 s 输入稳定输出：
  • 对单个 10 s 样本有 25 个有效 token
  • batch 张量输出为 spatial_embeddings [B, T_s_max, 768]
  • batch 张量输出为 llm_spatial_tokens [B, T_s_max, d_llm]
• 训练阶段能算出：
  • L_activity
  • L_azi
  • L_ele
  • L_dist
  • L_cls_aux
• loss 能稳定回传到 trunk

10. 一句话总结

当前最合理的文件级实现方式是：

• 不改原始 BEATs 主体文件
• 用新文件组装一个 FOA -> spatial features -> BEATs trunk -> 2.5Hz spatial embeddings -> fixed-slot supervision heads -> projector 的 Spatial-BEATs
• 训练时真正要训好的是前面的 spatial_embeddings
• 后面的 heads 只是为了把空间监督显式接出来