DAC-JAX: 基于JAX实现的Descript音频编码器

我们使用Google的JAX生态系统(Fjax、Optax、Orbax、AUX和CLU)实现了Descript音频编码器(DAC)。我们的 codebases 允许从原始PyTorch DAC中重用模型权重，我们确认在给定相同的输入情况下，两种实现会产生等效的标记序列和解码音频。我们提供了一个支持设备并行的训练和微调脚本，尽管我们仅使用小数据集进行了简短的训练运行进行验证。即使在有限的GPU内存下，原始DAC也可以通过将其处理为一系列重叠的“块”来压缩或解压缩长音频文件。我们在JAX中实现了该功能，并在两种类型的GPU上进行基准测试。在消费级GPU上，DAC-JAX在所有块大小下都优于原始DAC进行压缩和解压。在高性能、基于集群的GPU上，DAC-JAX对于小块大小优于原始DAC，但对于大块大小则表现更差。