# 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](/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`

推荐原则：

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

## 4. 现有文件如何复用

## 4.1 [BEATs.py](/apdcephfs_cq10/share_1603164/user/schmittzhu/code/unilm/beats/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](/apdcephfs_cq10/share_1603164/user/schmittzhu/code/unilm/beats/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，至少包含：

```python
{
    "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`

总损失：

```text
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 只是为了把空间监督显式接出来