BS-RoFormer，目前音频分离SOTA

Chirp
420 阅读3分钟

留下阅读 （2023）Music Source Separation with Band-Split RoPE Transformer 和（2023）Mel-Band RoFormer for Music Source Separation 的痕迹。

论文提出的模型针对音乐分离任务（Music source separation, MSS），用于分离音频中的元素，例如将音乐分离成伴奏和人声。

1. 总览

从 mvsep 的 leaderboard 来看，BS-RoFormer 和 Mel-Band RoFormer 性能领先，是目前的 SOTA（20250819）。两个模型都是字节跳动做出来的，Mel-Band RoFormer 是 BS-RoFormer 的变体。

BS-RoFormer 已经达成了 SOTA，但其频段分割方案没有理论支撑，是凭经验确定的。这个 Mel-Band RoFormer 就是来补足这个缺陷，根据 mel scale 设计出符合人类耳朵的频段分割方案。

数据集使用 MUSDB18HQ，验证方法有效性。

2. BS-RoFormer

2.1. 实现思路与损失函数

设 xx 代表输入音频，XX 是其复频域形式。

期望 XX 经过含有可学习参数 θ\theta 的网络层 fθf_\theta 预测出 cIRMs，

M^=fθ(X)\hat{M}=f_\theta(X)

借由这个 cIRMs 进行音频分离，最后就可以用逆短时傅里叶变换 iSTFT 获得预测结果 y^\hat{y}。

Y^=M^⊙X\hat{Y}=\hat{M}\odot X

什么是 cIRM？

理想比值掩蔽（Ideal Ratio Mask, IRM），可以理解为频谱的 mask。原频谱乘上 IRM 就能分离出想要的元素。

与之对应的是理想二值掩蔽（Ideal Binary Mask, IBM），这个 mask 就不是连续取值，而是只能取 0 和 1。

模型使用复数频谱，那就需要 complex IRMs，简称 cIRMs。

一句话总结：模型生成一个 mask，音频在复频域乘上这个 mask，获得推理结果。

关于损失函数，先看公式：

loss=∥y−y^∥+∑s=0S−1∥Y(s)−Y^(s)∥\text{loss}= \Vert y-\hat{y} \Vert + \sum^{S-1}_{s=0}\Vert Y^{(s)}-\hat{Y}^{(s)} \Vert

第一项很好理解，时域上的目标 yy 与预测结果 y^\hat{y} 的 MAE。

第二项是因为模型用到了多个不同窗口大小的 STFT，所以在频域上也计算 MAE。

模型的 STFT 窗口大小有 [4096, 2048, 1024, 512, 256]，步长 147，等价于每秒 300 帧。

2.2. 模型流程

输入时域音频 x∈RB×C×Lx\in\mathbb{R}^{B\times C\times L}。BB 是 batch，CC 是通道数，LL 是采样数。

首先，对 xx 进行分频与映射。

  1. 输入 xx 并 STFT 转换到复频域 X∈CB×C×T×FX\in\mathbb{C}^{B\times C\times T\times F}。TT 是帧数，FF 是频段数
  2. 对于 XX，分为 NN 个不均匀不重叠的子频段数据 Xn∈CB×C×T×FnX_n\in\mathbb{C}^{B\times C\times T\times F_n}
  3. 每个 XnX_n 经由一个由 RMSNorm 和 Linear 构成的 MLP，将 C×FnC\times F_n 维度映射为 DD，获得 Hn∈CB×T×DH_n\in\mathbb{C}^{B\times T\times D}
  4. 将所有 HnH_n 进行拼接，获得 H∈CB×T×N×DH\in\mathbb{C}^{B\times T\times N\times D}

记这个 HH 为 HlH^l（l∈[1,…,L]l\in[1, \dots, L]），然后在时间和频率两个维度分别使用 Transformer。

  1. HlH^l 变换维度到 (B×N)×T×D(B\times N)\times T\times D
  2. 施加 Transformer
  3. HlH^l 变换维度到 (B×T)×N×D(B\times T)\times N\times D
  4. 施加 Transformer
  5. 输出 Hl+1∈CB×T×N×DH^{l+1}\in\mathbb{C}^{B\times T\times N\times D}

最终获得 HLH^L，像开头对分频的处理一样进行映射，获得 cIRM。

  1. HLH^L 拥有分频 HnL∈CB×T×DH^L_n\in \mathbb{C}^{B\times T\times D}
  2. 每个 HnLH^L_n 经由一个由 RMSNorm、Linear、Tanh、Linear 和 GLU 构成的 MLP，将 DD 映射回 C×FnC\times F_n
  3. 拼接分频，获得 M^∈CB×C×T×F\hat{M}\in \mathbb{C}^{B\times C\times T\times F}

2.3. Transformer 设计

很标准的 Transformer 结构。

  • RMSNorm，采用 Pre-Norm 方案
  • 多头注意力机制
  • RoPE 位置编码
  • GELU 激活的 MLP

除此之外：

  • attention、attention 之后的 proj，以及 MLP 的上下投影，都有 dropout

2.4. 数据集与数据增强

Musdb18HQ 数据集有 100 首歌的训练集和 50 首歌的验证集，采样率 44.1kHz，每首歌分轨 vocal bass drum other 四轨。论文额外使用了自己独家的数据进行训练，有 450 首歌的训练集和 50 首歌的验证集。

评估指标为 SDR（signal-to-distortion ratio），其值越大越好。

SDR(y,y^)=10log⁡10∥y∥2∥y^−y∥2\text{SDR}(y, \hat{y})=10\log_{10}\frac{\Vert y\Vert^2}{\Vert \hat{y}-y\Vert^2}

SDR 有个变体叫做 SI-SDR（scale-invariant SDR），会在幅度上做归一化以避免 y^\hat{y} 音量变化导致分数急剧增大。

SI-SDR(y,y^)=10log⁡10∥αy∥2∥y^−αy∥2,α=⟨y^,y⟩∥y∥2\text{SI-SDR}(y, \hat{y})=10\log_{10}\frac{\Vert \alpha y\Vert^2}{\Vert \hat{y}-\alpha y\Vert^2}, \alpha=\frac{\langle\hat{y}, y\rangle}{\Vert y\Vert^2}

训练使用 8 秒的分段，每次取出所有 4 个分轨。各个分轨施加的增强技巧包含：

  • ±3 dB 的随机音量增益
  • 10% 概率替换为空波形
  • 分轨可能来自不同歌曲、同一歌曲的不同部分，xx 将会是分轨的线性和结果
  • 选取的分段音量必须大于 -50dB

2.5. 模型超参与其他细节

STFT 使用长度为 2048 的 Hann 窗口，间距 10ms。

分频总共分为了 64 段：

  • 1000Hz 以下分为 2 段
  • 1000-2000Hz 分为 4 段
  • 2000-4000Hz 分为 12 段
  • 4000-8000Hz 分为 24 段
  • 8000-16000Hz 分为 48 段
  • 16000Hz 以上分为 2 段

关于分频策略，Mel-Band RoFormer 提出了基于 mel-scale 的分段方案。见后文。

Transformer 层数 L=12L=12，8 个头，dropout=0.1，

模型超参：

  • L=12L=12，Transformer 层数
  • D=384D=384，隐空间维度
  • dropout=0.1\text{dropout}=0.1
  • MLP Scale=4\text{MLP Scale}=4

此时模型参数量为 93.4M。

训练细节：

  • 优化器 AdamW
  • 使用 EMA（Exponential Moving Averagin），decay 设置为 0.999
  • 学习率 5×10−45\times 10^{-4}，每 4 万步衰减到原来的 0.9 倍
  • 混合精度训练
  • STFT 操作保持 FP32
  • 每张 GPU 的 batch size 为 128
  • 16×A100 训练 4 周

推理时 8 秒分段，4 秒一间隔，相邻重合的 4 秒推理结果取平均。

3. Mel-Band RoFormer

从心理学上讲，人类对低频敏感对高频迟钝，换句话说就是低频部分需要更高的分辨率才更合适。BS-RoFormer 已经从此受益颇多，本文的 Mel-Band RoFormer 将在这方面做得更科学、更好。

论文使用 librosa 的 librosa.filters.mel 方法获得梅尔分频矩阵。

这个矩阵大概长这样：

[0.      0.00098 0.      0.      0.      0.      0.      0.      0.     ]
[0.      0.00025 0.00041 0.      0.      0.      0.      0.      0.     ]
[0.      0.      0.00024 0.00036 0.00012 0.      0.      0.      0.     ]
[0.      0.      0.      0.00006 0.0002  0.00023 0.00015 0.00008 0.     ]

每一行都是一个梅尔频谱滤波器。左低频右高频。

可见，滤波器都呈刺状，且低频的滤波器能量集中，高频的滤波器能量分散。

Mel-Band RoFormer 选择把这个矩阵二值化，转换为 mask。

4. 碎碎念

论文是两年前的了，模型放到现在依旧能打。很难不怀疑字节内部已经有更好的模型不想也没必要放出来，闷声发大财。恐怖如斯。

训练花了这么多资源和时间，看来我自己是没办法完全复现了。

这个方案简单直观，好评。但这么说要分 4 个轨道就需要训练 4 个不同的模型？那就不太优雅了。之后查阅代码实现时再看是怎么个情况。

看来纯频域的音频处理就能获得很好的效果，不需要像 HDemucs 那样时域频域混合，那问题简单了不少。需要注意 STFT 时保留虚部从而保留相位信息。

5. 参考来源

avatar
文章
阅读
粉丝