SNAC: Multi-Scale Neural Audio Codec
论文链接：https://arxiv.org/pdf/2410.14411
代码链接：https://github.com/hubertsiuzdak/snac
摘要
近年来，神经音频编解码器因其能够在极低比特率下高保真地表示音频信号而备受关注，这使得利用语言建模方法进行音频生成和理解成为可能。残差矢量量化 (RVQ) 已成为使用级联矢量量化 (VQ) 码本进行神经音频压缩的标准技术。本文提出了一种多尺度神经音频编解码器，它是 RVQ 的一个简单扩展，其中量化器可以以不同的时间分辨率运行。通过应用可变帧率的量化器层级结构，该编解码器能够适应跨多个时间尺度的音频结构。大量的客观和主观评估表明，这可以实现更高效的压缩。代码和模型权重已在 https://github.com/hubertsiuzdak/snac 开源。
1.介绍
近年来，神经音频压缩方法在保持音质的同时，实现了比传统编解码器更低的比特率。通过使用离散潜在变量来表示音频，这些方法已被证明在压缩之外也很有用，尤其是在生成模型中。这种方法已成功生成了自然流畅的语音和音乐连续片段。
然而，现有的音频 tokenizer 虽然能够实现高质量的重构，但其分词粒度过高，限制了它们捕捉长期结构的能力。具体而言，在 Transformer 架构中，处理扩展的注意力上下文窗口变得不切实际。理想情况下，离散的音频 token 也应该代表声音的更高层次特征。语义 token 由于其更高的信息密度而被广泛应用于音频生成系统中，这使得自回归模型能够更可靠地识别模式，这与自然语言处理类似。相比之下，高频声学 token 通常包含噪声，且模式化程度较低。
音频信号本质上涉及多个抽象层次。语音可以从局部声学或语音层面进行分析，也可以从韵律或语义层面进行分析。音乐在较长的时间尺度上也呈现出层级结构。对人类听觉皮层的研究也表明，声学信号是以层级方式处理的。
本文提出了一种名为 SNAC (Multi-Scale Neural Audio Codec), 的方法，它是对现有音频残差量化方法的简单扩展。该方法通过引入不同时间分辨率的量化，形成音频的多尺度离散表示。我们的实验（包括客观评价和主观评价）表明，所提出的方法能够实现更高效的压缩。此外，我们还通过引入噪声块、深度可分离卷积和局部窗口注意力机制，对 RVQGAN 框架进行了一些改进。
2. Related Work
Neural Audio Codecs。近年来，音频编解码器的发展利用深度学习来学习高效的音频表示，超越了传统的信号处理方法。矢量量化（VQ）将高维数据映射到离散的码本向量集，一直是压缩技术的基础工具。后来，它被应用于深度学习，并在 VQ-VAE 模型中实现了强大的离散潜在表示。VQ-VAE 在音频编码中的成功应用已得到证实。然而，由于码本大小的快速增长，VQ 在高比特率下效率会降低。这个问题早期就通过将量化器结构化为多个阶段来解决。如今，这种方法已有效地应用于现代端到端模型中，即残差矢量量化（RVQ），从而实现了具有可扩展比特率的高质量音频压缩。RVQ的进一步发展提高了神经音频压缩的效率和质量。
Multi-Scale Models。在生成模型中捕捉长期结构一直是一个挑战，针对这一问题，人们提出了多种方法。分层解码器被引入以应对长期音乐结构建模的难题。随后，人们探索了离散序列的多级模型。矢量量化码本的多尺度分层组织被提出，用于有效地对大型图像进行建模。此外，具有不同时间分辨率的独立 VQ-VAE 模型被用于表示音乐作品中更长的依赖关系。
3.Method
我们的模型基于 RVQGAN，RVQGAN 是一种编码器-解码器网络，其瓶颈部分采用了残差矢量量化 (RVQ)。它使用级联的 个矢量量化层，其中每一层将残差 映射到形状为 的 one-hot 向量序列，其中 表示帧数， 是编码器维度， 是码本维度。
Multi-Scale Residual Vector Quantization。我们的工作通过引入多尺度残差矢量量化（如图 1 所示）扩展了 RVQGAN。在每次迭代 中，我们通过因子 将残差下采样，执行码本查找，然后使用相同因子 上采样，以匹配 的原始时间分辨率 。在实际应用中，我们使用平均池化进行下采样，并使用最近邻插值进行上采样。
Noise Block。为了引入随机性并增强解码器的表达能力，我们在每个上采样层之后添加了一个噪声块。该噪声块通过更新输入来向激活值添加噪声：，其中 为高斯噪声， 表示逐元素乘法。这种机制允许模型注入与输入相关的噪声。我们发现，噪声块提高了重构质量，并能更好地利用码本。
Depthwise Convolution。深度可分离卷积的引入，旨在为视觉应用创建更轻量级的模型。通过对每个输入通道应用单个滤波器，该方法显著降低了计算量和模型大小。我们提出在生成器中使用深度可分离卷积，不仅可以减少参数数量，还可以稳定训练过程。基于生成对抗网络（GAN）的声码器稳定性极差，经常在训练初期出现梯度发散，这可能导致模型不稳定甚至崩溃。
Local Windowed Attention。在我们的模型中，编码器和解码器都包含一层时间分辨率最低的局部窗口注意力机制。这是因为注意力机制能够自适应地关注不同输入中的相关特征。此外，它还可以通过帮助捕获上下文表征来补充后续的平均池化操作。类似地，[13] 集成了 LSTM 层以更有效地建模时间依赖性。
4.Experiments
我们开展了两项研究：一项研究训练 SNAC 编解码器处理通用音频（包括音乐、音效和环境音），另一项研究则专注于语音。语音专项研究旨在以更低的比特率和更少的计算量保持音质。这种方法借鉴了传统音频编解码器，它们提供了专门用于语音编码的算法。
4.1 Model architecture
General Audio。我们的模型遵循文献 [2] 中描述的编码器-解码器架构，并在解码器中每个转置卷积之后添加了额外的噪声块。编码器和解码器均在最低时间分辨率下集成了局部窗口注意力层。除了嵌入层、输出投影层和上采样层之外，我们将大部分卷积替换为深度可分离卷积。编码器使用级联的下采样层，采样率分别为 ，解码器中相应的上采样层采样率分别为 。在 RVQ 中，我们使用 的下采样因子（步长），有效地将 44.1 kHz 的输入信号压缩成四个采样率分别为 14、29、57 和 115 Hz 的 token 序列。每个码本包含 4096 个条目（12位），总比特率为 2.6 kbps。该模型编码器包含 16M 参数，解码器包含 38.3M 参数，总计 54.5M 参数。我们采用相同的架构对 32 kHz 音频进行训练，得到的采样率分别为 10、21、42 和 83 Hz，总比特率为 1.9 kbps。
Speech。对于语音编解码器，我们通过将编码器（以及相应的解码器）中的下采样因子调整为 来修改架构。在残差矢量量化中，我们使用 的步长。该模型在 24 kHz 音频上进行训练，得到 12、23 和 47 Hz 的 token 序列，有效比特率为 984 比特/秒。此外，我们减少了卷积通道数，使得编码器中的参数数为 6.7M，解码器中的参数数为 13M，总共 19.8M 参数。我们在语音编解码器中省略了局部窗口注意力层，使架构成为纯卷积架构。
4.2 Training
我们的训练过程基本遵循文献 [2] 中描述的GAN设置，并采用了多周期判别器 [22] 和复杂的多尺度短时傅里叶变换（STFT）判别器 [23]。值得注意的是，由于使用了深度可分离卷积，我们能够为 AdamW 使用更高的初始学习率（6e-4），并以λ = 0.999994的速率逐次迭代衰减。我们没有使用梯度裁剪，因为训练过程保持稳定。
模型训练迭代次数为 80 万次，每次迭代包含 16 个时长约为 0.8 秒的音频片段。由此产生总计 2730 小时的独特音频样本，这些样本均从我们的内部数据集中无放回地抽取。我们使用通用音频数据集训练 SNAC 32 kHz 和 44 kHz 版本，其中音乐、音效和语音的采样权重分别为 80%、10% 和 10%。SNAC 24 kHz 版本则完全使用语音数据集进行训练。但需要注意的是，该数据集具有多样性，仍然可能使模型接触到非语音声音，例如播客片头音乐、广告歌曲或各种背景噪音和音效。