docs/gemini.md

Spatial-BEATs 最终实施指南 (Reference Implementation Guide)

本文档定义了 Spatial-BEATs 的模型架构、特征工程与训练流程的最终技术细节，作为代码实现的唯一参照。

1. 模型架构细节 (Architecture Specification)

1.1 输入前端 (Stem)

• 输入特征图: $7 \times 128 \times 1024$ (Channels $\times$ Mel-bins $\times$ Time-frames)。
• 通道定义:
  • [0:4]: W, X, Y, Z 的 Log-mel。
  • [4:7]: IVx, IVy, IVz (Intensity Vector)，按时间/频率对齐。
• Patch Embedding:
  • 结构: nn.Conv2d(7, embed_dim, kernel_size=16, stride=16)。
  • 初始化: 通道 0 (W) 复用 BEATs 预训练权重，通道 1-6 随机初始化。

1.2 空间 Token 提取 (Source Queries)

• Token 数量 ($K$): 4 个。
• 实现方式:
  • 定义 nn.Parameter(torch.randn(1, 4, embed_dim)) 作为 Source Queries。
  • 使用 2 层 Transformer Decoder 层。
  • Query: Source Queries。
  • Key/Value: BEATs Trunk 的输出序列 (Dense Patch Tokens)。
• 输出: 4 个维度为 embed_dim 的 Spatial Tokens。

1.3 预测头 (Prediction Heads)

每个 Spatial Token 独立连接以下 MLP 层：

• Objectness: Linear -> Sigmoid (1 unit)。
• Azimuth: Linear -> tanh (2 units: $\sin, \cos$)。计算角度使用 atan2。
• Elevation: Linear -> tanh (2 units: $\sin, \cos$)。
• Distance: Linear (1 unit, 单位：Centimeters)。
• Class: Linear -> Softmax (N units, 对应 FSD50k 类别)。

2. 坐标系与物理特征 (Spatial Physics)

2.1 坐标系 (DCASE Standard)

• 轴向: +x 前, +y 左, +z 上。
• 方位角 (Azimuth): $[-180, 180]$，逆时针增加。+90 度为左，-90 度为右。
• 仰角 (Elevation): $[-90, 90]$，向上增加。
• 距离 (Distance): 以 厘米 (cm) 为单位进行回归。

2.2 IV 计算 (Intensity Vector)

在特征提取阶段，按以下逻辑计算 IV：

• $I_x = \text{Re}{W^* \cdot X}$
• $I_y = \text{Re}{W^* \cdot Y}$
• $I_z = \text{Re}{W^* \cdot Z}$
• 所有的 $I$ 均通过 Mel 滤波器组进行映射，以匹配 Log-mel 的分辨率。

3. 训练策略 (Training Recipe)

3.1 损失函数 (Hungarian Loss)

• 匹配算法: 使用 scipy.optimize.linear_sum_assignment (Hungarian Matching) 匹配 4 个预测 Token 与 $N$ 个 GT 声源 ($N \le 4$)。
• 匹配代价 (Matching Cost): 综合位置误差 (Az/El/Dist)、类别误差和 Objectness 分数。
• 总损失:
  • 对匹配成功的 Token：计算 $L_{MSE}(pos) + L_{BCE}(obj) + L_{CrossEntropy}(cls)$。
  • 对未匹配的 Token：计算 $L_{BCE}(obj, 0)$。

3.2 训练阶段

1. Stage 1 (Stem & Head Warmup): 冻结 BEATs Trunk (Transformer 层)，仅训练新 Patch Embedding 和 Spatial Decoder/Heads。
2. Stage 2 (Joint Fine-tuning): 以 $1 \times 10^{-5}$ 的低学习率解冻整个 Trunk 进行微调。

4. LLM 接入接口 (LLM Interface)

• 提取后的 4 个 Spatial Tokens 将通过一个 Linear 投影层对齐到 LLM 的隐藏层空间。
• 在 Prompt 中，这 4 个 tokens 将按 object-wise 顺序排列，代表音频中的空间实体。