| # Spatial-BEATs 实现规格 |
|
|
| ## 1. 目标 |
|
|
| 本规格文档用于将前期讨论收敛为一个可以直接实施的 `Spatial-BEATs` 方案。 |
|
|
| 目标是构建一个独立的 `Spatial Encoder`： |
|
|
| - 输入为完整 `FOA` 音频及其派生空间特征 |
| - 完整的 `FOA` 特征经过 `BEATs backbone` |
| - 最大化复用 `BEATs` 预训练权重 |
| - 输出一组 `source-level spatial tokens` |
| - 这些 token 作为独立模态输入给 LLM |
| - 原有语义 audio encoder 保持不动 |
|
|
| 这里的关键原则是： |
|
|
| > 不是让 `W-only` 走主干，再外挂一个小空间 adapter；而是让完整 FOA 空间特征真正进入 BEATs 主干，并在主干之后产出结构化空间 token。 |
|
|
| ## 2. 最终任务定义 |
|
|
| ### 2.1 核心任务 |
|
|
| `Spatial-BEATs` 的主任务定义为： |
|
|
| - 给定一个多源 `FOA` 音频片段 |
| - 预测其中最多 `K` 个潜在声源的空间表示 |
| - 每个表示对应一个 `source token` |
|
|
| 每个 source token 至少承载： |
|
|
| - `objectness` |
| - `azimuth` |
| - `elevation` |
| - `distance` |
|
|
| 可选承载： |
|
|
| - `source class auxiliary logits` |
| - `source embedding` |
|
|
| ### 2.2 推荐监督形式 |
|
|
| 如果训练数据中每个源都有标注，则推荐采用： |
|
|
| - `set prediction` |
| - `K` 个预测 token 对 `N` 个 GT sources |
| - 用 `Hungarian matching` 做一一匹配 |
|
|
| 不建议采用： |
|
|
| - 单一 scene-level spatial token |
| - 仅回归整段音频的全局空间摘要 |
|
|
| 原因是这会损失多源结构，不利于后续 LLM 做关系推理。 |
|
|
| ## 3. 最终架构 |
|
|
| 推荐最终架构： |
|
|
| ```text |
| FOA waveform |
| -> SpatialBEATsPreprocessor |
| -> FOA feature map [B, C_foa, T, F] |
| -> FOA patch embedding |
| -> BEATs trunk |
| -> Spatial query decoder |
| -> K source tokens |
| -> Spatial prediction heads |
| -> LLM projector |
| ``` |
|
|
| 为了最大化复用 BEATs 主干，本方案尽量不改 trunk 内部的 Transformer 结构。 |
|
|
| ## 4. 输入特征定义 |
|
|
| ### 4.1 默认推荐特征 |
|
|
| 第一版推荐输入通道： |
|
|
| - `W_logmel` |
| - `X_logmel` |
| - `Y_logmel` |
| - `Z_logmel` |
| - `IVx` |
| - `IVy` |
| - `IVz` |
|
|
| 即： |
|
|
| - `C_foa = 7` |
|
|
| 这是默认推荐方案。 |
|
|
| ### 4.2 备选输入特征 |
|
|
| 若希望先降低复杂度，可以使用： |
|
|
| - `WXYZ logmel` |
|
|
| 即： |
|
|
| - `C_foa = 4` |
|
|
| 但这只适合最小原型。 |
| 如果目标是稳定学习空间方向与结构，优先使用 `WXYZ + IV`。 |
|
|
| ### 4.3 前端参数建议 |
|
|
| 为了最大化复用 BEATs 主干，推荐保持与 BEATs 接近的时频分辨率： |
|
|
| - sample rate：优先 `16k` |
| - mel bins：`128` |
| - frame length：`25 ms` |
| - frame shift：`10 ms` |
|
|
| 原因： |
|
|
| - 这能让 trunk 看到与原始 BEATs 更接近的 patch 几何结构 |
| - patch embedding 和后续序列长度更容易保持一致 |
| - 预训练权重复用更稳定 |
|
|
| ### 4.4 为什么不沿用 Spatial-AST 的 binaural 前端 |
|
|
| Spatial-AST 采用的是： |
|
|
| - 双耳 log-mel |
| - IPD |
|
|
| 这适合 binaural，不适合直接迁移到 FOA。 |
|
|
| FOA 下应优先利用： |
|
|
| - ambisonic 通道本身 |
| - intensity vector |
| - 其他 FOA 物理特征 |
|
|
| ## 5. 对 BEATs 具体修改哪些模块 |
|
|
| 下面按模块说明修改方案。 |
|
|
| ### 5.1 保留不动的模块 |
|
|
| 建议尽量保留： |
|
|
| - `TransformerEncoder` |
| - `TransformerSentenceEncoderLayer` |
| - `MultiheadAttention` |
| - `conv_pos` |
| - `LayerNorm` |
| - `FFN` |
| - `post_extract_proj` |
|
|
| 也就是 `backbone.py` 内的主干结构和 `BEATs.py` 中的 trunk 逻辑尽量不动。 |
|
|
| ### 5.2 必须修改的模块 |
|
|
| 必须重做： |
|
|
| 1. `preprocess` |
| 2. `patch_embedding` |
| 3. `extract_features` 输出头部逻辑 |
| 4. 下游 `predictor` |
|
|
| ### 5.3 推荐新增的模块 |
|
|
| 建议新增： |
|
|
| 1. `SpatialBEATsPreprocessor` |
| 2. `SpatialPatchEmbedding` |
| 3. `SpatialQueryDecoder` |
| 4. `SpatialPredictionHead` |
| 5. `SpatialTokenProjector` |
| 6. `HungarianMatcher` |
| 7. `SpatialSetCriterion` |
|
|
| ## 6. 代码级映射建议 |
|
|
| ### 6.1 现有文件建议 |
|
|
| 建议保留和复用： |
|
|
| - [BEATs.py](/apdcephfs_cq10/share_1603164/user/schmittzhu/code/unilm/beats/BEATs.py) |
| - [backbone.py](/apdcephfs_cq10/share_1603164/user/schmittzhu/code/unilm/beats/backbone.py) |
|
|
| 建议新增： |
|
|
| - `spatial_beats.py` |
| - `spatial_modules.py` |
| - `spatial_loss.py` |
| - `spatial_dataset.py` |
| - `train_spatial_beats.py` |
|
|
| ### 6.2 `spatial_beats.py` 建议包含 |
| |
| 建议实现： |
| |
| - `SpatialBEATsConfig` |
| - `SpatialBEATs` |
| - `SpatialBEATs.extract_spatial_tokens()` |
| - `SpatialBEATs.forward()` |
| |
| ### 6.3 `spatial_modules.py` 建议包含 |
|
|
| 建议实现： |
|
|
| - `SpatialBEATsPreprocessor` |
| - `SpatialPatchEmbedding` |
| - `SpatialQueryDecoder` |
| - `SpatialPredictionHead` |
| - `SpatialTokenProjector` |
|
|
| ### 6.4 `spatial_loss.py` 建议包含 |
| |
| 建议实现： |
| |
| - `HungarianMatcher` |
| - `SpatialSetCriterion` |
| |
| ## 7. 预训练权重如何复用 |
| |
| ## 7.1 默认推荐权重 |
| |
| 默认推荐： |
| |
| - `BEATs_iter3+ (AS2M) pre-trained` |
|
|
| 而不是： |
|
|
| - fine-tuned checkpoints |
|
|
| 原因： |
|
|
| - `pre-trained` 更适合作为 trunk 初始化 |
| - `fine-tuned` 更偏向 AudioSet 分类判别 |
| - 你这里的 spatial encoder 应与原语义 encoder 职责分离 |
|
|
| ### 7.2 必须直接加载的层 |
|
|
| 这些层建议直接加载原 BEATs checkpoint： |
|
|
| - `post_extract_proj` |
| - `encoder.pos_conv` |
| - `encoder.layers.*` |
| - `encoder.layer_norm` |
| - `layer_norm` |
|
|
| 即除了输入 stem 和输出头，主干参数都尽量继承。 |
|
|
| ### 7.3 需要特殊初始化的层 |
|
|
| 以下层因为 shape 不同，不能直接 strict load： |
|
|
| - `patch_embedding` |
| - 新增的 `query decoder` |
| - 新增的 `spatial heads` |
| - 新增的 `LLM projector` |
|
|
| ### 7.4 新 patch embedding 的初始化策略 |
|
|
| 原 BEATs stem 是： |
|
|
| - `Conv2d(1, embed_dim, kernel_size=patch, stride=patch)` |
|
|
| 新 stem 建议是： |
|
|
| - `Conv2d(C_foa, embed_dim, kernel_size=patch, stride=patch)` |
|
|
| 推荐初始化策略： |
|
|
| #### 方案 A：保守初始化，默认推荐 |
|
|
| - `W_logmel` 通道继承原 stem 权重 |
| - 其他空间通道初始化为 `0` 或较小随机值 |
|
|
| 优点： |
|
|
| - 最大程度保留原 BEATs 初始分布 |
| - trunk 适配更稳 |
|
|
| 缺点： |
|
|
| - 训练初期空间通道利用较慢 |
|
|
| #### 方案 B：通道 inflation |
|
|
| - 把原 stem 权重复制到全部输入通道 |
| - 再按通道数做归一化 |
|
|
| 优点： |
|
|
| - 所有通道一开始都能进入主干 |
|
|
| 缺点： |
|
|
| - 初始统计更可能偏离原 BEATs |
|
|
| 最终推荐： |
|
|
| - 第一版用 `方案 A` |
| - 后续做 ablation 再比较 `方案 B` |
|
|
| ## 8. Spatial token 模块的最终设计 |
|
|
| ### 8.1 为什么不用全局池化 |
|
|
| 原始 BEATs 的输出方式更接近： |
|
|
| - patch sequence |
| - mean pooling |
| - clip-level prediction |
|
|
| 这不适合多源空间任务。 |
|
|
| ### 8.2 最终推荐：Query Decoder |
|
|
| 在 trunk 输出后新增： |
|
|
| - `K` 个 learnable source queries |
| - 一个轻量 `cross-attention decoder` |
|
|
| 输入： |
|
|
| - encoder memory：`H in R^{B x T x D}` |
| - source queries：`Q in R^{B x K x D}` |
|
|
| 输出： |
|
|
| - `Z in R^{B x K x D}` |
|
|
| 这里的 `Z[:, i, :]` 即第 `i` 个 `source token` |
|
|
| ### 8.3 为什么 query decoder 是当前最优解 |
|
|
| 它的优点： |
|
|
| - 不改 trunk 内部结构 |
| - 仍然让完整 FOA 特征经过 backbone |
| - 适合多源 set prediction |
| - 最利于最大化复用 trunk 权重 |
|
|
| ## 9. 输出头设计 |
|
|
| 对每个 source token `z_i`，预测： |
|
|
| - `objectness` |
| - `azimuth` |
| - `elevation` |
| - `distance` |
| - 可选 `class_aux` |
|
|
| ### 9.1 离散还是连续 |
|
|
| 第一版推荐全部使用离散分类头： |
|
|
| - `azimuth`: 360 bins |
| - `elevation`: 180 bins |
| - `distance`: 按数据分桶，例如 `0.5m` 一档 |
|
|
| 原因： |
|
|
| - 与已有 Spatial-AST/BAT 经验一致 |
| - 分类头更稳 |
| - 更便于构造离散坐标 embedding |
|
|
| ### 9.2 objectness 头 |
|
|
| 推荐增加： |
|
|
| - `objectness_head: D -> 1` |
|
|
| 用于： |
|
|
| - 判断当前 token 是否对应真实声源 |
| - 作为 Hungarian matching 的一部分 |
| - 推理时做 token 保留/裁剪 |
|
|
| ### 9.3 类别头 |
|
|
| 类别头建议作为： |
|
|
| - `auxiliary head` |
|
|
| 而不是最终 LLM 的主要输入内容。 |
|
|
| 这样做的作用： |
|
|
| - 让 query token 更容易学会 source slot 对齐 |
| - 但不把 Spatial-BEATs 变成第二个强语义 encoder |
|
|
| ## 10. Loss 设计 |
|
|
| 推荐总损失： |
|
|
| ```text |
| L_total = |
| lambda_obj * L_obj |
| + lambda_azi * L_azi |
| + lambda_ele * L_ele |
| + lambda_dist * L_dist |
| + lambda_cls * L_cls_aux |
| ``` |
|
|
| ### 10.1 匹配方式 |
|
|
| 使用 `Hungarian matching`： |
|
|
| - 预测：`K` 个 token |
| - GT：`N` 个 sources |
| - 成本由以下项构成： |
| - objectness cost |
| - azimuth cost |
| - elevation cost |
| - distance cost |
| - optional class cost |
|
|
| ### 10.2 损失项定义 |
|
|
| 推荐： |
|
|
| - `L_obj`: BCE 或 focal loss |
| - `L_azi`: cross entropy |
| - `L_ele`: cross entropy |
| - `L_dist`: cross entropy |
| - `L_cls_aux`: cross entropy 或 BCE |
|
|
| ### 10.3 初始 loss 权重建议 |
|
|
| 第一版建议从以下权重起步： |
|
|
| ```text |
| lambda_obj = 1.0 |
| lambda_azi = 2.0 |
| lambda_ele = 2.0 |
| lambda_dist = 1.0 |
| lambda_cls = 0.25 |
| ``` |
|
|
| 解释： |
|
|
| - 方向任务通常更关键 |
| - 距离次之 |
| - objectness 必须稳定 |
| - 类别监督只作为辅助 |
|
|
| ### 10.4 不建议的做法 |
|
|
| 第一版不建议： |
|
|
| - 重分类损失压倒空间损失 |
| - 直接照搬 Spatial-AST 的 `1250 * cls` |
|
|
| 原因： |
|
|
| - Spatial-AST 的目标之一是保住 sound event detection |
| - 这里 `Spatial-BEATs` 的主要目标是空间 token |
| - 原项目已有独立语义 encoder |
|
|
| ## 11. 训练策略 |
|
|
| ### 11.1 第一阶段是否需要 SSL |
|
|
| 当前最终结论： |
|
|
| - 第一版 **不需要** 重新做 BEATs 式 SSL |
|
|
| 因为当前已经有： |
|
|
| - 多源监督 |
| - 每个源的空间标注 |
| - 可复用的 BEATs 主干预训练 |
|
|
| 所以第一阶段应优先做： |
|
|
| - `supervised multi-source spatial training` |
|
|
| ### 11.2 分阶段训练建议 |
|
|
| #### Stage A：Warmup |
|
|
| 冻结： |
|
|
| - 大部分 trunk |
|
|
| 只训练： |
|
|
| - FOA preprocessor |
| - patch embedding |
| - query decoder |
| - spatial heads |
| - LLM projector |
|
|
| 目的： |
|
|
| - 让新输入 stem 和新输出头稳定接入 trunk |
|
|
| #### Stage B：Upper-trunk finetune |
|
|
| 解冻： |
|
|
| - trunk 上层若干层 |
|
|
| 目的： |
|
|
| - 让主干逐步适应 FOA 空间任务 |
|
|
| #### Stage C：Near-full finetune |
|
|
| 进一步解冻： |
|
|
| - 更多 encoder layers |
|
|
| 目的： |
|
|
| - 提升空间表示上限 |
|
|
| ### 11.3 学习率建议 |
|
|
| 推荐： |
|
|
| - trunk：较小 lr |
| - 新模块：较大学习率 |
|
|
| 例如： |
|
|
| ```text |
| lr_trunk = 1e-5 ~ 5e-5 |
| lr_new = 1e-4 ~ 5e-4 |
| ``` |
|
|
| 并配合： |
|
|
| - layer-wise lr decay |
|
|
| ## 12. 最终输出给 LLM 的 spatial token 形式 |
|
|
| 这是本项目最关键的接口定义之一。 |
|
|
| ### 12.1 内部 token 形式 |
|
|
| `Spatial-BEATs` 内部输出： |
|
|
| - `Z in R^{B x K x D}` |
|
|
| 其中： |
|
|
| - `B`: batch size |
| - `K`: source token 数 |
| - `D`: Spatial-BEATs hidden dim，建议与 BEATs trunk 一致 |
|
|
| ### 12.2 不建议直接把 raw logits 喂给 LLM |
|
|
| 不建议直接给 LLM： |
|
|
| - azimuth logits |
| - elevation logits |
| - distance logits |
| - objectness logits |
|
|
| 这些是监督头，不是最终模态表示。 |
|
|
| ### 12.3 最终推荐的 LLM spatial token 形式 |
|
|
| 最终推荐送给 LLM 的每个 token 形式为： |
|
|
| ```text |
| s_i = Proj([z_i ; e_azi(i) ; e_ele(i) ; e_dist(i) ; e_obj(i)]) |
| ``` |
|
|
| 其中： |
|
|
| - `z_i`: query decoder 输出的 latent token |
| - `e_azi(i)`: 由预测 azimuth bin 查表得到的 embedding |
| - `e_ele(i)`: 由预测 elevation bin 查表得到的 embedding |
| - `e_dist(i)`: 由预测 distance bin 查表得到的 embedding |
| - `e_obj(i)`: 由 objectness/confidence 产生的 embedding |
| - `Proj`: 投影到 LLM hidden size 的 MLP/Linear |
|
|
| 最终： |
|
|
| - `s_i in R^{d_llm}` |
|
|
| ### 12.4 为什么采用“latent + structured embedding”的混合形式 |
|
|
| 原因： |
|
|
| 1. `z_i` 保留丰富的隐式空间结构信息 |
| 2. `坐标 embedding` 给 LLM 显式离散空间线索 |
| 3. `confidence` 有助于 LLM 区分可靠/不可靠 token |
|
|
| 这比单纯只传： |
|
|
| - raw latent token |
|
|
| 或者只传： |
|
|
| - 显式坐标 one-hot / scalar |
|
|
| 都更合适。 |
|
|
| ### 12.5 最终序列形式 |
|
|
| 送入 LLM 时推荐： |
|
|
| ```text |
| <SPATIAL_START>, s_1, s_2, ..., s_K, <SPATIAL_END> |
| ``` |
|
|
| 并且： |
|
|
| - 按 `objectness` 从高到低排序 |
| - 对低置信 token 可直接截断或 mask |
|
|
| ### 12.6 是否保留全部 K 个 token |
|
|
| 默认推荐： |
|
|
| - 训练时保留全部 `K` |
| - 推理时按 `objectness` 过滤 |
|
|
| 例如： |
|
|
| - 保留前 `K_keep` |
| - 或保留 `obj > threshold` 的 token |
|
|
| ## 13. 与原语义 audio encoder 的关系 |
|
|
| 为了避免“两个 encoder 在做同样的事”，推荐如下职责划分： |
|
|
| - 原语义 audio encoder：负责 `what` |
| - Spatial-BEATs：负责 `where / spatial structure / relations` |
|
|
| ### 13.1 是否允许 Spatial-BEATs 学类别 |
|
|
| 允许，但只作为辅助。 |
|
|
| 建议： |
|
|
| - 类别头只用于训练 |
| - 最终输入给 LLM 的空间 token 不直接暴露完整类别 logits |
|
|
| ### 13.2 是否需要和语义 encoder 做对齐 |
|
|
| 第一版不是必须。 |
|
|
| 若后续希望更强的 source grounding，可进一步加入： |
|
|
| - semantic distillation |
| - cross-encoder alignment |
| - source-wise contrastive loss |
|
|
| 但这些应放到第二阶段。 |
|
|
| ## 14. 第一版推荐配置 |
|
|
| 第一版默认建议： |
|
|
| - 输入特征：`WXYZ + IVxyz` |
| - `C_foa = 7` |
| - 采样率：`16k` |
| - mel bins：`128` |
| - patch 配置：与 BEATs 保持一致 |
| - 预训练权重：`BEATs_iter3+ AS2M pre-trained` |
| - trunk：最大化加载 |
| - patch stem：`W` 继承，其余通道小初始化 |
| - 输出：`K` 个 source tokens |
| - token 解码：轻量 query decoder |
| - 监督：Hungarian matching + 多头空间分类 |
| - LLM 输入：`latent + structured coordinate embedding` 的混合 token |
|
|
| ## 15. 实现优先级 |
|
|
| 推荐按如下优先级推进： |
|
|
| 1. 实现 `FOA preprocessor` |
| 2. 实现多通道 `patch embedding` |
| 3. 完成 trunk ckpt 加载 |
| 4. 实现 `query decoder` |
| 5. 实现 `objectness / azi / ele / dist` heads |
| 6. 实现 `Hungarian matcher + criterion` |
| 7. 实现 `LLM projector` |
| 8. 完成训练脚本 |
|
|
| ## 16. 当前仍需用户确认的问题 |
|
|
| 以下问题会直接影响第一版实现细节： |
|
|
| 1. `FOA` 数据当前主要采样率是多少？是 `16k`、`24k`、`32k` 还是 `48k`？ |
| 2. 每个样本中 `最大同时源数` 大概是多少？这会影响 `K` 的默认设定。 |
| 3. 每个源是否都有 `source-level class label`？如果有，类别头和匹配会更稳。 |
| 4. 你希望 `distance` 是离散分类还是连续回归？当前默认推荐离散分类。 |
| 5. 下游 LLM 的 hidden size 是多少？是否已有固定的 audio token projector？ |
| 6. 你是否希望 Spatial-BEATs 在第一版就具备一定的 source semantic 辅助能力，还是严格只做空间？ |
|
|
| ## 17. 结论 |
|
|
| 当前最终方案已经明确： |
|
|
| - **完整 FOA 特征进入 BEATs 主干** |
| - **最大化复用 trunk 预训练** |
| - **重做输入 stem** |
| - **重做输出为多源 spatial tokens** |
| - **第一版采用监督式 set prediction** |
| - **最终给 LLM 的不是 raw logits，而是融合 latent 与坐标 embedding 的 spatial tokens** |
|
|
| 这是当前最符合项目目标、也最稳妥的 `Spatial-BEATs` 方案。 |
|
|
